refactor: Add 'Retriever' type alias to field_typing/constants.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-20 11:46:27 -03:00
commit 09861ce504
2 changed files with 9 additions and 4 deletions

View file

@ -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",
]

View file

@ -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,
}