AG-UI Event-emitter for graphs. Compute state-deltas from from BaseModel pydantic state models

haiku_rag_slim/haiku/rag/agui/__init__.py

@@ -0,0 +1,39 @@
+"""Generic AG-UI protocol support for haiku.rag graphs."""
+
+from haiku.rag.agui.emitter import AGUIEmitter
+from haiku.rag.agui.events import (
+    AGUIEvent,
+    emit_activity,
+    emit_activity_delta,
+    emit_run_error,
+    emit_run_finished,
+    emit_run_started,
+    emit_state_delta,
+    emit_state_snapshot,
+    emit_step_finished,
+    emit_step_started,
+    emit_text_message,
+    emit_text_message_content,
+    emit_text_message_end,
+    emit_text_message_start,
+)
+from haiku.rag.agui.state import compute_state_delta
+
+__all__ = [
+    "AGUIEmitter",
+    "AGUIEvent",
+    "compute_state_delta",
+    "emit_activity",
+    "emit_activity_delta",
+    "emit_run_error",
+    "emit_run_finished",
+    "emit_run_started",
+    "emit_state_delta",
+    "emit_state_snapshot",
+    "emit_step_finished",
+    "emit_step_started",
+    "emit_text_message",
+    "emit_text_message_content",
+    "emit_text_message_end",
+    "emit_text_message_start",
+]

haiku_rag_slim/haiku/rag/agui/emitter.py

@@ -0,0 +1,189 @@
+"""Generic AG-UI event emitter for any graph execution."""
+
+import asyncio
+import hashlib
+from collections.abc import AsyncIterator
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+from haiku.rag.agui.events import (
+    AGUIEvent,
+    emit_activity,
+    emit_run_error,
+    emit_run_finished,
+    emit_run_started,
+    emit_state_delta,
+    emit_state_snapshot,
+    emit_step_finished,
+    emit_step_started,
+    emit_text_message,
+)
+
+
+class AGUIEmitter[StateT: BaseModel, ResultT]:
+    """Generic queue-backed AG-UI event emitter for any graph.
+
+    Manages the lifecycle of AG-UI events including:
+    - Run lifecycle (start, finish, error)
+    - Step lifecycle (start, finish)
+    - Text messages
+    - State synchronization (snapshots and deltas)
+    - Activity updates
+
+    Type parameters:
+        StateT: The Pydantic BaseModel type for graph state
+        ResultT: The result type returned by the graph
+    """
+
+    def __init__(self, thread_id: str | None = None, run_id: str | None = None):
+        """Initialize the emitter.
+
+        Args:
+            thread_id: Optional thread ID (generated from input hash if not provided)
+            run_id: Optional run ID (random UUID if not provided)
+        """
+        self._queue: asyncio.Queue[AGUIEvent | None] = asyncio.Queue()
+        self._closed = False
+        self._thread_id = thread_id or str(uuid4())
+        self._run_id = run_id or str(uuid4())
+        self._last_state: StateT | None = None
+        self._current_step: str | None = None
+
+    @property
+    def thread_id(self) -> str:
+        """Get the thread ID for this emitter."""
+        return self._thread_id
+
+    @property
+    def run_id(self) -> str:
+        """Get the run ID for this emitter."""
+        return self._run_id
+
+    def start_run(self, input_data: str, initial_state: StateT) -> None:
+        """Emit RunStarted and initial StateSnapshot.
+
+        Args:
+            input_data: The input that started the run
+            initial_state: The initial state of the graph
+        """
+        # If thread_id wasn't provided, generate from input hash
+        if not self._thread_id or self._thread_id == str(uuid4()):
+            self._thread_id = self._generate_thread_id(input_data)
+
+        self._emit(emit_run_started(self._thread_id, self._run_id, input_data))
+        self._emit(emit_state_snapshot(initial_state))
+        self._last_state = initial_state
+
+    def start_step(self, step_name: str) -> None:
+        """Emit StepStarted event.
+
+        Args:
+            step_name: Name of the step being started
+        """
+        self._current_step = step_name
+        self._emit(emit_step_started(step_name))
+
+    def finish_step(self) -> None:
+        """Emit StepFinished event for the current step."""
+        if self._current_step:
+            self._emit(emit_step_finished(self._current_step))
+            self._current_step = None
+
+    def log(self, message: str, role: str = "assistant") -> None:
+        """Emit a text message event.
+
+        Args:
+            message: The message content
+            role: The role of the sender (default: assistant)
+        """
+        self._emit(emit_text_message(message, role))
+
+    def update_state(self, new_state: StateT) -> None:
+        """Emit StateDelta for state change, or StateSnapshot if no previous state.
+
+        Args:
+            new_state: The updated state
+        """
+        if self._last_state:
+            # Emit delta if we have a previous state
+            delta_event = emit_state_delta(self._last_state, new_state)
+            # Only emit if there are actual changes
+            if delta_event.get("operations"):
+                self._emit(delta_event)
+        else:
+            # Emit snapshot if this is the first state update
+            self._emit(emit_state_snapshot(new_state))
+
+        self._last_state = new_state
+
+    def update_activity(
+        self, activity_type: str, content: str, message_id: str | None = None
+    ) -> None:
+        """Emit ActivitySnapshot event.
+
+        Args:
+            activity_type: Type of activity (e.g., "planning", "searching")
+            content: Description of the activity
+            message_id: Optional message ID to associate activity with
+        """
+        self._emit(emit_activity(activity_type, content, message_id))
+
+    def finish_run(self, result: ResultT) -> None:
+        """Emit RunFinished event.
+
+        Args:
+            result: The final result from the graph
+        """
+        self._emit(emit_run_finished(self._thread_id, self._run_id, result))
+
+    def error(self, error: Exception, code: str | None = None) -> None:
+        """Emit RunError event.
+
+        Args:
+            error: The exception that occurred
+            code: Optional error code
+        """
+        self._emit(emit_run_error(str(error), code))
+
+    def _emit(self, event: AGUIEvent) -> None:
+        """Put event in queue.
+
+        Args:
+            event: The event to emit
+        """
+        if not self._closed:
+            self._queue.put_nowait(event)
+
+    async def close(self) -> None:
+        """Close the emitter and stop event iteration."""
+        if self._closed:
+            return
+        self._closed = True
+        await self._queue.put(None)
+
+    def __aiter__(self) -> AsyncIterator[AGUIEvent]:
+        """Enable async iteration over events."""
+        return self._iter_events()
+
+    async def _iter_events(self) -> AsyncIterator[AGUIEvent]:
+        """Iterate over events from the queue."""
+        while True:
+            event = await self._queue.get()
+            if event is None:
+                break
+            yield event
+
+    @staticmethod
+    def _generate_thread_id(input_data: str) -> str:
+        """Generate a deterministic thread ID from input data.
+
+        Args:
+            input_data: The input data (e.g., question, prompt)
+
+        Returns:
+            A stable thread ID based on input hash
+        """
+        # Use hash of input for deterministic thread ID
+        hash_obj = hashlib.sha256(input_data.encode("utf-8"))
+        return hash_obj.hexdigest()[:16]

