Refactor AgentInitializerComponent to support optional memory parameter
This commit is contained in:
parent
8306a5bbd8
commit
f943ea26d6
1 changed files with 27 additions and 12 deletions
|
|
@ -1,7 +1,6 @@
|
|||
from typing import Callable, List, Union
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
from langchain.agents import AgentExecutor, AgentType, initialize_agent, types
|
||||
|
||||
from langflow import CustomComponent
|
||||
from langflow.field_typing import BaseChatMemory, BaseLanguageModel, Tool
|
||||
|
||||
|
|
@ -20,18 +19,34 @@ class AgentInitializerComponent(CustomComponent):
|
|||
"memory": {"display_name": "Memory"},
|
||||
"tools": {"display_name": "Tools"},
|
||||
"llm": {"display_name": "Language Model"},
|
||||
"code": {"advanced": True},
|
||||
}
|
||||
|
||||
def build(
|
||||
self, agent: str, llm: BaseLanguageModel, memory: BaseChatMemory, tools: List[Tool], max_iterations: int
|
||||
self,
|
||||
agent: str,
|
||||
llm: BaseLanguageModel,
|
||||
tools: List[Tool],
|
||||
max_iterations: int,
|
||||
memory: Optional[BaseChatMemory] = None,
|
||||
) -> Union[AgentExecutor, Callable]:
|
||||
agent = AgentType(agent)
|
||||
return initialize_agent(
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
agent=agent,
|
||||
memory=memory,
|
||||
return_intermediate_steps=True,
|
||||
handle_parsing_errors=True,
|
||||
max_iterations=max_iterations,
|
||||
)
|
||||
if memory:
|
||||
return initialize_agent(
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
agent=agent,
|
||||
memory=memory,
|
||||
return_intermediate_steps=True,
|
||||
handle_parsing_errors=True,
|
||||
max_iterations=max_iterations,
|
||||
)
|
||||
else:
|
||||
return initialize_agent(
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
agent=agent,
|
||||
return_intermediate_steps=True,
|
||||
handle_parsing_errors=True,
|
||||
max_iterations=max_iterations,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue