feat: added on_agent_final_answer-support to Agent callback_manager (#…

…5736) * chore: added on_agent_final_answer-support to Agent callback_manager * chore: format black * run pre-commit to format file * updated release notes * reverted sorted imports --------- Co-authored-by: Stefano Fiorucci <[email protected]>
deepset-ai · Oct 9, 2023 · 13fb7c5 · 13fb7c5
1 parent d0ff3fa
commit 13fb7c5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/haystack/agents/base.py b/haystack/agents/base.py
@@ -244,7 +244,9 @@ def __init__(
  self.max_steps = max_steps
  self.tm = tools_manager or ToolsManager()
  self.memory = memory or NoMemory()
- self.callback_manager = Events(("on_agent_start", "on_agent_step", "on_agent_finish", "on_new_token"))
+ self.callback_manager = Events(
+ ("on_agent_start", "on_agent_step", "on_agent_finish", "on_agent_final_answer", "on_new_token")
+ )
  self.prompt_node = prompt_node
  prompt_template = prompt_template or prompt_node.default_prompt_template or "zero-shot-react"
  resolved_prompt_template = prompt_node.get_prompt_template(prompt_template)
@@ -290,8 +292,12 @@ def on_agent_start(**kwargs: Any) -> None:
  agent_name = kwargs.pop("name", "react")
  print_text(f"\nAgent {agent_name} started with {kwargs}\n")
 
+ def on_agent_final_answer(final_answer: Dict[str, Any], **kwargs: Any) -> None:
+ pass
+
  self.tm.callback_manager.on_tool_finish += on_tool_finish
  self.callback_manager.on_agent_start += on_agent_start
+ self.callback_manager.on_agent_final_answer += on_agent_final_answer
 
  if streaming:
  self.callback_manager.on_new_token += lambda token, **kwargs: print_text(token, color=agent_color)
@@ -359,7 +365,9 @@ def run(
  agent_step = self._step(query, agent_step, params)
  finally:
  self.callback_manager.on_agent_finish(agent_step)
- return agent_step.final_answer(query=query)
+ final_answer = agent_step.final_answer(query=query)
+ self.callback_manager.on_agent_final_answer(final_answer)
+ return final_answer
 
  def _step(self, query: str, current_step: AgentStep, params: Optional[dict] = None):
  # plan next step using the LLM

diff --git a/releasenotes/notes/add-on_agent_final_answer_to_agent_base-7798ea8de2f43af0.yaml b/releasenotes/notes/add-on_agent_final_answer_to_agent_base-7798ea8de2f43af0.yaml
@@ -0,0 +1,4 @@
+---
+enhancements:
+ - |
+ added support for using `on_final_answer` trough `Agent` `callback_manager`