Add create_vectorstore_router_agent function and update VectorStoreRouterToolkitComponent

This commit is contained in:
anovazzi1 2024-01-16 18:04:14 -03:00
commit 758b4e443d
2 changed files with 6 additions and 4 deletions

View file

@ -2,7 +2,7 @@
from langflow import CustomComponent
from langchain_core.language_models.base import BaseLanguageModel
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreRouterToolkit
from langchain.agents import AgentExecutor
from langchain.agents import create_vectorstore_router_agent
from typing import Callable
class VectorStoreRouterAgentComponent(CustomComponent):
@ -20,4 +20,4 @@ class VectorStoreRouterAgentComponent(CustomComponent):
llm: BaseLanguageModel,
vectorstoreroutertoolkit: VectorStoreRouterToolkit
) -> Callable:
return AgentExecutor(llm=llm, toolkit=vectorstoreroutertoolkit)
return create_vectorstore_router_agent(llm=llm,toolkit=vectorstoreroutertoolkit)

View file

@ -1,6 +1,6 @@
from langflow import CustomComponent
from typing import List
from typing import List, Union
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
@ -19,5 +19,7 @@ class VectorStoreRouterToolkitComponent(CustomComponent):
self,
vectorstores: List[VectorStoreInfo],
llm: BaseLanguageModel
)->Tool:
)->Union[Tool,VectorStoreRouterToolkit]:
print("vectorstores",vectorstores)
print("llm",llm)
return VectorStoreRouterToolkit(vectorstores=vectorstores,llm=llm)