agentpool 2.1.9__py3-none-any.whl

agentpool/agents/agui_agent/agui_converters.py

@@ -0,0 +1,423 @@
+"""AG-UI to native event converters.
+
+This module provides conversion from AG-UI protocol events to native agentpool
+streaming events, enabling AGUIAgent to yield the same event types as native agents.
+
+Also provides conversion of native Tool objects to AG-UI Tool format for
+client-side tool execution.
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import TYPE_CHECKING, Any
+
+import anyenv
+from pydantic_ai import (
+    AudioUrl,
+    BinaryContent,
+    BinaryImage,
+    CachePoint,
+    DocumentUrl,
+    FileUrl,
+    ImageUrl,
+    PartDeltaEvent,
+    TextPartDelta,
+    ThinkingPartDelta,
+    VideoUrl,
+)
+
+from agentpool.agents.events import (
+    CustomEvent,
+    PlanUpdateEvent,
+    RunErrorEvent,
+    RunStartedEvent,
+    ToolCallProgressEvent,
+    ToolCallStartEvent as NativeToolCallStartEvent,
+)
+from agentpool.resource_providers.plan_provider import PlanEntry
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from ag_ui.core import BaseEvent, InputContent, Tool as AGUITool
+    from pydantic_ai import UserContent
+
+    from agentpool.agents.events import RichAgentStreamEvent
+    from agentpool.tools.base import Tool
+
+
+def agui_to_native_event(event: BaseEvent) -> RichAgentStreamEvent[Any] | None:  # noqa: PLR0911
+    """Convert AG-UI event to native streaming event.
+
+    Args:
+        event: AG-UI Event from SSE stream
+
+    Returns:
+        Corresponding native event, or None if no mapping exists
+    """
+    from ag_ui.core import (
+        ActivityDeltaEvent,
+        ActivitySnapshotEvent,
+        CustomEvent as AGUICustomEvent,
+        MessagesSnapshotEvent,
+        RawEvent,
+        RunErrorEvent as AGUIRunErrorEvent,
+        RunStartedEvent as AGUIRunStartedEvent,
+        StateDeltaEvent,
+        StateSnapshotEvent,
+        TextMessageChunkEvent,
+        TextMessageContentEvent,
+        TextMessageEndEvent,
+        TextMessageStartEvent,
+        ThinkingEndEvent,
+        ThinkingStartEvent,
+        ThinkingTextMessageContentEvent,
+        ThinkingTextMessageEndEvent,
+        ThinkingTextMessageStartEvent,
+        ToolCallArgsEvent,
+        ToolCallChunkEvent,
+        ToolCallEndEvent,
+        ToolCallResultEvent,
+        ToolCallStartEvent,
+    )
+
+    match event:
+        # === Lifecycle Events ===
+
+        case AGUIRunStartedEvent(thread_id=thread_id, run_id=run_id):
+            return RunStartedEvent(thread_id=thread_id, run_id=run_id)
+
+        case AGUIRunErrorEvent(message=message, code=code):
+            return RunErrorEvent(message=message, code=code)
+
+        # === Text Message Events ===
+
+        case TextMessageContentEvent(delta=delta):
+            return PartDeltaEvent(index=0, delta=TextPartDelta(content_delta=delta))
+
+        case TextMessageChunkEvent(delta=delta) if delta:
+            return PartDeltaEvent(index=0, delta=TextPartDelta(content_delta=delta))
+
+        case TextMessageStartEvent() | TextMessageEndEvent():
+            return None
+
+        # === Thinking/Reasoning Events ===
+
+        case ThinkingTextMessageContentEvent(delta=delta):
+            return PartDeltaEvent(index=0, delta=ThinkingPartDelta(content_delta=delta))
+
+        case ThinkingStartEvent() | ThinkingEndEvent():
+            # These mark thinking blocks but don't carry content
+            return None
+
+        case ThinkingTextMessageStartEvent() | ThinkingTextMessageEndEvent():
+            return None
+
+        # === Tool Call Events ===
+
+        case ToolCallStartEvent(tool_call_id=str() as tc_id, tool_call_name=name):
+            return NativeToolCallStartEvent(tool_call_id=tc_id, tool_name=name, title=name)
+
+        case ToolCallChunkEvent(tool_call_id=str() as tc_id, tool_call_name=str() as name):
+            return NativeToolCallStartEvent(tool_call_id=tc_id, tool_name=name, title=name)
+
+        case ToolCallArgsEvent(tool_call_id=tool_call_id, delta=_):
+            return ToolCallProgressEvent(tool_call_id=tool_call_id, status="in_progress")
+
+        case ToolCallResultEvent(tool_call_id=tc_id, content=content, message_id=_):
+            return ToolCallProgressEvent(tool_call_id=tc_id, status="completed", message=content)
+
+        case ToolCallEndEvent(tool_call_id=tool_call_id):
+            return ToolCallProgressEvent(tool_call_id=tool_call_id, status="completed")
+
+        # === Activity Events -> PlanUpdateEvent ===
+
+        case ActivitySnapshotEvent(
+            message_id=_, activity_type=activity_type, content=content, replace=_
+        ):
+            # Map activity content to plan entries if it looks like a plan
+            if activity_type.upper() == "PLAN" and isinstance(content, list):
+                entries = _content_to_plan_entries(content)
+                if entries:
+                    return PlanUpdateEvent(entries=entries)
+            # For other activity types, wrap as custom event
+            return CustomEvent(
+                event_data={"activity_type": activity_type, "content": content},
+                event_type=f"activity_{activity_type.lower()}",
+                source="ag-ui",
+            )
+
+        case ActivityDeltaEvent(message_id=_, activity_type=activity_type, patch=patch):
+            return CustomEvent(
+                event_data={"activity_type": activity_type, "patch": patch},
+                event_type=f"activity_delta_{activity_type.lower()}",
+                source="ag-ui",
+            )
+
+        # === State Management Events ===
+
+        case StateSnapshotEvent(snapshot=snapshot):
+            return CustomEvent(event_data=snapshot, event_type="state_snapshot", source="ag-ui")
+
+        case StateDeltaEvent(delta=delta):
+            return CustomEvent(event_data=delta, event_type="state_delta", source="ag-ui")
+
+        case MessagesSnapshotEvent(messages=messages):
+            data = [m.model_dump() if hasattr(m, "model_dump") else m for m in messages]
+            return CustomEvent(event_data=data, event_type="messages_snapshot", source="ag-ui")
+
+        # === Special Events ===
+
+        case RawEvent(event=raw_event, source=source):
+            return CustomEvent(event_data=raw_event, event_type="raw", source=source or "ag-ui")
+
+        case AGUICustomEvent(name=name, value=value):
+            return CustomEvent(event_data=value, event_type=name, source="ag-ui")
+
+        case _:
+            return None
+
+
+def _content_to_plan_entries(content: list[Any]) -> list[PlanEntry]:
+    """Convert AG-UI activity content to PlanEntry list.
+
+    Args:
+        content: List of plan items from ActivitySnapshotEvent
+
+    Returns:
+        List of PlanEntry objects
+    """
+    entries: list[PlanEntry] = []
+    for item in content:
+        if isinstance(item, dict):
+            # Try to extract plan entry fields
+            entry_content = item.get("content") or item.get("text") or item.get("description", "")
+            priority = item.get("priority", "medium")
+            status = item.get("status", "pending")
+
+            # Normalize values
+            if priority not in ("high", "medium", "low"):
+                priority = "medium"
+            if status not in ("pending", "in_progress", "completed"):
+                status = "pending"
+
+            if entry_content:
+                entry = PlanEntry(content=str(entry_content), priority=priority, status=status)
+                entries.append(entry)
+        elif isinstance(item, str):
+            entries.append(PlanEntry(content=item, priority="medium", status="pending"))
+    return entries
+
+
+def to_agui_input_content(
+    parts: UserContent | Sequence[UserContent] | None,
+) -> list[InputContent]:
+    """Convert pydantic-ai UserContent parts to AG-UI InputContent format.
+
+    Args:
+        parts: UserContent part(s) to convert (str, ImageUrl, BinaryContent, etc.)
+
+    Returns:
+        List of AG-UI InputContent items
+    """
+    from ag_ui.core import BinaryInputContent, TextInputContent
+
+    if parts is None:
+        return []
+
+    # Normalize to list
+    part_list = (
+        [parts] if isinstance(parts, str | FileUrl | BinaryContent | CachePoint) else list(parts)
+    )
+
+    result: list[InputContent] = []
+    for part in part_list:
+        match part:
+            case str() as text:
+                result.append(TextInputContent(text=text))
+
+            case ImageUrl(url=url, media_type=media_type):
+                result.append(BinaryInputContent(url=str(url), mime_type=media_type))
+
+            case AudioUrl(url=url, media_type=media_type):
+                result.append(BinaryInputContent(url=str(url), mime_type=media_type))
+
+            case DocumentUrl(url=url, media_type=media_type):
+                result.append(BinaryInputContent(url=str(url), mime_type=media_type))
+
+            case VideoUrl(url=url, media_type=media_type):
+                result.append(BinaryInputContent(url=str(url), mime_type=media_type))
+
+            case FileUrl(url=url, media_type=media_type):
+                # Generic FileUrl fallback
+                result.append(BinaryInputContent(url=str(url), mime_type=media_type))
+
+            case BinaryImage(data=data, media_type=media_type):
+                encoded = base64.b64encode(data).decode("utf-8")
+                result.append(BinaryInputContent(data=encoded, mime_type=media_type))
+
+            case BinaryContent(data=data, media_type=media_type):
+                encoded = base64.b64encode(data).decode("utf-8")
+                result.append(BinaryInputContent(data=encoded, mime_type=media_type))
+
+            case CachePoint():
+                # Cache points are markers, not actual content - skip
+                pass
+
+            case _:
+                # Fallback: convert to string
+                result.append(TextInputContent(text=str(part)))
+
+    return result
+
+
+def to_agui_tool(tool: Tool) -> AGUITool:
+    """Convert native Tool to AG-UI Tool format.
+
+    Args:
+        tool: native Tool instance
+
+    Returns:
+        AG-UI Tool with JSON Schema parameters
+    """
+    from ag_ui.core import Tool as AGUITool
+
+    schema = tool.schema
+    func_schema = schema["function"]
+    return AGUITool(
+        name=func_schema["name"],
+        description=func_schema.get("description", ""),
+        parameters=func_schema.get("parameters", {"type": "object", "properties": {}}),
+    )
+
+
+def _repair_partial_json(buffer: str) -> str:
+    """Attempt to repair truncated JSON for preview purposes.
+
+    Handles common truncation cases:
+    - Unclosed strings
+    - Missing closing braces/brackets
+    - Trailing commas
+
+    Args:
+        buffer: Potentially incomplete JSON string
+
+    Returns:
+        Repaired JSON string (may still be invalid in edge cases)
+    """
+    if not buffer:
+        return "{}"
+
+    result = buffer.rstrip()
+
+    # Check if we're in the middle of a string by counting unescaped quotes
+    in_string = False
+    i = 0
+    while i < len(result):
+        char = result[i]
+        if char == "\\" and i + 1 < len(result):
+            i += 2  # Skip escaped character
+            continue
+        if char == '"':
+            in_string = not in_string
+        i += 1
+
+    # Close unclosed string
+    if in_string:
+        result += '"'
+
+    # Remove trailing comma (invalid JSON)
+    result = result.rstrip()
+    if result.endswith(","):
+        result = result[:-1]
+
+    # Balance braces and brackets
+    open_braces = result.count("{") - result.count("}")
+    open_brackets = result.count("[") - result.count("]")
+
+    result += "]" * max(0, open_brackets)
+    result += "}" * max(0, open_braces)
+
+    return result
+
+
+class ToolCallAccumulator:
+    """Accumulates streamed tool call arguments.
+
+    AG-UI streams tool call arguments as deltas, this class accumulates them
+    and provides the complete arguments when the tool call ends.
+    """
+
+    def __init__(self) -> None:
+        self._calls: dict[str, dict[str, Any]] = {}
+
+    def start(self, tool_call_id: str, tool_name: str) -> None:
+        """Start tracking a new tool call."""
+        self._calls[tool_call_id] = {"name": tool_name, "args_buffer": ""}
+
+    def add_args(self, tool_call_id: str, delta: str) -> None:
+        """Add argument delta to a tool call."""
+        if tool_call_id in self._calls:
+            self._calls[tool_call_id]["args_buffer"] += delta
+
+    def complete(self, tool_call_id: str) -> tuple[str, dict[str, Any]] | None:
+        """Complete a tool call and return (tool_name, parsed_args).
+
+        Returns:
+            Tuple of (tool_name, args_dict) or None if call not found
+        """
+        if tool_call_id not in self._calls:
+            return None
+
+        call_data = self._calls.pop(tool_call_id)
+        args_str = call_data["args_buffer"]
+        try:
+            args = anyenv.load_json(args_str) if args_str else {}
+        except anyenv.JsonLoadError:
+            args = {"raw": args_str}
+        return call_data["name"], args
+
+    def get_pending(self, tool_call_id: str) -> tuple[str, str] | None:
+        """Get pending call data (tool_name, args_buffer) without completing."""
+        if tool_call_id not in self._calls:
+            return None
+        data = self._calls[tool_call_id]
+        return data["name"], data["args_buffer"]
+
+    def get_partial_args(self, tool_call_id: str) -> dict[str, Any]:
+        """Get best-effort parsed args from incomplete JSON stream.
+
+        Uses heuristics to complete truncated JSON for preview purposes.
+        Handles unclosed strings, missing braces/brackets, and trailing commas.
+
+        Args:
+            tool_call_id: Tool call ID
+
+        Returns:
+            Partially parsed arguments or empty dict
+        """
+        if tool_call_id not in self._calls:
+            return {}
+
+        buffer = self._calls[tool_call_id]["args_buffer"]
+        if not buffer:
+            return {}
+
+        # Try direct parse first
+        try:
+            return anyenv.load_json(buffer)
+        except anyenv.JsonLoadError:
+            pass
+
+        # Try to repair truncated JSON
+        try:
+            repaired = _repair_partial_json(buffer)
+            return anyenv.load_json(repaired)
+        except anyenv.JsonLoadError:
+            return {}
+
+    def clear(self) -> None:
+        """Clear all pending tool calls."""
+        self._calls.clear()

agentpool/agents/agui_agent/chunk_transformer.py

@@ -0,0 +1,204 @@
+"""Chunk transformer for AG-UI events.
+
+Transforms TEXT_MESSAGE_CHUNK and TOOL_CALL_CHUNK events into proper
+START/CONTENT/END sequences for easier downstream processing.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from ag_ui.core import BaseEvent, TextMessageChunkEvent, ToolCallChunkEvent
+
+
+logger = get_logger(__name__)
+
+
+class ChunkTransformer:
+    """Transforms CHUNK events into proper START/CONTENT/END sequences.
+
+    AG-UI supports two streaming modes:
+    1. Explicit: TEXT_MESSAGE_START -> TEXT_MESSAGE_CONTENT* -> TEXT_MESSAGE_END
+    2. Chunk: TEXT_MESSAGE_CHUNK* (implicit start/end based on message_id changes)
+
+    This transformer normalizes chunk mode into explicit mode for consistent handling.
+    Same applies for TOOL_CALL_CHUNK events.
+    """
+
+    def __init__(self) -> None:
+        """Initialize transformer."""
+        # Track active text message: message_id -> role
+        self._active_text: dict[str, str] | None = None
+        # Track active tool call: tool_call_id -> (name, parent_message_id)
+        self._active_tool: dict[str, tuple[str, str | None]] | None = None
+
+    def transform(self, event: BaseEvent) -> list[BaseEvent]:
+        """Transform a single event, potentially expanding chunks.
+
+        Args:
+            event: Input event
+
+        Returns:
+            List of output events (may be empty, single, or multiple)
+        """
+        from ag_ui.core import EventType
+
+        match event.type:
+            case EventType.TEXT_MESSAGE_CHUNK:
+                return self._handle_text_chunk(event)  # type: ignore[arg-type]
+
+            case EventType.TOOL_CALL_CHUNK:
+                return self._handle_tool_chunk(event)  # type: ignore[arg-type]
+
+            # These events close any pending chunks
+            case (
+                EventType.TEXT_MESSAGE_START
+                | EventType.TEXT_MESSAGE_END
+                | EventType.TOOL_CALL_START
+                | EventType.TOOL_CALL_END
+                | EventType.RUN_FINISHED
+                | EventType.RUN_ERROR
+            ):
+                close_events = self._close_all_pending()
+                return [*close_events, event]
+
+            case _:
+                # Pass through other events unchanged
+                return [event]
+
+    def _handle_text_chunk(self, event: TextMessageChunkEvent) -> list[BaseEvent]:
+        """Handle TEXT_MESSAGE_CHUNK event."""
+        from ag_ui.core import EventType, TextMessageContentEvent, TextMessageStartEvent
+
+        result: list[BaseEvent] = []
+        message_id = event.message_id
+        role = event.role or "assistant"
+        delta = event.delta
+
+        # Check if we need to close current text message (different ID)
+        if self._active_text is not None:
+            current_id = next(iter(self._active_text.keys()))
+            if message_id and message_id != current_id:
+                result.extend(self._close_text_message())
+
+        # Start new message if needed
+        if self._active_text is None and message_id:
+            self._active_text = {message_id: role}
+            start_event = TextMessageStartEvent(
+                type=EventType.TEXT_MESSAGE_START,
+                message_id=message_id,
+                role=role,
+            )
+            result.append(start_event)
+
+        # Emit content if we have delta and active message
+        if delta and self._active_text:
+            current_id = next(iter(self._active_text.keys()))
+            content_event = TextMessageContentEvent(
+                type=EventType.TEXT_MESSAGE_CONTENT,
+                message_id=current_id,
+                delta=delta,
+            )
+            result.append(content_event)
+
+        return result
+
+    def _handle_tool_chunk(self, event: ToolCallChunkEvent) -> list[BaseEvent]:
+        """Handle TOOL_CALL_CHUNK event."""
+        from ag_ui.core import EventType, ToolCallArgsEvent, ToolCallStartEvent
+
+        result: list[BaseEvent] = []
+        tool_call_id = event.tool_call_id
+        tool_name = event.tool_call_name
+        parent_id = event.parent_message_id
+        delta = event.delta
+
+        # Check if we need to close current tool call (different ID)
+        if self._active_tool is not None:
+            current_id = next(iter(self._active_tool.keys()))
+            if tool_call_id and tool_call_id != current_id:
+                result.extend(self._close_tool_call())
+
+        # Start new tool call if needed
+        if self._active_tool is None and tool_call_id and tool_name:
+            self._active_tool = {tool_call_id: (tool_name, parent_id)}
+            start_event = ToolCallStartEvent(
+                type=EventType.TOOL_CALL_START,
+                tool_call_id=tool_call_id,
+                tool_call_name=tool_name,
+                parent_message_id=parent_id,
+            )
+            result.append(start_event)
+
+        # Emit args if we have delta and active tool call
+        if delta and self._active_tool:
+            current_id = next(iter(self._active_tool.keys()))
+            args_event = ToolCallArgsEvent(
+                type=EventType.TOOL_CALL_ARGS,
+                tool_call_id=current_id,
+                delta=delta,
+            )
+            result.append(args_event)
+
+        return result
+
+    def _close_text_message(self) -> list[BaseEvent]:
+        """Close active text message."""
+        from ag_ui.core import EventType, TextMessageEndEvent
+
+        if self._active_text is None:
+            return []
+
+        message_id = next(iter(self._active_text.keys()))
+        self._active_text = None
+
+        end_event = TextMessageEndEvent(
+            type=EventType.TEXT_MESSAGE_END,
+            message_id=message_id,
+        )
+        logger.debug("Chunk transformer: TEXT_MESSAGE_END", message_id=message_id)
+        return [end_event]
+
+    def _close_tool_call(self) -> list[BaseEvent]:
+        """Close active tool call."""
+        from ag_ui.core import EventType, ToolCallEndEvent
+
+        if self._active_tool is None:
+            return []
+
+        tool_call_id = next(iter(self._active_tool.keys()))
+        self._active_tool = None
+
+        end_event = ToolCallEndEvent(
+            type=EventType.TOOL_CALL_END,
+            tool_call_id=tool_call_id,
+        )
+        logger.debug("Chunk transformer: TOOL_CALL_END", tool_call_id=tool_call_id)
+        return [end_event]
+
+    def _close_all_pending(self) -> list[BaseEvent]:
+        """Close all pending chunks (text and tool)."""
+        result: list[BaseEvent] = []
+        result.extend(self._close_text_message())
+        result.extend(self._close_tool_call())
+        return result
+
+    def flush(self) -> list[BaseEvent]:
+        """Flush any pending events at end of stream.
+
+        Call this when the stream ends to ensure all pending
+        messages/tool calls are properly closed.
+
+        Returns:
+            List of closing events
+        """
+        return self._close_all_pending()
+
+    def reset(self) -> None:
+        """Reset transformer state."""
+        self._active_text = None
+        self._active_tool = None

