From c05ab078fae2330d13aa668386e3aa7dea96be4b Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 17:28:08 -0300 Subject: [PATCH] Refactor VectorStoreInfoComponent and VectorStoreToolkitComponent This commit refactors the VectorStoreInfoComponent and VectorStoreToolkitComponent classes. The changes include importing the VectorStoreInfo class from the correct module and adding a new parameter, llm, to the build method of the VectorStoreToolkitComponent class. --- .../langflow/components/toolkits/VectorStoreInfo.py | 13 ++----------- .../components/toolkits/VectorStoreToolkit.py | 7 ++++++- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/src/backend/langflow/components/toolkits/VectorStoreInfo.py b/src/backend/langflow/components/toolkits/VectorStoreInfo.py index 3018d975e..27a5b3792 100644 --- a/src/backend/langflow/components/toolkits/VectorStoreInfo.py +++ b/src/backend/langflow/components/toolkits/VectorStoreInfo.py @@ -2,7 +2,7 @@ from langflow import CustomComponent from langchain.vectorstores import VectorStore from typing import Union, Callable -from langflow.field_typing import Chain +from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo class VectorStoreInfoComponent(CustomComponent): display_name = "VectorStoreInfo" @@ -20,14 +20,5 @@ class VectorStoreInfoComponent(CustomComponent): vectorstore: VectorStore, description: str, name: str, - ) -> Union[Chain, Callable]: - # Since the actual implementation of VectorStoreInfo is not provided, this is a placeholder - # Replace VectorStoreInfo with the actual class that should be instantiated - # This is a hypothetical class, actual implementation may vary - class VectorStoreInfo: - def __init__(self, vectorstore, description, name): - self.vectorstore = vectorstore - self.description = description - self.name = name - + ) -> Union[VectorStoreInfo, Callable]: return VectorStoreInfo(vectorstore=vectorstore, description=description, name=name) diff --git a/src/backend/langflow/components/toolkits/VectorStoreToolkit.py b/src/backend/langflow/components/toolkits/VectorStoreToolkit.py index 236e2146f..118a06fd9 100644 --- a/src/backend/langflow/components/toolkits/VectorStoreToolkit.py +++ b/src/backend/langflow/components/toolkits/VectorStoreToolkit.py @@ -2,6 +2,9 @@ from langflow import CustomComponent from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreToolkit from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo +from langflow.field_typing import ( + BaseLanguageModel, +) from langflow.field_typing import ( Tool, ) @@ -13,10 +16,12 @@ class VectorStoreToolkitComponent(CustomComponent): def build_config(self): return { "vectorstore_info": {"display_name": "Vector Store Info"}, + "llm": {"display_name": "LLM"}, } def build( self, vectorstore_info: VectorStoreInfo, + llm: BaseLanguageModel, ) -> Tool: - return VectorStoreToolkit(vectorstore_info=vectorstore_info) + return VectorStoreToolkit(vectorstore_info=vectorstore_info,llm=llm)