AGUI starlette server

haiku_rag_slim/haiku/rag/agui/__init__.py

@@ -18,6 +18,12 @@ from haiku.rag.agui.events import (
     emit_text_message_end,
     emit_text_message_start,
 )
+from haiku.rag.agui.server import (
+    RunAgentInput,
+    create_agui_app,
+    create_research_server,
+    format_sse_event,
+)
 from haiku.rag.agui.state import compute_state_delta
 from haiku.rag.agui.stream import stream_graph
 
@@ -25,7 +31,10 @@ __all__ = [
     "AGUIConsoleRenderer",
     "AGUIEmitter",
     "AGUIEvent",
+    "RunAgentInput",
     "compute_state_delta",
+    "create_agui_app",
+    "create_research_server",
     "emit_activity",
     "emit_activity_delta",
     "emit_run_error",
@@ -39,5 +48,6 @@ __all__ = [
     "emit_text_message_content",
     "emit_text_message_end",
     "emit_text_message_start",
+    "format_sse_event",
     "stream_graph",
 ]

haiku_rag_slim/haiku/rag/agui/server.py

@@ -0,0 +1,209 @@
+"""AG-UI HTTP server implementation for graph execution."""
+
+import json
+from collections.abc import AsyncIterator, Callable
+from typing import Any, Protocol
+
+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.agui.events import AGUIEvent
+from haiku.rag.agui.stream import stream_graph
+from haiku.rag.config.models import AGUIConfig
+
+
+class GraphDeps(Protocol):
+    """Protocol for graph dependencies that support AG-UI emission."""
+
+    agui_emitter: 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_research_server(config: Any, db_path: Any | None = None) -> Starlette:
+    """Create AG-UI server for research graph.
+
+    This is a convenience function specifically for the research graph.
+
+    Args:
+        config: Application config with research settings
+        db_path: Optional database path override
+
+    Returns:
+        Starlette app configured for research graph
+    """
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.research.dependencies import ResearchContext
+    from haiku.rag.research.graph import build_research_graph
+    from haiku.rag.research.state import ResearchDeps, ResearchState
+
+    # Store client reference for proper lifecycle management
+    _client_cache: dict[str, HaikuRAG] = {}
+
+    def graph_factory() -> Graph:
+        """Create research graph instance."""
+        return build_research_graph(config)
+
+    def state_factory(input_state: dict[str, Any]) -> ResearchState:
+        """Create research state from input."""
+        # Extract question from input state or messages
+        question = input_state.get("question", "")
+        if not question:
+            # Try to get from first message if available
+            messages = input_state.get("messages", [])
+            if messages:
+                question = messages[0].get("content", "")
+
+        # Create context and state
+        context = ResearchContext(original_question=question)
+        return ResearchState.from_config(context=context, config=config)
+
+    def deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
+        """Create research dependencies."""
+        # Use provided db_path or fallback to config
+        effective_db_path = (
+            db_path
+            or input_config.get("db_path")
+            or config.storage.data_dir / "haiku.rag.lancedb"
+        )
+
+        # Reuse existing client if available
+        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 ResearchDeps(client=_client_cache[path_key])
+
+    # Use AG-UI config from app config
+    return create_agui_app(
+        graph_factory=graph_factory,
+        state_factory=state_factory,
+        deps_factory=deps_factory,
+        config=config.agui,
+    )

haiku_rag_slim/haiku/rag/app.py

@@ -445,6 +445,7 @@ class HaikuRAGApp:
         enable_mcp: bool = True,
         mcp_transport: str | None = None,
         mcp_port: int = 8001,
+        enable_agui: bool = False,
     ):
         """Start the server with selected services."""
         async with HaikuRAG(self.db_path, config=self.config) as client:
@@ -472,6 +473,30 @@ class HaikuRAGApp:
                 mcp_task = asyncio.create_task(run_mcp())
                 tasks.append(mcp_task)
 
+            # Start AG-UI server if enabled
+            if enable_agui:
+
+                async def run_agui():
+                    import uvicorn
+
+                    from haiku.rag.agui import create_research_server
+
+                    logger.info(
+                        f"Starting AG-UI server on {self.config.agui.host}:{self.config.agui.port}"
+                    )
+                    app = create_research_server(self.config, db_path=self.db_path)
+                    config = uvicorn.Config(
+                        app=app,
+                        host=self.config.agui.host,
+                        port=self.config.agui.port,
+                        log_level="info",
+                    )
+                    server = uvicorn.Server(config)
+                    await server.serve()
+
+                agui_task = asyncio.create_task(run_agui())
+                tasks.append(agui_task)
+
             if not tasks:
                 logger.warning("No services enabled")
                 return

haiku_rag_slim/haiku/rag/cli.py

@@ -396,7 +396,7 @@ def download_models_cmd():
 
 @cli.command(
     "serve",
-    help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
+    help="Start haiku.rag server. Use --monitor, --mcp, and/or --agui to enable services.",
 )
 def serve(
     db: Path | None = typer.Option(
@@ -424,12 +424,17 @@ def serve(
         "--mcp-port",
         help="Port to bind MCP server to (ignored with --stdio)",
     ),
+    agui: bool = typer.Option(
+        False,
+        "--agui",
+        help="Enable AG-UI HTTP server for graph streaming",
+    ),
 ) -> None:
     """Start the server with selected services."""
     # Require at least one service flag
-    if not (monitor or mcp):
+    if not (monitor or mcp or agui):
         typer.echo(
-            "Error: At least one service flag (--monitor or --mcp) must be specified"
+            "Error: At least one service flag (--monitor, --mcp, or --agui) must be specified"
         )
         raise typer.Exit(1)
 
@@ -447,6 +452,7 @@ def serve(
             enable_mcp=mcp,
             mcp_transport=transport,
             mcp_port=mcp_port,
+            enable_agui=agui,
         )
     )
 

