Clean up

examples/ag-ui-research/README.md

@@ -1,47 +1,78 @@
-# Haiku.rag Interactive Research Assistant
+# Interactive Research Assistant
 
-Interactive research assistant powered by **Haiku.rag**, **Pydantic AI**, and **AG-UI** protocol. Ask complex questions and watch the multi-agent research process unfold in real-time with synchronized state between backend and frontend.
+Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic AI](https://ai.pydantic.dev/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
+
+## Features
+
+- **Multi-step research workflow**: Question decomposition, search, analysis, and synthesis
+- **Human-in-the-loop**: Approve or revise research plans before execution
+- **Live state synchronization**: Real-time updates of research progress between backend and frontend
+- **Context expansion**: Automatically expands top search results for better context
+- **Rich reporting**: Generates structured reports with findings, conclusions, and citations
 
 ## Quick Start
 
 ### Prerequisites
 
 - Docker and Docker Compose
-- Ollama running on host (or configure another QA provider)
+- A haiku.rag database with indexed documents
+- Ollama (or configure another LLM provider)
 
 ### Setup
 
-1. **Clone the repository**
+1. **Prepare your knowledge base**
    ```bash
-   git clone <repository-url>
-   cd haiku.rag/examples/ag-ui-research
-   ```
-
-2. **Configure environment** (optional, defaults to Ollama with gpt-oss:latest)
-
-   ```bash
-   cp .env.example .env
-   # Edit .env to customize provider/model or add API keys
-   ```
-
-   See [haiku.rag configuration docs](https://ggozad.github.io/haiku.rag/configuration/) for provider setup.
-
-3. **Prepare your knowledge base**
-
-   Create and populate a haiku.rag database:
-   ```bash
-   # Create a data directory
    mkdir -p data
-   # Add documents (requires haiku-rag installed locally)
    haiku-rag add "Your documents here" --db data/haiku_rag.lancedb
    # Or add from files
    haiku-rag add-src document.pdf --db data/haiku_rag.lancedb
    ```
 
-4. **Start the application**
+2. **Configure environment** (optional)
+   ```bash
+   cp .env.example .env
+   # Edit .env to customize provider/model
+   ```
+   See [haiku.rag configuration](https://ggozad.github.io/haiku.rag/configuration/) for details.
+
+3. **Start the application**
    ```bash
    docker compose up --build
    ```
 
-5. **Open the application**
+4. **Access the interface**
    - Frontend: http://localhost:3000
+   - Backend health: http://localhost:8000/health
+
+## How It Works
+
+1. **Ask a question**: Type your research question in the chat
+2. **Review the plan**: The agent decomposes your question into 3 sub-questions
+3. **Approve or revise**: Choose to approve the plan or request changes
+4. **Watch it work**: The agent automatically:
+   - Searches the knowledge base for each sub-question
+   - Extracts key insights from search results
+   - Evaluates overall confidence in findings
+5. **Get your report**: Receive a structured research report with citations
+
+## Architecture
+
+- **Backend** (Python): Pydantic AI agent with haiku.rag integration
+  - `agent.py`: Research agent with tool definitions
+  - `main.py`: Starlette app serving AG-UI protocol
+
+- **Frontend** (Next.js): CopilotKit/AG-UI interface
+  - Real-time state synchronization with backend
+  - Interactive approval workflow
+  - Collapsible research plan and insights display
+
+## Configuration
+
+Environment variables (see `.env.example`):
+
+- `DB_PATH`: Path to haiku.rag database (default: `haiku_rag.lancedb`)
+- `QA_PROVIDER`: LLM provider (default: `ollama`)
+- `QA_MODEL`: Model name (default: `gpt-oss:latest`)
+- `OLLAMA_BASE_URL`: Ollama endpoint (default: `http://host.docker.internal:11434`)
+
+For other providers (OpenAI, Anthropic, etc.), see [haiku.rag configuration docs](https://ggozad.github.io/haiku.rag/configuration/).

examples/ag-ui-research/backend/agent.py

@@ -1,7 +1,4 @@
-"""Pydantic AI research agent for haiku.rag with AG-UI protocol."""
+import json
-
-from __future__ import annotations
-
 from dataclasses import dataclass
 
 from ag_ui.core import EventType, StateSnapshotEvent
@@ -18,33 +15,15 @@ class ResearchState(BaseModel):
     """Shared state between research agent and frontend."""
 
     question: str = ""
-    phase: str = "idle"  # idle|planning|searching|analyzing|evaluating|done
+    phase: str = "idle"
-    status: str = ""  # Human-readable message
+    status: str = ""
-
+    plan: list[dict] = []
-    # Research plan with embedded search results
-    plan: list[
-        dict
-    ] = []  # [{id, question, status: pending|searching|done, search_results: {type, results: [...]}}]
     current_question_index: int = 0
-
+    insights: list[dict] = []
-    # Accumulated findings
+    document_registry: dict[str, dict] = {}
-    insights: list[
+    current_document: dict | None = None
-        dict
-    ] = []  # [{summary, confidence, source_refs: [{chunk_id, document_uri, document_title, chunk_position}]}]
-
-    # Document registry - tracks all referenced documents
-    document_registry: dict[
-        str, dict
-    ] = {}  # {doc_uri: {title, chunks_referenced: [chunk_id]}}
-
-    # Document viewer state
-    current_document: dict | None = None  # {uri, title, content, total_chunks}
-
-    # Final output
     confidence: float = 0.0
-    final_report: dict | None = (
+    final_report: dict | None = None
-        None  # {title, summary, findings, conclusions, citations: [{document_uri, document_title, chunk_ids}]}
-    )
 
 
 @dataclass
@@ -55,7 +34,6 @@ class ResearchDeps(StateDeps[ResearchState]):
 
 
 def _as_state_snapshot(ctx: RunContext[ResearchDeps]) -> StateSnapshotEvent:
-    """Helper to create state snapshot event for AG-UI synchronization."""
     return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)
 
 
@@ -112,20 +90,11 @@ Remember: Call tools ONE AT A TIME in sequence. Each tool must complete before c
     async def propose_research_plan(
         ctx: RunContext[ResearchDeps], question: str
     ) -> StateSnapshotEvent:
-        """Propose a research plan by decomposing the question into sub-questions.
+        """Propose a research plan by decomposing the question into sub-questions."""
-
-        Args:
-            question: The main research question to decompose
-        """
-        # Update state with the question
         ctx.deps.state.question = question
         ctx.deps.state.phase = "planning"
         ctx.deps.state.status = "Decomposing question into sub-questions..."
-        print(
-            f"[AGENT] Updated state: phase={ctx.deps.state.phase}, question={question}"
-        )
 
-        # Use LLM to decompose the question
         decompose_prompt = f"""Break down this research question into exactly 3 specific sub-questions that would help answer it comprehensively.
 
 Research Question: {question}
@@ -134,20 +103,15 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
 
         response = await ctx.deps.client.ask(decompose_prompt)
 
-        # Parse the response (simplified - assume it returns reasonable sub-questions)
-        import json
-
         try:
             sub_questions = json.loads(response)
         except json.JSONDecodeError:
-            # Fallback: split by newlines and clean up
             sub_questions = [
                 q.strip().lstrip("0123456789.-) ")
                 for q in response.split("\n")
                 if q.strip()
             ][:3]
 
-        # Create plan
         plan = [
             {"id": i, "question": q, "status": "pending"}
             for i, q in enumerate(sub_questions)
@@ -156,9 +120,6 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
         ctx.deps.state.plan = plan
         ctx.deps.state.current_question_index = 0
         ctx.deps.state.status = f"Proposed plan with {len(plan)} sub-questions"
-        print(f"[AGENT] Plan created with {len(plan)} sub-questions")
-        print("[AGENT] Sending state snapshot to frontend")
-        print("[AGENT] *** NEXT STEP: Agent should call approve_research_plan ***")
 
         return _as_state_snapshot(ctx)
 
@@ -168,49 +129,32 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
         question_id: int,
         search_type: str = "hybrid",
     ) -> StateSnapshotEvent:
-        """Execute search for a specific sub-question.
+        """Execute search for a specific sub-question."""
-
-        Args:
-            question_id: ID of the sub-question from the plan
-            search_type: Type of search (hybrid, vector, or fts)
-        """
-        # Get the question from plan
         plan = ctx.deps.state.plan
         if question_id >= len(plan):
             raise ValueError(f"Question ID {question_id} not found in plan")
 
         question = plan[question_id]["question"]
-
-        # Update state
         ctx.deps.state.phase = "searching"
         ctx.deps.state.current_question_index = question_id
         ctx.deps.state.status = f"Searching: {question}"
         plan[question_id]["status"] = "searching"
 
-        # Execute search
         search_results = await ctx.deps.client.search(
             question, limit=5, search_type=search_type
         )
 
-        # Expand context for top 3 results
+        expanded_map = {}
-        if len(search_results) > 0:
+        if search_results:
-            # Get top 3 for context expansion
-            top_results = search_results[:3]
             expanded_results = await ctx.deps.client.expand_context(
-                top_results, radius=2
+                search_results[:3], radius=2
             )
-
-            # Create a map of expanded chunks
             expanded_map = {
                 chunk.id: (chunk, score) for chunk, score in expanded_results
             }
-        else:
-            expanded_map = {}
 
-        # Process results and update document registry
         results = []
         for chunk, score in search_results:
-            # Update document registry
             doc_uri = chunk.document_uri or "unknown"
             doc_title = chunk.document_title or chunk.document_uri or "Unknown"
 
@@ -228,41 +172,27 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
                     chunk.id
                 )
 
-            # Check if this chunk was expanded
+            expanded_chunk, _ = (
-            if chunk.id in expanded_map:
+                expanded_map[chunk.id] if chunk.id in expanded_map else (chunk, score)
-                expanded_chunk, _ = expanded_map[chunk.id]
+            )
-                result_data = {
+            result_data = {
-                    "chunk": expanded_chunk.content[:500],  # Truncate for display
+                "chunk": expanded_chunk.content[:500],
-                    "chunk_id": chunk.id,
+                "chunk_id": chunk.id,
-                    "document_uri": doc_uri,
+                "document_uri": doc_uri,
-                    "document_title": doc_title,
+                "document_title": doc_title,
-                    "chunk_position": chunk.order,
+                "chunk_position": chunk.order,
-                    "full_chunk_content": expanded_chunk.content,
+                "full_chunk_content": expanded_chunk.content,
-                    "score": round(score, 3),
+                "score": round(score, 3),
-                    "expanded": True,
+                "expanded": chunk.id in expanded_map,
-                }
+            }
-            else:
-                result_data = {
-                    "chunk": chunk.content[:500],  # Truncate for display
-                    "chunk_id": chunk.id,
-                    "document_uri": doc_uri,
-                    "document_title": doc_title,
-                    "chunk_position": chunk.order,
-                    "full_chunk_content": chunk.content,
-                    "score": round(score, 3),
-                    "expanded": False,
-                }
-
             results.append(result_data)
 
-        # Store search results in the plan item
         plan[question_id]["search_results"] = {
             "type": search_type,
             "results": results,
         }
         plan[question_id]["status"] = "searched"
         ctx.deps.state.status = f"Found {len(results)} results"
-        print("[AGENT] Search complete, sending state snapshot")
 
         return _as_state_snapshot(ctx)
 
@@ -271,14 +201,7 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
         ctx: RunContext[ResearchDeps],
         question_id: int,
     ) -> StateSnapshotEvent:
-        """Extract key insights from search results for a specific question.
+        """Extract key insights from search results for a specific question."""
-
-        IMPORTANT: You must call search_question for this question_id BEFORE calling this tool.
-        This tool requires that search results already exist for the given question.
-
-        Args:
-            question_id: ID of the question whose results to analyze
-        """
         plan = ctx.deps.state.plan
         if question_id >= len(plan):
             raise ValueError(f"Question ID {question_id} not found in plan")
@@ -287,27 +210,19 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
         if "search_results" not in question_item:
             raise ValueError(
                 f"No search results found for question ID {question_id}. "
-                f"You must call search_question(question_id={question_id}) first before extracting insights."
+                f"You must call search_question(question_id={question_id}) first."
             )
 
         search_results = question_item["search_results"]
-
-        # Update state
         ctx.deps.state.phase = "analyzing"
         ctx.deps.state.status = "Extracting insights from results..."
 
-        # Build context from results with chunk IDs for reference
+        context_parts = [
-        context_parts = []
+            f"[Result {idx}] [Source: {r['document_title']}] {r['full_chunk_content']}"
-        for idx, r in enumerate(search_results["results"]):
+            for idx, r in enumerate(search_results["results"])
-            context_parts.append(
+        ]
-                f"[Result {idx}] [Source: {r['document_title']}] {r['full_chunk_content']}"
-            )
         context = "\n\n".join(context_parts)
 
-        # Use LLM to extract insights with structured output
-        from pydantic import BaseModel
-        from pydantic_ai import Agent
-
         class InsightResult(BaseModel):
             summary: str
             confidence: float
@@ -324,7 +239,6 @@ Search Results:
 
 For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
 
-        # Create a temporary agent with structured output using the same model
         insight_agent: Agent[None, InsightsList] = Agent(
             ctx.model,
             output_type=InsightsList,
@@ -339,15 +253,11 @@ For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
             }
             for insight in result.output.insights
         ]
-        print(f"[AGENT] Extracted {len(raw_insights)} insights using structured output")
 
-        # Convert result indices to structured source references
         new_insights = []
         for insight in raw_insights:
-            result_indices = insight.get("result_indices", [])
             source_refs = []
-
+            for idx in insight.get("result_indices", []):
-            for idx in result_indices:
                 if 0 <= idx < len(search_results["results"]):
                     result = search_results["results"][idx]
                     source_refs.append(
@@ -367,13 +277,9 @@ For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
                 }
             )
 
-        # Add to accumulated insights
         ctx.deps.state.insights.extend(new_insights)
-
-        # Mark question as fully done (searched + analyzed)
         plan[question_id]["status"] = "done"
         ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
-        print("[AGENT] Insights extracted, sending state snapshot")
 
         return _as_state_snapshot(ctx)
 
@@ -386,15 +292,12 @@ For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
         if not insights:
             raise ValueError("No insights collected yet")
 
-        # Update state
         ctx.deps.state.phase = "evaluating"
         ctx.deps.state.status = "Evaluating research confidence..."
 
-        # Calculate confidence (simple average of insight confidences)
         confidences = [i.get("confidence", 0.5) for i in insights]
         overall_confidence = sum(confidences) / len(confidences) if confidences else 0
 
-        # Use LLM to evaluate completeness
         eval_prompt = f"""Evaluate if these insights provide a confident answer to: "{ctx.deps.state.question}"
 
 Insights collected:
@@ -409,25 +312,14 @@ Return JSON: {{"confidence": 0.0-1.0, "gaps": ["gap1", "gap2"], "recommendation"
 
         response = await ctx.deps.client.ask(eval_prompt)
 
-        # Parse evaluation
-        import json
-
         try:
             evaluation = json.loads(response)
             overall_confidence = evaluation.get("confidence", overall_confidence)
         except json.JSONDecodeError:
-            evaluation = {
+            pass
-                "confidence": overall_confidence,
-                "gaps": [],
-                "recommendation": "finalize"
-                if overall_confidence > 0.7
-                else "continue",
-            }
 
-        # Update state
         ctx.deps.state.confidence = overall_confidence
         ctx.deps.state.status = f"Confidence: {overall_confidence:.0%}"
-        print("[AGENT] Confidence evaluated, sending state snapshot")
 
         return _as_state_snapshot(ctx)
 
@@ -440,22 +332,17 @@ Return JSON: {{"confidence": 0.0-1.0, "gaps": ["gap1", "gap2"], "recommendation"
         if not insights:
             raise ValueError("No insights to synthesize")
 
-        # Update state
         ctx.deps.state.phase = "synthesizing"
         ctx.deps.state.status = "Generating final report..."
 
-        # Build summary of insights with source information
         insights_summary = []
         for i in insights:
             source_titles = [ref["document_title"] for ref in i.get("source_refs", [])]
-            unique_sources = list(
+            unique_sources = list(dict.fromkeys(source_titles))
-                dict.fromkeys(source_titles)
-            )  # Preserve order, remove duplicates
             insights_summary.append(
                 f"- {i['summary']} (sources: {', '.join(unique_sources[:2])})"
             )
 
-        # Build report prompt
         report_prompt = f"""Generate a comprehensive research report answering: "{ctx.deps.state.question}"
 
 Based on these insights:
@@ -478,13 +365,9 @@ Return JSON with format:
 
         response = await ctx.deps.client.ask(report_prompt)
 
-        # Parse report
-        import json
-
         try:
             report = json.loads(response)
         except json.JSONDecodeError:
-            # Fallback report
             report = {
                 "title": ctx.deps.state.question,
                 "summary": response[:300],
@@ -493,25 +376,19 @@ Return JSON with format:
                 "sources": [],
             }
 
-        # Build structured citations from document registry
+        citations = [
-        citations = []
+            {
-        for doc_uri, doc_info in ctx.deps.state.document_registry.items():
+                "document_uri": doc_uri,
-            citations.append(
+                "document_title": doc_info["title"],
-                {
+                "chunk_ids": doc_info["chunks_referenced"],
-                    "document_uri": doc_uri,
+            }
-                    "document_title": doc_info["title"],
+            for doc_uri, doc_info in ctx.deps.state.document_registry.items()
-                    "chunk_ids": doc_info["chunks_referenced"],
+        ]
-                }
-            )
-
-        # Add citations to report
         report["citations"] = citations
 
-        # Update state
         ctx.deps.state.final_report = report
         ctx.deps.state.phase = "done"
         ctx.deps.state.status = "Research complete"
-        print("[AGENT] Report complete, sending state snapshot")
 
         return _as_state_snapshot(ctx)
 
@@ -520,15 +397,8 @@ Return JSON with format:
         ctx: RunContext[ResearchDeps],
         document_uri: str,
     ) -> StateSnapshotEvent:
-        """Retrieve and display the full content of a document by its URI.
+        """Retrieve and display the full content of a document by its URI."""
-
-        Args:
-            document_uri: The URI identifier of the document to retrieve
-        """
-        # Update state
         ctx.deps.state.status = f"Retrieving document: {document_uri}"
-
-        # Get document from haiku.rag
         document = await ctx.deps.client.get_document_by_uri(document_uri)
 
         if document is None:
@@ -536,15 +406,12 @@ Return JSON with format:
             ctx.deps.state.current_document = {
                 "uri": document_uri,
                 "title": "Not Found",
-                "content": f"Document with URI '{document_uri}' was not found in the database.",
+                "content": f"Document with URI '{document_uri}' was not found.",
                 "total_chunks": 0,
             }
         else:
-            # Get all chunks for this document to count them
             all_chunks = await ctx.deps.client.search(
-                query="",  # Empty query to get all chunks
+                query="", limit=1000, search_type="fts"
-                limit=1000,
-                search_type="fts",
             )
             chunks_for_doc = [
                 c for c, _ in all_chunks if c.document_uri == document_uri
@@ -557,11 +424,7 @@ Return JSON with format:
                 "total_chunks": len(chunks_for_doc),
                 "metadata": document.metadata,
             }
-            ctx.deps.state.status = (
+            ctx.deps.state.status = f"Retrieved: {document.title or document_uri}"
-                f"Retrieved document: {document.title or document_uri}"
-            )
-
-        print(f"[AGENT] Document retrieved: {document_uri}")
 
         return _as_state_snapshot(ctx)
 

examples/ag-ui-research/backend/main.py

@@ -1,5 +1,3 @@
-"""Main entry point for the haiku.rag AG-UI research assistant backend."""
-
 import logging
 import os
 from contextlib import asynccontextmanager
@@ -17,47 +15,37 @@ from haiku.rag.config import Config
 
 logger = logging.getLogger(__name__)
 
-# Global client instance
 client: HaikuRAG | None = None
+ag_ui_app = None
 
 
 @asynccontextmanager
 async def lifespan(app):
-    """Manage HaikuRAG client lifecycle."""
     global client
-
-    # Get database path from environment or use default
     db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
     db_path = Path(db_path_str)
 
     if not db_path.exists():
-        logger.error(
+        logger.error(f"Database not found at {db_path}")
-            f"Database not found at {db_path}. Please initialize haiku.rag first."
-        )
         logger.error("Run: haiku-rag add <path-to-documents>")
         raise RuntimeError(f"Database not found: {db_path}")
 
     logger.info(f"Initializing HaikuRAG client with database: {db_path}")
     client = HaikuRAG(db_path)
-
     logger.info("Research assistant backend ready")
-    logger.info(f"QA Provider: {Config.QA_PROVIDER}")
+    logger.info(f"QA Provider: {Config.QA_PROVIDER}, Model: {Config.QA_MODEL}")
-    logger.info(f"QA Model: {Config.QA_MODEL}")
 
     yield
 
-    # Cleanup
     if client:
         logger.info("Closing HaikuRAG client")
         client.close()
 
 
-# Create research agent instance
 agent = create_agent()
 
 
 async def health(request):
-    """Health check endpoint."""
     db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
     return JSONResponse(
         {
@@ -72,28 +60,16 @@ async def health(request):
     )
 
 
-# Create AG-UI app once with the client
-# State will be managed per-session by AG-UI
-ag_ui_app = None
-
-
 def get_ag_ui_app():
-    """Get or create AG-UI app."""
     global ag_ui_app
     if ag_ui_app is None and client is not None:
-        if client is None:
-            raise RuntimeError("Client not initialized")
-
-        # Create deps with shared client but new state per session
         research_deps = ResearchDeps(client=client, state=ResearchState())
-        logger.info("Creating AG-UI app with initial state")
+        logger.info("Creating AG-UI app")
         ag_ui_app = agent.to_ag_ui(deps=research_deps)
     return ag_ui_app
 
 
-# Initialize AG-UI app after client is ready in lifespan
 async def agent_endpoint(scope, receive, send):
-    """Proxy requests to AG-UI app."""
     app = get_ag_ui_app()
     if app is None:
         response = JSONResponse({"error": "Client not initialized"}, status_code=503)
@@ -102,7 +78,6 @@ async def agent_endpoint(scope, receive, send):
     await app(scope, receive, send)
 
 
-# Mount the AG-UI app at /agent and add health endpoint
 app = Starlette(
     routes=[
         Route("/health", health),

examples/ag-ui-research/frontend/components/Agent.tsx

@@ -1,6 +1,5 @@
 "use client";
 
-import React, { useState } from "react";
 import {
 	CopilotKit,
 	useCoAgent,
@@ -76,7 +75,6 @@ interface ResearchState {
 }
 
 function AgentContent() {
-	// Use useCoAgent to sync state with the backend research agent
 	const { state } = useCoAgent<ResearchState>({
 		name: "research_agent",
 		initialState: {
@@ -93,132 +91,104 @@ function AgentContent() {
 		},
 	});
 
-	// Log state changes
-	console.log("[FRONTEND] Current state:", state);
-
-	// Human-in-the-loop: Request approval for research plan
-	console.log("[FRONTEND] Registering approve_research_plan action");
 	useCopilotAction({
 		name: "approve_research_plan",
 		description:
 			"Request user approval for the research plan. Returns 'APPROVED' if approved or 'REVISE' if user wants to revise.",
 		parameters: [],
-		renderAndWaitForResponse: ({ respond, status }) => {
+		renderAndWaitForResponse: ({ respond, status }) => (
-			console.log(
+			<div
-				"[FRONTEND ACTION] renderAndWaitForResponse called",
+				style={{
-				{ status }
+					padding: "1.5rem",
-			);
+					background: "white",
-
+					borderRadius: "8px",
-			return (
+					border: "2px solid #4299e1",
-				<div
+					marginBottom: "1rem",
+					boxShadow: "0 2px 8px rgba(0,0,0,0.1)",
+				}}
+			>
+				<h3
 					style={{
-						padding: "1.5rem",
+						fontSize: "1.25rem",
-						background: "white",
+						fontWeight: "bold",
-						borderRadius: "8px",
-						border: "2px solid #4299e1",
 						marginBottom: "1rem",
-						boxShadow: "0 2px 8px rgba(0,0,0,0.1)",
+						color: "#2d3748",
 					}}
 				>
-					<h3
+					Research Plan Approval
-						style={{
+				</h3>
-							fontSize: "1.25rem",
+				<p
-							fontWeight: "bold",
+					style={{
-							marginBottom: "1rem",
+						fontSize: "0.875rem",
-							color: "#2d3748",
+						color: "#4a5568",
-						}}
+						marginBottom: "1rem",
-					>
+					}}
-						Research Plan Approval
+				>
-					</h3>
+					Please review the research plan in the right pane.
-					<p
+				</p>
-						style={{
-							fontSize: "0.875rem",
-							color: "#4a5568",
-							marginBottom: "1rem",
-						}}
-					>
-						Please review the research plan in the right pane.
-					</p>
 
-					<div
+				<div
+					style={{
+						display: "flex",
+						gap: "1rem",
+					}}
+					className={status !== "executing" ? "hidden" : ""}
+				>
+					<button
+						type="button"
+						onClick={() => respond?.("REVISE")}
+						disabled={status !== "executing"}
 						style={{
-							display: "flex",
+							flex: 1,
-							gap: "1rem",
+							padding: "0.75rem",
+							background: "white",
+							border: "2px solid #e2e8f0",
+							borderRadius: "6px",
+							fontSize: "0.875rem",
+							fontWeight: "600",
+							cursor: status === "executing" ? "pointer" : "not-allowed",
+							opacity: status === "executing" ? 1 : 0.5,
 						}}
-						className={status !== "executing" ? "hidden" : ""}
 					>
-						<button
+						Revise Plan
-							type="button"
+					</button>
-							onClick={() => respond?.("REVISE")}
+					<button
-							disabled={status !== "executing"}
+						type="button"
-							style={{
+						onClick={() => respond?.("APPROVED")}
-								flex: 1,
+						disabled={status !== "executing"}
-								padding: "0.75rem",
+						style={{
-								background: "white",
+							flex: 1,
-								border: "2px solid #e2e8f0",
+							padding: "0.75rem",
-								borderRadius: "6px",
+							background: "#4299e1",
-								fontSize: "0.875rem",
+							color: "white",
-								fontWeight: "600",
+							border: "none",
-								cursor: status === "executing" ? "pointer" : "not-allowed",
+							borderRadius: "6px",
-								opacity: status === "executing" ? 1 : 0.5,
+							fontSize: "0.875rem",
-							}}
+							fontWeight: "600",
-						>
+							cursor: status === "executing" ? "pointer" : "not-allowed",
-							Revise Plan
+							opacity: status === "executing" ? 1 : 0.5,
-						</button>
+						}}
-						<button
+					>
-							type="button"
+						Approve & Start Research
-							onClick={() => respond?.("APPROVED")}
+					</button>
-							disabled={status !== "executing"}
-							style={{
-								flex: 1,
-								padding: "0.75rem",
-								background: "#4299e1",
-								color: "white",
-								border: "none",
-								borderRadius: "6px",
-								fontSize: "0.875rem",
-								fontWeight: "600",
-								cursor: status === "executing" ? "pointer" : "not-allowed",
-								opacity: status === "executing" ? 1 : 0.5,
-							}}
-						>
-							Approve & Start Research
-						</button>
-					</div>
 				</div>
-			);
+			</div>
-		},
+		),
 	});
 
-	// Render state updates from the research agent
 	useCoAgentStateRender<ResearchState>({
 		name: "research_agent",
 		render: ({ state: newState }) => {
-			console.log("[FRONTEND] State render update:", newState);
+			const phaseMessages: Record<string, string> = {
-			// Show different messages based on phase
+				planning: "Planning research...",
-			let phaseMessage = "";
+				searching: "Searching...",
-			switch (newState.phase) {
+				analyzing: "Extracting insights...",
-				case "planning":
+				evaluating: `Evaluating confidence: ${(newState.confidence * 100).toFixed(0)}%`,
-					phaseMessage = "Planning research...";
+				synthesizing: "Generating final report...",
-					break;
+				done: "Research complete!",
-				case "searching":
+			};
-					phaseMessage = "Searching...";
+			const phaseMessage =
-					break;
+				phaseMessages[newState.phase] || newState.status || "Ready";
-				case "analyzing":
-					phaseMessage = "Extracting insights...";
-					break;
-				case "evaluating":
-					phaseMessage = `Evaluating confidence: ${(newState.confidence * 100).toFixed(0)}%`;
-					break;
-				case "synthesizing":
-					phaseMessage = "Generating final report...";
-					break;
-				case "done":
-					phaseMessage = "Research complete!";
-					break;
-				default:
-					phaseMessage = newState.status || "Ready";
-			}
 
 			return (
 				<div

examples/ag-ui-research/frontend/components/StateDisplay.tsx

@@ -100,7 +100,9 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 	};
 
 	// Calculate research progress
-	const completedQuestions = state.plan.filter((q) => q.status === "done").length;
+	const completedQuestions = state.plan.filter(
+		(q) => q.status === "done",
+	).length;
 	const totalQuestions = state.plan.length;
 	const researchProgress =
 		totalQuestions > 0 ? (completedQuestions / totalQuestions) * 100 : 0;
@@ -393,7 +395,8 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 												color:
 													item.status === "done"
 														? "#48bb78"
-														: item.status === "searching" || item.status === "searched"
+														: item.status === "searching" ||
+																item.status === "searched"
 															? "#4299e1"
 															: "#a0aec0",
 												flexShrink: 0,