Add first_tools_rejected_callback

Author: Bobholamovic

erniebot-agent/src/erniebot_agent/agents/function_agent.py

@@ -16,7 +16,17 @@
 import logging
 from collections import deque
 from dataclasses import dataclass, replace
-from typing import Deque, Final, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import (
+    Callable,
+    Deque,
+    Final,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 from erniebot_agent.agents.agent import Agent
 from erniebot_agent.agents.callback.callback_manager import CallbackManager
@@ -82,6 +92,7 @@ def __init__(
         plugins: Optional[List[str]] = None,
         max_steps: Optional[int] = None,
         first_tools: Optional[Sequence[BaseTool]] = [],
+        first_tools_rejected_callback: Optional[Callable[[BaseTool, List[Message], AgentStep], None]] = None,
     ) -> None:
         """Initialize a function agent.
 
@@ -133,6 +144,7 @@ def __init__(
                     raise RuntimeError("The tool in `first_tools` must be in the tools list.")
         else:
             self._first_tools = []
+        self._first_tools_rejected_callback = first_tools_rejected_callback
         self._snapshots: Deque[FunctionAgentRunSnapshot] = deque(maxlen=5)
         self._snapshot_for_curr_run: Optional[FunctionAgentRunSnapshot] = None
 
@@ -170,6 +182,8 @@ async def _run(self, prompt: str, files: Optional[Sequence[File]] = None) -> Age
             else:
                 # If tool choice not work, skip this round
                 _logger.warning(f"Selected tool [{tool.tool_name}] not work")
+                if self._first_tools_rejected_callback is not None:
+                    self._first_tools_rejected_callback(tool, new_messages, curr_step)
 
         while num_steps_taken < self.max_steps:
             curr_step, new_messages = await self._step(chat_history)