Remove custom ag-ui from app backend.

app/backend/agent.py

@@ -1,17 +1,14 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
 
+from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic import BaseModel
-from pydantic_ai import Agent, RunContext, format_as_xml
+from pydantic_ai import Agent, RunContext, ToolReturn, format_as_xml
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import get_model
 
-if TYPE_CHECKING:
-    from haiku.rag.graph.agui.emitter import AGUIEmitter
-
 
 class CitationInfo(BaseModel):
     """Citation info for frontend display."""
@@ -74,7 +71,6 @@ class ChatDeps:
 
     client: HaikuRAG
     config: AppConfig
-    agui_emitter: "AGUIEmitter | None" = None
     search_results: list[SearchResult] | None = None
     session_state: ChatSessionState | None = None
 
@@ -135,7 +131,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         ctx: RunContext[ChatDeps],
         query: str,
         document_name: str | None = None,
-    ) -> str:
+    ) -> ToolReturn:
         """Search the knowledge base for relevant documents.
 
         Use this when you need to find documents or explore the knowledge base.
@@ -147,12 +143,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         """
         from search_agent import SearchAgent
 
-        if ctx.deps.agui_emitter:
-            msg = f"Searching: {query}"
-            if document_name:
-                msg += f" (in {document_name})"
-            ctx.deps.agui_emitter.log(msg)
-
         # Build context from conversation history
         context = None
         if ctx.deps.session_state and ctx.deps.session_state.qa_history:
@@ -169,7 +159,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         ctx.deps.search_results = results
 
         if not results:
-            return "No results found."
+            return ToolReturn(return_value="No results found.")
 
         # Build citation infos for frontend display
         citation_infos = [
@@ -186,23 +176,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             for i, r in enumerate(results)
         ]
 
-        # Emit search results as citations
+        # Build new state with citations
-        if ctx.deps.agui_emitter:
+        new_state = ChatSessionState(
-            ctx.deps.agui_emitter.update_state(
+            session_id=(
-                ChatSessionState(
+                ctx.deps.session_state.session_id if ctx.deps.session_state else ""
-                    session_id=(
+            ),
-                        ctx.deps.session_state.session_id
+            citations=citation_infos,
-                        if ctx.deps.session_state
+            qa_history=(
-                        else ""
+                ctx.deps.session_state.qa_history if ctx.deps.session_state else []
-                    ),
+            ),
-                    citations=citation_infos,
+        )
-                    qa_history=(
-                        ctx.deps.session_state.qa_history
-                        if ctx.deps.session_state
-                        else []
-                    ),
-                )
-            )
 
         # Return detailed results for the agent to present
         result_lines = []
@@ -219,14 +202,23 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             line += f"\n    {snippet}"
             result_lines.append(line)
 
