Skip to content

Commit

Permalink
update comments
Browse files Browse the repository at this point in the history
  • Loading branch information
Southpika committed Jan 9, 2024
1 parent 867ddec commit 5cb3c86
Showing 1 changed file with 8 additions and 8 deletions.
16 changes: 8 additions & 8 deletions erniebot-agent/src/erniebot_agent/agents/function_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ def __init__(
file_manager: Optional[FileManager] = None,
plugins: Optional[List[str]] = None,
max_steps: Optional[int] = None,
first_tools: Optional[List[BaseTool]] = [],
first_tools: Optional[Sequence[BaseTool]] = [],
) -> None:
"""Initialize a function agent.
Expand All @@ -96,8 +96,8 @@ def __init__(
`plugins` to `[]` to disable the use of plugins.
max_steps: The maximum number of steps in each agent run. If `None`,
use a default value.
first_tools: A list of tools that will be used first when running
the agent.
first_tools: Tools scheduled to be called sequentially at the
beginning of each agent run.
"""
super().__init__(
llm=llm,
Expand Down Expand Up @@ -135,7 +135,7 @@ async def _run(self, prompt: str, files: Optional[Sequence[File]] = None) -> Age
chat_history.append(run_input)

for tool in self._first_tools:
curr_step, new_messages = await self._step(chat_history, use_tool=tool)
curr_step, new_messages = await self._step(chat_history, selected_tool=tool)
chat_history.extend(new_messages)
num_steps_taken += 1
if not isinstance(curr_step, NoActionStep):
Expand All @@ -161,15 +161,15 @@ async def _run(self, prompt: str, files: Optional[Sequence[File]] = None) -> Age
return response

async def _step(
self, chat_history: List[Message], use_tool: Optional[BaseTool] = None
self, chat_history: List[Message], selected_tool: Optional[BaseTool] = None
) -> Tuple[AgentStep, List[Message]]:
new_messages: List[Message] = []
input_messages = self.memory.get_messages() + chat_history
if use_tool is not None:
tool_choice = {"type": "function", "function": {"name": use_tool.tool_name}}
if selected_tool is not None:
tool_choice = {"type": "function", "function": {"name": selected_tool.tool_name}}
llm_resp = await self.run_llm(
messages=input_messages,
functions=[use_tool.function_call_schema()], # only regist one tool
functions=[selected_tool.function_call_schema()], # only regist one tool
tool_choice=tool_choice,
)
else:
Expand Down

0 comments on commit 5cb3c86

Please sign in to comment.