feat: send NewCheckpoint msgs on tool use generations
What does this merge request do and why?
This reverts this MR: chore: send only non-empty new-checkpoints gprc (!2938 - merged)
We don't perform a round-trip on NewCheckpoint anymore, so it should be safe to send it on token generation to keep the channel busy.
When a tool-use is streamed, we can send NewCheckpoint with the snapshot of the state made before the tool-use started the generation.
Combined with fine-grained-tool-streaming-2025-05-14 beta feature:
- It'll help to keep the channel busy while the tool-use is generated
- It the first step towards tool-use message streaming
Potential tool streaming
- Cherry-pick feat: enable fine-grained-tool-streaming-2025-0... (!3459 - merged)
- Apply the diff below
- Ask Agentic Chat to generate large tool-use message
- The message is streamed
Click to expand
diff --git a/duo_workflow_service/checkpointer/notifier.py b/duo_workflow_service/checkpointer/notifier.py
index 078fa4b48..2a00d0b6c 100644
--- a/duo_workflow_service/checkpointer/notifier.py
+++ b/duo_workflow_service/checkpointer/notifier.py
@@ -14,6 +14,7 @@ from duo_workflow_service.checkpointer.gitlab_workflow import (
)
from duo_workflow_service.entities.state import (
MessageTypeEnum,
+ ToolInfo,
UiChatLog,
WorkflowStatusEnum,
)
@@ -33,6 +34,7 @@ class UserInterface:
self.ui_chat_log: list[UiChatLog] = []
self.status = WorkflowStatusEnum.NOT_STARTED
self.steps: list[dict] = []
+ self.last_message_chunk = None
async def send_event(
self,
@@ -95,15 +97,13 @@ class UserInterface:
Args:
message (BaseMessage): The message chunk to be processed and added to the log.
"""
- content = StrOutputParser().invoke(message) or ""
- if not content:
- return
if (
not self.ui_chat_log
- or self.ui_chat_log[-1]["message_type"] != MessageTypeEnum.AGENT
+ or not self.ui_chat_log[-1]["message_type"] in [MessageTypeEnum.AGENT, MessageTypeEnum.TOOL]
or self.ui_chat_log[-1]["status"]
):
+ self.last_message_chunk = message
last_message = UiChatLog(
status=None,
correlation_id=None,
@@ -116,6 +116,12 @@ class UserInterface:
)
self.ui_chat_log.append(last_message)
else:
+ self.last_message_chunk += message
last_message = self.ui_chat_log[-1]
- last_message["content"] = last_message["content"] + content
+ last_message["content"] = StrOutputParser().invoke(self.last_message_chunk) or ""
+
+ if self.last_message_chunk.tool_calls:
+ tool_call = self.last_message_chunk.tool_calls[0]
+ last_message["tool_info"] = ToolInfo(name=tool_call["name"], args=tool_call["args"])
+ last_message["message_type"] = MessageTypeEnum.TOOL
Related to gitlab-org/gitlab#570575 (closed)
Edited by Dylan Griffith