refactor(agents): change tool_names from set to list

This commit changes the `tool_names` variable from a set to a list in the `JsonAgent`, `CSVAgent`, `VectorStoreAgent`, `SQLAgent`, `VectorStoreRouterAgent`, `MalfoyAgent`, and `load_agent_executor_from_config` classes. This is done to ensure that the `tool_names` variable is always a list, even when there is only one tool.
This commit is contained in:
Gabriel Almeida 2023-05-02 15:11:32 -03:00
commit 7dddd641cf
3 changed files with 8 additions and 8 deletions

View file

@ -51,7 +51,7 @@ class JsonAgent(AgentExecutor):
@classmethod @classmethod
def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel): def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel):
tools = toolkit.get_tools() tools = toolkit.get_tools()
tool_names = {tool.name for tool in tools} tool_names = [tool.name for tool in tools]
prompt = ZeroShotAgent.create_prompt( prompt = ZeroShotAgent.create_prompt(
tools, tools,
prefix=JSON_PREFIX, prefix=JSON_PREFIX,
@ -109,7 +109,7 @@ class CSVAgent(AgentExecutor):
llm=llm, llm=llm,
prompt=partial_prompt, prompt=partial_prompt,
) )
tool_names = {tool.name for tool in tools} tool_names = [tool.name for tool in tools]
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
return cls.from_agent_and_tools(agent=agent, tools=tools, verbose=True) return cls.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
@ -146,7 +146,7 @@ class VectorStoreAgent(AgentExecutor):
llm=llm, llm=llm,
prompt=prompt, prompt=prompt,
) )
tool_names = {tool.name for tool in tools} tool_names = [tool.name for tool in tools]
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
return AgentExecutor.from_agent_and_tools( return AgentExecutor.from_agent_and_tools(
agent=agent, tools=tools, verbose=True agent=agent, tools=tools, verbose=True
@ -212,7 +212,7 @@ class SQLAgent(AgentExecutor):
llm=llm, llm=llm,
prompt=prompt, prompt=prompt,
) )
tool_names = {tool.name for tool in tools} # type: ignore tool_names = [tool.name for tool in tools] # type: ignore
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
return AgentExecutor.from_agent_and_tools( return AgentExecutor.from_agent_and_tools(
agent=agent, agent=agent,
@ -255,7 +255,7 @@ class VectorStoreRouterAgent(AgentExecutor):
llm=llm, llm=llm,
prompt=prompt, prompt=prompt,
) )
tool_names = {tool.name for tool in tools} tool_names = [tool.name for tool in tools]
agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
return AgentExecutor.from_agent_and_tools( return AgentExecutor.from_agent_and_tools(
agent=agent, tools=tools, verbose=True agent=agent, tools=tools, verbose=True

View file

@ -21,7 +21,7 @@ class MalfoyAgent(AgentExecutor):
@classmethod @classmethod
def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel): def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel):
tools = toolkit.get_tools() tools = toolkit.get_tools()
tool_names = {tool.name for tool in tools} tool_names = [tool.name for tool in tools]
prompt = ZeroShotAgent.create_prompt( prompt = ZeroShotAgent.create_prompt(
tools, tools,
prefix=JSON_PREFIX, prefix=JSON_PREFIX,

View file

@ -1,5 +1,5 @@
import json import json
from typing import Any, Callable, Dict, Iterable, Optional from typing import Any, Callable, Dict, Optional
from langchain.agents import ZeroShotAgent from langchain.agents import ZeroShotAgent
from langchain.agents import agent as agent_module from langchain.agents import agent as agent_module
@ -180,7 +180,7 @@ def load_agent_executor_from_config(
**kwargs: Any, **kwargs: Any,
): ):
tools = load_tools_from_config(config["allowed_tools"]) tools = load_tools_from_config(config["allowed_tools"])
config["allowed_tools"] = {tool.name for tool in tools} if tools else [] config["allowed_tools"] = [tool.name for tool in tools] if tools else []
agent_obj = load_agent_from_config(config, llm, tools, **kwargs) agent_obj = load_agent_from_config(config, llm, tools, **kwargs)
return AgentExecutor.from_agent_and_tools( return AgentExecutor.from_agent_and_tools(