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 langflow import CustomComponent
from langchain_core.language_models.base import BaseLanguageModel from langchain_core.language_models.base import BaseLanguageModel
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreRouterToolkit 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 from typing import Callable
class VectorStoreRouterAgentComponent(CustomComponent): class VectorStoreRouterAgentComponent(CustomComponent):
@ -20,4 +20,4 @@ class VectorStoreRouterAgentComponent(CustomComponent):
llm: BaseLanguageModel, llm: BaseLanguageModel,
vectorstoreroutertoolkit: VectorStoreRouterToolkit vectorstoreroutertoolkit: VectorStoreRouterToolkit
) -> Callable: ) -> 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 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 VectorStoreRouterToolkit
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo
from langflow.field_typing import BaseLanguageModel,Tool from langflow.field_typing import BaseLanguageModel,Tool
@ -19,5 +19,7 @@ class VectorStoreRouterToolkitComponent(CustomComponent):
self, self,
vectorstores: List[VectorStoreInfo], vectorstores: List[VectorStoreInfo],
llm: BaseLanguageModel llm: BaseLanguageModel
)->Tool: )->Union[Tool,VectorStoreRouterToolkit]:
print("vectorstores",vectorstores)
print("llm",llm)
return VectorStoreRouterToolkit(vectorstores=vectorstores,llm=llm) return VectorStoreRouterToolkit(vectorstores=vectorstores,llm=llm)