X-iZhang
542 lines
22 KiB
Python
542 lines
22 KiB
Python
"""Stream event generator and chunk processing helpers.
|
|
|
|
Async generator that streams events from an agent graph,
|
|
plus helpers for processing AI message chunks and tool results.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import mimetypes
|
|
import os
|
|
from typing import Any, AsyncIterator
|
|
|
|
from langchain_core.messages import AIMessage, AIMessageChunk # type: ignore[import-untyped]
|
|
|
|
from .emitter import StreamEventEmitter
|
|
from .tracker import ToolCallTracker
|
|
from .utils import DisplayLimits, is_success
|
|
|
|
# Image media types returned by DeepAgents read_file
|
|
_IMAGE_MEDIA_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp", "image/bmp", "image/svg+xml"}
|
|
|
|
|
|
def _extract_tool_content(msg) -> tuple[str, bool]:
|
|
"""Extract display-safe content from a ToolMessage.
|
|
|
|
DeepAgents ``read_file`` returns image content as
|
|
``ToolMessage(content=[ImageContentBlock])`` with
|
|
``additional_kwargs["read_file_media_type"]`` set.
|
|
Stringifying that would dump huge base64 data into the display.
|
|
|
|
Returns:
|
|
(content_string, is_image) — a short summary for images,
|
|
or the raw string content for normal results.
|
|
"""
|
|
additional = getattr(msg, "additional_kwargs", None) or {}
|
|
media_type = additional.get("read_file_media_type", "")
|
|
if media_type and media_type in _IMAGE_MEDIA_TYPES:
|
|
# Extract path from the tool call args if available
|
|
file_path = additional.get("read_file_path", "")
|
|
if not file_path:
|
|
file_path = getattr(msg, "name", "image")
|
|
return f"[OK] Image displayed: {file_path} ({media_type})", True
|
|
|
|
content = getattr(msg, "content", "")
|
|
# Guard against list-type content (image content blocks without metadata)
|
|
if isinstance(content, list):
|
|
# Check if any block looks like image data
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if block.get("type") == "image" or "base64" in block:
|
|
return "[OK] Image displayed", True
|
|
# Non-image list content — join text blocks
|
|
parts = []
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
text = block.get("text", "")
|
|
if text:
|
|
parts.append(text)
|
|
elif isinstance(block, str):
|
|
parts.append(block)
|
|
return "\n".join(parts) if parts else str(content), False
|
|
|
|
return str(content), False
|
|
|
|
|
|
async def stream_agent_events(
|
|
agent: Any,
|
|
message: str,
|
|
thread_id: str,
|
|
metadata: dict | None = None,
|
|
media: list[str] | None = None,
|
|
) -> AsyncIterator[dict]:
|
|
"""Stream events from the agent graph using async iteration.
|
|
|
|
Uses agent.astream() with subgraphs=True to see sub-agent activity.
|
|
|
|
Args:
|
|
agent: Compiled state graph from create_deep_agent()
|
|
message: User message
|
|
thread_id: Thread ID for conversation persistence
|
|
metadata: Optional metadata dict merged into the LangGraph config
|
|
(e.g. agent_name, updated_at for checkpoint persistence).
|
|
media: Optional list of local file paths for attachments.
|
|
|
|
Yields:
|
|
Event dicts: thinking, text, tool_call, tool_result,
|
|
subagent_start, subagent_tool_call, subagent_tool_result, subagent_end,
|
|
done, error
|
|
"""
|
|
config: dict[str, Any] = {"configurable": {"thread_id": thread_id}}
|
|
if metadata:
|
|
config["metadata"] = metadata
|
|
emitter = StreamEventEmitter()
|
|
main_tracker = ToolCallTracker()
|
|
full_response = ""
|
|
|
|
# Track sub-agent names
|
|
_key_to_name: dict[str, str] = {} # subagent_key -> display name (cache)
|
|
_announced_names: list[str] = [] # ordered queue of announced task names
|
|
_assigned_names: set[str] = set() # names already assigned to a namespace
|
|
_announced_task_ids: list[str] = [] # ordered task tool_call_ids
|
|
_task_id_to_name: dict[str, str] = {} # tool_call_id -> sub-agent name
|
|
_subagent_trackers: dict[str, ToolCallTracker] = {} # namespace_key -> tracker
|
|
|
|
def _register_task_tool_call(tc_data: dict) -> str | None:
|
|
"""Register or update a task tool call, return subagent name if started/updated."""
|
|
tool_id = tc_data.get("id", "")
|
|
if not tool_id:
|
|
return None
|
|
args = tc_data.get("args", {}) or {}
|
|
desc = str(args.get("description", "")).strip()
|
|
sa_name = str(args.get("subagent_type", "")).strip()
|
|
if not sa_name:
|
|
# Fallback to description snippet (may be empty during streaming)
|
|
sa_name = desc.split("\n")[0].strip()
|
|
sa_name = sa_name[:30] + "\u2026" if len(sa_name) > 30 else sa_name
|
|
if not sa_name:
|
|
sa_name = "sub-agent"
|
|
|
|
if tool_id not in _announced_task_ids:
|
|
_announced_task_ids.append(tool_id)
|
|
_announced_names.append(sa_name)
|
|
_task_id_to_name[tool_id] = sa_name
|
|
return sa_name
|
|
|
|
# Update mapping if we learned a better name later
|
|
current = _task_id_to_name.get(tool_id, "sub-agent")
|
|
if sa_name != "sub-agent" and current != sa_name:
|
|
_task_id_to_name[tool_id] = sa_name
|
|
try:
|
|
idx = _announced_task_ids.index(tool_id)
|
|
if idx < len(_announced_names):
|
|
_announced_names[idx] = sa_name
|
|
except ValueError:
|
|
pass
|
|
return sa_name
|
|
return None
|
|
|
|
def _extract_task_id(namespace: tuple) -> tuple[str | None, str | None]:
|
|
"""Extract task tool_call_id from namespace if present.
|
|
|
|
Returns (task_id, task_ns_element) or (None, None).
|
|
"""
|
|
for part in namespace:
|
|
part_str = str(part)
|
|
if "task:" in part_str:
|
|
tail = part_str.split("task:", 1)[1]
|
|
task_id = tail.split(":", 1)[0] if tail else ""
|
|
if task_id:
|
|
return task_id, part_str
|
|
return None, None
|
|
|
|
def _next_announced_name() -> str | None:
|
|
"""Get next announced name that hasn't been assigned yet."""
|
|
for announced in _announced_names:
|
|
if announced not in _assigned_names:
|
|
_assigned_names.add(announced)
|
|
return announced
|
|
return None
|
|
|
|
def _find_task_id_from_metadata(metadata: dict | None) -> str | None:
|
|
"""Try to find a task tool_call_id in metadata."""
|
|
if not metadata:
|
|
return None
|
|
candidates = (
|
|
"tool_call_id",
|
|
"task_id",
|
|
"parent_run_id",
|
|
"root_run_id",
|
|
"run_id",
|
|
)
|
|
for key in candidates:
|
|
val = metadata.get(key)
|
|
if val and val in _task_id_to_name:
|
|
return val
|
|
return None
|
|
|
|
def _get_subagent_key(namespace: tuple, metadata: dict | None) -> str | None:
|
|
"""Stable key for tracker/mapping per sub-agent namespace."""
|
|
if not namespace:
|
|
return None
|
|
task_id, task_ns = _extract_task_id(namespace)
|
|
if task_ns:
|
|
return task_ns
|
|
meta_task_id = _find_task_id_from_metadata(metadata)
|
|
if meta_task_id:
|
|
return f"task:{meta_task_id}"
|
|
if metadata:
|
|
for key in ("parent_run_id", "root_run_id", "run_id", "graph_id", "node_id"):
|
|
val = metadata.get(key)
|
|
if val:
|
|
return f"{key}:{val}"
|
|
return str(namespace)
|
|
|
|
def _get_subagent_name(namespace: tuple, metadata: dict | None) -> str | None:
|
|
"""Resolve sub-agent name from namespace, or None if main agent.
|
|
|
|
Priority:
|
|
0) metadata["lc_agent_name"] -- most reliable, set by DeepAgents framework.
|
|
1) Match task_id embedded in namespace to announced tool_call_id.
|
|
2) Use cached key mapping (only real names, never "sub-agent").
|
|
3) Queue-based: assign next announced name to this key.
|
|
4) Fallback: return "sub-agent" WITHOUT caching.
|
|
"""
|
|
if not namespace:
|
|
return None
|
|
|
|
key = _get_subagent_key(namespace, metadata) or str(namespace)
|
|
|
|
# 0) lc_agent_name from metadata -- the REAL sub-agent name
|
|
# set by the DeepAgents framework on every namespace event.
|
|
if metadata:
|
|
lc_name = metadata.get("lc_agent_name", "")
|
|
if isinstance(lc_name, str):
|
|
lc_name = lc_name.strip()
|
|
# Filter out generic/framework names
|
|
if lc_name and lc_name not in (
|
|
"sub-agent", "agent", "tools", "EvoScientist",
|
|
"LangGraph", "",
|
|
):
|
|
_key_to_name[key] = lc_name
|
|
return lc_name
|
|
|
|
# 1) Resolve by task_id if present in namespace
|
|
task_id, _task_ns = _extract_task_id(namespace)
|
|
if task_id and task_id in _task_id_to_name:
|
|
name = _task_id_to_name[task_id]
|
|
if name and name != "sub-agent":
|
|
_assigned_names.add(name)
|
|
_key_to_name[key] = name
|
|
return name
|
|
|
|
meta_task_id = _find_task_id_from_metadata(metadata)
|
|
if meta_task_id and meta_task_id in _task_id_to_name:
|
|
name = _task_id_to_name[meta_task_id]
|
|
if name and name != "sub-agent":
|
|
_assigned_names.add(name)
|
|
_key_to_name[key] = name
|
|
return name
|
|
|
|
# 2) Cached real name for this key (skip if it's "sub-agent")
|
|
cached = _key_to_name.get(key)
|
|
if cached and cached != "sub-agent":
|
|
return cached
|
|
|
|
# 3) Assign next announced name from queue (skip "sub-agent" entries)
|
|
for announced in _announced_names:
|
|
if announced not in _assigned_names and announced != "sub-agent":
|
|
_assigned_names.add(announced)
|
|
_key_to_name[key] = announced
|
|
return announced
|
|
|
|
# 4) No real names available yet -- return generic WITHOUT caching
|
|
return "sub-agent"
|
|
|
|
# Build user message content: text + inline images + file path references
|
|
user_content: str | list[dict[str, Any]] = message
|
|
if media:
|
|
_IMAGE_EXTS = frozenset({".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"})
|
|
_MAX_INLINE_SIZE = 5 * 1024 * 1024 # 5 MB
|
|
content_blocks: list[dict[str, Any]] = []
|
|
if message:
|
|
content_blocks.append({"type": "text", "text": message})
|
|
def _read_file_b64(path: str) -> str:
|
|
with open(path, "rb") as fh:
|
|
return base64.b64encode(fh.read()).decode("ascii")
|
|
|
|
file_refs: list[str] = []
|
|
for path in media:
|
|
ext = os.path.splitext(path)[1].lower()
|
|
is_image = ext in _IMAGE_EXTS and await asyncio.to_thread(os.path.isfile, path)
|
|
if is_image:
|
|
fsize = await asyncio.to_thread(os.path.getsize, path)
|
|
if fsize <= _MAX_INLINE_SIZE:
|
|
mime = mimetypes.guess_type(path)[0] or "image/png"
|
|
b64 = await asyncio.to_thread(_read_file_b64, path)
|
|
content_blocks.append({"type": "image_url", "image_url": {
|
|
"url": f"data:{mime};base64,{b64}",
|
|
}})
|
|
else:
|
|
file_refs.append(path)
|
|
else:
|
|
file_refs.append(path)
|
|
if file_refs:
|
|
ref_text = "\n".join(
|
|
f"[attached file: {os.path.basename(p)}] path: {p}" for p in file_refs
|
|
)
|
|
content_blocks.append({"type": "text", "text": ref_text})
|
|
if content_blocks:
|
|
user_content = content_blocks
|
|
|
|
try:
|
|
async for chunk in agent.astream(
|
|
{"messages": [{"role": "user", "content": user_content}]},
|
|
config=config,
|
|
stream_mode=["messages", "updates"],
|
|
subgraphs=True,
|
|
):
|
|
# Multi-mode + subgraphs: 3-tuple (namespace, mode, data)
|
|
# Single-mode + subgraphs: 2-tuple (namespace, data) — fallback
|
|
if not isinstance(chunk, tuple):
|
|
continue
|
|
|
|
namespace: tuple = ()
|
|
data: Any
|
|
mode_str: str
|
|
|
|
if len(chunk) == 3:
|
|
namespace, mode_str, data = chunk
|
|
if not isinstance(namespace, tuple):
|
|
namespace = ()
|
|
elif len(chunk) == 2:
|
|
first = chunk[0]
|
|
if isinstance(first, tuple):
|
|
namespace = first
|
|
data = chunk[1]
|
|
else:
|
|
data = chunk
|
|
mode_str = "messages"
|
|
else:
|
|
continue
|
|
|
|
# Skip non-messages modes (updates, etc.)
|
|
if mode_str == "updates":
|
|
continue
|
|
if mode_str != "messages":
|
|
continue
|
|
|
|
# Unpack message + metadata from data
|
|
msg: Any
|
|
metadata: dict = {}
|
|
if isinstance(data, tuple) and len(data) >= 2:
|
|
msg = data[0]
|
|
metadata = data[1] or {}
|
|
else:
|
|
msg = data
|
|
|
|
# Filter summarization middleware synthetic messages
|
|
if isinstance(metadata, dict) and metadata.get("lc_source") == "summarization":
|
|
continue
|
|
|
|
subagent = _get_subagent_name(namespace, metadata)
|
|
subagent_tracker = None
|
|
if subagent:
|
|
tracker_key = _get_subagent_key(namespace, metadata) or str(namespace)
|
|
subagent_tracker = _subagent_trackers.setdefault(tracker_key, ToolCallTracker())
|
|
|
|
# Extract token usage from main-agent AIMessages
|
|
if isinstance(msg, (AIMessageChunk, AIMessage)) and not subagent:
|
|
usage = getattr(msg, "usage_metadata", None)
|
|
if usage:
|
|
inp = usage.get("input_tokens", 0) if isinstance(usage, dict) else getattr(usage, "input_tokens", 0)
|
|
out = usage.get("output_tokens", 0) if isinstance(usage, dict) else getattr(usage, "output_tokens", 0)
|
|
if inp or out:
|
|
yield emitter.usage_stats(inp, out).data
|
|
|
|
# Process AIMessageChunk / AIMessage
|
|
if isinstance(msg, (AIMessageChunk, AIMessage)):
|
|
if subagent:
|
|
# Sub-agent content -- emit sub-agent events
|
|
for ev in _process_chunk_content(msg, emitter, subagent_tracker):
|
|
if ev.type == "tool_call":
|
|
yield emitter.subagent_tool_call(
|
|
subagent, ev.data["name"], ev.data["args"], ev.data.get("id", "")
|
|
).data
|
|
# Skip text/thinking from sub-agents (too noisy)
|
|
|
|
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
|
for tc in msg.tool_calls:
|
|
name = tc.get("name", "")
|
|
args = tc.get("args", {})
|
|
tool_id = tc.get("id", "")
|
|
# Skip empty-name chunks (incomplete streaming fragments)
|
|
if not name and not tool_id:
|
|
continue
|
|
yield emitter.subagent_tool_call(
|
|
subagent, name, args if isinstance(args, dict) else {}, tool_id
|
|
).data
|
|
else:
|
|
# Main agent content
|
|
for ev in _process_chunk_content(msg, emitter, main_tracker):
|
|
if ev.type == "text":
|
|
full_response += ev.data.get("content", "")
|
|
yield ev.data
|
|
|
|
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
|
for ev in _process_tool_calls(msg.tool_calls, emitter, main_tracker):
|
|
yield ev.data
|
|
# Detect task tool calls -> announce sub-agent
|
|
tc_data = ev.data
|
|
if tc_data.get("name") == "task":
|
|
started_name = _register_task_tool_call(tc_data)
|
|
if started_name:
|
|
desc = str(tc_data.get("args", {}).get("description", "")).strip()
|
|
yield emitter.subagent_start(started_name, desc).data
|
|
|
|
# Process ToolMessage (tool execution result)
|
|
elif hasattr(msg, "type") and msg.type == "tool":
|
|
if subagent:
|
|
if subagent_tracker:
|
|
subagent_tracker.finalize_all()
|
|
for info in subagent_tracker.emit_all_pending():
|
|
yield emitter.subagent_tool_call(
|
|
subagent,
|
|
info.name,
|
|
info.args,
|
|
info.id,
|
|
).data
|
|
name = getattr(msg, "name", "unknown")
|
|
raw_content, _is_img = _extract_tool_content(msg)
|
|
content = raw_content[:DisplayLimits.TOOL_RESULT_MAX]
|
|
success = is_success(content)
|
|
yield emitter.subagent_tool_result(subagent, name, content, success).data
|
|
else:
|
|
for ev in _process_tool_result(msg, emitter, main_tracker):
|
|
yield ev.data
|
|
# Tool result can re-emit tool_call with full args; update task mapping
|
|
if ev.type == "tool_call" and ev.data.get("name") == "task":
|
|
started_name = _register_task_tool_call(ev.data)
|
|
if started_name:
|
|
desc = str(ev.data.get("args", {}).get("description", "")).strip()
|
|
yield emitter.subagent_start(started_name, desc).data
|
|
# Check if this is a task result -> sub-agent ended
|
|
name = getattr(msg, "name", "")
|
|
if name == "task":
|
|
tool_call_id = getattr(msg, "tool_call_id", "")
|
|
# Find the sub-agent name via tool_call_id map
|
|
sa_name = _task_id_to_name.get(tool_call_id, "sub-agent")
|
|
yield emitter.subagent_end(sa_name).data
|
|
|
|
except Exception as e:
|
|
yield emitter.error(str(e)).data
|
|
raise
|
|
|
|
yield emitter.done(full_response).data
|
|
|
|
|
|
def _process_chunk_content(chunk, emitter: StreamEventEmitter, tracker: ToolCallTracker):
|
|
"""Process content blocks from an AI message chunk."""
|
|
content = chunk.content
|
|
|
|
if isinstance(content, str):
|
|
if content:
|
|
yield emitter.text(content)
|
|
return
|
|
|
|
blocks = None
|
|
if hasattr(chunk, "content_blocks"):
|
|
try:
|
|
blocks = chunk.content_blocks
|
|
except Exception:
|
|
blocks = None
|
|
|
|
if blocks is None:
|
|
if isinstance(content, dict):
|
|
blocks = [content]
|
|
elif isinstance(content, list):
|
|
blocks = content
|
|
else:
|
|
return
|
|
|
|
for raw_block in blocks:
|
|
block = raw_block
|
|
if not isinstance(block, dict):
|
|
if hasattr(block, "model_dump"):
|
|
block = block.model_dump()
|
|
elif hasattr(block, "dict"):
|
|
block = block.dict()
|
|
else:
|
|
continue
|
|
|
|
block_type = block.get("type")
|
|
|
|
if block_type in ("thinking", "reasoning"):
|
|
thinking_text = block.get("thinking") or block.get("reasoning") or ""
|
|
if thinking_text:
|
|
yield emitter.thinking(thinking_text)
|
|
|
|
elif block_type == "text":
|
|
text = block.get("text") or block.get("content") or ""
|
|
if text:
|
|
yield emitter.text(text)
|
|
|
|
elif block_type in ("tool_use", "tool_call"):
|
|
tool_id = block.get("id", "")
|
|
name = block.get("name", "")
|
|
args = block.get("input") if block_type == "tool_use" else block.get("args")
|
|
args_payload = args if isinstance(args, dict) else {}
|
|
|
|
if tool_id:
|
|
tracker.update(tool_id, name=name, args=args_payload)
|
|
if tracker.is_ready(tool_id):
|
|
tracker.mark_emitted(tool_id)
|
|
yield emitter.tool_call(name, args_payload, tool_id)
|
|
|
|
elif block_type == "input_json_delta":
|
|
partial_json = block.get("partial_json", "")
|
|
if partial_json:
|
|
tracker.append_json_delta(partial_json, block.get("index", 0))
|
|
|
|
elif block_type == "tool_call_chunk":
|
|
tool_id = block.get("id", "")
|
|
name = block.get("name", "")
|
|
if tool_id:
|
|
tracker.update(tool_id, name=name)
|
|
partial_args = block.get("args", "")
|
|
if isinstance(partial_args, str) and partial_args:
|
|
tracker.append_json_delta(partial_args, block.get("index", 0))
|
|
|
|
|
|
def _process_tool_calls(tool_calls: list, emitter: StreamEventEmitter, tracker: ToolCallTracker):
|
|
"""Process tool_calls from chunk.tool_calls attribute."""
|
|
for tc in tool_calls:
|
|
tool_id = tc.get("id", "")
|
|
if tool_id:
|
|
name = tc.get("name", "")
|
|
args = tc.get("args", {})
|
|
args_payload = args if isinstance(args, dict) else {}
|
|
|
|
tracker.update(tool_id, name=name, args=args_payload)
|
|
if tracker.is_ready(tool_id):
|
|
tracker.mark_emitted(tool_id)
|
|
yield emitter.tool_call(name, args_payload, tool_id)
|
|
|
|
|
|
def _process_tool_result(chunk, emitter: StreamEventEmitter, tracker: ToolCallTracker):
|
|
"""Process a ToolMessage result."""
|
|
tracker.finalize_all()
|
|
|
|
# Re-emit all tool calls with complete args
|
|
for info in tracker.get_all():
|
|
yield emitter.tool_call(info.name, info.args, info.id)
|
|
|
|
name = getattr(chunk, "name", "unknown")
|
|
raw_content, _is_img = _extract_tool_content(chunk)
|
|
content = raw_content[:DisplayLimits.TOOL_RESULT_MAX]
|
|
if len(raw_content) > DisplayLimits.TOOL_RESULT_MAX:
|
|
content += "\n... (truncated)"
|
|
|
|
success = is_success(content)
|
|
yield emitter.tool_result(name, content, success)
|