From 4636776ac3f172f7a46b8e2f863c780ee4f32503 Mon Sep 17 00:00:00 2001
From: Andrei Fajardo <92402603+nerdai@users.noreply.github.com>
Date: Fri, 22 Mar 2024 23:12:42 -0400
Subject: [PATCH] fix: bring back callback_manager to AgentRunner stream_chat
(#12189)
* bring back cb stream_chat
* rm storage
---
.../llama_index/core/agent/runner/base.py | 17 ++++++++++-------
1 file changed, 10 insertions(+), 7 deletions(-)
diff --git a/llama-index-core/llama_index/core/agent/runner/base.py b/llama-index-core/llama_index/core/agent/runner/base.py
index 6cc721e84c..ee0da35300 100644
--- a/llama-index-core/llama_index/core/agent/runner/base.py
+++ b/llama-index-core/llama_index/core/agent/runner/base.py
@@ -628,6 +628,7 @@ class AgentRunner(BaseAgentRunner):
return chat_response
@dispatcher.span
+ @trace_method("chat")
def stream_chat(
self,
message: str,
@@ -637,13 +638,15 @@ class AgentRunner(BaseAgentRunner):
# override tool choice is provided as input.
if tool_choice is None:
tool_choice = self.default_tool_choice
- chat_response = self._chat(
- message=message,
- chat_history=chat_history,
- tool_choice=tool_choice,
- mode=ChatResponseMode.STREAM,
- )
- assert isinstance(chat_response, StreamingAgentChatResponse)
+ with self.callback_manager.event(
+ CBEventType.AGENT_STEP,
+ payload={EventPayload.MESSAGES: [message]},
+ ) as e:
+ chat_response = self._chat(
+ message, chat_history, tool_choice, mode=ChatResponseMode.STREAM
+ )
+ assert isinstance(chat_response, StreamingAgentChatResponse)
+ e.on_end(payload={EventPayload.RESPONSE: chat_response})
return chat_response
@trace_method("chat")
--
GitLab