haiku_rag_slim/haiku/rag/agui/events.py

@@ -0,0 +1,248 @@
+"""Generic AG-UI event creation utilities for any graph."""
+
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+from haiku.rag.agui.state import compute_state_delta
+
+# Type aliases for AG-UI events (actual types from ag_ui.core will be used at runtime)
+AGUIEvent = dict[str, Any]
+
+
+def emit_run_started(thread_id: str, run_id: str, input_data: str) -> dict[str, Any]:
+    """Create a RunStarted event.
+
+    Args:
+        thread_id: Unique identifier for the conversation thread
+        run_id: Unique identifier for this run
+        input_data: The input that started the run
+
+    Returns:
+        RunStarted event dict
+    """
+    return {
+        "type": "RUN_STARTED",
+        "threadId": thread_id,
+        "runId": run_id,
+        "input": input_data,
+    }
+
+
+def emit_run_finished(thread_id: str, run_id: str, result: Any) -> dict[str, Any]:
+    """Create a RunFinished event.
+
+    Args:
+        thread_id: Unique identifier for the conversation thread
+        run_id: Unique identifier for this run
+        result: The final result of the run
+
+    Returns:
+        RunFinished event dict
+    """
+    # Convert result to dict if it's a Pydantic model
+    if hasattr(result, "model_dump"):
+        result = result.model_dump()
+
+    return {
+        "type": "RUN_FINISHED",
+        "threadId": thread_id,
+        "runId": run_id,
+        "result": result,
+    }
+
+
+def emit_run_error(message: str, code: str | None = None) -> dict[str, Any]:
+    """Create a RunError event.
+
+    Args:
+        message: Error message
+        code: Optional error code
+
+    Returns:
+        RunError event dict
+    """
+    event: dict[str, Any] = {
+        "type": "RUN_ERROR",
+        "message": message,
+    }
+    if code:
+        event["code"] = code
+    return event
+
+
+def emit_step_started(step_name: str) -> dict[str, Any]:
+    """Create a StepStarted event.
+
+    Args:
+        step_name: Name of the step being started
+
+    Returns:
+        StepStarted event dict
+    """
+    return {
+        "type": "STEP_STARTED",
+        "stepName": step_name,
+    }
+
+
+def emit_step_finished(step_name: str) -> dict[str, Any]:
+    """Create a StepFinished event.
+
+    Args:
+        step_name: Name of the step that finished
+
+    Returns:
+        StepFinished event dict
+    """
+    return {
+        "type": "STEP_FINISHED",
+        "stepName": step_name,
+    }
+
+
+def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]:
+    """Create a TextMessageChunk event (convenience wrapper).
+
+    This creates a complete text message in one event.
+
+    Args:
+        content: The message content
+        role: The role of the sender (default: assistant)
+
+    Returns:
+        TextMessageChunk event dict
+    """
+    message_id = str(uuid4())
+    return {
+        "type": "TEXT_MESSAGE_CHUNK",
+        "messageId": message_id,
+        "role": role,
+        "content": content,
+    }
+
+
+def emit_text_message_start(message_id: str, role: str = "assistant") -> dict[str, Any]:
+    """Create a TextMessageStart event.
+
+    Args:
+        message_id: Unique identifier for this message
+        role: The role of the sender
+
+    Returns:
+        TextMessageStart event dict
+    """
+    return {
+        "type": "TEXT_MESSAGE_START",
+        "messageId": message_id,
+        "role": role,
+    }
+
+
+def emit_text_message_content(message_id: str, delta: str) -> dict[str, Any]:
+    """Create a TextMessageContent event.
+
+    Args:
+        message_id: Identifier for the message being streamed
+        delta: Content chunk to append
+
+    Returns:
+        TextMessageContent event dict
+    """
+    return {
+        "type": "TEXT_MESSAGE_CONTENT",
+        "messageId": message_id,
+        "delta": delta,
+    }
+
+
+def emit_text_message_end(message_id: str) -> dict[str, Any]:
+    """Create a TextMessageEnd event.
+
+    Args:
+        message_id: Identifier for the message being completed
+
+    Returns:
+        TextMessageEnd event dict
+    """
+    return {
+        "type": "TEXT_MESSAGE_END",
+        "messageId": message_id,
+    }
+
+
+def emit_state_snapshot(state: BaseModel) -> dict[str, Any]:
+    """Create a StateSnapshot event.
+
+    Args:
+        state: The complete state to snapshot (any Pydantic BaseModel)
+
+    Returns:
+        StateSnapshot event dict
+    """
+    return {
+        "type": "STATE_SNAPSHOT",
+        "snapshot": state.model_dump(),
+    }
+
+
+def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, Any]:
+    """Create a StateDelta event with JSON Patch operations.
+
+    Args:
+        old_state: Previous state (any Pydantic BaseModel)
+        new_state: Current state (same type as old_state)
+
+    Returns:
+        StateDelta event dict
+    """
+    operations = compute_state_delta(old_state, new_state)
+    return {
+        "type": "STATE_DELTA",
+        "operations": operations,
+    }
+
+
+def emit_activity(
+    activity_type: str,
+    content: str,
+    message_id: str | None = None,
+) -> dict[str, Any]:
+    """Create an ActivitySnapshot event.
+
+    Args:
+        activity_type: Type of activity (e.g., "planning", "searching")
+        content: Description of the activity
+        message_id: Optional message ID to associate activity with
+
+    Returns:
+        ActivitySnapshot event dict
+    """
+    event: dict[str, Any] = {
+        "type": "ACTIVITY_SNAPSHOT",
+        "activityType": activity_type,
+        "content": content,
+    }
+    if message_id:
+        event["messageId"] = message_id
+    return event
+
+
+def emit_activity_delta(
+    message_id: str, operations: list[dict[str, Any]]
+) -> dict[str, Any]:
+    """Create an ActivityDelta event with JSON Patch operations.
+
+    Args:
+        message_id: Message ID of the activity being updated
+        operations: JSON Patch operations to apply
+
+    Returns:
+        ActivityDelta event dict
+    """
+    return {
+        "type": "ACTIVITY_DELTA",
+        "messageId": message_id,
+        "operations": operations,
+    }

