From 217225467f5f35395d3d9fc2627024213d6b340d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 25 Jun 2023 18:57:35 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A8=20refactor(loading.py):=20add=20ty?= =?UTF-8?q?pe=20hints=20to=20function=20parameters=20and=20return=20types?= =?UTF-8?q?=20This=20commit=20adds=20type=20hints=20to=20the=20function=20?= =?UTF-8?q?parameters=20and=20return=20types=20in=20the=20loading.py=20fil?= =?UTF-8?q?e.=20This=20improves=20the=20readability=20and=20maintainabilit?= =?UTF-8?q?y=20of=20the=20codebase=20by=20making=20it=20easier=20to=20unde?= =?UTF-8?q?rstand=20the=20expected=20types=20of=20the=20parameters=20and?= =?UTF-8?q?=20return=20values=20of=20the=20functions.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/interface/initialize/loading.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/src/backend/langflow/interface/initialize/loading.py b/src/backend/langflow/interface/initialize/loading.py index 88b981f9d..4acc21383 100644 --- a/src/backend/langflow/interface/initialize/loading.py +++ b/src/backend/langflow/interface/initialize/loading.py @@ -1,5 +1,5 @@ import json -from typing import Any, Callable, Dict, Sequence +from typing import Any, Callable, Dict, Sequence, Type from langchain.agents import ZeroShotAgent from langchain.agents import agent as agent_module @@ -16,6 +16,12 @@ from langflow.interface.toolkits.base import toolkits_creator from langflow.interface.chains.base import chain_creator from langflow.interface.utils import load_file_into_dict from langflow.utils import validate +from langchain.text_splitter import TextSplitter +from langchain.chains.base import Chain +from langchain.vectorstores.base import VectorStore +from langchain.document_loaders.base import BaseLoader +from langchain.embeddings.base import Embeddings +from langchain.prompts.base import BasePromptTemplate def instantiate_class(node_type: str, base_type: str, params: Dict) -> Any: @@ -76,7 +82,7 @@ def instantiate_based_on_type(class_object, base_type, node_type, params): return class_object(**params) -def instantiate_chains(node_type, class_object, params): +def instantiate_chains(node_type, class_object: Type[Chain], params: Dict): if "retriever" in params and hasattr(params["retriever"], "as_retriever"): params["retriever"] = params["retriever"].as_retriever() if node_type in chain_creator.from_method_nodes: @@ -88,11 +94,11 @@ def instantiate_chains(node_type, class_object, params): return class_object(**params) -def instantiate_agent(class_object, params): +def instantiate_agent(class_object: Type[Chain], params: Dict): return load_agent_executor(class_object, params) -def instantiate_prompt(node_type, class_object, params): +def instantiate_prompt(node_type, class_object: Type[BasePromptTemplate], params: Dict): if node_type == "ZeroShotPrompt": if "tools" not in params: params["tools"] = [] @@ -100,7 +106,7 @@ def instantiate_prompt(node_type, class_object, params): return class_object(**params) -def instantiate_tool(node_type, class_object, params): +def instantiate_tool(node_type, class_object: Type[BaseTool], params: Dict): if node_type == "JsonSpec": params["dict_"] = load_file_into_dict(params.pop("path")) return class_object(**params) @@ -118,7 +124,7 @@ def instantiate_tool(node_type, class_object, params): return class_object(**params) -def instantiate_toolkit(node_type, class_object, params): +def instantiate_toolkit(node_type, class_object: Type[BaseToolkit], params: Dict): loaded_toolkit = class_object(**params) # Commenting this out for now to use toolkits as normal tools # if toolkits_creator.has_create_function(node_type): @@ -128,7 +134,7 @@ def instantiate_toolkit(node_type, class_object, params): return loaded_toolkit -def instantiate_embedding(class_object, params): +def instantiate_embedding(class_object: Type[Embeddings], params: Dict): params.pop("model", None) params.pop("headers", None) try: @@ -142,7 +148,7 @@ def instantiate_embedding(class_object, params): return class_object(**params) -def instantiate_vectorstore(class_object, params): +def instantiate_vectorstore(class_object: Type[VectorStore], params: Dict): search_kwargs = params.pop("search_kwargs", {}) if initializer := vecstore_initializer.get(class_object.__name__): vecstore = initializer(class_object, params) @@ -158,7 +164,7 @@ def instantiate_vectorstore(class_object, params): return vecstore -def instantiate_documentloader(class_object, params): +def instantiate_documentloader(class_object: Type[BaseLoader], params: Dict): if "file_filter" in params: # file_filter will be a string but we need a function # that will be used to filter the files using file_filter @@ -187,7 +193,7 @@ def instantiate_documentloader(class_object, params): return docs -def instantiate_textsplitter(class_object, params): +def instantiate_textsplitter(class_object: Type[TextSplitter], params: Dict): try: documents = params.pop("documents") except KeyError as e: @@ -195,11 +201,17 @@ def instantiate_textsplitter(class_object, params): "The source you provided did not load correctly or was empty." "Try changing the chunk_size of the Text Splitter." ) from e - text_splitter = class_object(**params) + if "separator_type" in params and params["separator_type"] == "Text": + text_splitter = class_object(**params) + else: + params["language"] = params.pop("separator_type", None) + params.pop("separators", None) + text_splitter = class_object.from_language(**params) + return text_splitter.split_documents(documents) -def instantiate_utility(node_type, class_object, params): +def instantiate_utility(node_type, class_object, params: Dict): if node_type == "SQLDatabase": return class_object.from_uri(params.pop("uri")) return class_object(**params)