diff --git a/CHANGELOG.md b/CHANGELOG.md index cbb5acf2..04664605 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,10 +5,11 @@ - **AG-UI Protocol Support**: Full AG-UI (Agent-UI) protocol implementation for graph execution with event streaming - New `AGUIEmitter` class for emitting AG-UI events from graphs - - Support for all AG-UI event types: lifecycle events (`RUN_STARTED`, `RUN_FINISHED`, `RUN_ERROR`), step events (`STEP_STARTED`, `STEP_FINISHED`), state updates (`STATE_SNAPSHOT`), activity narration (`ACTIVITY_SNAPSHOT`), and text messages (`TEXT_MESSAGE_CHUNK`) + - Support for all AG-UI event types: lifecycle events (`RUN_STARTED`, `RUN_FINISHED`, `RUN_ERROR`), step events (`STEP_STARTED`, `STEP_FINISHED`), state updates (`STATE_SNAPSHOT`, `STATE_DELTA`), activity narration (`ACTIVITY_SNAPSHOT`), and text messages (`TEXT_MESSAGE_CHUNK`) - `AGUIConsoleRenderer` for rendering AG-UI event streams to terminal with Rich formatting - `stream_graph()` utility function for executing graphs with AG-UI event emission - State diff computation for efficient state synchronization + - **Delta State Updates**: AG-UI emitter now supports incremental state updates via JSON Patch operations (`STATE_DELTA` events) to reduce bandwidth, configurable via `use_deltas` parameter (enabled by default) - **AG-UI Server**: Starlette-based HTTP server for serving graphs via AG-UI protocol - Server-Sent Events (SSE) streaming endpoint at `/v1/agent/stream` - Health check endpoint at `/health` @@ -19,6 +20,7 @@ - Step-by-step execution visibility via AG-UI events - **CLI AG-UI Flag**: New `--agui` flag for `serve` command to start AG-UI server - **Graph Module**: New unified `haiku.rag.graph` module containing all graph-related functionality +- **Common Graph Nodes**: New factory functions (`create_plan_node`, `create_search_node`) in `haiku.rag.graph.common.nodes` for reusable graph components ### Changed diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index e7f7a829..b569540d 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -2,7 +2,11 @@ import json from collections.abc import AsyncIterator, Callable -from typing import Any, Protocol +from pathlib import Path +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + from haiku.rag.config.models import AppConfig from pydantic import BaseModel, Field from pydantic_graph.beta import Graph @@ -14,6 +18,7 @@ from starlette.responses import JSONResponse, StreamingResponse from starlette.routing import Route from haiku.rag.config.models import AGUIConfig +from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.agui.events import AGUIEvent from haiku.rag.graph.agui.stream import stream_graph @@ -21,7 +26,7 @@ from haiku.rag.graph.agui.stream import stream_graph class GraphDeps(Protocol): """Protocol for graph dependencies that support AG-UI emission.""" - agui_emitter: Any | None + agui_emitter: AGUIEmitter[Any, Any] | None class RunAgentInput(BaseModel): @@ -146,7 +151,7 @@ def format_sse_event(event: AGUIEvent) -> str: return f"data: {event_json}\n\n" -def create_agui_server(config: Any, db_path: Any | None = None) -> Starlette: +def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Starlette: """Create AG-UI server with both research and deep ask endpoints. Args: @@ -167,7 +172,7 @@ def create_agui_server(config: Any, db_path: Any | None = None) -> Starlette: # Store client reference for proper lifecycle management _client_cache: dict[str, HaikuRAG] = {} - def get_client(effective_db_path: Any) -> HaikuRAG: + def get_client(effective_db_path: Path) -> HaikuRAG: """Get or create cached client.""" path_key = str(effective_db_path) if path_key not in _client_cache: diff --git a/haiku_rag_slim/haiku/rag/graph/agui/stream.py b/haiku_rag_slim/haiku/rag/graph/agui/stream.py index ee6f1ac5..acde9298 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/stream.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/stream.py @@ -3,24 +3,28 @@ import asyncio from collections.abc import AsyncIterator from contextlib import suppress -from typing import Any, Protocol +from typing import Protocol, TypeVar from pydantic import BaseModel +from pydantic_graph.beta import Graph from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.agui.events import AGUIEvent +StateT = TypeVar("StateT", bound=BaseModel) +ResultT = TypeVar("ResultT") -class GraphDeps(Protocol): + +class GraphDeps[StateT: BaseModel, ResultT](Protocol): """Protocol for graph dependencies that support AG-UI emission.""" - agui_emitter: AGUIEmitter[Any, Any] | None + agui_emitter: AGUIEmitter[StateT, ResultT] | None -async def stream_graph( - graph: Any, - state: BaseModel, - deps: GraphDeps, +async def stream_graph[StateT: BaseModel, DepsT: GraphDeps, ResultT]( + graph: Graph[StateT, DepsT, None, ResultT], + state: StateT, + deps: DepsT, use_deltas: bool = True, ) -> AsyncIterator[AGUIEvent]: """Run a graph and yield AG-UI events as they occur. @@ -48,8 +52,8 @@ async def stream_graph( raise TypeError("deps must have an 'agui_emitter' attribute") # Create AG-UI emitter - emitter: AGUIEmitter[Any, Any] = AGUIEmitter(use_deltas=use_deltas) - deps.agui_emitter = emitter + emitter: AGUIEmitter[StateT, ResultT] = AGUIEmitter(use_deltas=use_deltas) + deps.agui_emitter = emitter # type: ignore[assignment] async def _execute() -> None: try: diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py index 179542a6..ee600874 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/state.py @@ -1,20 +1,22 @@ import asyncio from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from pydantic import BaseModel, Field from haiku.rag.client import HaikuRAG from haiku.rag.graph.deep_qa.dependencies import DeepQAContext +from haiku.rag.graph.deep_qa.models import DeepQAAnswer if TYPE_CHECKING: from haiku.rag.config.models import AppConfig + from haiku.rag.graph.agui.emitter import AGUIEmitter @dataclass class DeepQADeps: client: HaikuRAG - agui_emitter: Any | None = None + agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None semaphore: asyncio.Semaphore | None = None