🚀 feat(custom.py): add AgentType enum to improve readability and type safety
The `AgentType` enum is added to the `langchain.agents.custom` module to improve readability and type safety. The `InitializeAgent` class now uses the `AgentType` enum to ensure that the `agent` parameter is a valid value from the enum.
This commit is contained in:
parent
2b08969c0d
commit
d3a2367932
1 changed files with 4 additions and 0 deletions
|
|
@ -6,6 +6,7 @@ from langchain.agents import (
|
|||
Tool,
|
||||
ZeroShotAgent,
|
||||
initialize_agent,
|
||||
AgentType,
|
||||
)
|
||||
from langchain.agents.agent_toolkits import (
|
||||
SQLDatabaseToolkit,
|
||||
|
|
@ -297,6 +298,9 @@ class InitializeAgent(CustomAgentExecutor):
|
|||
agent: str,
|
||||
memory: Optional[BaseChatMemory] = None,
|
||||
):
|
||||
# Find which value in the AgentType enum corresponds to the string
|
||||
# passed in as agent
|
||||
agent = AgentType(agent)
|
||||
return initialize_agent(
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue