my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import importlib
#2	import os
#3	from typing import Any, Optional
#4
#5	from langchain_cohere import ChatCohere
#6
#7	from embedchain.config import BaseLlmConfig
#8	from embedchain.helpers.json_serializable import register_deserializable
#9	from embedchain.llm.base import BaseLlm
#10
#11
#12	@register_deserializable
#13	class CohereLlm(BaseLlm):
#14	def __init__(self, config: Optional[BaseLlmConfig] = None):
#15	try:
#16	importlib.import_module("cohere")
#17	except ModuleNotFoundError:
#18	raise ModuleNotFoundError(
#19	"The required dependencies for Cohere are not installed."
#20	"Please install with `pip install langchain_cohere==1.16.0`"
#21	) from None
#22
#23	super().__init__(config=config)
#24	if not self.config.api_key and "COHERE_API_KEY" not in os.environ:
#25	raise ValueError("Please set the COHERE_API_KEY environment variable or pass it in the config.")
#26
#27	def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
#28	if self.config.system_prompt:
#29	raise ValueError("CohereLlm does not support `system_prompt`")
#30
#31	if self.config.token_usage:
#32	response, token_info = self._get_answer(prompt, self.config)
#33	model_name = "cohere/" + self.config.model
#34	if model_name not in self.config.model_pricing_map:
#35	raise ValueError(
#36	f"Model {model_name} not found in `model_prices_and_context_window.json`. \
#37	You can disable token usage by setting `token_usage` to False."
#38	)
#39	total_cost = (
#40	self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["input_tokens"]
#41	) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["output_tokens"]
#42	response_token_info = {
#43	"prompt_tokens": token_info["input_tokens"],
#44	"completion_tokens": token_info["output_tokens"],
#45	"total_tokens": token_info["input_tokens"] + token_info["output_tokens"],
#46	"total_cost": round(total_cost, 10),
#47	"cost_currency": "USD",
#48	}
#49	return response, response_token_info
#50	return self._get_answer(prompt, self.config)
#51
#52	@staticmethod
#53	def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
#54	api_key = config.api_key or os.environ["COHERE_API_KEY"]
#55	kwargs = {
#56	"model_name": config.model or "command-r",
#57	"temperature": config.temperature,
#58	"max_tokens": config.max_tokens,
#59	"together_api_key": api_key,
#60	}
#61
#62	chat = ChatCohere(**kwargs)
#63	chat_response = chat.invoke(prompt)
#64	if config.token_usage:
#65	return chat_response.content, chat_response.response_metadata["token_count"]
#66	return chat_response.content
#67

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public