From 758b4e443dd38c1d02bf2fa7a910d8d0bf041341 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 18:04:14 -0300 Subject: [PATCH] Add create_vectorstore_router_agent function and update VectorStoreRouterToolkitComponent --- .../langflow/components/agents/VectorStoreRouterAgent.py | 4 ++-- .../components/toolkits/VectorStoreRouterToolkit.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/components/agents/VectorStoreRouterAgent.py b/src/backend/langflow/components/agents/VectorStoreRouterAgent.py index d613f3080..8a3c3dc6a 100644 --- a/src/backend/langflow/components/agents/VectorStoreRouterAgent.py +++ b/src/backend/langflow/components/agents/VectorStoreRouterAgent.py @@ -2,7 +2,7 @@ from langflow import CustomComponent from langchain_core.language_models.base import BaseLanguageModel from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreRouterToolkit -from langchain.agents import AgentExecutor +from langchain.agents import create_vectorstore_router_agent from typing import Callable class VectorStoreRouterAgentComponent(CustomComponent): @@ -20,4 +20,4 @@ class VectorStoreRouterAgentComponent(CustomComponent): llm: BaseLanguageModel, vectorstoreroutertoolkit: VectorStoreRouterToolkit ) -> Callable: - return AgentExecutor(llm=llm, toolkit=vectorstoreroutertoolkit) + return create_vectorstore_router_agent(llm=llm,toolkit=vectorstoreroutertoolkit) diff --git a/src/backend/langflow/components/toolkits/VectorStoreRouterToolkit.py b/src/backend/langflow/components/toolkits/VectorStoreRouterToolkit.py index 7264f1328..fb1c66ef4 100644 --- a/src/backend/langflow/components/toolkits/VectorStoreRouterToolkit.py +++ b/src/backend/langflow/components/toolkits/VectorStoreRouterToolkit.py @@ -1,6 +1,6 @@ from langflow import CustomComponent -from typing import List +from typing import List, Union from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreRouterToolkit from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo from langflow.field_typing import BaseLanguageModel,Tool @@ -19,5 +19,7 @@ class VectorStoreRouterToolkitComponent(CustomComponent): self, vectorstores: List[VectorStoreInfo], llm: BaseLanguageModel - )->Tool: + )->Union[Tool,VectorStoreRouterToolkit]: + print("vectorstores",vectorstores) + print("llm",llm) return VectorStoreRouterToolkit(vectorstores=vectorstores,llm=llm)