diff --git a/haiku_rag_slim/haiku/rag/agui/__init__.py b/haiku_rag_slim/haiku/rag/agui/__init__.py index 470d1b74..29ecd62d 100644 --- a/haiku_rag_slim/haiku/rag/agui/__init__.py +++ b/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", ] diff --git a/haiku_rag_slim/haiku/rag/agui/server.py b/haiku_rag_slim/haiku/rag/agui/server.py new file mode 100644 index 00000000..e0f94765 --- /dev/null +++ b/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: \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, + ) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 6b1fe382..bac4c6f9 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index e92dbe01..ab6a1914 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/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, ) ) diff --git a/haiku_rag_slim/haiku/rag/config/__init__.py b/haiku_rag_slim/haiku/rag/config/__init__.py index 970ba4f2..b6325451 100644 --- a/haiku_rag_slim/haiku/rag/config/__init__.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/config/loader.py b/haiku_rag_slim/haiku/rag/config/loader.py index df21f0bb..d6584351 100644 --- a/haiku_rag_slim/haiku/rag/config/loader.py +++ b/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": ["*"], + }, } diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 2c432e63..4883859d 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/graph/nodes/__init__.py b/haiku_rag_slim/haiku/rag/graph/nodes/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/haiku_rag_slim/test_agui_server.py b/haiku_rag_slim/test_agui_server.py new file mode 100755 index 00000000..bf231267 --- /dev/null +++ b/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()