diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py
index e2da2fb..0ee6428 100644
--- a/app/services/grok/services/chat.py
+++ b/app/services/grok/services/chat.py
@@ -33,7 +33,7 @@ _CHAT_SEMAPHORE = None
_CHAT_SEM_VALUE = None
-def extract_tool_text(raw: str) -> str:
+def extract_tool_text(raw: str, rollout_id: str = "") -> str:
if not raw:
return ""
name_match = re.search(
@@ -60,12 +60,13 @@ def extract_tool_text(raw: str) -> str:
label = name
text = args
+ prefix = f"[{rollout_id}]" if rollout_id else ""
if name == "web_search":
- label = "[WebSearch]"
+ label = f"{prefix}[WebSearch]"
if isinstance(payload, dict):
text = payload.get("query") or payload.get("q") or ""
elif name == "search_images":
- label = "[SearchImage]"
+ label = f"{prefix}[SearchImage]"
if isinstance(payload, dict):
text = (
payload.get("image_description")
@@ -74,7 +75,7 @@ def extract_tool_text(raw: str) -> str:
or ""
)
elif name == "chatroom_send":
- label = "[AgentThink]"
+ label = f"{prefix}[AgentThink]"
if isinstance(payload, dict):
text = payload.get("message") or ""
@@ -415,6 +416,7 @@ class StreamProcessor(proc_base.BaseProcessor):
def __init__(self, model: str, token: str = "", show_think: bool = None):
super().__init__(model, token)
self.response_id: str = None
+ self.rollout_id: str = ""
self.fingerprint: str = ""
self.think_opened: bool = False
self.role_sent: bool = False
@@ -444,7 +446,7 @@ class StreamProcessor(proc_base.BaseProcessor):
return "".join(output_parts)
end_pos = end_idx + len(end_tag)
self._tool_usage_buffer += rest[:end_pos]
- line = extract_tool_text(self._tool_usage_buffer)
+ line = extract_tool_text(self._tool_usage_buffer, self.rollout_id)
if line:
if output_parts and not output_parts[-1].endswith("\n"):
output_parts[-1] += "\n"
@@ -470,7 +472,7 @@ class StreamProcessor(proc_base.BaseProcessor):
end_pos = end_idx + len(end_tag)
raw_card = rest[start_idx:end_pos]
- line = extract_tool_text(raw_card)
+ line = extract_tool_text(raw_card, self.rollout_id)
if line:
if output_parts and not output_parts[-1].endswith("\n"):
output_parts[-1] += "\n"
@@ -554,6 +556,8 @@ class StreamProcessor(proc_base.BaseProcessor):
self.fingerprint = llm.get("modelHash", "")
if rid := resp.get("responseId"):
self.response_id = rid
+ if rid := resp.get("rolloutId"):
+ self.rollout_id = str(rid)
if not self.role_sent:
yield self._sse(role="assistant")
@@ -697,11 +701,17 @@ class CollectProcessor(proc_base.BaseProcessor):
result = content
if "xai:tool_usage_card" in self.filter_tags:
+ rollout_id = ""
+ rollout_match = re.search(
+ r"${title}