Add agent step callback type (run-llama#7652)

Arize-ai · Sep 12, 2023 · 9ad04f2 · 9ad04f2
1 parent fd5c65f
commit 9ad04f2
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 16 deletions.
diff --git a/llama_index/agent/openai_agent.py b/llama_index/agent/openai_agent.py
@@ -337,10 +337,15 @@ def chat(
         chat_history: Optional[List[ChatMessage]] = None,
         function_call: Union[str, dict] = "auto",
     ) -> AgentChatResponse:
-        chat_response = self._chat(
-            message, chat_history, function_call, mode=ChatResponseMode.WAIT
-        )
-        assert isinstance(chat_response, AgentChatResponse)
+        with self.callback_manager.event(
+            CBEventType.AGENT_STEP,
+            payload={EventPayload.MESSAGES: [message]},
+        ) as e:
+            chat_response = self._chat(
+                message, chat_history, function_call, mode=ChatResponseMode.WAIT
+            )
+            assert isinstance(chat_response, AgentChatResponse)
+            e.on_end(payload={EventPayload.RESPONSE: chat_response})
         return chat_response
 
     @trace_method("chat")
@@ -350,10 +355,15 @@ async def achat(
         chat_history: Optional[List[ChatMessage]] = None,
         function_call: Union[str, dict] = "auto",
     ) -> AgentChatResponse:
-        chat_response = await self._achat(
-            message, chat_history, function_call, mode=ChatResponseMode.WAIT
-        )
-        assert isinstance(chat_response, AgentChatResponse)
+        with self.callback_manager.event(
+            CBEventType.AGENT_STEP,
+            payload={EventPayload.MESSAGES: [message]},
+        ) as e:
+            chat_response = await self._achat(
+                message, chat_history, function_call, mode=ChatResponseMode.WAIT
+            )
+            assert isinstance(chat_response, AgentChatResponse)
+            e.on_end(payload={EventPayload.RESPONSE: chat_response})
         return chat_response
 
     @trace_method("chat")
@@ -363,10 +373,15 @@ def stream_chat(
         chat_history: Optional[List[ChatMessage]] = None,
         function_call: Union[str, dict] = "auto",
     ) -> StreamingAgentChatResponse:
-        chat_response = self._chat(
-            message, chat_history, function_call, 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, function_call, mode=ChatResponseMode.STREAM
+            )
+            assert isinstance(chat_response, StreamingAgentChatResponse)
+            e.on_end(payload={EventPayload.RESPONSE: chat_response})
         return chat_response
 
     @trace_method("chat")
@@ -376,10 +391,15 @@ async def astream_chat(
         chat_history: Optional[List[ChatMessage]] = None,
         function_call: Union[str, dict] = "auto",
     ) -> StreamingAgentChatResponse:
-        chat_response = await self._achat(
-            message, chat_history, function_call, mode=ChatResponseMode.STREAM
-        )
-        assert isinstance(chat_response, StreamingAgentChatResponse)
+        with self.callback_manager.event(
+            CBEventType.AGENT_STEP,
+            payload={EventPayload.MESSAGES: [message]},
+        ) as e:
+            chat_response = await self._achat(
+                message, chat_history, function_call, mode=ChatResponseMode.STREAM
+            )
+            assert isinstance(chat_response, StreamingAgentChatResponse)
+            e.on_end(payload={EventPayload.RESPONSE: chat_response})
         return chat_response
 
 

diff --git a/llama_index/callbacks/schema.py b/llama_index/callbacks/schema.py
@@ -40,6 +40,7 @@ class CBEventType(str, Enum):
     FUNCTION_CALL = "function_call"
     RERANKING = "reranking"
     EXCEPTION = "exception"
+    AGENT_STEP = "agent_step"
 
 
 class EventPayload(str, Enum):

diff --git a/llama_index/callbacks/wandb_callback.py b/llama_index/callbacks/wandb_callback.py
@@ -387,6 +387,8 @@ def _map_event_type_to_span_kind(
             span_kind = self._trace_tree.SpanKind.LLM
         elif event_type == CBEventType.QUERY:
             span_kind = self._trace_tree.SpanKind.AGENT
+        elif event_type == CBEventType.AGENT_STEP:
+            span_kind = self._trace_tree.SpanKind.AGENT
         elif event_type == CBEventType.RETRIEVE:
             span_kind = self._trace_tree.SpanKind.TOOL
         elif event_type == CBEventType.SYNTHESIZE: