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 committed by Gabriel Luiz Freitas Almeida
commit 4ee891711c
3 changed files with 8 additions and 8 deletions

View file

@ -51,7 +51,7 @@ class JsonAgent(AgentExecutor):
@classmethod
def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel):
tools = toolkit.get_tools()
tool_names = {tool.name for tool in tools}
tool_names = [tool.name for tool in tools]
prompt = ZeroShotAgent.create_prompt(
tools,
prefix=JSON_PREFIX,
@ -109,7 +109,7 @@ class CSVAgent(AgentExecutor):
llm=llm,
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)
return cls.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
@ -146,7 +146,7 @@ class VectorStoreAgent(AgentExecutor):
llm=llm,
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)
return AgentExecutor.from_agent_and_tools(
agent=agent, tools=tools, verbose=True
@ -212,7 +212,7 @@ class SQLAgent(AgentExecutor):
llm=llm,
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)
return AgentExecutor.from_agent_and_tools(
agent=agent,
@ -255,7 +255,7 @@ class VectorStoreRouterAgent(AgentExecutor):
llm=llm,
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)
return AgentExecutor.from_agent_and_tools(
agent=agent, tools=tools, verbose=True

View file

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

View file

@ -1,5 +1,5 @@
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 agent as agent_module
@ -180,7 +180,7 @@ def load_agent_executor_from_config(
**kwargs: Any,
):
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)
return AgentExecutor.from_agent_and_tools(