From f83120bcece9f5c2ee55c18e3bd309f6428c66fb Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 17:56:42 -0300 Subject: [PATCH] Update VectorStoreAgent and VectorStoreToolkit --- .../langflow/components/agents/VectorStoreAgent.py | 10 +++++----- .../langflow/components/toolkits/VectorStoreToolkit.py | 3 ++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/components/agents/VectorStoreAgent.py b/src/backend/langflow/components/agents/VectorStoreAgent.py index 339e3f713..99424a797 100644 --- a/src/backend/langflow/components/agents/VectorStoreAgent.py +++ b/src/backend/langflow/components/agents/VectorStoreAgent.py @@ -1,7 +1,7 @@ from langflow import CustomComponent -from langchain.agents import AgentExecutor -from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo +from langchain.agents import AgentExecutor, create_vectorstore_agent +from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreToolkit from typing import Union, Callable from langflow.field_typing import BaseLanguageModel @@ -12,12 +12,12 @@ class VectorStoreAgentComponent(CustomComponent): def build_config(self): return { "llm": {"display_name": "LLM"}, - "vectorstoreinfo": {"display_name": "Vector Store Info"}, + "vector_store_toolkit": {"display_name": "Vector Store Info"}, } def build( self, llm: BaseLanguageModel, - vectorstoreinfo: VectorStoreInfo, + vector_store_toolkit: VectorStoreToolkit, ) -> Union[AgentExecutor, Callable]: - return AgentExecutor(llm=llm, vectorstore=vectorstoreinfo) + return create_vectorstore_agent(llm=llm,toolkit=vector_store_toolkit) diff --git a/src/backend/langflow/components/toolkits/VectorStoreToolkit.py b/src/backend/langflow/components/toolkits/VectorStoreToolkit.py index 118a06fd9..58529b373 100644 --- a/src/backend/langflow/components/toolkits/VectorStoreToolkit.py +++ b/src/backend/langflow/components/toolkits/VectorStoreToolkit.py @@ -8,6 +8,7 @@ from langflow.field_typing import ( from langflow.field_typing import ( Tool, ) +from typing import Union class VectorStoreToolkitComponent(CustomComponent): display_name = "VectorStoreToolkit" @@ -23,5 +24,5 @@ class VectorStoreToolkitComponent(CustomComponent): self, vectorstore_info: VectorStoreInfo, llm: BaseLanguageModel, - ) -> Tool: + ) -> Union[Tool,VectorStoreToolkit]: return VectorStoreToolkit(vectorstore_info=vectorstore_info,llm=llm)