refactor(agents, loading): change tool_names from list to set

This commit changes the `tool_names` variable from a list to a set in several
places in the codebase. This is done to improve performance and readability
since the `tool_names` variable is used as a set in the code. The affected
files are `custom.py`, `prebuilt.py`, and `loading.py`.
This commit is contained in:
Gabriel Almeida 2023-05-02 10:56:01 -03:00
commit 4eea69c563
3 changed files with 7 additions and 7 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

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