my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Any, Optional
#3
#4	from embedchain.config import BaseLlmConfig
#5	from embedchain.helpers.json_serializable import register_deserializable
#6	from embedchain.llm.base import BaseLlm
#7
#8
#9	@register_deserializable
#10	class MistralAILlm(BaseLlm):
#11	def __init__(self, config: Optional[BaseLlmConfig] = None):
#12	super().__init__(config)
#13	if not self.config.api_key and "MISTRAL_API_KEY" not in os.environ:
#14	raise ValueError("Please set the MISTRAL_API_KEY environment variable or pass it in the config.")
#15
#16	def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
#17	if self.config.token_usage:
#18	response, token_info = self._get_answer(prompt, self.config)
#19	model_name = "mistralai/" + self.config.model
#20	if model_name not in self.config.model_pricing_map:
#21	raise ValueError(
#22	f"Model {model_name} not found in `model_prices_and_context_window.json`. \
#23	You can disable token usage by setting `token_usage` to False."
#24	)
#25	total_cost = (
#26	self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_tokens"]
#27	) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["completion_tokens"]
#28	response_token_info = {
#29	"prompt_tokens": token_info["prompt_tokens"],
#30	"completion_tokens": token_info["completion_tokens"],
#31	"total_tokens": token_info["prompt_tokens"] + token_info["completion_tokens"],
#32	"total_cost": round(total_cost, 10),
#33	"cost_currency": "USD",
#34	}
#35	return response, response_token_info
#36	return self._get_answer(prompt, self.config)
#37
#38	@staticmethod
#39	def _get_answer(prompt: str, config: BaseLlmConfig):
#40	try:
#41	from langchain_core.messages import HumanMessage, SystemMessage
#42	from langchain_mistralai.chat_models import ChatMistralAI
#43	except ModuleNotFoundError:
#44	raise ModuleNotFoundError(
#45	"The required dependencies for MistralAI are not installed."
#46	'Please install with `pip install --upgrade "embedchain[mistralai]"`'
#47	) from None
#48
#49	api_key = config.api_key or os.getenv("MISTRAL_API_KEY")
#50	client = ChatMistralAI(mistral_api_key=api_key)
#51	messages = []
#52	if config.system_prompt:
#53	messages.append(SystemMessage(content=config.system_prompt))
#54	messages.append(HumanMessage(content=prompt))
#55	kwargs = {
#56	"model": config.model or "mistral-tiny",
#57	"temperature": config.temperature,
#58	"max_tokens": config.max_tokens,
#59	"top_p": config.top_p,
#60	}
#61
#62	# TODO: Add support for streaming
#63	if config.stream:
#64	answer = ""
#65	for chunk in client.stream(**kwargs, input=messages):
#66	answer += chunk.content
#67	return answer
#68	else:
#69	chat_response = client.invoke(**kwargs, input=messages)
#70	if config.token_usage:
#71	return chat_response.content, chat_response.response_metadata["token_usage"]
#72	return chat_response.content
#73

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public