Spaces:
Running
Running
| from langchain_openai import ChatOpenAI | |
| from pydantic.v1 import SecretStr | |
| from typing_extensions import override | |
| from langflow.base.models.aiml_constants import AimlModels | |
| from langflow.base.models.model import LCModelComponent | |
| from langflow.field_typing import LanguageModel | |
| from langflow.field_typing.range_spec import RangeSpec | |
| from langflow.inputs import DictInput, DropdownInput, FloatInput, IntInput, SecretStrInput, StrInput | |
| from langflow.inputs.inputs import HandleInput | |
| class AIMLModelComponent(LCModelComponent): | |
| display_name = "AIML" | |
| description = "Generates text using AIML LLMs." | |
| icon = "AIML" | |
| name = "AIMLModel" | |
| documentation = "https://docs.aimlapi.com/api-reference" | |
| 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.", | |
| range_spec=RangeSpec(min=0, max=128000), | |
| ), | |
| DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True), | |
| DropdownInput( | |
| name="model_name", | |
| display_name="Model Name", | |
| advanced=False, | |
| options=[], | |
| refresh_button=True, | |
| ), | |
| StrInput( | |
| name="aiml_api_base", | |
| display_name="AIML API Base", | |
| advanced=True, | |
| info="The base URL of the OpenAI API. Defaults to https://api.aimlapi.com . " | |
| "You can change this to use other APIs like JinaChat, LocalAI and Prem.", | |
| ), | |
| SecretStrInput( | |
| name="api_key", | |
| display_name="AIML API Key", | |
| info="The AIML API Key to use for the OpenAI model.", | |
| advanced=False, | |
| value="AIML_API_KEY", | |
| ), | |
| FloatInput(name="temperature", display_name="Temperature", value=0.1), | |
| 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 update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None): | |
| if field_name in {"api_key", "aiml_api_base", "model_name"}: | |
| aiml = AimlModels() | |
| aiml.get_aiml_models() | |
| build_config["model_name"]["options"] = aiml.chat_models | |
| return build_config | |
| def build_model(self) -> LanguageModel: # type: ignore[type-var] | |
| aiml_api_key = self.api_key | |
| temperature = self.temperature | |
| model_name: str = self.model_name | |
| max_tokens = self.max_tokens | |
| model_kwargs = self.model_kwargs or {} | |
| aiml_api_base = self.aiml_api_base or "https://api.aimlapi.com/v2" | |
| openai_api_key = aiml_api_key.get_secret_value() if isinstance(aiml_api_key, SecretStr) else aiml_api_key | |
| # TODO: Once OpenAI fixes their o1 models, this part will need to be removed | |
| # to work correctly with o1 temperature settings. | |
| if "o1" in model_name: | |
| temperature = 1 | |
| return ChatOpenAI( | |
| model=model_name, | |
| temperature=temperature, | |
| api_key=openai_api_key, | |
| base_url=aiml_api_base, | |
| max_tokens=max_tokens or None, | |
| **model_kwargs, | |
| ) | |
| def _get_exception_message(self, e: Exception): | |
| """Get a message from an OpenAI exception. | |
| Args: | |
| e (Exception): The exception to get the message from. | |
| Returns: | |
| str: The message from the exception. | |
| """ | |
| try: | |
| from openai.error import BadRequestError | |
| except ImportError: | |
| return None | |
| if isinstance(e, BadRequestError): | |
| message = e.json_body.get("error", {}).get("message", "") | |
| if message: | |
| return message | |
| return None | |