agentpool/agents/agui_agent/event_types.py

@@ -0,0 +1,83 @@
+"""Extended AG-UI event types.
+
+The upstream ag_ui.core.Event union is missing thinking events.
+This module provides a complete Event type until the SDK is fixed.
+
+TODO: Remove this workaround once upstream SDK is fixed.
+See: https://github.com/ag-ui-protocol/ag-ui/pull/753
+"""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from ag_ui.core import (
+    ActivityDeltaEvent,
+    ActivitySnapshotEvent,
+    CustomEvent,
+    MessagesSnapshotEvent,
+    RawEvent,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    StepFinishedEvent,
+    StepStartedEvent,
+    TextMessageChunkEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ThinkingTextMessageContentEvent,
+    ThinkingTextMessageEndEvent,
+    ThinkingTextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallChunkEvent,
+    ToolCallEndEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+from pydantic import Field
+
+
+# Complete Event union including thinking events (missing from upstream SDK)
+Event = Annotated[
+    # Text message events
+    TextMessageStartEvent
+    | TextMessageContentEvent
+    | TextMessageEndEvent
+    | TextMessageChunkEvent
+    # Thinking events (missing from upstream Event union)
+    | ThinkingStartEvent
+    | ThinkingEndEvent
+    | ThinkingTextMessageStartEvent
+    | ThinkingTextMessageContentEvent
+    | ThinkingTextMessageEndEvent
+    # Tool call events
+    | ToolCallStartEvent
+    | ToolCallArgsEvent
+    | ToolCallEndEvent
+    | ToolCallChunkEvent
+    | ToolCallResultEvent
+    # State events
+    | StateSnapshotEvent
+    | StateDeltaEvent
+    | MessagesSnapshotEvent
+    # Activity events
+    | ActivitySnapshotEvent
+    | ActivityDeltaEvent
+    # Lifecycle events
+    | RunStartedEvent
+    | RunFinishedEvent
+    | RunErrorEvent
+    | StepStartedEvent
+    | StepFinishedEvent
+    # Special events
+    | RawEvent
+    | CustomEvent,
+    Field(discriminator="type"),
+]
+
+__all__ = ["Event"]