Merge pull request #45 from rongxinzy/techdebt/general-agent-mainline

Refactor general agent event streaming

Author: hawkli-1994

swarmmind/agents/general_agent.py

@@ -9,6 +9,7 @@
 import logging
 import uuid
 from collections.abc import AsyncGenerator, Generator
+from dataclasses import dataclass, field
 from typing import Any
 
 import deerflow.client as deerflow_client_module
@@ -26,6 +27,18 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _StreamCaptureState:
+    """Mutable capture state for a single async DeerFlow turn."""
+
+    current_chunk_msg_id: str | None = None
+    accumulated_reasoning: str = ""
+    accumulated_content: str = ""
+    final_text: str = ""
+    tool_results: list[str] = field(default_factory=list)
+    seen_ids: set[str] = field(default_factory=set)
+
+
 class SwarmMindDeerFlowClient(DeerFlowClient):
     """DeerFlow client wrapper that injects SwarmMind product identity."""
 
@@ -364,14 +377,7 @@ async def _astream_events(
         state: dict[str, Any] = {"messages": [HumanMessage(content=goal, id=current_user_message_id)]}
         runtime_context = {"thread_id": thread_id}
 
-        seen_ids: set[str] = set()
-        final_text = ""
-        tool_results: list[str] = []
-
-        # Token-level streaming accumulators (reset per LLM invocation)
-        current_chunk_msg_id: str | None = None
-        accumulated_reasoning = ""
-        accumulated_content = ""
+        capture_state = _StreamCaptureState()
 
         async for mode_tag, chunk in self._client._agent.astream(
             state,
@@ -381,133 +387,27 @@ async def _astream_events(
         ):
             if mode_tag == "messages":
                 msg_chunk, _metadata = chunk
-                if not isinstance(msg_chunk, AIMessageChunk):
-                    continue
-
-                chunk_id = getattr(msg_chunk, "id", None)
-                if chunk_id and chunk_id != current_chunk_msg_id:
-                    # New LLM invocation started; reset accumulators
-                    current_chunk_msg_id = chunk_id
-                    accumulated_reasoning = ""
-                    accumulated_content = ""
-
-                if not current_chunk_msg_id:
-                    current_chunk_msg_id = str(uuid.uuid4())
-
-                # Stream reasoning tokens
-                reasoning_delta = self._extract_reasoning_delta(msg_chunk)
-                if reasoning_delta:
-                    accumulated_reasoning += reasoning_delta
-                    yield {
-                        "type": "assistant_reasoning",
-                        "message_id": current_chunk_msg_id,
-                        "content": accumulated_reasoning,
-                    }
-
-                # Stream content tokens
-                content_delta = self._extract_content_delta(msg_chunk)
-                if content_delta:
-                    accumulated_content += content_delta
-                    yield {
-                        "type": "assistant_message",
-                        "message_id": current_chunk_msg_id,
-                        "content": accumulated_content,
-                    }
+                for event in self._process_messages_mode_chunk(msg_chunk, capture_state):
+                    yield event
 
             elif mode_tag == "custom":
-                # Handle custom events from task_tool (task_started, task_running, task_completed, task_failed)
-                event = chunk
-                logger.debug("Custom event received: %s", event)
-                if isinstance(event, dict) and event.get("type") in (
-                    "task_started",
-                    "task_running",
-                    "task_completed",
-                    "task_failed",
-                ):
-                    logger.info("Task event: type=%s, task_id=%s", event.get("type"), event.get("task_id"))
-                    yield {
-                        "type": "custom_event",
-                        "event_type": event["type"],
-                        "task_id": event.get("task_id"),
-                        "description": event.get("description"),
-                        "message": event.get("message"),
-                        "result": event.get("result"),
-                        "error": event.get("error"),
-                    }
+                event = self._process_custom_mode_chunk(chunk)
+                if event is not None:
+                    yield event
 
             elif mode_tag == "values":
                 messages = chunk.get("messages", [])
-                turn_anchor_index = next(
-                    (
-                        index
-                        for index, message in enumerate(messages)
-                        if isinstance(message, HumanMessage) and getattr(message, "id", None) == current_user_message_id
-                    ),
-                    -1,
-                )
-
-                if turn_anchor_index == -1:
-                    continue
-
-                for msg in messages[turn_anchor_index + 1 :]:
-                    if isinstance(msg, HumanMessage):
-                        continue
-
-                    msg_id = getattr(msg, "id", None)
-                    if msg_id and msg_id in seen_ids:
-                        continue
-                    if msg_id:
-                        seen_ids.add(msg_id)
-
-                    if isinstance(msg, AIMessage):
-                        # Tool calls (only from values mode for completeness)
-                        if msg.tool_calls:
-                            tool_names = [tc.get("name") for tc in msg.tool_calls]
-                            logger.info("AI tool calls: %s", tool_names)
-                            yield {
-                                "type": "assistant_tool_calls",
-                                "message_id": msg_id,
-                                "tool_calls": [
-                                    {
-                                        "name": tool_call.get("name"),
-                                        "args": tool_call.get("args", {}),
-                                        "id": tool_call.get("id"),
-                                    }
-                                    for tool_call in msg.tool_calls
-                                ],
-                            }
-
-                        # Track final text from complete messages
-                        content = self._client._extract_text(msg.content)
-                        if content:
-                            final_text = content
-
-                    elif isinstance(msg, ToolMessage):
-                        tool_name = getattr(msg, "name", None) or "unknown"
-                        tool_content = self._client._extract_text(msg.content)
-                        logger.info(
-                            "Tool result: name=%s, content_preview=%s",
-                            tool_name,
-                            tool_content[:100] if tool_content else "(empty)",
-                        )
-                        if tool_content:
-                            tool_results.append(f"[{tool_name}]: {tool_content[:200]}")
-
-                        yield {
-                            "type": "tool_result",
-                            "message_id": msg_id,
-                            "tool_name": tool_name,
-                            "tool_call_id": getattr(msg, "tool_call_id", None),
-                            "content": tool_content,
-                        }
+                for msg in self._iter_new_turn_messages(messages, current_user_message_id, capture_state.seen_ids):
+                    for event in self._process_values_mode_message(msg, capture_state):
+                        yield event
 
         # Fallback: if messages mode captured content but values mode didn't
-        if not final_text and accumulated_content:
-            final_text = accumulated_content
+        if not capture_state.final_text and capture_state.accumulated_content:
+            capture_state.final_text = capture_state.accumulated_content
 
         # Store results for the caller to retrieve
-        self._last_final_text = final_text
-        self._last_tool_results = tool_results
+        self._last_final_text = capture_state.final_text
+        self._last_tool_results = capture_state.tool_results
 
     def stream_events(
         self,
@@ -552,6 +452,157 @@ def _run_deerflow_turn(
 
         return final_text, tool_results
 
+    def _process_messages_mode_chunk(
+        self,
+        msg_chunk: object,
+        capture_state: _StreamCaptureState,
+    ) -> list[dict[str, Any]]:
+        """Convert a streaming AI chunk into accumulated reasoning/content events."""
+        if not isinstance(msg_chunk, AIMessageChunk):
+            return []
+
+        chunk_id = getattr(msg_chunk, "id", None)
+        if chunk_id and chunk_id != capture_state.current_chunk_msg_id:
+            capture_state.current_chunk_msg_id = chunk_id
+            capture_state.accumulated_reasoning = ""
+            capture_state.accumulated_content = ""
+
+        if not capture_state.current_chunk_msg_id:
+            capture_state.current_chunk_msg_id = str(uuid.uuid4())
+
+        events: list[dict[str, Any]] = []
+        reasoning_delta = self._extract_reasoning_delta(msg_chunk)
+        if reasoning_delta:
+            capture_state.accumulated_reasoning += reasoning_delta
+            events.append(
+                {
+                    "type": "assistant_reasoning",
+                    "message_id": capture_state.current_chunk_msg_id,
+                    "content": capture_state.accumulated_reasoning,
+                }
+            )
+
+        content_delta = self._extract_content_delta(msg_chunk)
+        if content_delta:
+            capture_state.accumulated_content += content_delta
+            events.append(
+                {
+                    "type": "assistant_message",
+                    "message_id": capture_state.current_chunk_msg_id,
+                    "content": capture_state.accumulated_content,
+                }
+            )
+
+        return events
+
+    @staticmethod
+    def _process_custom_mode_chunk(event: object) -> dict[str, Any] | None:
+        """Normalize supported custom task events from DeerFlow."""
+        logger.debug("Custom event received: %s", event)
+        if not isinstance(event, dict) or event.get("type") not in {
+            "task_started",
+            "task_running",
+            "task_completed",
+            "task_failed",
+        }:
+            return None
+
+        logger.info("Task event: type=%s, task_id=%s", event.get("type"), event.get("task_id"))
+        return {
+            "type": "custom_event",
+            "event_type": event["type"],
+            "task_id": event.get("task_id"),
+            "description": event.get("description"),
+            "message": event.get("message"),
+            "result": event.get("result"),
+            "error": event.get("error"),
+        }
+
+    @staticmethod
+    def _iter_new_turn_messages(
+        messages: list[object],
+        current_user_message_id: str,
+        seen_ids: set[str],
+    ) -> Generator[object, None, None]:
+        """Yield unseen non-user messages after the current turn anchor."""
+        turn_anchor_index = next(
+            (
+                index
+                for index, message in enumerate(messages)
+                if isinstance(message, HumanMessage) and getattr(message, "id", None) == current_user_message_id
+            ),
+            -1,
+        )
+        if turn_anchor_index == -1:
+            return
+
+        for msg in messages[turn_anchor_index + 1 :]:
+            if isinstance(msg, HumanMessage):
+                continue
+
+            msg_id = getattr(msg, "id", None)
+            if msg_id and msg_id in seen_ids:
+                continue
+            if msg_id:
+                seen_ids.add(msg_id)
+            yield msg
+
+    def _process_values_mode_message(
+        self,
+        msg: object,
+        capture_state: _StreamCaptureState,
+    ) -> list[dict[str, Any]]:
+        """Convert full values-mode messages into runtime events and summaries."""
+        msg_id = getattr(msg, "id", None)
+
+        if isinstance(msg, AIMessage):
+            events: list[dict[str, Any]] = []
+            if msg.tool_calls:
+                tool_names = [tc.get("name") for tc in msg.tool_calls]
+                logger.info("AI tool calls: %s", tool_names)
+                events.append(
+                    {
+                        "type": "assistant_tool_calls",
+                        "message_id": msg_id,
+                        "tool_calls": [
+                            {
+                                "name": tool_call.get("name"),
+                                "args": tool_call.get("args", {}),
+                                "id": tool_call.get("id"),
+                            }
+                            for tool_call in msg.tool_calls
+                        ],
+                    }
+                )
+
+            content = self._client._extract_text(msg.content)
+            if content:
+                capture_state.final_text = content
+            return events
+
+        if isinstance(msg, ToolMessage):
+            tool_name = getattr(msg, "name", None) or "unknown"
+            tool_content = self._client._extract_text(msg.content)
+            logger.info(
+                "Tool result: name=%s, content_preview=%s",
+                tool_name,
+                tool_content[:100] if tool_content else "(empty)",
+            )
+            if tool_content:
+                capture_state.tool_results.append(f"[{tool_name}]: {tool_content[:200]}")
+
+            return [
+                {
+                    "type": "tool_result",
+                    "message_id": msg_id,
+                    "tool_name": tool_name,
+                    "tool_call_id": getattr(msg, "tool_call_id", None),
+                    "content": tool_content,
+                }
+            ]
+
+        return []
+
     def _resolve_runtime_options(
         self,
         runtime_options: ConversationRuntimeOptions | None = None,