🚀 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,
|
Tool,
|
||||||
ZeroShotAgent,
|
ZeroShotAgent,
|
||||||
initialize_agent,
|
initialize_agent,
|
||||||
|
AgentType,
|
||||||
)
|
)
|
||||||
from langchain.agents.agent_toolkits import (
|
from langchain.agents.agent_toolkits import (
|
||||||
SQLDatabaseToolkit,
|
SQLDatabaseToolkit,
|
||||||
|
|
@ -297,6 +298,9 @@ class InitializeAgent(CustomAgentExecutor):
|
||||||
agent: str,
|
agent: str,
|
||||||
memory: Optional[BaseChatMemory] = None,
|
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(
|
return initialize_agent(
|
||||||
tools=tools,
|
tools=tools,
|
||||||
llm=llm,
|
llm=llm,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue