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)