state synchronization between backend agent and frontend

examples/ag-ui-research/.env.example

@@ -11,8 +11,8 @@ OLLAMA_BASE_URL=http://host.docker.internal:11434
 
 # Path to the LanceDB database
 # For Docker: /app/data/haiku_rag.lancedb
-# For local development: ./haiku_rag.lancedb
+# For local development: Use absolute path to existing database
-DB_PATH=haiku_rag.lancedb
+DB_PATH=~/SOME_FOLDER/haiku.rag.lancedb
 
 # API keys (set as needed for your QA provider)
 # OPENAI_API_KEY=your-key-here

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

@@ -1,10 +1,15 @@
 """Pydantic AI research agent for haiku.rag with AG-UI protocol."""
 
+from __future__ import annotations
+
+from dataclasses import dataclass
+
 from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic import BaseModel
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.ag_ui import StateDeps
 
+from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.graph.common import get_model
 
@@ -13,23 +18,41 @@ class ResearchState(BaseModel):
     """Shared state between research agent and frontend."""
 
     question: str = ""
-    status: str = "idle"
+    phase: str = "idle"  # idle|planning|searching|analyzing|evaluating|done
-    current_iteration: int = 0
+    status: str = ""  # Human-readable message
-    max_iterations: int = 2
+
+    # Research plan
+    plan: list[dict] = []  # [{id, question, status: pending|searching|done}]
+    current_question_index: int = 0
+
+    # Search results (live updates)
+    current_search: dict | None = (
+        None  # {query, type, results: [{chunk, score, expanded}]}
+    )
+
+    # Accumulated findings
+    insights: list[dict] = []  # [{summary, confidence, sources}]
+
+    # Final output
     confidence: float = 0.0
-    plan: list[dict] = []
-    findings: list[dict] = []
     final_report: dict | None = None
 
 
-def _as_state_snapshot(ctx: RunContext[StateDeps[ResearchState]]) -> StateSnapshotEvent:
+@dataclass
-    """Helper to create a state snapshot event for AG-UI."""
+class ResearchDeps(StateDeps[ResearchState]):
+    """Dependencies for the research agent with HaikuRAG client."""
+
+    client: HaikuRAG
+
+
+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)
 
 
 def create_agent(
     qa_provider: str = Config.QA_PROVIDER, qa_model: str = Config.QA_MODEL
-) -> Agent[StateDeps[ResearchState], str]:
+) -> Agent[ResearchDeps, str]:
     """Create and configure the research agent.
 
     Args:
@@ -38,33 +61,327 @@ def create_agent(
     """
     agent = Agent(
         model=get_model(qa_provider, qa_model),
-        deps_type=StateDeps[ResearchState],
+        deps_type=ResearchDeps,
-        instructions="""You are a research assistant powered by haiku.rag.
+        instructions="""You are a research co-pilot powered by haiku.rag.
 
-You help users conduct deep research on complex questions by:
+You work step-by-step with the user to conduct deep research on complex questions.
-- Breaking down questions into sub-questions
-- Searching through a knowledge base
-- Evaluating findings for completeness and confidence
-- Synthesizing comprehensive reports with citations
 
-The state is shared with the frontend application, showing research progress in real-time.
+Your workflow:
+1. When user asks a question, propose a research plan (3-5 sub-questions)
+2. Wait for user approval before proceeding
+3. For each sub-question:
+   - Announce what you're searching for
+   - Execute search and show results with scores
+   - Extract insights from the results
+   - Ask user if they want to continue to next question
+4. Evaluate overall confidence in your findings
+5. Ask user if confident enough or should search more
+6. Synthesize final report with citations
 
-Currently, tools are placeholder stubs. Full integration with haiku.rag research pipeline
+Be transparent: always announce what you're doing before you do it.
-will be implemented in the next phase.""",
+Show search scores, explain your reasoning, and involve the user in decisions.
+""",
     )
 
     @agent.tool
-    async def get_research_status(ctx: RunContext[StateDeps[ResearchState]]) -> dict:
+    async def propose_research_plan(
-        """Get the current research state and progress."""
+        ctx: RunContext[ResearchDeps], question: str
-        return {
+    ) -> StateSnapshotEvent:
-            "question": ctx.deps.state.question,
+        """Propose a research plan by decomposing the question into sub-questions.
-            "status": ctx.deps.state.status,
+
-            "iteration": ctx.deps.state.current_iteration,
+        Args:
-            "max_iterations": ctx.deps.state.max_iterations,
+            question: The main research question to decompose
-            "confidence": ctx.deps.state.confidence,
+        """
-            "has_plan": len(ctx.deps.state.plan) > 0,
+        # Update state with the question
-            "findings_count": len(ctx.deps.state.findings),
+        ctx.deps.state.question = question
-            "has_report": ctx.deps.state.final_report is not None,
+        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 3-5 specific sub-questions that would help answer it comprehensively.
+
+Research Question: {question}
+
+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()
+            ][:5]
+
+        # Create plan
+        plan = [
+            {"id": i, "question": q, "status": "pending"}
+            for i, q in enumerate(sub_questions)
+        ]
+
+        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")
+
+        return _as_state_snapshot(ctx)
+
+    @agent.tool
+    async def search_question(
+        ctx: RunContext[ResearchDeps],
+        question_id: int,
+        search_type: str = "hybrid",
+    ) -> StateSnapshotEvent:
+        """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
+        if len(search_results) > 0:
+            # Get top 3 for context expansion
+            top_results = search_results[:3]
+            expanded_results = await ctx.deps.client.expand_context(
+                top_results, 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
+        results = []
+        for chunk, score in search_results:
+            # Check if this chunk was expanded
+            if chunk.id in expanded_map:
+                expanded_chunk, _ = expanded_map[chunk.id]
+                result_data = {
+                    "chunk": expanded_chunk.content[:500],  # Truncate for display
+                    "score": round(score, 3),
+                    "source": chunk.document_title or chunk.document_uri or "Unknown",
+                    "expanded": True,
+                }
+            else:
+                result_data = {
+                    "chunk": chunk.content[:500],  # Truncate for display
+                    "score": round(score, 3),
+                    "source": chunk.document_title or chunk.document_uri or "Unknown",
+                    "expanded": False,
+                }
+
+            results.append(result_data)
+
+        # Update state
+        ctx.deps.state.current_search = {
+            "query": question,
+            "type": search_type,
+            "results": results,
         }
+        plan[question_id]["status"] = "done"
+        ctx.deps.state.status = f"Found {len(results)} results"
+        print("[AGENT] Search complete, sending state snapshot")
+
+        return _as_state_snapshot(ctx)
+
+    @agent.tool
+    async def extract_insights_from_results(
+        ctx: RunContext[ResearchDeps],
+    ) -> StateSnapshotEvent:
+        """Extract key insights from current search results."""
+        current_search = ctx.deps.state.current_search
+        if not current_search:
+            raise ValueError("No current search results to analyze")
+
+        # Update state
+        ctx.deps.state.phase = "analyzing"
+        ctx.deps.state.status = "Extracting insights from results..."
+
+        # Build context from results
+        context = "\n\n".join(
+            [f"[Source: {r['source']}] {r['chunk']}" for r in current_search["results"]]
+        )
+
+        # Use LLM to extract insights
+        extract_prompt = f"""Analyze these search results and extract 1-3 key insights that help answer the question: "{current_search["query"]}"
+
+Search Results:
+{context}
+
+Return a JSON array of insights with format:
+[{{"summary": "brief insight", "confidence": 0.0-1.0, "sources": ["source1", "source2"]}}]"""
+
+        response = await ctx.deps.client.ask(extract_prompt)
+
+        # Parse insights
+        import json
+
+        try:
+            new_insights = json.loads(response)
+        except json.JSONDecodeError:
+            # Fallback: create simple insight
+            new_insights = [
+                {
+                    "summary": response[:200],
+                    "confidence": 0.7,
+                    "sources": [r["source"] for r in current_search["results"][:3]],
+                }
+            ]
+
+        # Add to accumulated insights
+        ctx.deps.state.insights.extend(new_insights)
+
+        # Clear current search
+        ctx.deps.state.current_search = None
+        ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
+        print("[AGENT] Insights extracted, sending state snapshot")
+
+        return _as_state_snapshot(ctx)
+
+    @agent.tool
+    async def evaluate_research_confidence(
+        ctx: RunContext[ResearchDeps],
+    ) -> StateSnapshotEvent:
+        """Evaluate overall confidence in the research findings."""
+        insights = ctx.deps.state.insights
+        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:
+{chr(10).join([f"- {i['summary']}" for i in insights])}
+
+Assess:
+1. Do we have enough information to answer the question?
+2. What gaps remain?
+3. Overall confidence (0.0-1.0)
+
+Return JSON: {{"confidence": 0.0-1.0, "gaps": ["gap1", "gap2"], "recommendation": "continue" or "finalize"}}"""
+
+        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 = {
+                "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)
+
+    @agent.tool
+    async def synthesize_final_report(
+        ctx: RunContext[ResearchDeps],
+    ) -> StateSnapshotEvent:
+        """Generate final research report with citations."""
+        insights = ctx.deps.state.insights
+        if not insights:
+            raise ValueError("No insights to synthesize")
+
+        # Update state
+        ctx.deps.state.phase = "synthesizing"
+        ctx.deps.state.status = "Generating final report..."
+
+        # Build report prompt
+        report_prompt = f"""Generate a comprehensive research report answering: "{ctx.deps.state.question}"
+
+Based on these insights:
+{chr(10).join([f"- {i['summary']} (sources: {', '.join(i.get('sources', [])[:2])})" for i in insights])}
+
+Create a structured report with:
+- Executive Summary (2-3 sentences)
+- Main Findings (bullet points)
+- Conclusions
+- Sources
+
+Return JSON with format:
+{{
+    "title": "...",
+    "summary": "...",
+    "findings": ["finding1", "finding2", ...],
+    "conclusions": ["conclusion1", ...],
+    "sources": ["source1", "source2", ...]
+}}"""
+
+        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],
+                "findings": [i["summary"] for i in insights],
+                "conclusions": ["See findings above"],
+                "sources": list(
+                    set([s for i in insights for s in i.get("sources", [])])
+                ),
+            }
+
+        # 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)
 
     return agent

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

@@ -1,21 +1,64 @@
 """Main entry point for the haiku.rag AG-UI research assistant backend."""
 
-from agent import ResearchState, create_agent
+import logging
-from pydantic_ai.ag_ui import StateDeps
+import os
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from agent import ResearchDeps, ResearchState, create_agent
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import JSONResponse
 from starlette.routing import Mount, Route
 
+from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 
-# Create research agent instance using haiku.rag config
+logger = logging.getLogger(__name__)
+
+# Global client instance
+client: HaikuRAG | None = 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(
+            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 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(
         {
             "status": "healthy",
@@ -23,18 +66,47 @@ async def health(request):
             "qa_provider": Config.QA_PROVIDER,
             "qa_model": Config.QA_MODEL,
             "ollama_base_url": Config.OLLAMA_BASE_URL,
+            "db_path": db_path_str,
+            "db_exists": Path(db_path_str).exists(),
         }
     )
 
 
-# Convert PydanticAI agent to AG-UI compatible ASGI app
+# Create AG-UI app once with the client
-ag_ui_app = agent.to_ag_ui(deps=StateDeps(ResearchState()))  # type: ignore[arg-type]
+# 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")
+        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)
+        await response(scope, receive, send)
+        return
+    await app(scope, receive, send)
+
 
 # Mount the AG-UI app at /agent and add health endpoint
 app = Starlette(
     routes=[
         Route("/health", health),
-        Mount("/agent", ag_ui_app),
+        Mount("/agent", agent_endpoint),
     ],
     middleware=[
         Middleware(
@@ -45,6 +117,7 @@ app = Starlette(
             allow_headers=["*"],
         )
     ],
+    lifespan=lifespan,
 )
 
 if __name__ == "__main__":

examples/ag-ui-research/docker-compose.yml

@@ -8,14 +8,14 @@ services:
     environment:
       - QA_PROVIDER=${QA_PROVIDER:-ollama}
       - QA_MODEL=${QA_MODEL:-gpt-oss:latest}
-      - DB_PATH=${DB_PATH:-/app/data/haiku_rag.lancedb}
+      - DB_PATH=/app/data/haiku.rag.lancedb
       - OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
     volumes:
       - ./backend:/app
       - /app/.venv
-      - ./data:/app/data
+      - ${DB_PATH}:/app/data/haiku.rag.lancedb
     networks:
       - ag-ui-network
     extra_hosts:

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

@@ -11,13 +11,37 @@ import StateDisplay from "./StateDisplay";
 
 interface ResearchState {
 	question: string;
+	phase: string; // idle|planning|searching|analyzing|evaluating|done
 	status: string;
-	current_iteration: number;
+	plan: Array<{
-	max_iterations: number;
+		id: number;
+		question: string;
+		status: string; // pending|searching|done
+	}>;
+	current_question_index: number;
+	current_search: {
+		query: string;
+		type: string;
+		results?: Array<{
+			chunk: string;
+			score: number;
+			source: string;
+			expanded: boolean;
+		}>;
+	} | null;
+	insights: Array<{
+		summary: string;
+		confidence: number;
+		sources: string[];
+	}>;
 	confidence: number;
-	plan: Array<Record<string, unknown>>;
+	final_report: {
-	findings: Array<Record<string, unknown>>;
+		title: string;
-	final_report: Record<string, unknown> | null;
+		summary: string;
+		findings: string[];
+		conclusions: string[];
+		sources: string[];
+	} | null;
 }
 
 function AgentContent() {
@@ -26,20 +50,52 @@ function AgentContent() {
 		name: "research_agent",
 		initialState: {
 			question: "",
-			status: "idle",
+			phase: "idle",
-			current_iteration: 0,
+			status: "",
-			max_iterations: 2,
-			confidence: 0.0,
 			plan: [],
-			findings: [],
+			current_question_index: 0,
+			current_search: null,
+			insights: [],
+			confidence: 0.0,
 			final_report: null,
 		},
 	});
 
+	// Log state changes
+	console.log("[FRONTEND] Current state:", state);
+
 	// Render state updates from the research agent
 	useCoAgentStateRender<ResearchState>({
 		name: "research_agent",
 		render: ({ state: newState }) => {
+			console.log("[FRONTEND] State render update:", newState);
+			// Show different messages based on phase
+			let phaseMessage = "";
+			switch (newState.phase) {
+				case "planning":
+					phaseMessage = "Planning research...";
+					break;
+				case "searching":
+					phaseMessage = newState.current_search
+						? `Searching: ${newState.current_search.query}`
+						: "Searching...";
+					break;
+				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
 					style={{
@@ -50,9 +106,7 @@ function AgentContent() {
 						border: "1px solid #91d5ff",
 					}}
 				>
-					<strong>Research Update:</strong> Status: {newState.status},
+					<strong>Research Update:</strong> {phaseMessage}
-					Iteration: {newState.current_iteration}/{newState.max_iterations},
-					Confidence: {(newState.confidence * 100).toFixed(0)}%
 				</div>
 			);
 		},

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

@@ -1,14 +1,40 @@
 "use client";
 
+import { useState } from "react";
+
 interface ResearchState {
 	question: string;
+	phase: string;
 	status: string;
-	current_iteration: number;
+	plan: Array<{
-	max_iterations: number;
+		id: number;
+		question: string;
+		status: string;
+	}>;
+	current_question_index: number;
+	current_search: {
+		query: string;
+		type: string;
+		results?: Array<{
+			chunk: string;
+			score: number;
+			source: string;
+			expanded: boolean;
+		}>;
+	} | null;
+	insights: Array<{
+		summary: string;
+		confidence: number;
+		sources: string[];
+	}>;
 	confidence: number;
-	plan: Array<Record<string, unknown>>;
+	final_report: {
-	findings: Array<Record<string, unknown>>;
+		title: string;
-	final_report: Record<string, unknown> | null;
+		summary: string;
+		findings: string[];
+		conclusions: string[];
+		sources: string[];
+	} | null;
 }
 
 interface StateDisplayProps {
@@ -16,6 +42,33 @@ interface StateDisplayProps {
 }
 
 export default function StateDisplay({ state }: StateDisplayProps) {
+	const [expandedSections, setExpandedSections] = useState<
+		Record<string, boolean>
+	>({
+		plan: true,
+		search: true,
+		insights: true,
+		report: true,
+	});
+
+	const toggleSection = (section: string) => {
+		setExpandedSections((prev) => ({
+			...prev,
+			[section]: !prev[section],
+		}));
+	};
+
+	// Phase indicator
+	const phases = [
+		"idle",
+		"planning",
+		"searching",
+		"analyzing",
+		"evaluating",
+		"done",
+	];
+	const currentPhaseIndex = phases.indexOf(state.phase);
+
 	return (
 		<div
 			style={{
@@ -36,31 +89,66 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 			>
 				Research State
 			</h2>
-			<p
-				style={{
-					color: "#4a5568",
-					lineHeight: "1.6",
-					marginBottom: "1rem",
-					fontSize: "0.875rem",
-				}}
-			>
-				This state is shared between the research agent and the frontend via the
-				AG-UI protocol.
-			</p>
 
-			<div
+			{/* Phase Progress */}
-				style={{
+			<div style={{ marginBottom: "2rem" }}>
-					display: "grid",
+				<div
-					gap: "1rem",
+					style={{
-					marginTop: "1rem",
+						fontSize: "0.875rem",
-				}}
+						color: "#718096",
-			>
+						marginBottom: "0.5rem",
+					}}
+				>
+					Progress
+				</div>
+				<div style={{ display: "flex", gap: "0.5rem", alignItems: "center" }}>
+					{phases.slice(1).map((phase, idx) => (
+						<div key={phase} style={{ display: "flex", alignItems: "center" }}>
+							<div
+								style={{
+									padding: "0.25rem 0.75rem",
+									background:
+										idx < currentPhaseIndex
+											? "#48bb78"
+											: idx === currentPhaseIndex
+												? "#4299e1"
+												: "#e2e8f0",
+									color:
+										idx < currentPhaseIndex || idx === currentPhaseIndex
+											? "white"
+											: "#718096",
+									borderRadius: "4px",
+									fontSize: "0.75rem",
+									fontWeight: "600",
+									textTransform: "capitalize",
+								}}
+							>
+								{phase}
+							</div>
+							{idx < phases.length - 2 && (
+								<div
+									style={{
+										width: "1rem",
+										height: "2px",
+										background: idx < currentPhaseIndex ? "#48bb78" : "#e2e8f0",
+										margin: "0 0.25rem",
+									}}
+								/>
+							)}
+						</div>
+					))}
+				</div>
+			</div>
+
+			{/* Question */}
+			{state.question && (
 				<div
 					style={{
 						padding: "1rem",
 						background: "#f7fafc",
 						borderRadius: "4px",
 						border: "1px solid #e2e8f0",
+						marginBottom: "1rem",
 					}}
 				>
 					<div
@@ -79,61 +167,54 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 							color: "#2d3748",
 						}}
 					>
-						{state.question || "No question yet"}
+						{state.question}
 					</div>
 				</div>
+			)}
 
+			{/* Confidence Meter */}
+			{state.confidence > 0 && (
 				<div
 					style={{
-						display: "grid",
+						padding: "1rem",
-						gridTemplateColumns: "repeat(2, 1fr)",
+						background: "#f7fafc",
-						gap: "1rem",
+						borderRadius: "4px",
+						border: "1px solid #e2e8f0",
+						marginBottom: "1rem",
 					}}
 				>
 					<div
 						style={{
-							padding: "1rem",
+							fontSize: "0.875rem",
-							background: "#f7fafc",
+							color: "#718096",
-							borderRadius: "4px",
+							marginBottom: "0.5rem",
-							border: "1px solid #e2e8f0",
 						}}
 					>
-						<div
+						Confidence
-							style={{
-								fontSize: "0.875rem",
-								color: "#718096",
-								marginBottom: "0.25rem",
-							}}
-						>
-							Status
-						</div>
-						<div
-							style={{
-								fontSize: "1.125rem",
-								fontWeight: "bold",
-								color: state.status === "idle" ? "#718096" : "#38a169",
-							}}
-						>
-							{state.status}
-						</div>
 					</div>
-
+					<div style={{ display: "flex", alignItems: "center", gap: "1rem" }}>
-					<div
-						style={{
-							padding: "1rem",
-							background: "#f7fafc",
-							borderRadius: "4px",
-							border: "1px solid #e2e8f0",
-						}}
-					>
 						<div
 							style={{
-								fontSize: "0.875rem",
+								flex: 1,
-								color: "#718096",
+								height: "1rem",
-								marginBottom: "0.25rem",
+								background: "#e2e8f0",
+								borderRadius: "4px",
+								overflow: "hidden",
 							}}
 						>
-							Confidence
+							<div
+								style={{
+									width: `${state.confidence * 100}%`,
+									height: "100%",
+									background:
+										state.confidence > 0.8
+											? "#48bb78"
+											: state.confidence > 0.5
+												? "#ed8936"
+												: "#f56565",
+									transition: "width 0.3s ease",
+								}}
+							/>
 						</div>
 						<div
 							style={{
@@ -141,137 +222,456 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 								fontWeight: "bold",
 								color:
 									state.confidence > 0.8
-										? "#38a169"
+										? "#48bb78"
 										: state.confidence > 0.5
-											? "#d69e2e"
+											? "#ed8936"
-											: "#e53e3e",
+											: "#f56565",
 							}}
 						>
 							{(state.confidence * 100).toFixed(0)}%
 						</div>
 					</div>
 				</div>
+			)}
 
-				<div
+			{/* Research Plan */}
-					style={{
+			{state.plan.length > 0 && (
-						display: "grid",
+				<div style={{ marginBottom: "1rem" }}>
-						gridTemplateColumns: "repeat(3, 1fr)",
+					<button
-						gap: "1rem",
+						type="button"
-					}}
+						onClick={() => toggleSection("plan")}
-				>
-					<div
 						style={{
-							padding: "1rem",
+							width: "100%",
-							background: "#f7fafc",
+							display: "flex",
-							borderRadius: "4px",
+							justifyContent: "space-between",
+							alignItems: "center",
+							padding: "0.75rem",
+							background: "#edf2f7",
 							border: "1px solid #e2e8f0",
+							borderRadius: "4px",
+							cursor: "pointer",
+							fontSize: "1rem",
+							fontWeight: "600",
+							color: "#2d3748",
 						}}
 					>
+						<span>Research Plan ({state.plan.length} questions)</span>
+						<span>{expandedSections.plan ? "▼" : "▶"}</span>
+					</button>
+					{expandedSections.plan && (
 						<div
 							style={{
-								fontSize: "0.875rem",
+								padding: "1rem",
-								color: "#718096",
+								background: "#f7fafc",
-								marginBottom: "0.25rem",
+								border: "1px solid #e2e8f0",
+								borderTop: "none",
+								borderRadius: "0 0 4px 4px",
 							}}
 						>
-							Progress
+							{state.plan.map((item) => (
+								<div
+									key={item.id}
+									style={{
+										padding: "0.75rem",
+										background: "white",
+										borderRadius: "4px",
+										marginBottom: "0.5rem",
+										border: "1px solid #e2e8f0",
+										display: "flex",
+										gap: "0.75rem",
+									}}
+								>
+									<div
+										style={{
+											fontSize: "1.25rem",
+											color:
+												item.status === "done"
+													? "#48bb78"
+													: item.status === "searching"
+														? "#4299e1"
+														: "#a0aec0",
+										}}
+									>
+										{item.status === "done"
+											? "✓"
+											: item.status === "searching"
+												? "🔍"
+												: "⏳"}
+									</div>
+									<div style={{ flex: 1 }}>
+										<div
+											style={{
+												fontSize: "0.875rem",
+												color: "#4a5568",
+											}}
+										>
+											{item.question}
+										</div>
+									</div>
+								</div>
+							))}
 						</div>
-						<div
+					)}
-							style={{
-								fontSize: "1.125rem",
-								fontWeight: "bold",
-								color: "#2d3748",
-							}}
-						>
-							{state.current_iteration} / {state.max_iterations}
-						</div>
-					</div>
-
-					<div
-						style={{
-							padding: "1rem",
-							background: "#f7fafc",
-							borderRadius: "4px",
-							border: "1px solid #e2e8f0",
-						}}
-					>
-						<div
-							style={{
-								fontSize: "0.875rem",
-								color: "#718096",
-								marginBottom: "0.25rem",
-							}}
-						>
-							Plan Items
-						</div>
-						<div
-							style={{
-								fontSize: "1.125rem",
-								fontWeight: "bold",
-								color: "#2d3748",
-							}}
-						>
-							{state.plan.length}
-						</div>
-					</div>
-
-					<div
-						style={{
-							padding: "1rem",
-							background: "#f7fafc",
-							borderRadius: "4px",
-							border: "1px solid #e2e8f0",
-						}}
-					>
-						<div
-							style={{
-								fontSize: "0.875rem",
-								color: "#718096",
-								marginBottom: "0.25rem",
-							}}
-						>
-							Findings
-						</div>
-						<div
-							style={{
-								fontSize: "1.125rem",
-								fontWeight: "bold",
-								color: "#2d3748",
-							}}
-						>
-							{state.findings.length}
-						</div>
-					</div>
 				</div>
+			)}
 
-				<div
+			{/* Current Search Results */}
-					style={{
+			{state.current_search && (
-						padding: "1rem",
+				<div style={{ marginBottom: "1rem" }}>
-						background: "#f7fafc",
+					<button
-						borderRadius: "4px",
+						type="button"
-						border: "1px solid #e2e8f0",
+						onClick={() => toggleSection("search")}
-					}}
-				>
-					<div
 						style={{
-							fontSize: "0.875rem",
+							width: "100%",
-							color: "#718096",
+							display: "flex",
-							marginBottom: "0.25rem",
+							justifyContent: "space-between",
+							alignItems: "center",
+							padding: "0.75rem",
+							background: "#edf2f7",
+							border: "1px solid #e2e8f0",
+							borderRadius: "4px",
+							cursor: "pointer",
+							fontSize: "1rem",
+							fontWeight: "600",
+							color: "#2d3748",
 						}}
 					>
-						Final Report
+						<span>
-					</div>
+							Search Results: {state.current_search.query.substring(0, 50)}...
-					<div
+						</span>
-						style={{
+						<span>{expandedSections.search ? "▼" : "▶"}</span>
-							fontSize: "1.125rem",
+					</button>
-							fontWeight: "bold",
+					{expandedSections.search && (
-							color: state.final_report ? "#38a169" : "#a0aec0",
+						<div
-						}}
+							style={{
-					>
+								padding: "1rem",
-						{state.final_report ? "Ready" : "Not ready"}
+								background: "#f7fafc",
-					</div>
+								border: "1px solid #e2e8f0",
+								borderTop: "none",
+								borderRadius: "0 0 4px 4px",
+							}}
+						>
+							{state.current_search.results && (
+								<div>
+									<div
+										style={{
+											fontSize: "0.875rem",
+											color: "#718096",
+											marginBottom: "0.5rem",
+										}}
+									>
+										Type: {state.current_search.type} |{" "}
+										{state.current_search.results.length} results
+									</div>
+									{state.current_search.results.map((result, idx) => (
+										<div
+											key={`${result.source}-${idx}`}
+											style={{
+												padding: "0.75rem",
+												background: "white",
+												borderRadius: "4px",
+												marginBottom: "0.5rem",
+												border: "1px solid #e2e8f0",
+											}}
+										>
+											<div
+												style={{
+													display: "flex",
+													justifyContent: "space-between",
+													marginBottom: "0.5rem",
+												}}
+											>
+												<span
+													style={{
+														fontSize: "0.875rem",
+														fontWeight: "600",
+														color: "#4a5568",
+													}}
+												>
+													{result.source}
+												</span>
+												<div style={{ display: "flex", gap: "0.5rem" }}>
+													{result.expanded && (
+														<span
+															style={{
+																fontSize: "0.75rem",
+																padding: "0.125rem 0.5rem",
+																background: "#bee3f8",
+																color: "#2c5282",
+																borderRadius: "4px",
+															}}
+														>
+															Expanded
+														</span>
+													)}
+													<span
+														style={{
+															fontSize: "0.875rem",
+															fontWeight: "bold",
+															color:
+																result.score > 0.8
+																	? "#48bb78"
+																	: result.score > 0.6
+																		? "#ed8936"
+																		: "#a0aec0",
+														}}
+													>
+														{result.score.toFixed(2)}
+													</span>
+												</div>
+											</div>
+											<div
+												style={{
+													fontSize: "0.875rem",
+													color: "#718096",
+													lineHeight: "1.4",
+												}}
+											>
+												{result.chunk}...
+											</div>
+										</div>
+									))}
+								</div>
+							)}
+						</div>
+					)}
 				</div>
-			</div>
+			)}
+
+			{/* Insights */}
+			{state.insights.length > 0 && (
+				<div style={{ marginBottom: "1rem" }}>
+					<button
+						type="button"
+						onClick={() => toggleSection("insights")}
+						style={{
+							width: "100%",
+							display: "flex",
+							justifyContent: "space-between",
+							alignItems: "center",
+							padding: "0.75rem",
+							background: "#edf2f7",
+							border: "1px solid #e2e8f0",
+							borderRadius: "4px",
+							cursor: "pointer",
+							fontSize: "1rem",
+							fontWeight: "600",
+							color: "#2d3748",
+						}}
+					>
+						<span>Key Insights ({state.insights.length})</span>
+						<span>{expandedSections.insights ? "▼" : "▶"}</span>
+					</button>
+					{expandedSections.insights && (
+						<div
+							style={{
+								padding: "1rem",
+								background: "#f7fafc",
+								border: "1px solid #e2e8f0",
+								borderTop: "none",
+								borderRadius: "0 0 4px 4px",
+							}}
+						>
+							{state.insights.map((insight, idx) => (
+								<div
+									key={`${insight.summary.substring(0, 30)}-${idx}`}
+									style={{
+										padding: "0.75rem",
+										background: "white",
+										borderRadius: "4px",
+										marginBottom: "0.5rem",
+										border: "1px solid #e2e8f0",
+									}}
+								>
+									<div
+										style={{
+											display: "flex",
+											justifyContent: "space-between",
+											marginBottom: "0.5rem",
+										}}
+									>
+										<span
+											style={{
+												fontSize: "0.75rem",
+												padding: "0.125rem 0.5rem",
+												background: "#c6f6d5",
+												color: "#22543d",
+												borderRadius: "4px",
+											}}
+										>
+											{(insight.confidence * 100).toFixed(0)}% confidence
+										</span>
+										<span
+											style={{
+												fontSize: "0.75rem",
+												color: "#718096",
+											}}
+										>
+											{insight.sources.length} sources
+										</span>
+									</div>
+									<div
+										style={{
+											fontSize: "0.875rem",
+											color: "#2d3748",
+											lineHeight: "1.5",
+										}}
+									>
+										{insight.summary}
+									</div>
+								</div>
+							))}
+						</div>
+					)}
+				</div>
+			)}
+
+			{/* Final Report */}
+			{state.final_report && (
+				<div>
+					<button
+						type="button"
+						onClick={() => toggleSection("report")}
+						style={{
+							width: "100%",
+							display: "flex",
+							justifyContent: "space-between",
+							alignItems: "center",
+							padding: "0.75rem",
+							background: "#edf2f7",
+							border: "1px solid #e2e8f0",
+							borderRadius: "4px",
+							cursor: "pointer",
+							fontSize: "1rem",
+							fontWeight: "600",
+							color: "#2d3748",
+						}}
+					>
+						<span>Final Report</span>
+						<span>{expandedSections.report ? "▼" : "▶"}</span>
+					</button>
+					{expandedSections.report && (
+						<div
+							style={{
+								padding: "1.5rem",
+								background: "white",
+								border: "1px solid #e2e8f0",
+								borderTop: "none",
+								borderRadius: "0 0 4px 4px",
+							}}
+						>
+							<h3
+								style={{
+									fontSize: "1.25rem",
+									fontWeight: "600",
+									marginBottom: "1rem",
+									color: "#2d3748",
+								}}
+							>
+								{state.final_report.title}
+							</h3>
+							<div style={{ marginBottom: "1.5rem" }}>
+								<h4
+									style={{
+										fontSize: "0.875rem",
+										fontWeight: "600",
+										color: "#718096",
+										marginBottom: "0.5rem",
+									}}
+								>
+									Executive Summary
+								</h4>
+								<p
+									style={{
+										fontSize: "0.875rem",
+										color: "#4a5568",
+										lineHeight: "1.6",
+									}}
+								>
+									{state.final_report.summary}
+								</p>
+							</div>
+							<div style={{ marginBottom: "1.5rem" }}>
+								<h4
+									style={{
+										fontSize: "0.875rem",
+										fontWeight: "600",
+										color: "#718096",
+										marginBottom: "0.5rem",
+									}}
+								>
+									Main Findings
+								</h4>
+								<ul
+									style={{
+										paddingLeft: "1.5rem",
+										fontSize: "0.875rem",
+										color: "#4a5568",
+										lineHeight: "1.6",
+									}}
+								>
+									{state.final_report.findings.map((finding) => (
+										<li key={finding} style={{ marginBottom: "0.5rem" }}>
+											{finding}
+										</li>
+									))}
+								</ul>
+							</div>
+							<div style={{ marginBottom: "1.5rem" }}>
+								<h4
+									style={{
+										fontSize: "0.875rem",
+										fontWeight: "600",
+										color: "#718096",
+										marginBottom: "0.5rem",
+									}}
+								>
+									Conclusions
+								</h4>
+								<ul
+									style={{
+										paddingLeft: "1.5rem",
+										fontSize: "0.875rem",
+										color: "#4a5568",
+										lineHeight: "1.6",
+									}}
+								>
+									{state.final_report.conclusions.map((conclusion) => (
+										<li key={conclusion} style={{ marginBottom: "0.5rem" }}>
+											{conclusion}
+										</li>
+									))}
+								</ul>
+							</div>
+							<div>
+								<h4
+									style={{
+										fontSize: "0.875rem",
+										fontWeight: "600",
+										color: "#718096",
+										marginBottom: "0.5rem",
+									}}
+								>
+									Sources
+								</h4>
+								<div
+									style={{
+										fontSize: "0.75rem",
+										color: "#718096",
+										lineHeight: "1.4",
+									}}
+								>
+									{state.final_report.sources.map((source) => (
+										<div key={source} style={{ marginBottom: "0.25rem" }}>
+											{source}
+										</div>
+									))}
+								</div>
+							</div>
+						</div>
+					)}
+				</div>
+			)}
 		</div>
 	);
 }