Spaces:
Sleeping
Sleeping
| from langchain_mistralai import ChatMistralAI | |
| from pydantic.v1 import SecretStr | |
| from langflow.base.models.model import LCModelComponent | |
| from langflow.field_typing import LanguageModel | |
| from langflow.inputs.inputs import HandleInput | |
| from langflow.io import BoolInput, DropdownInput, FloatInput, IntInput, SecretStrInput, StrInput | |
| class MistralAIModelComponent(LCModelComponent): | |
| display_name = "MistralAI" | |
| description = "Generates text using MistralAI LLMs." | |
| icon = "MistralAI" | |
| name = "MistralModel" | |
| inputs = [ | |
| *LCModelComponent._base_inputs, | |
| IntInput( | |
| name="max_tokens", | |
| display_name="Max Tokens", | |
| advanced=True, | |
| info="The maximum number of tokens to generate. Set to 0 for unlimited tokens.", | |
| ), | |
| DropdownInput( | |
| name="model_name", | |
| display_name="Model Name", | |
| advanced=False, | |
| options=[ | |
| "open-mixtral-8x7b", | |
| "open-mixtral-8x22b", | |
| "mistral-small-latest", | |
| "mistral-medium-latest", | |
| "mistral-large-latest", | |
| "codestral-latest", | |
| ], | |
| value="codestral-latest", | |
| ), | |
| StrInput( | |
| name="mistral_api_base", | |
| display_name="Mistral API Base", | |
| advanced=True, | |
| info="The base URL of the Mistral API. Defaults to https://api.mistral.ai/v1. " | |
| "You can change this to use other APIs like JinaChat, LocalAI and Prem.", | |
| ), | |
| SecretStrInput( | |
| name="api_key", | |
| display_name="Mistral API Key", | |
| info="The Mistral API Key to use for the Mistral model.", | |
| advanced=False, | |
| ), | |
| FloatInput(name="temperature", display_name="Temperature", advanced=False, value=0.5), | |
| IntInput(name="max_retries", display_name="Max Retries", advanced=True, value=5), | |
| IntInput(name="timeout", display_name="Timeout", advanced=True, value=60), | |
| IntInput(name="max_concurrent_requests", display_name="Max Concurrent Requests", advanced=True, value=3), | |
| FloatInput(name="top_p", display_name="Top P", advanced=True, value=1), | |
| IntInput(name="random_seed", display_name="Random Seed", value=1, advanced=True), | |
| BoolInput(name="safe_mode", display_name="Safe Mode", advanced=True), | |
| HandleInput( | |
| name="output_parser", | |
| display_name="Output Parser", | |
| info="The parser to use to parse the output of the model", | |
| advanced=True, | |
| input_types=["OutputParser"], | |
| ), | |
| ] | |
| def build_model(self) -> LanguageModel: # type: ignore[type-var] | |
| mistral_api_key = self.api_key | |
| temperature = self.temperature | |
| model_name = self.model_name | |
| max_tokens = self.max_tokens | |
| mistral_api_base = self.mistral_api_base or "https://api.mistral.ai/v1" | |
| max_retries = self.max_retries | |
| timeout = self.timeout | |
| max_concurrent_requests = self.max_concurrent_requests | |
| top_p = self.top_p | |
| random_seed = self.random_seed | |
| safe_mode = self.safe_mode | |
| api_key = SecretStr(mistral_api_key).get_secret_value() if mistral_api_key else None | |
| return ChatMistralAI( | |
| max_tokens=max_tokens or None, | |
| model_name=model_name, | |
| endpoint=mistral_api_base, | |
| api_key=api_key, | |
| temperature=temperature, | |
| max_retries=max_retries, | |
| timeout=timeout, | |
| max_concurrent_requests=max_concurrent_requests, | |
| top_p=top_p, | |
| random_seed=random_seed, | |
| safe_mode=safe_mode, | |
| ) | |