-        return f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
+        return ToolReturn(
+            return_value=f"Found {len(results)} results:\n\n"
+            + "\n\n".join(result_lines),
+            metadata=[
+                StateSnapshotEvent(
+                    type=EventType.STATE_SNAPSHOT,
+                    snapshot=new_state.model_dump(),
+                )
+            ],
+        )
 
     @agent.tool
     async def ask(
         ctx: RunContext[ChatDeps],
         question: str,
         document_name: str | None = None,
-    ) -> str:
+    ) -> ToolReturn:
         """Answer a specific question using the knowledge base.
 
         Use this for direct questions that need a focused answer with citations.
@@ -241,12 +233,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         from haiku.rag.graph.research.models import Citation, SearchAnswer
         from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 
-        if ctx.deps.agui_emitter:
-            msg = f"Answering: {question}"
-            if document_name:
-                msg += f" (in {document_name})"
-            ctx.deps.agui_emitter.log(msg)
-
         # Build filter from document_name
         doc_filter = build_document_filter(document_name) if document_name else None
 
@@ -290,8 +276,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             search_filter=doc_filter,
             max_concurrency=ctx.deps.config.research.max_concurrency,
         )
-        # Don't pass agui_emitter to research graph - its state model differs from ChatSessionState
-        # The ask tool handles final state emission with citations
         deps = ResearchDeps(
             client=ctx.deps.client,
         )
@@ -323,23 +307,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             )
             ctx.deps.session_state.qa_history.append(qa_response)
 
-        # Emit updated state with citations AND accumulated qa_history
+        # Build new state with citations AND accumulated qa_history
-        if ctx.deps.agui_emitter:
+        new_state = ChatSessionState(
-            ctx.deps.agui_emitter.update_state(
+            session_id=(
-                ChatSessionState(
+                ctx.deps.session_state.session_id if ctx.deps.session_state else ""
-                    session_id=(
+            ),
-                        ctx.deps.session_state.session_id
+            citations=citation_infos,
-                        if ctx.deps.session_state
+            qa_history=(
-                        else ""
+                ctx.deps.session_state.qa_history if ctx.deps.session_state else []
-                    ),
+            ),
-                    citations=citation_infos,
+        )
-                    qa_history=(
-                        ctx.deps.session_state.qa_history
-                        if ctx.deps.session_state
-                        else []
-                    ),
-                )
-            )
 
         # Format answer with citation references and confidence
         answer_text = result.answer
@@ -347,7 +324,15 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos)))
             answer_text = f"{answer_text}\n\nSources: {citation_refs}"
 
-        return answer_text
+        return ToolReturn(
+            return_value=answer_text,
+            metadata=[
+                StateSnapshotEvent(
+                    type=EventType.STATE_SNAPSHOT,
+                    snapshot=new_state.model_dump(),
+                )
+            ],
+        )
 
     @agent.tool
     async def get_document(
@@ -361,9 +346,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         Args:
             query: The document title or URI to look up
         """
-        if ctx.deps.agui_emitter:
-            ctx.deps.agui_emitter.log(f"Fetching document: {query}")
-
         # Try exact URI match first
         doc = await ctx.deps.client.get_document_by_uri(query)
 

app/backend/main.py

@@ -3,57 +3,19 @@ import os
 from pathlib import Path
 
 from agent import ChatDeps, ChatSessionState, QAResponse, create_chat_agent
-from anyio import (
-    EndOfStream,
-    create_memory_object_stream,
-    create_task_group,
-    move_on_after,
-)
-from anyio.streams.memory import MemoryObjectSendStream
 from dotenv import load_dotenv
-from pydantic_ai.messages import (
+from pydantic_ai.ui import SSE_CONTENT_TYPE
-    ModelMessage,
+from pydantic_ai.ui.ag_ui import AGUIAdapter
-    ModelRequest,
-    ModelResponse,
-    TextPart,
-    UserPromptPart,
-)
 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.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import load_yaml_config
 from haiku.rag.config.models import AppConfig
-from haiku.rag.graph.agui.emitter import AGUIEmitter
-from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
-
-
-def convert_messages_to_history(
-    messages: list[dict[str, str]],
-) -> list[ModelMessage]:
-    """Convert AG-UI/CopilotKit messages to pydantic-ai message history.
-
-    Skips the last message since it will be passed as user_prompt to agent.run().
-    """
-    history: list[ModelMessage] = []
-
-    # Skip the last message - it will be the current user prompt
-    for msg in messages[:-1]:
-        role = msg.get("role", "")
-        content = msg.get("content", "")
-
-        if role == "user":
-            history.append(ModelRequest(parts=[UserPromptPart(content=content)]))
-        elif role == "assistant":
-            history.append(ModelResponse(parts=[TextPart(content=content)]))
-        # Skip other roles (system, tool, etc.) for now
-
-    return history
-
 
 load_dotenv()
 
@@ -103,109 +65,40 @@ def get_client(effective_db_path: Path) -> HaikuRAG:
     return _client_cache[path_key]
 
 
-async def stream_chat(request: Request) -> StreamingResponse:
+async def stream_chat(request: Request) -> Response:
     """Chat streaming endpoint with AG-UI protocol."""
-    body = await request.json()
+    body = await request.body()
-    logger.info(f"Received request: {list(body.keys())}")
+    logger.info("Received chat request")
-    input_data = RunAgentInput(**body)
 
