EvoScientist 0.0.1.dev2__py3-none-any.whl → 0.0.1.dev3__py3-none-any.whl

EvoScientist/stream/events.py

@@ -0,0 +1,415 @@
+"""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.
+"""
+
+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
+
+
+async def stream_agent_events(agent: Any, message: str, thread_id: str) -> 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
+
+    Yields:
+        Event dicts: thinking, text, tool_call, tool_result,
+                     subagent_start, subagent_tool_call, subagent_tool_result, subagent_end,
+                     done, error
+    """
+    config = {"configurable": {"thread_id": thread_id}}
+    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"
+
+    try:
+        async for chunk in agent.astream(
+            {"messages": [{"role": "user", "content": message}]},
+            config=config,
+            stream_mode="messages",
+            subgraphs=True,
+        ):
+            # With subgraphs=True, event is (namespace, (message, metadata))
+            namespace: tuple = ()
+            data: Any = chunk
+
+            if isinstance(chunk, tuple) and len(chunk) >= 2:
+                first = chunk[0]
+                if isinstance(first, tuple):
+                    # (namespace_tuple, (message, metadata))
+                    namespace = first
+                    data = chunk[1]
+                else:
+                    # (message, metadata) -- no namespace
+                    data = chunk
+
+            # 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
+
+            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())
+
+            # 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 = str(getattr(msg, "content", ""))
+                    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 = str(getattr(chunk, "content", ""))
+    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)