diff --git a/src/backend/base/langflow/field_typing/__init__.py b/src/backend/base/langflow/field_typing/__init__.py index 57eba2f47..f967116f5 100644 --- a/src/backend/base/langflow/field_typing/__init__.py +++ b/src/backend/base/langflow/field_typing/__init__.py @@ -25,6 +25,7 @@ from .constants import ( TextSplitter, Tool, VectorStore, + Retriever, ) from .range_spec import RangeSpec @@ -82,4 +83,5 @@ __all__ = [ "Tool", "VectorStore", "BaseChatModel", + "Retriever", ] diff --git a/src/backend/base/langflow/field_typing/constants.py b/src/backend/base/langflow/field_typing/constants.py index 0fd15e4f0..feae7880d 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 typing import Callable, Dict, Text, TypeAlias, Union + from langchain.agents.agent import AgentExecutor from langchain.chains.base import Chain from langchain.memory.chat_memory import BaseChatMemory @@ -7,6 +7,7 @@ from langchain_core.document_loaders import BaseLoader from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.language_models import BaseLanguageModel, BaseLLM +from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.memory import BaseMemory from langchain_core.output_parsers import BaseOutputParser from langchain_core.prompts import BasePromptTemplate, ChatPromptTemplate, PromptTemplate @@ -16,8 +17,9 @@ from langchain_core.vectorstores import VectorStore from langchain_text_splitters import TextSplitter -NestedDict = Dict[str, Union[str, Dict]] -LanguageModel = Union[BaseLanguageModel, BaseLLM, BaseChatModel] +NestedDict: TypeAlias = Dict[str, Union[str, Dict]] +LanguageModel: TypeAlias = Union[BaseLanguageModel, BaseLLM, BaseChatModel] +Retriever: TypeAlias = BaseRetriever class Object: @@ -61,4 +63,5 @@ CUSTOM_COMPONENT_SUPPORTED_TYPES = { "Object": Object, "Callable": Callable, "LanguageModel": LanguageModel, + "Retriever": Retriever, }