Better typing

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
 

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:

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:

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