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 <edwin.jose@datastax.com>
This commit is contained in:
Rodrigo Nader
GitHub
parent
1c974b6c60
commit
174468a281
3 changed files with 245 additions and 0 deletions
|
|
@ -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",
|
||||
|
|
|
|||
115
src/backend/base/langflow/components/models/language_model.py
Normal file
115
src/backend/base/langflow/components/models/language_model.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue