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.
This commit is contained in:
anovazzi1 2024-01-16 17:28:08 -03:00
commit c05ab078fa
2 changed files with 8 additions and 12 deletions

View file

@ -2,7 +2,7 @@
from langflow import CustomComponent from langflow import CustomComponent
from langchain.vectorstores import VectorStore from langchain.vectorstores import VectorStore
from typing import Union, Callable from typing import Union, Callable
from langflow.field_typing import Chain from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo
class VectorStoreInfoComponent(CustomComponent): class VectorStoreInfoComponent(CustomComponent):
display_name = "VectorStoreInfo" display_name = "VectorStoreInfo"
@ -20,14 +20,5 @@ class VectorStoreInfoComponent(CustomComponent):
vectorstore: VectorStore, vectorstore: VectorStore,
description: str, description: str,
name: str, name: str,
) -> Union[Chain, Callable]: ) -> Union[VectorStoreInfo, 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
return VectorStoreInfo(vectorstore=vectorstore, description=description, name=name) return VectorStoreInfo(vectorstore=vectorstore, description=description, name=name)

View file

@ -2,6 +2,9 @@
from langflow import CustomComponent from langflow import CustomComponent
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreToolkit from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreToolkit
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo
from langflow.field_typing import (
BaseLanguageModel,
)
from langflow.field_typing import ( from langflow.field_typing import (
Tool, Tool,
) )
@ -13,10 +16,12 @@ class VectorStoreToolkitComponent(CustomComponent):
def build_config(self): def build_config(self):
return { return {
"vectorstore_info": {"display_name": "Vector Store Info"}, "vectorstore_info": {"display_name": "Vector Store Info"},
"llm": {"display_name": "LLM"},
} }
def build( def build(
self, self,
vectorstore_info: VectorStoreInfo, vectorstore_info: VectorStoreInfo,
llm: BaseLanguageModel,
) -> Tool: ) -> Tool:
return VectorStoreToolkit(vectorstore_info=vectorstore_info) return VectorStoreToolkit(vectorstore_info=vectorstore_info,llm=llm)