diff --git a/app/backend/agent.py b/app/backend/agent.py index c653d688..9bd00891 100644 --- a/app/backend/agent.py +++ b/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 - if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.update_state( - ChatSessionState( - session_id=( - ctx.deps.session_state.session_id - if ctx.deps.session_state - else "" - ), - citations=citation_infos, - qa_history=( - ctx.deps.session_state.qa_history - if ctx.deps.session_state - else [] - ), - ) - ) + # Build new state with citations + new_state = ChatSessionState( + session_id=( + ctx.deps.session_state.session_id 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 - if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.update_state( - ChatSessionState( - session_id=( - ctx.deps.session_state.session_id - if ctx.deps.session_state - else "" - ), - citations=citation_infos, - qa_history=( - ctx.deps.session_state.qa_history - if ctx.deps.session_state - else [] - ), - ) - ) + # Build new state with citations AND accumulated qa_history + new_state = ChatSessionState( + session_id=( + ctx.deps.session_state.session_id 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) diff --git a/app/backend/main.py b/app/backend/main.py index 6a5d52ba..d3c4ebf2 100644 --- a/app/backend/main.py +++ b/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 ( - ModelMessage, - ModelRequest, - ModelResponse, - TextPart, - UserPromptPart, -) +from pydantic_ai.ui import SSE_CONTENT_TYPE +from pydantic_ai.ui.ag_ui import AGUIAdapter 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() - logger.info(f"Received request: {list(body.keys())}") - input_data = RunAgentInput(**body) + body = await request.body() + logger.info("Received chat request") - user_message = "" - message_history: list[ModelMessage] = [] - if input_data.messages: - user_message = input_data.messages[-1].get("content", "") - message_history = convert_messages_to_history(input_data.messages) + # Parse request to build run_input + accept = request.headers.get("accept", SSE_CONTENT_TYPE) + run_input = AGUIAdapter.build_run_input(body) - 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( - send_stream: MemoryObjectSendStream[str], - ) -> None: - """Execute agent and forward events to stream.""" - async with send_stream: - try: - # Create emitter for streaming - emitter: AGUIEmitter = AGUIEmitter( - thread_id=input_data.thread_id, - run_id=input_data.run_id, - use_deltas=True, - ) + # Build deps with session state + thread_id = getattr(run_input, "thread_id", None) + deps = ChatDeps( + client=get_client(db_path), + config=Config, + session_state=ChatSessionState( + session_id=thread_id or "", + qa_history=initial_qa_history, + ), + ) - # Get client - effective_db_path = db_path - if input_data.config and input_data.config.get("db_path"): - effective_db_path = Path(input_data.config["db_path"]) - 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 + # Use AGUIAdapter for streaming + adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept) + event_stream = adapter.run_stream(deps=deps) + sse_event_stream = adapter.encode_stream(event_stream) return StreamingResponse( - event_generator(), - media_type="text/event-stream", + sse_event_stream, + media_type=accept, headers={ "Cache-Control": "no-cache", "Connection": "keep-alive",