-    user_message = ""
+    # Parse request to build run_input
-    message_history: list[ModelMessage] = []
+    accept = request.headers.get("accept", SSE_CONTENT_TYPE)
-    if input_data.messages:
+    run_input = AGUIAdapter.build_run_input(body)
-        user_message = input_data.messages[-1].get("content", "")
-        message_history = convert_messages_to_history(input_data.messages)
 
-    send_stream, receive_stream = create_memory_object_stream[str]()
+    # Restore qa_history from incoming state
+    initial_qa_history: list[QAResponse] = []
+    state = getattr(run_input, "state", None)
+    if state and "qa_history" in state:
+        initial_qa_history = [QAResponse(**qa) for qa in state.get("qa_history", [])]
 
-    async def run_agent_with_streaming(
+    # Build deps with session state
-        send_stream: MemoryObjectSendStream[str],
+    thread_id = getattr(run_input, "thread_id", None)
-    ) -> None:
+    deps = ChatDeps(
-        """Execute agent and forward events to stream."""
+        client=get_client(db_path),
-        async with send_stream:
+        config=Config,
-            try:
+        session_state=ChatSessionState(
-                # Create emitter for streaming
+            session_id=thread_id or "",
-                emitter: AGUIEmitter = AGUIEmitter(
+            qa_history=initial_qa_history,
-                    thread_id=input_data.thread_id,
+        ),
-                    run_id=input_data.run_id,
+    )
-                    use_deltas=True,
-                )
 
-                # Get client
+    # Use AGUIAdapter for streaming
-                effective_db_path = db_path
+    adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
-                if input_data.config and input_data.config.get("db_path"):
+    event_stream = adapter.run_stream(deps=deps)
-                    effective_db_path = Path(input_data.config["db_path"])
+    sse_event_stream = adapter.encode_stream(event_stream)
-                client = get_client(effective_db_path)
-
-                # Parse incoming state to restore qa_history
-                initial_qa_history: list[QAResponse] = []
-                if input_data.state and "qa_history" in input_data.state:
-                    initial_qa_history = [
-                        QAResponse(**qa)
-                        for qa in input_data.state.get("qa_history", [])
-                    ]
-
-                # Create initial state with restored history
-                initial_state = ChatSessionState(
-                    session_id=input_data.thread_id or "",
-                    qa_history=initial_qa_history,
-                )
-
-                # Create deps with session state
-                deps = ChatDeps(
-                    client=client,
-                    config=Config,
-                    agui_emitter=emitter,
-                    session_state=initial_state,
-                )
-                emitter.start_run(initial_state=initial_state)
-
-                # Forward events
-                async def forward_events():
-                    async for event in emitter:
-                        event_type = event.get("type")
-                        logger.debug(f"AG-UI event: {event_type}")
-                        await send_stream.send(format_sse_event(event))
-
-                # Run agent and forward concurrently
-                async with create_task_group() as tg:
-                    tg.start_soon(forward_events)
-
-                    result = await chat_agent.run(
-                        user_message, deps=deps, message_history=message_history
-                    )
-                    emitter.log(result.output)
-                    emitter.finish_run(result.output)
-                    await emitter.close()
-
-            except Exception as e:
-                logger.exception("Error executing agent")
-                try:
-                    await send_stream.send(
-                        format_sse_event({"type": "RUN_ERROR", "message": str(e)})
-                    )
-                except Exception:
-                    pass
-
-    async def event_generator():
-        """Generate SSE events with heartbeat to keep connection alive."""
-        async with create_task_group() as tg:
-            tg.start_soon(run_agent_with_streaming, send_stream)
-            async with receive_stream:
-                while True:
-                    try:
-                        # Wait for event with timeout, send heartbeat if nothing received
-                        with move_on_after(15):  # 15 second timeout
-                            event_str = await receive_stream.receive()
-                            yield event_str
-                            continue
-                        # No event received within timeout - send SSE comment as heartbeat
-                        yield ": heartbeat\n\n"
-                    except EndOfStream:
-                        break
 
     return StreamingResponse(
-        event_generator(),
+        sse_event_stream,
-        media_type="text/event-stream",
+        media_type=accept,
         headers={
             "Cache-Control": "no-cache",
             "Connection": "keep-alive",