diff --git a/src/backend/base/langflow/field_typing/__init__.py b/src/backend/base/langflow/field_typing/__init__.py index 1038e438e..57eba2f47 100644 --- a/src/backend/base/langflow/field_typing/__init__.py +++ b/src/backend/base/langflow/field_typing/__init__.py @@ -7,6 +7,7 @@ from .constants import ( BaseLLM, BaseLoader, BaseMemory, + BaseChatModel, BaseOutputParser, BasePromptTemplate, BaseRetriever, @@ -80,4 +81,5 @@ __all__ = [ "TextSplitter", "Tool", "VectorStore", + "BaseChatModel", ] diff --git a/src/backend/base/langflow/field_typing/constants.py b/src/backend/base/langflow/field_typing/constants.py index 38846ba77..3e9ca973f 100644 --- a/src/backend/base/langflow/field_typing/constants.py +++ b/src/backend/base/langflow/field_typing/constants.py @@ -1,5 +1,5 @@ from typing import Callable, Dict, Text, Union - +from langchain_core.language_models.chat_models import BaseChatModel from langchain.agents.agent import AgentExecutor from langchain.chains.base import Chain from langchain.memory.chat_memory import BaseChatMemory @@ -48,6 +48,7 @@ LANGCHAIN_BASE_TYPES = { "BaseOutputParser": BaseOutputParser, "BaseMemory": BaseMemory, "BaseChatMemory": BaseChatMemory, + "BaseChatModel": BaseChatModel, } # Langchain base types plus Python base types CUSTOM_COMPONENT_SUPPORTED_TYPES = {