From 5e94bfb8b27ad5d1f06200c0bd379106e2bad57e Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 9 Jan 2024 19:16:55 -0300 Subject: [PATCH] Add VectorStoreInfo import to VectorStoreAgent.py --- src/backend/langflow/components/agents/VectorStoreAgent.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/components/agents/VectorStoreAgent.py b/src/backend/langflow/components/agents/VectorStoreAgent.py index e554e0d96..339e3f713 100644 --- a/src/backend/langflow/components/agents/VectorStoreAgent.py +++ b/src/backend/langflow/components/agents/VectorStoreAgent.py @@ -1,8 +1,9 @@ from langflow import CustomComponent from langchain.agents import AgentExecutor +from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo from typing import Union, Callable -from langflow.field_typing import BaseLanguageModel, VectorStore +from langflow.field_typing import BaseLanguageModel class VectorStoreAgentComponent(CustomComponent): display_name = "VectorStoreAgent" @@ -17,6 +18,6 @@ class VectorStoreAgentComponent(CustomComponent): def build( self, llm: BaseLanguageModel, - vectorstoreinfo: VectorStore, + vectorstoreinfo: VectorStoreInfo, ) -> Union[AgentExecutor, Callable]: return AgentExecutor(llm=llm, vectorstore=vectorstoreinfo)