haiku_rag_slim/haiku/rag/agui/state.py

@@ -0,0 +1,34 @@
+"""Generic AG-UI state utilities for any Pydantic BaseModel."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+def compute_state_delta(
+    old_state: BaseModel, new_state: BaseModel
+) -> list[dict[str, Any]]:
+    """Compute JSON Patch (RFC 6902) operations from old state to new state.
+
+    Args:
+        old_state: Previous state (any Pydantic BaseModel)
+        new_state: Current state (same type as old_state)
+
+    Returns:
+        List of JSON Patch operations
+    """
+    operations: list[dict[str, Any]] = []
+
+    # Convert states to dicts for comparison
+    old_dict = old_state.model_dump()
+    new_dict = new_state.model_dump()
+
+    # Compare each field and generate patches
+    for key, new_value in new_dict.items():
+        old_value = old_dict.get(key)
+
+        if old_value != new_value:
+            # Simple replace operation
+            operations.append({"op": "replace", "path": f"/{key}", "value": new_value})
+
+    return operations

haiku_rag_slim/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "lancedb==0.25.2",
     "pathspec>=0.12.1",
     "pydantic>=2.12.3",
-    "pydantic-ai-slim[openai,fastmcp,logfire]>=1.11.1",
+    "pydantic-ai-slim[openai,fastmcp,logfire,ag-ui]>=1.11.1",
     "python-dotenv>=1.2.1",
     "pyyaml>=6.0.3",
     "rich>=14.2.0",

uv.lock

@@ -38,6 +38,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5f/a0/d9ef19f780f319c21ee90ecfef4431cbeeca95bec7f14071785c17b6029b/accelerate-1.10.1-py3-none-any.whl", hash = "sha256:3621cff60b9a27ce798857ece05e2b9f56fcc71631cfb31ccf71f0359c311f11", size = 374909, upload-time = "2025-08-25T13:57:04.55Z" },
 ]
 
+[[package]]
+name = "ag-ui-protocol"
+version = "0.1.10"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pydantic" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/67/bb/5a5ec893eea5805fb9a3db76a9888c3429710dfb6f24bbb37568f2cf7320/ag_ui_protocol-0.1.10.tar.gz", hash = "sha256:3213991c6b2eb24bb1a8c362ee270c16705a07a4c5962267a083d0959ed894f4", size = 6945, upload-time = "2025-11-06T15:17:17.068Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8f/78/eb55fabaab41abc53f52c0918a9a8c0f747807e5306273f51120fd695957/ag_ui_protocol-0.1.10-py3-none-any.whl", hash = "sha256:c81e6981f30aabdf97a7ee312bfd4df0cd38e718d9fc10019c7d438128b93ab5", size = 7889, upload-time = "2025-11-06T15:17:15.325Z" },
+]
+
 [[package]]
 name = "aiohappyeyeballs"
 version = "2.6.1"
@@ -1188,7 +1200,7 @@ dependencies = [
     { name = "lancedb" },
     { name = "pathspec" },
     { name = "pydantic" },
-    { name = "pydantic-ai-slim", extra = ["fastmcp", "logfire", "openai"] },
+    { name = "pydantic-ai-slim", extra = ["ag-ui", "fastmcp", "logfire", "openai"] },
     { name = "python-dotenv" },
     { name = "pyyaml" },
     { name = "rich" },
@@ -1247,7 +1259,7 @@ requires-dist = [
     { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'" },
     { name = "pydantic-ai-slim", extras = ["groq"], marker = "extra == 'groq'" },
     { name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" },
-    { name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire"], specifier = ">=1.11.1" },
+    { name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire", "ag-ui"], specifier = ">=1.11.1" },
     { name = "pydantic-ai-slim", extras = ["vertexai"], marker = "extra == 'vertexai'" },
     { name = "python-dotenv", specifier = ">=1.2.1" },
     { name = "pyyaml", specifier = ">=6.0.3" },
@@ -3019,6 +3031,10 @@ wheels = [
 ]
 
 [package.optional-dependencies]
+ag-ui = [
+    { name = "ag-ui-protocol" },
+    { name = "starlette" },
+]
 anthropic = [
     { name = "anthropic" },
 ]