Add llm parameter to build method in VectorStoreRouterToolkitComponent

This commit is contained in:
anovazzi1 2024-01-16 17:33:39 -03:00
commit 00d80a4b9b

View file

@ -3,6 +3,7 @@ from langflow import CustomComponent
from typing import List
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
class VectorStoreRouterToolkitComponent(CustomComponent):
display_name = "VectorStoreRouterToolkit"
@ -11,14 +12,12 @@ class VectorStoreRouterToolkitComponent(CustomComponent):
def build_config(self):
return {
"vectorstores": {"display_name": "Vector Stores"},
"llm": {"display_name": "LLM"},
}
def build(
self,
vectorstores: List[VectorStoreInfo],
):
# Assuming the class `VectorStoreRouterToolkit` exists within a module, but since there
# is no further information provided about the module structure, I will assume it is
# accessible from the current context. If it's in `langchain.vectorstores`, it should be
# imported from there.
return VectorStoreRouterToolkit(vectorstores=vectorstores)
llm: BaseLanguageModel
)->Tool:
return VectorStoreRouterToolkit(vectorstores=vectorstores,llm=llm)