haiku_rag_slim/haiku/rag/config/__init__.py

@@ -6,6 +6,7 @@ from haiku.rag.config.loader import (
     load_yaml_config,
 )
 from haiku.rag.config.models import (
+    AGUIConfig,
     AppConfig,
     EmbeddingsConfig,
     LanceDBConfig,
@@ -22,6 +23,7 @@ from haiku.rag.config.models import (
 
 __all__ = [
     "Config",
+    "AGUIConfig",
     "AppConfig",
     "StorageConfig",
     "MonitorConfig",

haiku_rag_slim/haiku/rag/config/loader.py

@@ -84,4 +84,12 @@ def generate_default_config() -> dict:
                 "research_base_url": "",
             },
         },
+        "agui": {
+            "host": "0.0.0.0",
+            "port": 8000,
+            "cors_origins": ["*"],
+            "cors_credentials": True,
+            "cors_methods": ["GET", "POST", "OPTIONS"],
+            "cors_headers": ["*"],
+        },
     }

haiku_rag_slim/haiku/rag/config/models.py

@@ -77,6 +77,15 @@ class ProvidersConfig(BaseModel):
     vllm: VLLMConfig = Field(default_factory=VLLMConfig)
 
 
+class AGUIConfig(BaseModel):
+    host: str = "0.0.0.0"
+    port: int = 8000
+    cors_origins: list[str] = ["*"]
+    cors_credentials: bool = True
+    cors_methods: list[str] = ["GET", "POST", "OPTIONS"]
+    cors_headers: list[str] = ["*"]
+
+
 class AppConfig(BaseModel):
     environment: str = "production"
     storage: StorageConfig = Field(default_factory=StorageConfig)
@@ -88,3 +97,4 @@ class AppConfig(BaseModel):
     research: ResearchConfig = Field(default_factory=ResearchConfig)
     processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
     providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
+    agui: AGUIConfig = Field(default_factory=AGUIConfig)

haiku_rag_slim/test_agui_server.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""Test script for AG-UI server functionality."""
+
+import asyncio
+import json
+
+import httpx
+
+
+async def test_agui_server():
+    """Test the AG-UI server endpoint."""
+    base_url = "http://localhost:8000"
+
+    # Test health check
+    async with httpx.AsyncClient() as client:
+        print("Testing health check...")
+        response = await client.get(f"{base_url}/health")
+        print(f"Health check response: {response.status_code}")
+        print(f"Health check data: {response.json()}")
+
+        # Test AG-UI streaming endpoint
+        print("\nTesting AG-UI stream endpoint...")
+        request_data = {
+            "threadId": "test-thread-1",
+            "runId": "test-run-1",
+            "state": {"question": "What is pydantic-graph?"},
+            "messages": [],
+            "config": {},
+        }
+
+        print(f"Request data: {json.dumps(request_data, indent=2)}")
+
+        # Send request and stream response
+        async with client.stream(
+            "POST",
+            f"{base_url}/v1/agent/stream",
+            json=request_data,
+            timeout=120.0,
+        ) as response:
+            print(f"Response status: {response.status_code}")
+            print("Streaming events...\n")
+
+            event_count = 0
+            async for line in response.aiter_lines():
+                if line.startswith("data: "):
+                    event_data = line[6:]  # Remove "data: " prefix
+                    try:
+                        event = json.loads(event_data)
+                        event_type = event.get("type", "UNKNOWN")
+                        print(f"Event {event_count}: {event_type}")
+
+                        # Show specific event details
+                        if event_type == "RUN_STARTED":
+                            print(f"  Thread ID: {event.get('threadId')}")
+                            print(f"  Run ID: {event.get('runId')}")
+                        elif event_type == "STEP_STARTED":
+                            print(f"  Step: {event.get('stepName')}")
+                        elif event_type == "ACTIVITY_SNAPSHOT":
+                            print(f"  Activity: {event.get('content')}")
+                        elif event_type == "STATE_SNAPSHOT":
+                            state = event.get("snapshot", {})
+                            if "context" in state:
+                                context = state["context"]
+                                if "sub_questions" in context:
+                                    num_questions = len(context["sub_questions"])
+                                    print(f"  Sub-questions: {num_questions}")
+                        elif event_type == "RUN_FINISHED":
+                            result = event.get("result", {})
+                            if "title" in result:
+                                print(f"  Report Title: {result['title']}")
+                        elif event_type == "RUN_ERROR":
+                            print(f"  Error: {event.get('message')}")
+
+                        event_count += 1
+                    except json.JSONDecodeError as e:
+                        print(f"Failed to parse event: {e}")
+                        print(f"Raw line: {line}")
+
+            print(f"\nTotal events received: {event_count}")
+
+
+if __name__ == "__main__":
+    print("AG-UI Server Test")
+    print("=" * 50)
+    print("Make sure to start the server first with:")
+    print("  haiku-rag serve --agui --agui-port 8000")
+    print("=" * 50)
+    print()
+
+    try:
+        asyncio.run(test_agui_server())
+    except httpx.ConnectError:
+        print("ERROR: Could not connect to server at http://localhost:8000")
+        print("Make sure the AG-UI server is running.")
+    except KeyboardInterrupt:
+        print("\nTest interrupted by user")
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()