haiku.rag-slim 0.16.0__py3-none-any.whl

haiku/rag/graph/agui/server.py

@@ -0,0 +1,310 @@
+"""AG-UI HTTP server implementation for graph execution."""
+
+import json
+from collections.abc import AsyncIterator, Callable
+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
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.cors import CORSMiddleware
+from starlette.requests import Request
+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
+
+
+class GraphDeps(Protocol):
+    """Protocol for graph dependencies that support AG-UI emission."""
+
+    agui_emitter: AGUIEmitter[Any, Any] | None
+
+
+class RunAgentInput(BaseModel):
+    """AG-UI protocol run agent input.
+
+    See: https://docs.ag-ui.com/concepts/agents#runagentinput
+    """
+
+    thread_id: str | None = Field(None, alias="threadId")
+    run_id: str | None = Field(None, alias="runId")
+    state: dict[str, Any] = Field(default_factory=dict)
+    messages: list[dict[str, Any]] = Field(default_factory=list)
+    config: dict[str, Any] = Field(default_factory=dict)
+
+
+def create_agui_app(
+    graph_factory: Callable[[], Graph],
+    state_factory: Callable[[dict[str, Any]], BaseModel],
+    deps_factory: Callable[[dict[str, Any]], GraphDeps],
+    config: AGUIConfig,
+) -> Starlette:
+    """Create Starlette app with AG-UI endpoint.
+
+    Args:
+        graph_factory: Factory function to create graph instance
+        state_factory: Factory to create initial state from input
+        deps_factory: Factory to create graph dependencies
+        config: AG-UI server configuration
+
+    Returns:
+        Starlette application with AG-UI endpoints
+    """
+
+    async def event_stream(
+        input_data: RunAgentInput,
+    ) -> AsyncIterator[str]:
+        """Generate SSE event stream from graph execution.
+
+        Yields:
+            Server-Sent Events formatted strings
+        """
+        # Create graph, state, and dependencies
+        graph = graph_factory()
+
+        # Create initial state from input
+        initial_state = state_factory(input_data.state)
+
+        # Create dependencies (may use config from input)
+        deps = deps_factory(input_data.config)
+
+        # Execute graph and stream events
+        async for event in stream_graph(graph, initial_state, deps):
+            # Format as SSE event
+            event_data = format_sse_event(event)
+            yield event_data
+
+    async def stream_agent(request: Request) -> StreamingResponse:
+        """AG-UI agent stream endpoint.
+
+        Accepts AG-UI RunAgentInput and streams events via SSE.
+        """
+        # Parse request body
+        body = await request.json()
+        input_data = RunAgentInput(**body)
+
+        # Return SSE stream
+        return StreamingResponse(
+            event_stream(input_data),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",  # Disable buffering in nginx
+            },
+        )
+
+    async def health_check(_: Request) -> JSONResponse:
+        """Health check endpoint."""
+        return JSONResponse({"status": "healthy"})
+
+    # Define routes
+    routes = [
+        Route("/v1/agent/stream", stream_agent, methods=["POST"]),
+        Route("/health", health_check, methods=["GET"]),
+    ]
+
+    # Configure CORS middleware
+    middleware = [
+        Middleware(
+            CORSMiddleware,
+            allow_origins=config.cors_origins,
+            allow_credentials=config.cors_credentials,
+            allow_methods=config.cors_methods,
+            allow_headers=config.cors_headers,
+        )
+    ]
+
+    # Create Starlette app
+    app = Starlette(
+        routes=routes,
+        middleware=middleware,
+        debug=False,
+    )
+
+    return app
+
+
+def format_sse_event(event: AGUIEvent) -> str:
+    """Format AG-UI event as Server-Sent Event.
+
+    Args:
+        event: AG-UI event dictionary
+
+    Returns:
+        SSE formatted string with event data
+    """
+    # Convert event to JSON
+    event_json = json.dumps(event, ensure_ascii=False)
+
+    # Format as SSE
+    # Each event is: data: <json>\n\n
+    return f"data: {event_json}\n\n"
+
+
+def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Starlette:
+    """Create AG-UI server with both research and deep ask endpoints.
+
+    Args:
+        config: Application config with research and qa settings
+        db_path: Optional database path override
+
+    Returns:
+        Starlette app with research and deep ask endpoints
+    """
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
+    from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
+    from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
+    from haiku.rag.graph.research.dependencies import ResearchContext
+    from haiku.rag.graph.research.graph import build_research_graph
+    from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+
+    # Store client reference for proper lifecycle management
+    _client_cache: dict[str, 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:
+            _client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=config)
+        return _client_cache[path_key]
+
+    # Research graph factories
+    def research_graph_factory() -> Graph:
+        return build_research_graph(config)
+
+    def research_state_factory(input_state: dict[str, Any]) -> ResearchState:
+        question = input_state.get("question", "")
+        if not question:
+            messages = input_state.get("messages", [])
+            if messages:
+                question = messages[0].get("content", "")
+        context = ResearchContext(original_question=question)
+        return ResearchState.from_config(context=context, config=config)
+
+    def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
+        effective_db_path = (
+            db_path
+            or input_config.get("db_path")
+            or config.storage.data_dir / "haiku.rag.lancedb"
+        )
+        return ResearchDeps(client=get_client(effective_db_path))
+
+    # Deep ask graph factories
+    def deep_ask_graph_factory() -> Graph:
+        return build_deep_qa_graph(config)
+
+    def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState:
+        question = input_state.get("question", "")
+        if not question:
+            messages = input_state.get("messages", [])
+            if messages:
+                question = messages[0].get("content", "")
+        use_citations = input_state.get("use_citations", False)
+        context = DeepQAContext(original_question=question, use_citations=use_citations)
+        return DeepQAState.from_config(context=context, config=config)
+
+    def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps:
+        effective_db_path = (
+            db_path
+            or input_config.get("db_path")
+            or config.storage.data_dir / "haiku.rag.lancedb"
+        )
+        return DeepQADeps(client=get_client(effective_db_path))
+
+    # Create event stream functions for each graph type
+    async def research_event_stream(
+        input_data: RunAgentInput,
+    ) -> AsyncIterator[str]:
+        """Generate SSE event stream from research graph execution."""
+        graph = research_graph_factory()
+        initial_state = research_state_factory(input_data.state)
+        deps = research_deps_factory(input_data.config)
+
+        async for event in stream_graph(graph, initial_state, deps):
+            event_data = format_sse_event(event)
+            yield event_data
+
+    async def deep_ask_event_stream(
+        input_data: RunAgentInput,
+    ) -> AsyncIterator[str]:
+        """Generate SSE event stream from deep ask graph execution."""
+        graph = deep_ask_graph_factory()
+        initial_state = deep_ask_state_factory(input_data.state)
+        deps = deep_ask_deps_factory(input_data.config)
+
+        async for event in stream_graph(graph, initial_state, deps):
+            event_data = format_sse_event(event)
+            yield event_data
+
+    # Endpoint handlers
+    async def stream_research(request: Request) -> StreamingResponse:
+        """Research graph streaming endpoint."""
+        body = await request.json()
+        input_data = RunAgentInput(**body)
+
+        return StreamingResponse(
+            research_event_stream(input_data),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+    async def stream_deep_ask(request: Request) -> StreamingResponse:
+        """Deep ask graph streaming endpoint."""
+        body = await request.json()
+        input_data = RunAgentInput(**body)
+
+        return StreamingResponse(
+            deep_ask_event_stream(input_data),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+    async def health_check(_: Request) -> JSONResponse:
+        """Health check endpoint."""
+        return JSONResponse({"status": "healthy"})
+
+    # Define routes
+    routes = [
+        Route("/v1/research/stream", stream_research, methods=["POST"]),
+        Route("/v1/deep-ask/stream", stream_deep_ask, methods=["POST"]),
+        Route("/health", health_check, methods=["GET"]),
+    ]
+
+    # Configure CORS middleware
+    middleware = [
+        Middleware(
+            CORSMiddleware,
+            allow_origins=config.agui.cors_origins,
+            allow_credentials=config.agui.cors_credentials,
+            allow_methods=config.agui.cors_methods,
+            allow_headers=config.agui.cors_headers,
+        )
+    ]
+
+    # Create Starlette app
+    app = Starlette(
+        routes=routes,
+        middleware=middleware,
+        debug=False,
+    )
+
+    return app

haiku/rag/graph/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/graph/agui/stream.py

@@ -0,0 +1,86 @@
+"""Generic graph streaming with AG-UI events."""
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import suppress
+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[StateT: BaseModel, ResultT](Protocol):
+    """Protocol for graph dependencies that support AG-UI emission."""
+
+    agui_emitter: AGUIEmitter[StateT, ResultT] | None
+
+
+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.
+
+    This is a generic streaming function that works with any pydantic-graph
+    that follows the AG-UI pattern:
+    - State must be a Pydantic BaseModel
+    - Deps must have an optional agui_emitter attribute
+    - Graph must be a pydantic-graph Graph instance
+
+    Args:
+        graph: The pydantic-graph Graph to execute
+        state: Initial state (Pydantic BaseModel)
+        deps: Graph dependencies with agui_emitter support
+        use_deltas: Whether to emit state deltas instead of full snapshots (default: True)
+
+    Yields:
+        AG-UI event dictionaries
+
+    Raises:
+        TypeError: If deps doesn't support agui_emitter
+        RuntimeError: If graph doesn't produce a result
+    """
+    if not hasattr(deps, "agui_emitter"):
+        raise TypeError("deps must have an 'agui_emitter' attribute")
+
+    # Create AG-UI emitter
+    emitter: AGUIEmitter[StateT, ResultT] = AGUIEmitter(use_deltas=use_deltas)
+    deps.agui_emitter = emitter  # type: ignore[assignment]
+
+    async def _execute() -> None:
+        try:
+            # Start the run with initial state
+            emitter.start_run(initial_state=state)
+
+            # Execute the graph
+            result = await graph.run(state=state, deps=deps)
+
+            if result is None:
+                raise RuntimeError("Graph did not produce a result")
+
+            # Finish the run with the result
+            emitter.finish_run(result)
+        except Exception as exc:
+            # Emit error event
+            emitter.error(exc)
+        finally:
+            await emitter.close()
+
+    runner = asyncio.create_task(_execute())
+
+    try:
+        async for event in emitter:
+            yield event
+    finally:
+        if not runner.done():
+            runner.cancel()
+        with suppress(asyncio.CancelledError):
+            await runner

haiku/rag/graph/common/__init__.py

@@ -0,0 +1,5 @@
+"""Common utilities for graph implementations."""
+
+from haiku.rag.graph.common.utils import get_model
+
+__all__ = ["get_model"]

haiku/rag/graph/common/models.py

@@ -0,0 +1,42 @@
+"""Common models used across different graph implementations."""
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class ResearchPlan(BaseModel):
+    """A structured research plan with sub-questions to explore."""
+
+    sub_questions: list[str] = Field(
+        ...,
+        description="Specific questions to research, phrased as complete questions",
+    )
+
+    @field_validator("sub_questions")
+    @classmethod
+    def validate_sub_questions(cls, v: list[str]) -> list[str]:
+        if len(v) < 1:
+            raise ValueError("Must have at least 1 sub-question")
+        if len(v) > 12:
+            raise ValueError("Cannot have more than 12 sub-questions")
+        return v
+
+
+class SearchAnswer(BaseModel):
+    """Answer from a search operation with sources."""
+
+    query: str = Field(..., description="The question that was answered")
+    answer: str = Field(..., description="The comprehensive answer to the question")
+    context: list[str] = Field(
+        default_factory=list,
+        description="Relevant snippets that directly support the answer",
+    )
+    sources: list[str] = Field(
+        default_factory=list,
+        description="Source URIs or titles that contributed to this answer",
+    )
+    confidence: float = Field(
+        default=1.0,
+        description="Confidence score for this answer (0-1)",
+        ge=0.0,
+        le=1.0,
+    )