Skip to content

Add error event type of v2 chat agent

Shinya Maeda requested to merge add-error-type-to-v2-chat-agent into master

This is a high priority MR for Switch to Chat Agent V2 (gitlab-org#13533 - closed). Please prioritize the review and merge.

What does this MR do and why?

This MR adds a new agent event class Error to capture the error events during the streaming.

Closes Return error event from v2/chat/agent (gitlab-org/modelops/applied-ml/code-suggestions/ai-assist#652 - closed)

AI Gateway counter-part: feat(chat): return error event from v2/chat/agent (gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!1376 - merged) (We can merge in any order)

MR acceptance checklist

Please evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Screenshots or screen recordings

Peek_2024-09-25_17-27

NOTE: The duplicate error message is a separate issue.

How to set up and validate locally

Apply the following patch in AI Gateway:

diff --git a/ai_gateway/chat/agents/react.py b/ai_gateway/chat/agents/react.py
index a64d0ebd..2ba3d52f 100644
--- a/ai_gateway/chat/agents/react.py
+++ b/ai_gateway/chat/agents/react.py
@@ -235,6 +235,7 @@ class ReActAgent(Prompt[ReActAgentInputs, TypeAgentEvent]):
         astream = super().astream(input, config=config, **kwargs)
         len_final_answer = 0
 
+        counter = 0
         try:
             async for event in astream:
                 if is_feature_enabled(FeatureFlag.EXPANDED_AI_LOGGING):
@@ -249,6 +250,11 @@ class ReActAgent(Prompt[ReActAgentInputs, TypeAgentEvent]):
 
                     len_final_answer = len(event.text)
 
+                    if counter > 10:
+                        yield AgentError(message="dummy")
+
+                    counter += 1
+
                 events.append(event)
         except Exception as e:
             yield AgentError(message=str(e))
Edited by Shinya Maeda

Merge request reports

Loading