Update VectorStoreAgent and VectorStoreToolkit
This commit is contained in:
parent
00d80a4b9b
commit
f83120bcec
2 changed files with 7 additions and 6 deletions
|
|
@ -1,7 +1,7 @@
|
|||
|
||||
from langflow import CustomComponent
|
||||
from langchain.agents import AgentExecutor
|
||||
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreInfo
|
||||
from langchain.agents import AgentExecutor, create_vectorstore_agent
|
||||
from langchain.agents.agent_toolkits.vectorstore.toolkit import VectorStoreToolkit
|
||||
from typing import Union, Callable
|
||||
from langflow.field_typing import BaseLanguageModel
|
||||
|
||||
|
|
@ -12,12 +12,12 @@ class VectorStoreAgentComponent(CustomComponent):
|
|||
def build_config(self):
|
||||
return {
|
||||
"llm": {"display_name": "LLM"},
|
||||
"vectorstoreinfo": {"display_name": "Vector Store Info"},
|
||||
"vector_store_toolkit": {"display_name": "Vector Store Info"},
|
||||
}
|
||||
|
||||
def build(
|
||||
self,
|
||||
llm: BaseLanguageModel,
|
||||
vectorstoreinfo: VectorStoreInfo,
|
||||
vector_store_toolkit: VectorStoreToolkit,
|
||||
) -> Union[AgentExecutor, Callable]:
|
||||
return AgentExecutor(llm=llm, vectorstore=vectorstoreinfo)
|
||||
return create_vectorstore_agent(llm=llm,toolkit=vector_store_toolkit)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from langflow.field_typing import (
|
|||
from langflow.field_typing import (
|
||||
Tool,
|
||||
)
|
||||
from typing import Union
|
||||
|
||||
class VectorStoreToolkitComponent(CustomComponent):
|
||||
display_name = "VectorStoreToolkit"
|
||||
|
|
@ -23,5 +24,5 @@ class VectorStoreToolkitComponent(CustomComponent):
|
|||
self,
|
||||
vectorstore_info: VectorStoreInfo,
|
||||
llm: BaseLanguageModel,
|
||||
) -> Tool:
|
||||
) -> Union[Tool,VectorStoreToolkit]:
|
||||
return VectorStoreToolkit(vectorstore_info=vectorstore_info,llm=llm)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue