From 9c766846c05282d173cd87b878e19e1872bbb9fe Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Dec 2025 17:26:45 +0200 Subject: [PATCH] Use proper ag-ui tool calls in ag-ui-example. Fuck copilotkit --- CHANGELOG.md | 6 +- docs/agents.md | 9 +- .../a2a-server/haiku_rag_a2a/a2a/__init__.py | 2 +- examples/ag-ui-research/README.md | 23 +- examples/ag-ui-research/backend/agent.py | 3 + examples/ag-ui-research/backend/main.py | 98 ++-- examples/ag-ui-research/frontend/Dockerfile | 2 +- .../frontend/components/Agent.tsx | 542 ++++++++++-------- .../frontend/components/StateDisplay.tsx | 4 +- examples/ag-ui-research/frontend/package.json | 6 +- 10 files changed, 395 insertions(+), 300 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ffb6b76..acb50d7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,8 +12,10 @@ - New `human_decide` graph node emits AG-UI tool calls (`TOOL_CALL_START/ARGS/END`) for frontend integration - New `emit_tool_call_start()`, `emit_tool_call_args()`, `emit_tool_call_end()` AG-UI event helpers - New `AGUIEmitter.emit()` method for direct event emission -- **AG-UI Research Example**: Updated with interactive decision UI - - Decision panel with question editing (add/remove) at each decision point +- **AG-UI Research Example**: Human-in-the-loop research with client-side tool calling + - Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events + - Tool results sent directly to backend `/v1/research/stream` endpoint + - Backend queues decisions and continues the research graph - **HotpotQA Evaluation**: Added HotpotQA dataset adapter for multi-hop QA benchmarks - Extracts unique documents from validation set context paragraphs - Uses MAP for retrieval evaluation (multiple supporting documents per question) diff --git a/docs/agents.md b/docs/agents.md index 69bc0e82..e5caf19f 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -219,8 +219,9 @@ In interactive mode, you can: - Execute searches and review collected answers - Continue researching or synthesize when ready -For a web-based interactive experience with visual decision UI, see the [AG-UI Research Example](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research). The web interface provides: +For a web-based interactive experience, see the [AG-UI Research Example](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research). The example demonstrates AG-UI client-side tool calling: -- Question editing panel to add/remove sub-questions at decision points -- Search and Generate Report buttons for controlling research flow -- Live state display showing answers, confidence, and progress +- Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events +- Decision UI rendered inline in the chat at each decision point +- Question editing (add/remove) and action buttons (Search, Generate Report) +- Tool results sent directly to the backend endpoint which queues decisions and continues the graph diff --git a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py index 745c9903..11ab5cac 100644 --- a/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py +++ b/examples/a2a-server/haiku_rag_a2a/a2a/__init__.py @@ -6,7 +6,7 @@ import logfire from pydantic_ai import Agent, RunContext from haiku.rag.config import AppConfig, Config -from haiku.rag.graph.common import get_model +from haiku.rag.utils import get_model from .context import load_message_history, save_message_history from .models import A2AConfig, AgentDependencies, SearchResult diff --git a/examples/ag-ui-research/README.md b/examples/ag-ui-research/README.md index 1fa202c9..3925fba1 100644 --- a/examples/ag-ui-research/README.md +++ b/examples/ag-ui-research/README.md @@ -93,7 +93,7 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), ### Agent + Graph Pattern -This example demonstrates the **agent+graph** architecture pattern: +This example demonstrates the **agent+graph** architecture with AG-UI client-side tool calls: 1. **Conversational Agent** (`agent.py`): - Pydantic AI agent handles user conversations @@ -102,13 +102,14 @@ This example demonstrates the **agent+graph** architecture pattern: 2. **Interactive Research Graph** (haiku.rag): - Multi-step research workflow invoked by the agent's tool - - Pauses at decision points waiting for human input via async queue - - Emits AG-UI events for real-time progress tracking + - At decision points, emits AG-UI `TOOL_CALL_START/ARGS/END` events for `human_decision` + - Waits for tool result via async queue before continuing -3. **Decision Endpoint** (`main.py`): - - `/v1/research/decide` receives human decisions from frontend - - Forwards decisions to the waiting graph via `HumanDecision` queue - - Supports actions: `search`, `synthesize`, `modify_questions` +3. **Client-Side Tool Handling** (AG-UI pattern): + - Frontend listens for `human_decision` tool calls via AG-UI events + - Renders decision UI inline in chat when tool call is received + - User decision sent directly to backend `/v1/research/stream` endpoint + - Backend extracts tool result from messages and routes to waiting graph via async queue 4. **Shared Event Stream**: - `AGUIEmitter` is shared between agent and graph @@ -120,14 +121,14 @@ This example demonstrates the **agent+graph** architecture pattern: - **Backend** (Python): - Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base - `agent.py`: Pydantic AI agent with `run_research` tool, manages `ActiveResearch` registry - - `main.py`: Custom AG-UI streaming endpoint, decision endpoint for human input + - `main.py`: Custom AG-UI streaming endpoint, extracts tool results from messages - Real-time event forwarding from emitter to SSE stream - **Frontend** (Next.js/React): - - CopilotKit for AG-UI protocol integration + - AG-UI protocol integration for real-time streaming + - Handles `human_decision` tool calls with inline decision UI - Split-pane UI: chat on left, live research state on right - - Decision UI: question editor with add/remove, search and generate report buttons - - Real-time state synchronization via Server-Sent Events (SSE) + - Tool results sent directly to backend endpoint ## Configuration diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index fe15557d..c30e17c2 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -50,6 +50,7 @@ class AgentDeps: agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None search_filter: str | None = None thread_id: str | None = None + research_result: "ResearchReport | None" = None model = get_model(Config.research.model, Config) @@ -119,6 +120,8 @@ async def run_research(ctx: RunContext[AgentDeps], question: str) -> str: if ctx.deps.agui_emitter: ctx.deps.agui_emitter.log("Research complete!") + # Store result for main.py to emit RUN_FINISHED after agent completes + ctx.deps.research_result = result return f"""Research completed successfully! diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 6254c53e..a5211c4b 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -63,11 +63,54 @@ def get_client(effective_db_path: Path) -> HaikuRAG: return _client_cache[path_key] +def extract_tool_result(messages: list[dict]) -> dict | None: + """Extract human_decision tool result from messages if present.""" + for msg in reversed(messages): + # Check for tool result message (CopilotKit sends role="tool") + if msg.get("role") == "tool": + content = msg.get("content") + # Content may be a string (JSON) or dict + if isinstance(content, str): + try: + content = json.loads(content) + except json.JSONDecodeError: + continue + if isinstance(content, dict) and "action" in content: + return content + return None + + async def stream_research_agent(request: Request) -> StreamingResponse: """Agent streaming endpoint with research graph integration.""" body = await request.json() + logger.info(f"Received request body keys: {list(body.keys())}") + if "tools" in body: + logger.info(f"Frontend tools received: {body['tools']}") input_data = RunAgentInput(**body) + thread_id = input_data.thread_id + active_research = _active_research.get(thread_id) if thread_id else None + + # Check if this is a tool result for active research + if active_research and input_data.messages: + tool_result = extract_tool_result(input_data.messages) + if tool_result: + logger.info(f"Received tool result: {tool_result}") + action = tool_result.get("action", "search") + questions = tool_result.get("questions") + + decision = HumanDecision( + action=action, + questions=questions, + ) + await active_research.queue.put(decision) + + # Return acknowledgment - the original stream will continue + return StreamingResponse( + iter([format_sse_event({"type": "TOOL_RESULT_RECEIVED"})]), + media_type="text/event-stream", + ) + user_message = "" if input_data.messages: user_message = input_data.messages[-1].get("content", "") @@ -100,8 +143,6 @@ async def stream_research_agent(request: Request) -> StreamingResponse: ids_str = ", ".join(f"'{id}'" for id in document_ids) search_filter = f"id IN ({ids_str})" - thread_id = input_data.thread_id - # Create agent dependencies with shared emitter agent_deps = AgentDeps( client=client, @@ -121,9 +162,12 @@ async def stream_research_agent(request: Request) -> StreamingResponse: # Forward emitter events to stream async def forward_events(): async for event in emitter: - # Log events for debugging - logger.info(f"AG-UI Event: {event}") event_type = event.get("type") + logger.info(f"AG-UI event: {event_type}") + + # Log tool call events for debugging + if event_type and event_type.startswith("TOOL_CALL"): + logger.info(f"Tool call event: {event}") # Convert ACTIVITY_SNAPSHOT to STATE_DELTA for CopilotKit # As CopilotKit does not handle ACTIVITY_SNAPSHOT events @@ -133,7 +177,6 @@ async def stream_research_agent(request: Request) -> StreamingResponse: message = content.get("message", "") # Emit STATE_DELTA to patch activity info into state - # Use "add" op which creates or replaces the value delta_event = { "type": "STATE_DELTA", "delta": [ @@ -152,22 +195,6 @@ async def stream_research_agent(request: Request) -> StreamingResponse: await send_stream.send(format_sse_event(delta_event)) continue - # When human_decision tool starts, set awaiting_decision flag - if event_type == "TOOL_CALL_START": - tool_name = event.get("toolCallName") - if tool_name == "human_decision": - delta_event = { - "type": "STATE_DELTA", - "delta": [ - { - "op": "add", - "path": "/awaiting_decision", - "value": True, - } - ], - } - await send_stream.send(format_sse_event(delta_event)) - # Sync state to ActiveResearch when human_decision tool call if event_type == "TOOL_CALL_ARGS" and thread_id: delta = event.get("delta", "{}") @@ -191,6 +218,9 @@ async def stream_research_agent(request: Request) -> StreamingResponse: result = await agent.run(user_message, deps=agent_deps) emitter.log(result.output) + # Emit RUN_FINISHED with research result if available + if agent_deps.research_result: + emitter.finish_run(agent_deps.research_result) await emitter.close() except Exception as e: @@ -283,36 +313,10 @@ async def visualize_chunk(request: Request) -> JSONResponse: ) -async def research_decide(request: Request) -> JSONResponse: - """Endpoint to receive human decisions for active research.""" - body = await request.json() - action = body.get("action", "search") - questions = body.get("questions", []) - - # Get first active research (single-user example) - active = next(iter(_active_research.values()), None) - if not active: - return JSONResponse({"error": "No active research found"}, status_code=404) - - # When "search" action is sent with questions, use "modify_questions" to update them - effective_action = ( - "modify_questions" if action == "search" and questions else action - ) - - decision = HumanDecision( - action=effective_action, - questions=questions or None, - ) - await active.queue.put(decision) - - return JSONResponse({"status": "ok", "action": effective_action}) - - # Create Starlette app app = Starlette( routes=[ Route("/v1/research/stream", stream_research_agent, methods=["POST"]), - Route("/v1/research/decide", research_decide, methods=["POST"]), Route("/api/documents", list_documents, methods=["GET"]), Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]), Route("/health", health_check, methods=["GET"]), diff --git a/examples/ag-ui-research/frontend/Dockerfile b/examples/ag-ui-research/frontend/Dockerfile index b5da8b25..0959f3ef 100644 --- a/examples/ag-ui-research/frontend/Dockerfile +++ b/examples/ag-ui-research/frontend/Dockerfile @@ -4,7 +4,7 @@ FROM node:22-alpine WORKDIR /app COPY package.json package-lock.json* ./ -RUN npm ci +RUN npm install --legacy-peer-deps COPY . . EXPOSE 3000 diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx index 4cc1c731..2f1ebc92 100644 --- a/examples/ag-ui-research/frontend/components/Agent.tsx +++ b/examples/ag-ui-research/frontend/components/Agent.tsx @@ -1,66 +1,282 @@ "use client"; -import { CopilotKit, useCoAgent } from "@copilotkit/react-core"; +import { + CopilotKit, + useCoAgent, + useCopilotAction, + useCopilotContext, +} from "@copilotkit/react-core"; import { CopilotChat } from "@copilotkit/react-ui"; import "@copilotkit/react-ui/styles.css"; -import { useState, useEffect } from "react"; +import { useState } from "react"; import DocumentSelector from "./DocumentSelector"; import StateDisplay from "./StateDisplay"; -interface Citation { - document_id: string; - chunk_id: string; - document_uri: string; - document_title?: string; - page_numbers: number[]; - headings?: string[]; - content: string; -} - interface SearchAnswer { query: string; answer: string; confidence: number; cited_chunks: string[]; - citations: Citation[]; + citations: { + document_id: string; + chunk_id: string; + document_uri: string; + document_title?: string; + page_numbers: number[]; + headings?: string[]; + content: string; + }[]; } -interface ResearchContext { +interface ResearchState { + context: { + original_question: string; + sub_questions: string[]; + qa_responses: SearchAnswer[]; + }; + iterations: number; + max_iterations: number; + confidence_threshold: number; + max_concurrency: number; + last_eval: { + new_questions: string[]; + confidence_score: number; + is_sufficient: boolean; + reasoning: string; + } | null; + result?: { + title: string; + executive_summary: string; + main_findings: string[]; + conclusions: string[]; + limitations: string[]; + recommendations: string[]; + sources_summary: string; + }; + current_activity?: string; + current_activity_message?: string; + documentFilter?: string[]; +} + +interface DecisionArgs { original_question: string; sub_questions: string[]; qa_responses: SearchAnswer[]; } -interface EvaluationResult { - new_questions: string[]; - confidence_score: number; - is_sufficient: boolean; - reasoning: string; +type DecisionAction = "search" | "synthesize" | "modify_questions"; + +interface DecisionResult { + action: DecisionAction; + questions?: string[]; } -interface ResearchReport { - title: string; - executive_summary: string; - main_findings: string[]; - conclusions: string[]; - limitations: string[]; - recommendations: string[]; - sources_summary: string; +function DecisionUI({ + args, + onResolve, +}: { + args: DecisionArgs; + onResolve: (result: DecisionResult) => void | Promise; +}) { + const [editableQuestions, setEditableQuestions] = useState( + args.sub_questions || [], + ); + const [newQuestion, setNewQuestion] = useState(""); + const [submitting, setSubmitting] = useState(false); + + const qaCount = args.qa_responses?.length || 0; + const hasQuestions = editableQuestions.length > 0; + const canSearch = hasQuestions && !submitting; + const canSynthesize = qaCount > 0 && !submitting; + + const questionsModified = + editableQuestions.length !== args.sub_questions.length || + editableQuestions.some((q, i) => q !== args.sub_questions[i]); + + const handleSubmit = (action: DecisionAction, questions?: string[]) => { + setSubmitting(true); + onResolve({ action, questions }); + }; + + const handleSearch = () => { + handleSubmit( + questionsModified ? "modify_questions" : "search", + editableQuestions, + ); + }; + + const handleSynthesize = () => { + handleSubmit("synthesize"); + }; + + const handleRemoveQuestion = (index: number) => { + if (submitting) return; + setEditableQuestions(editableQuestions.filter((_, i) => i !== index)); + }; + + const handleAddQuestion = () => { + if (submitting || !newQuestion.trim()) return; + setEditableQuestions([...editableQuestions, newQuestion.trim()]); + setNewQuestion(""); + }; + + if (submitting) { + return null; + } + + return ( +
+
+ Research Decision Point +
+ +
+ {qaCount} answers collected +
+ +
+
+ Pending Questions ({editableQuestions.length}): +
+ {editableQuestions.map((q, idx) => ( +
+ {q} + +
+ ))} +
+ +
+ setNewQuestion(e.target.value)} + placeholder="Add a new question..." + style={{ + flex: 1, + padding: "0.5rem", + border: "1px solid #cbd5e1", + borderRadius: "4px", + fontSize: "0.85rem", + }} + onKeyDown={(e) => { + if (e.key === "Enter") handleAddQuestion(); + }} + /> + +
+ +
+ + +
+
+ ); } -interface ResearchState { - context: ResearchContext; - iterations: number; - max_iterations: number; - confidence_threshold: number; - max_concurrency: number; - last_eval: EvaluationResult | null; - result?: ResearchReport; - current_activity?: string; - current_activity_message?: string; - documentFilter?: string[]; - awaiting_decision?: boolean; -} +const BACKEND_URL = + process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000"; function AgentContent() { const { state, setState, running } = useCoAgent({ @@ -77,61 +293,79 @@ function AgentContent() { max_concurrency: 1, last_eval: null, documentFilter: [], - awaiting_decision: false, }, }); - const [editableQuestions, setEditableQuestions] = useState([]); - const [newQuestion, setNewQuestion] = useState(""); - const [submitting, setSubmitting] = useState(false); - - // Sync editable questions when state changes - useEffect(() => { - if (state.awaiting_decision && state.context.sub_questions) { - setEditableQuestions([...state.context.sub_questions]); - } - }, [state.awaiting_decision, state.context.sub_questions]); + const { threadId } = useCopilotContext(); const handleDocumentFilterChange = (ids: string[]) => { setState({ ...state, documentFilter: ids }); }; - const handleRemoveQuestion = (index: number) => { - setEditableQuestions(editableQuestions.filter((_, i) => i !== index)); - }; - - const handleAddQuestion = () => { - if (newQuestion.trim()) { - setEditableQuestions([...editableQuestions, newQuestion.trim()]); - setNewQuestion(""); + const sendToolResult = async (result: DecisionResult) => { + if (!threadId) { + console.error("No threadId available to send tool result"); + return; } - }; - const handleDecision = async (action: "search" | "synthesize") => { - setSubmitting(true); try { - const response = await fetch( - `${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/v1/research/decide`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - thread_id: state.context.original_question, // Use as identifier - action, - questions: editableQuestions, - }), - } - ); - if (response.ok) { - setState({ ...state, awaiting_decision: false }); + const response = await fetch(`${BACKEND_URL}/v1/research/stream`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + threadId, + messages: [ + { + id: crypto.randomUUID(), + role: "tool", + content: JSON.stringify(result), + }, + ], + }), + }); + + if (!response.ok) { + console.error("Failed to send tool result:", response.status); } } catch (error) { - console.error("Failed to send decision:", error); - } finally { - setSubmitting(false); + console.error("Error sending tool result:", error); } }; + useCopilotAction({ + name: "human_decision", + description: "Pause for human decision on research direction", + parameters: [ + { + name: "original_question", + type: "string", + description: "The original research question", + }, + { + name: "sub_questions", + type: "string[]", + description: "Pending sub-questions to search", + }, + { + name: "qa_responses", + type: "object[]", + description: "Answers collected so far", + }, + ], + renderAndWaitForResponse: ({ args, status }) => { + if (status === "complete") { + return null; + } + + return ( + + ); + }, + }); + return ( <>
- {/* Chat on the left */}
- {/* State display on the right */}
- {/* Document filter - hidden when research is running */} {!running && (
)} - {/* Decision UI when awaiting human input - hidden when report exists, submitting, or not running */} - {state.awaiting_decision && !submitting && !state.result && running && ( -
-
- Research Decision Point -
- -
- {state.context.qa_responses?.length || 0} answers collected | Iteration {state.iterations || 0} -
- - {/* Questions list */} -
-
- Pending Questions ({editableQuestions.length}): -
- {editableQuestions.map((q, idx) => ( -
- {q} - -
- ))} -
- - {/* Add question input */} -
- setNewQuestion(e.target.value)} - placeholder="Add a new question..." - disabled={submitting} - style={{ - flex: 1, - padding: "0.5rem", - border: "1px solid #cbd5e1", - borderRadius: "4px", - fontSize: "0.85rem", - }} - onKeyDown={(e) => { - if (e.key === "Enter") { - handleAddQuestion(); - } - }} - /> - -
- - {/* Action buttons */} -
- - -
-
- )} -
diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx index aa7ee112..3c9e129f 100644 --- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx +++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx @@ -162,7 +162,7 @@ export default function StateDisplay({ state }: StateDisplayProps) { }} > {/* Question */} - {state.context.original_question && ( + {state.context?.original_question && (
0 && ( + {state.context?.qa_responses && state.context.qa_responses.length > 0 && (