🚀 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:
Gabriel Luiz Freitas Almeida 2023-06-26 13:27:34 -03:00
commit d3a2367932

View file

@ -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,