forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
"""Load agent."""
|
|
from typing import Any, List
|
|
|
|
from langchain.agents.agent import AgentExecutor
|
|
from langchain.agents.mrkl.base import ZeroShotAgent
|
|
from langchain.agents.react.base import ReActDocstoreAgent
|
|
from langchain.agents.self_ask_with_search.base import SelfAskWithSearchAgent
|
|
from langchain.agents.tools import Tool
|
|
from langchain.llms.base import BaseLLM
|
|
|
|
AGENT_TO_CLASS = {
|
|
"zero-shot-react-description": ZeroShotAgent,
|
|
"react-docstore": ReActDocstoreAgent,
|
|
"self-ask-with-search": SelfAskWithSearchAgent,
|
|
}
|
|
|
|
|
|
def initialize_agent(
|
|
tools: List[Tool],
|
|
llm: BaseLLM,
|
|
agent: str = "zero-shot-react-description",
|
|
**kwargs: Any,
|
|
) -> AgentExecutor:
|
|
"""Load agent given tools and LLM.
|
|
|
|
Args:
|
|
tools: List of tools this agent has access to.
|
|
llm: Language model to use as the agent.
|
|
agent: The agent to use. Valid options are:
|
|
`zero-shot-react-description`, `react-docstore`, `self-ask-with-search`.
|
|
**kwargs: Additional key word arguments to pass to the agent.
|
|
|
|
Returns:
|
|
An agent.
|
|
"""
|
|
if agent not in AGENT_TO_CLASS:
|
|
raise ValueError(
|
|
f"Got unknown agent type: {agent}. "
|
|
f"Valid types are: {AGENT_TO_CLASS.keys()}."
|
|
)
|
|
agent_cls = AGENT_TO_CLASS[agent]
|
|
agent_obj = agent_cls.from_llm_and_tools(llm, tools)
|
|
return AgentExecutor.from_agent_and_tools(agent=agent_obj, tools=tools, **kwargs)
|