From 174468a281b8c56ff8ed6f02a399faa9e5806e04 Mon Sep 17 00:00:00 2001 From: Rodrigo Nader Date: Thu, 13 Mar 2025 02:03:55 -0300 Subject: [PATCH] feat: add a unified language model component. (#6994) * add a unified language model component with a few providers * [autofix.ci] apply automated fixes * fix errors and add tests * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Edwin Jose --- .../langflow/components/models/__init__.py | 2 + .../components/models/language_model.py | 115 ++++++++++++++++ .../models/test_language_model_component.py | 128 ++++++++++++++++++ 3 files changed, 245 insertions(+) create mode 100644 src/backend/base/langflow/components/models/language_model.py create mode 100644 src/backend/tests/unit/components/models/test_language_model_component.py diff --git a/src/backend/base/langflow/components/models/__init__.py b/src/backend/base/langflow/components/models/__init__.py index c96410e9e..d9bbc391d 100644 --- a/src/backend/base/langflow/components/models/__init__.py +++ b/src/backend/base/langflow/components/models/__init__.py @@ -8,6 +8,7 @@ from .deepseek import DeepSeekModelComponent from .google_generative_ai import GoogleGenerativeAIComponent from .groq import GroqModel from .huggingface import HuggingFaceEndpointsComponent +from .language_model import LanguageModelComponent from .lmstudiomodel import LMStudioModelComponent from .maritalk import MaritalkModelComponent from .mistral import MistralAIModelComponent @@ -34,6 +35,7 @@ __all__ = [ "GroqModel", "HuggingFaceEndpointsComponent", "LMStudioModelComponent", + "LanguageModelComponent", "MaritalkModelComponent", "MistralAIModelComponent", "NVIDIAModelComponent", diff --git a/src/backend/base/langflow/components/models/language_model.py b/src/backend/base/langflow/components/models/language_model.py new file mode 100644 index 000000000..c9e0f8185 --- /dev/null +++ b/src/backend/base/langflow/components/models/language_model.py @@ -0,0 +1,115 @@ +from typing import Any + +from langchain_anthropic import ChatAnthropic +from langchain_openai import ChatOpenAI + +from langflow.base.models.anthropic_constants import ANTHROPIC_MODELS +from langflow.base.models.model import LCModelComponent +from langflow.base.models.openai_constants import OPENAI_MODEL_NAMES +from langflow.field_typing import LanguageModel +from langflow.field_typing.range_spec import RangeSpec +from langflow.inputs.inputs import BoolInput +from langflow.io import DropdownInput, MessageTextInput, SecretStrInput, SliderInput +from langflow.schema.dotdict import dotdict + + +class LanguageModelComponent(LCModelComponent): + display_name = "Language Model" + description = "Runs a language model given a specified provider. " + icon = "brain-circuit" + name = "LanguageModel" + category = "models" + + inputs = [ + DropdownInput( + name="provider", + display_name="Model Provider", + options=["OpenAI", "Anthropic"], + value="OpenAI", + info="Select the model provider", + real_time_refresh=True, + options_metadata=[{"icon": "OpenAI"}, {"icon": "Anthropic"}], + ), + DropdownInput( + name="model_name", + display_name="Model Name", + options=OPENAI_MODEL_NAMES, + value=OPENAI_MODEL_NAMES[0], + info="Select the model to use", + ), + SecretStrInput( + name="api_key", + display_name="OpenAI API Key", + info="Model Provider API key", + required=False, + show=True, + real_time_refresh=True, + ), + MessageTextInput( + name="input_value", + display_name="Input", + info="The input text to send to the model", + ), + MessageTextInput( + name="system_message", + display_name="System Message", + info="A system message that helps set the behavior of the assistant", + advanced=True, + ), + BoolInput( + name="stream", + display_name="Stream", + info="Whether to stream the response", + value=False, + advanced=True, + ), + SliderInput( + name="temperature", + display_name="Temperature", + value=0.1, + info="Controls randomness in responses", + range_spec=RangeSpec(min=0, max=1, step=0.01), + advanced=True, + ), + ] + + def build_model(self) -> LanguageModel: + provider = self.provider + model_name = self.model_name + temperature = self.temperature + stream = self.stream + + if provider == "OpenAI": + if not self.api_key: + msg = "OpenAI API key is required when using OpenAI provider" + raise ValueError(msg) + return ChatOpenAI( + model_name=model_name, + temperature=temperature, + streaming=stream, + openai_api_key=self.api_key, + ) + if provider == "Anthropic": + if not self.api_key: + msg = "Anthropic API key is required when using Anthropic provider" + raise ValueError(msg) + return ChatAnthropic( + model=model_name, + temperature=temperature, + streaming=stream, + anthropic_api_key=self.api_key, + ) + msg = f"Unknown provider: {provider}" + raise ValueError(msg) + + def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None) -> dotdict: + if field_name == "provider": + if field_value == "OpenAI": + build_config["model_name"]["options"] = OPENAI_MODEL_NAMES + build_config["model_name"]["value"] = OPENAI_MODEL_NAMES[0] + build_config["api_key"]["display_name"] = "OpenAI API Key" + elif field_value == "Anthropic": + build_config["model_name"]["options"] = ANTHROPIC_MODELS + build_config["model_name"]["value"] = ANTHROPIC_MODELS[0] + build_config["api_key"]["display_name"] = "Anthropic API Key" + return build_config diff --git a/src/backend/tests/unit/components/models/test_language_model_component.py b/src/backend/tests/unit/components/models/test_language_model_component.py new file mode 100644 index 000000000..f03be6b15 --- /dev/null +++ b/src/backend/tests/unit/components/models/test_language_model_component.py @@ -0,0 +1,128 @@ +from unittest.mock import MagicMock, patch + +import pytest +from langflow.base.models.anthropic_constants import ANTHROPIC_MODELS +from langflow.base.models.openai_constants import OPENAI_MODEL_NAMES +from langflow.components.models import LanguageModelComponent + +from tests.base import ComponentTestBaseWithClient + + +@pytest.mark.usefixtures("client") +class TestLanguageModelComponent(ComponentTestBaseWithClient): + @pytest.fixture + def component_class(self): + return LanguageModelComponent + + @pytest.fixture + def default_kwargs(self): + return { + "provider": "OpenAI", + "model_name": "gpt-3.5-turbo", + "api_key": "test-api-key", + "temperature": 0.1, + "system_message": "You are a helpful assistant.", + "input_value": "Hello, how are you?", + "stream": False, + } + + @pytest.fixture + def file_names_mapping(self): + """Return the file names mapping for version-specific files.""" + + async def test_update_build_config_openai(self, component_class, default_kwargs): + component = component_class(**default_kwargs) + build_config = { + "model_name": {"options": [], "value": ""}, + "api_key": {"display_name": "API Key"}, + } + updated_config = component.update_build_config(build_config, "OpenAI", "provider") + assert updated_config["model_name"]["options"] == OPENAI_MODEL_NAMES + assert updated_config["model_name"]["value"] == OPENAI_MODEL_NAMES[0] + assert updated_config["api_key"]["display_name"] == "OpenAI API Key" + + async def test_update_build_config_anthropic(self, component_class, default_kwargs): + component = component_class(**default_kwargs) + build_config = { + "model_name": {"options": [], "value": ""}, + "api_key": {"display_name": "API Key"}, + } + updated_config = component.update_build_config(build_config, "Anthropic", "provider") + assert updated_config["model_name"]["options"] == ANTHROPIC_MODELS + assert updated_config["model_name"]["value"] == ANTHROPIC_MODELS[0] + assert updated_config["api_key"]["display_name"] == "Anthropic API Key" + + @patch("langflow.components.models.language_model.ChatOpenAI") + async def test_build_model_openai(self, mock_chat_openai, component_class, default_kwargs): + # Setup mock + mock_instance = MagicMock() + mock_chat_openai.return_value = mock_instance + + # Create and configure the component + component = component_class(**default_kwargs) + component.provider = "OpenAI" + component.model_name = "gpt-3.5-turbo" + component.api_key = "test-key" + component.temperature = 0.5 + component.stream = False + + # Build the model + model = component.build_model() + + # Verify the ChatOpenAI was called with the correct parameters + mock_chat_openai.assert_called_once_with( + model_name="gpt-3.5-turbo", + temperature=0.5, + streaming=False, + openai_api_key="test-key", + ) + assert model == mock_instance + + @patch("langflow.components.models.language_model.ChatAnthropic") + async def test_build_model_anthropic(self, mock_chat_anthropic, component_class, default_kwargs): + # Setup mock + mock_instance = MagicMock() + mock_chat_anthropic.return_value = mock_instance + + # Create and configure the component + component = component_class(**default_kwargs) + component.provider = "Anthropic" + component.model_name = ANTHROPIC_MODELS[0] # Use the first model from the constants + component.api_key = "test-key" + component.temperature = 0.7 + component.stream = False + + # Build the model + model = component.build_model() + + # Verify the ChatAnthropic was called with the correct parameters + mock_chat_anthropic.assert_called_once_with( + model=ANTHROPIC_MODELS[0], + temperature=0.7, + streaming=False, + anthropic_api_key="test-key", + ) + assert model == mock_instance + + async def test_build_model_openai_missing_api_key(self, component_class, default_kwargs): + component = component_class(**default_kwargs) + component.provider = "OpenAI" + component.api_key = None + + with pytest.raises(ValueError, match="OpenAI API key is required when using OpenAI provider"): + component.build_model() + + async def test_build_model_anthropic_missing_api_key(self, component_class, default_kwargs): + component = component_class(**default_kwargs) + component.provider = "Anthropic" + component.api_key = None + + with pytest.raises(ValueError, match="Anthropic API key is required when using Anthropic provider"): + component.build_model() + + async def test_build_model_unknown_provider(self, component_class, default_kwargs): + component = component_class(**default_kwargs) + component.provider = "Unknown" + + with pytest.raises(ValueError, match="Unknown provider: Unknown"): + component.build_model()