diff --git a/haiku_rag_slim/haiku/rag/agui/__init__.py b/haiku_rag_slim/haiku/rag/agui/__init__.py new file mode 100644 index 00000000..fe38fc94 --- /dev/null +++ b/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", +] diff --git a/haiku_rag_slim/haiku/rag/agui/emitter.py b/haiku_rag_slim/haiku/rag/agui/emitter.py new file mode 100644 index 00000000..4f9e0794 --- /dev/null +++ b/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] diff --git a/haiku_rag_slim/haiku/rag/agui/events.py b/haiku_rag_slim/haiku/rag/agui/events.py new file mode 100644 index 00000000..fca360b1 --- /dev/null +++ b/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, + } diff --git a/haiku_rag_slim/haiku/rag/agui/state.py b/haiku_rag_slim/haiku/rag/agui/state.py new file mode 100644 index 00000000..42ee34c9 --- /dev/null +++ b/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 diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index d17cdf34..37baa321 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/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", diff --git a/uv.lock b/uv.lock index 4e692b4a..0d4f35c6 100644 --- a/uv.lock +++ b/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" }, ]