Refactor ag-ui-research. Drop human-in-the-loop, use MemoryObjectSendStream to merge the graph and agent streams together

CHANGELOG.md

@@ -21,6 +21,12 @@
 - **CLI AG-UI Flag**: New `--agui` flag for `serve` command to start AG-UI server
 - **Graph Module**: New unified `haiku.rag.graph` module containing all graph-related functionality
 - **Common Graph Nodes**: New factory functions (`create_plan_node`, `create_search_node`) in `haiku.rag.graph.common.nodes` for reusable graph components
+- **AG-UI Research Example**: New full-stack example (`examples/ag-ui-research`) demonstrating agent+graph architecture with CopilotKit frontend
+  - Pydantic AI agent with research tool that invokes the research graph
+  - Custom AG-UI streaming endpoint with anyio memory streams
+  - React/Next.js frontend with split-pane UI showing live research state
+  - Real-time progress tracking of questions, answers, insights, and gaps
+  - Docker Compose setup for easy local development
 
 ### Changed
 

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

@@ -3,6 +3,11 @@
 # Must be an absolute path to an existing database created with haiku-rag
 DB_PATH=/absolute/path/to/your/haiku.rag.lancedb
 
+# Ollama API base URL (if using Ollama for local models)
+# If running Ollama on your host machine, use your machine's IP address
+# that the Docker container can reach (not localhost)
+OLLAMA_BASE_URL=http://host.docker.internal:11434
+
 # API keys (set as needed for your QA provider)
 # OPENAI_API_KEY=your-key-here
 # ANTHROPIC_API_KEY=your-key-here

examples/ag-ui-research/README.md

@@ -1,16 +1,16 @@
 # Interactive Research Assistant
 
-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.
+Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
 
 [Watch demo video](https://vimeo.com/1128874386)
 
 ## Features
 
-- **Multi-step research workflow**: Question decomposition, search, analysis, and synthesis
+- **Multi-iteration research graph**: Automated question decomposition, search, insight extraction, and gap analysis
-- **Human-in-the-loop**: Approve or revise research plans before execution
+- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered
-- **Live state synchronization**: Real-time updates of research progress between backend and frontend
+- **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol
-- **Context expansion**: Automatically expands top search results for better context
+- **Insight & gap tracking**: Structured insights with provenance and automatic gap identification
-- **Rich reporting**: Generates structured reports with findings, conclusions, and citations
+- **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
 
 ## Quick Start
 
@@ -57,32 +57,67 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
 ## 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
+2. **Plan phase**: The research graph automatically:
-3. **Approve or revise**: Choose to approve the plan or request changes
+   - Decomposes your question into targeted sub-questions
-4. **Watch it work**: The agent automatically:
+   - Gathers initial context about the topic
-   - Searches the knowledge base for each sub-question
+3. **Research iterations**: The graph autonomously:
-   - Extracts key insights from search results
+   - Searches the knowledge base for each sub-question in parallel
-   - Evaluates overall confidence in findings
+   - Extracts structured insights with source provenance
-5. **Get your report**: Receive a structured research report with citations
+   - Identifies information gaps and assesses confidence
+   - Generates new follow-up questions for gaps
+   - Iterates until confidence threshold is met or max iterations reached
+4. **Synthesis**: Generates a comprehensive research report with:
+   - Executive summary
+   - Main findings with supporting evidence
+   - Conclusions and recommendations
+   - Source citations
 
 ## Architecture
 
-- **Backend** (Python): Pydantic AI agent with haiku.rag integration
+### Agent + Graph Pattern
-  - Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
-  - `agent.py`: Research agent with tool definitions
-  - `main.py`: Starlette app serving AG-UI protocol
 
-- **Frontend** (Next.js): CopilotKit/AG-UI interface
+This example demonstrates the **agent+graph** architecture pattern:
-  - Real-time state synchronization with backend
+
-  - Interactive approval workflow
+1. **Conversational Agent** (`agent.py`):
-  - Collapsible research plan and insights display
+   - Pydantic AI agent handles user conversations
+   - Decides when to invoke the research tool based on user intent
+   - Responds directly to greetings/casual chat without tools
+   - Formats research results for the user
+
+2. **Research Graph** (haiku.rag):
+   - Multi-step research workflow invoked by the agent's tool
+   - Autonomous execution with plan → search → analyze → decide → synthesize flow
+   - Emits AG-UI events for real-time progress tracking
+
+3. **Shared Event Stream**:
+   - `AGUIEmitter` is shared between agent and graph
+   - Events from both flow through a single stream to the frontend
+   - Custom streaming endpoint (`main.py`) uses anyio memory streams for proper async handling
+
+### Components
+
+- **Backend** (Python):
+  - Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
+  - `agent.py`: Pydantic AI agent with `run_research` tool
+  - `main.py`: Custom AG-UI streaming endpoint with anyio memory object streams
+  - Real-time event forwarding from emitter to SSE stream
+  - Filters out `ACTIVITY_SNAPSHOT` events (not yet supported by CopilotKit)
+
+- **Frontend** (Next.js/React):
+  - CopilotKit for AG-UI protocol integration
+  - Split-pane UI: chat on left, live research state on right
+  - Real-time state synchronization via Server-Sent Events (SSE)
+  - `StateDisplay` component with collapsible sections for questions, insights, and gaps
 
 ## Configuration
 
 Configuration is done through `haiku.rag.yaml` (see `haiku.rag.yaml.example`):
 
-- `qa.provider`: LLM provider (default: `ollama`)
+- `research.provider`: LLM provider (default: `ollama`)
-- `qa.model`: Model name (default: `gpt-oss:latest`)
+- `research.model`: Model name (default: `gpt-oss:latest`)
+- `research.max_iterations`: Maximum research iterations (default: `3`)
+- `research.confidence_threshold`: Confidence threshold for completion (default: `0.8`)
+- `research.max_concurrency`: Parallel sub-question processing (default: `1`)
 - `providers.ollama.base_url`: Ollama endpoint (default: `http://host.docker.internal:11434`)
 
 Environment variables (see `.env.example`):

examples/ag-ui-research/backend/Dockerfile

@@ -3,15 +3,13 @@ FROM ghcr.io/ggozad/haiku.rag:latest
 WORKDIR /app
 
 # Copy backend application files
-COPY agent.py main.py ./
+COPY main.py agent.py ./
-COPY pyproject.toml ./
 
-# Install backend dependencies
+# Install additional dependencies for the example
-# Note: haiku-rag is already installed in the base image
+# Note: haiku-rag-slim is already installed in the base image
 RUN pip install --no-cache-dir \
     starlette>=0.45.2 \
     uvicorn[standard]>=0.34.2 \
-    pydantic-ai-slim[ag-ui,openai]>=1.1.0 \
     python-dotenv>=1.0.1
 
 EXPOSE 8000

examples/ag-ui-research/backend/README.md

@@ -1,6 +1,6 @@
 # Haiku.rag Research Assistant Backend
 
-FastAPI backend for the haiku.rag interactive research assistant, using Pydantic AI with AG-UI protocol support.
+Starlette backend for the haiku.rag interactive research assistant, using the research graph with AG-UI protocol support.
 
 ## Setup
 
@@ -11,7 +11,17 @@ uv run python main.py
 
 The server starts on `http://localhost:8000` and uses [haiku.rag configuration](https://ggozad.github.io/haiku.rag/configuration/).
 
+## Architecture
+
+The backend uses `create_agui_server()` from `haiku.rag.graph.agui.server` which provides:
+
+- **Research graph execution**: Multi-iteration research workflow with insight/gap tracking
+- **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
+- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
+- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa`
+
 ## Endpoints
 
-- `GET /health` - Health check
+- `GET /health` - Health check with configuration info
-- `POST /agent` - AG-UI protocol endpoint
+- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol)
+- `POST /agent/deep_qa/stream` - Deep QA graph streaming endpoint (AG-UI protocol)

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

@@ -1,432 +1,111 @@
-import json
+"""Research assistant agent with graph integration."""
-from dataclasses import dataclass
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
 
-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.config import load_yaml_config
+from haiku.rag.config.models import AppConfig
 from haiku.rag.graph.common import get_model
+from haiku.rag.graph.research.dependencies import ResearchContext
+from haiku.rag.graph.research.graph import build_research_graph
+from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 
+if TYPE_CHECKING:
+    from haiku.rag.graph.agui.emitter import AGUIEmitter
+    from haiku.rag.graph.research.models import ResearchReport
 
-class ResearchState(BaseModel):
+# Load config
-    """Shared state between research agent and frontend."""
+config_path = Path("/app/haiku.rag.yaml")
-
+Config = (
-    question: str = ""
+    AppConfig.model_validate(load_yaml_config(config_path))
-    phase: str = "idle"
+    if config_path.exists()
-    status: str = ""
+    else AppConfig()
-    plan: list[dict] = []
+)
-    current_question_index: int = 0
-    insights: list[dict] = []
-    document_registry: dict[str, dict] = {}
-    current_document: dict | None = None
-    confidence: float = 0.0
-    final_report: dict | None = None
 
 
 @dataclass
-class ResearchDeps(StateDeps[ResearchState]):
+class AgentDeps:
-    """Dependencies for the research agent with HaikuRAG client."""
+    """Dependencies for research agent."""
 
     client: HaikuRAG
+    agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
 
 
-def _as_state_snapshot(ctx: RunContext[ResearchDeps]) -> StateSnapshotEvent:
+model = get_model(Config.research.provider, Config.research.model)
-    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)
 
-
+agent = Agent(
-def create_agent(
+    model,
-    qa_provider: str = Config.qa.provider, qa_model: str = Config.qa.model
+    deps_type=AgentDeps,
-) -> Agent[ResearchDeps, str]:
+    system_prompt="""You are an advanced research assistant powered by haiku.rag.
-    """Create and configure the research agent.
-
-    Args:
-        qa_provider: QA provider for the agent (default: from Config.qa.provider)
-        qa_model: Model name to use (default: from Config.qa.model)
-    """
-    print(f"[AGENT SETUP] Creating agent with provider={qa_provider}, model={qa_model}")
-    agent = Agent(
-        model=get_model(qa_provider, qa_model),
-        deps_type=ResearchDeps,
-        instructions="""You are a research co-pilot powered by haiku.rag.
-
-Your workflow MUST follow these exact steps in order:
-1. Call propose_research_plan with the user's question
-2. After propose_research_plan completes, IMMEDIATELY call approve_research_plan (with no arguments)
-3. WAIT for approve_research_plan to return:
-   - If it returns "APPROVED", proceed to step 4
-   - If it returns "REVISE", ask the user "How would you like me to revise the research plan?" and wait for their response
-   - Once you receive their revision feedback, revise the plan and go back to step 1
-4. Once approved, process questions ONE AT A TIME:
-   - Call search_question(question_id=0) and WAIT for it to complete
-   - Then call extract_insights_from_results(question_id=0) and WAIT for it to complete
-   - Then call search_question(question_id=1) and WAIT for it to complete
-   - Then call extract_insights_from_results(question_id=1) and WAIT for it to complete
-   - Then call search_question(question_id=2) and WAIT for it to complete
-   - Then call extract_insights_from_results(question_id=2) and WAIT for it to complete
-5. After all questions are processed, call evaluate_research_confidence
-6. Ask user if they want to finalize or continue researching
-7. When user approves, call synthesize_final_report
 
 CRITICAL RULES:
-- MANDATORY: Call approve_research_plan immediately after propose_research_plan - NO EXCEPTIONS
+1. For greetings (hi, hello, hey, etc) or casual chat: respond directly WITHOUT using any tools
-- If approve_research_plan returns "REVISE", ask the user for revision feedback naturally in chat
+2. For questions about yourself or the system: respond directly WITHOUT using any tools
-- Call ONE tool at a time - wait for each tool to return before calling the next
+3. For substantive questions requiring information: ALWAYS use the run_research tool
-- NEVER call extract_insights_from_results until search_question has completed and returned results
+4. NEVER answer substantive questions from your own knowledge - always use the tool
-- DO NOT explain what you're about to do - just call the tool
-- The state updates will show the user what's happening - you don't need to narrate
-- Process all 3 questions automatically without asking for approval between them
 
-Document Viewing:
+How to decide:
-- When user asks to "show document X", call get_full_document with the document_uri
+- "Hi" / "Hello" / "How are you?" → Respond directly, NO tools
+- "What can you do?" → Respond directly, NO tools
+- "How does X work in the codebase?" → Use run_research tool
+- "Tell me about Y" → Use run_research tool
 
-Remember: Call tools ONE AT A TIME in sequence. Each tool must complete before calling the next.
+When you use run_research, the graph will decompose questions, search the knowledge base,
-""",
+extract insights, and generate a comprehensive report.
+
+Be friendly and conversational in all responses.""",
+)
+
+
+@agent.tool
+async def run_research(ctx: RunContext[AgentDeps], question: str) -> str:
+    """Execute research graph on a substantive question.
+
+    Use for questions requiring knowledge base search.
+    DO NOT use for greetings or casual conversation.
+    """
+    if ctx.deps.agui_emitter:
+        ctx.deps.agui_emitter.log(f"🔍 Starting research on: {question}")
+
+    graph = build_research_graph(Config)
+    context = ResearchContext(original_question=question)
+    state = ResearchState.from_config(context=context, config=Config)
+
+    graph_deps = ResearchDeps(
+        client=ctx.deps.client,
+        agui_emitter=ctx.deps.agui_emitter,
     )
 
-    @agent.tool
+    try:
-    async def propose_research_plan(
+        result = await graph.run(state=state, deps=graph_deps)
-        ctx: RunContext[ResearchDeps], question: str
-    ) -> StateSnapshotEvent:
-        """Propose a research plan by decomposing the question into sub-questions."""
-        ctx.deps.state.question = question
-        ctx.deps.state.phase = "planning"
-        ctx.deps.state.status = "Decomposing question into sub-questions..."
 
-        decompose_prompt = f"""Break down this research question into exactly 3 specific sub-questions that would help answer it comprehensively.
+        if ctx.deps.agui_emitter:
+            ctx.deps.agui_emitter.log("✅ Research complete!")
 
-Research Question: {question}
+        return f"""Research completed successfully!
 
-Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?", "Question 3?"]"""
+Question: {question}
 
-        response = await ctx.deps.client.ask(decompose_prompt)
+Executive Summary: {result.executive_summary}
 
-        try:
+Main Findings:
-            sub_questions = json.loads(response)
+{chr(10).join(f"- {finding}" for finding in result.main_findings[:3])}
-        except json.JSONDecodeError:
-            sub_questions = [
-                q.strip().lstrip("0123456789.-) ")
-                for q in response.split("\n")
-                if q.strip()
-            ][:3]
 
-        plan = [
+Conclusions:
-            {"id": i, "question": q, "status": "pending"}
+{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])}
-            for i, q in enumerate(sub_questions)
-        ]
 
-        ctx.deps.state.plan = plan
+Total insights gathered: {len(state.context.insights)}
-        ctx.deps.state.current_question_index = 0
+Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
-        ctx.deps.state.status = f"Proposed plan with {len(plan)} sub-questions"
+Iterations completed: {state.iterations}
 
-        return _as_state_snapshot(ctx)
+The full research report with all citations has been provided to the user.
+"""
 
-    @agent.tool
+    except Exception as e:
-    async def search_question(
+        if ctx.deps.agui_emitter:
-        ctx: RunContext[ResearchDeps],
+            ctx.deps.agui_emitter.log(f"❌ Research error: {str(e)}")
-        question_id: int,
+        return f"I encountered an error while researching: {str(e)}"
-        search_type: str = "hybrid",
-    ) -> StateSnapshotEvent:
-        """Execute search for a specific sub-question."""
-        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"]
-        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"
-
-        search_results = await ctx.deps.client.search(
-            question, limit=5, search_type=search_type
-        )
-
-        expanded_map = {}
-        if search_results:
-            expanded_results = await ctx.deps.client.expand_context(
-                search_results[:3], radius=2
-            )
-            expanded_map = {
-                chunk.id: (chunk, score) for chunk, score in expanded_results
-            }
-
-        results = []
-        for chunk, score in search_results:
-            doc_uri = chunk.document_uri or "unknown"
-            doc_title = chunk.document_title or chunk.document_uri or "Unknown"
-
-            if doc_uri not in ctx.deps.state.document_registry:
-                ctx.deps.state.document_registry[doc_uri] = {
-                    "title": doc_title,
-                    "chunks_referenced": [],
-                }
-
-            if (
-                chunk.id
-                not in ctx.deps.state.document_registry[doc_uri]["chunks_referenced"]
-            ):
-                ctx.deps.state.document_registry[doc_uri]["chunks_referenced"].append(
-                    chunk.id
-                )
-
-            expanded_chunk, _ = (
-                expanded_map[chunk.id] if chunk.id in expanded_map else (chunk, score)
-            )
-            result_data = {
-                "chunk": expanded_chunk.content[:500],
-                "chunk_id": chunk.id,
-                "document_uri": doc_uri,
-                "document_title": doc_title,
-                "chunk_position": chunk.order,
-                "full_chunk_content": expanded_chunk.content,
-                "score": round(score, 3),
-                "expanded": chunk.id in expanded_map,
-            }
-            results.append(result_data)
-
-        plan[question_id]["search_results"] = {
-            "type": search_type,
-            "results": results,
-        }
-        plan[question_id]["status"] = "searched"
-        ctx.deps.state.status = f"Found {len(results)} results"
-
-        return _as_state_snapshot(ctx)
-
-    @agent.tool
-    async def extract_insights_from_results(
-        ctx: RunContext[ResearchDeps],
-        question_id: int,
-    ) -> StateSnapshotEvent:
-        """Extract key insights from search results for a specific question."""
-        plan = ctx.deps.state.plan
-        if question_id >= len(plan):
-            raise ValueError(f"Question ID {question_id} not found in plan")
-
-        question_item = plan[question_id]
-        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."
-            )
-
-        search_results = question_item["search_results"]
-        ctx.deps.state.phase = "analyzing"
-        ctx.deps.state.status = "Extracting insights from results..."
-
-        context_parts = [
-            f"[Result {idx}] [Source: {r['document_title']}] {r['full_chunk_content']}"
-            for idx, r in enumerate(search_results["results"])
-        ]
-        context = "\n\n".join(context_parts)
-
-        class InsightResult(BaseModel):
-            summary: str
-            confidence: float
-            result_indices: list[int]
-
-        class InsightsList(BaseModel):
-            insights: list[InsightResult]
-
-        question_text = question_item["question"]
-        extract_prompt = f"""Analyze these search results and extract 1-3 key insights that help answer the question: "{question_text}"
-
-Search Results:
-{context}
-
-For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
-
-        insight_agent: Agent[None, InsightsList] = Agent(
-            ctx.model,
-            output_type=InsightsList,
-            retries=3,
-        )
-
-        result = await insight_agent.run(extract_prompt)
-        raw_insights = [
-            {
-                "summary": insight.summary,
-                "confidence": insight.confidence,
-                "result_indices": insight.result_indices,
-            }
-            for insight in result.output.insights
-        ]
-
-        new_insights = []
-        for insight in raw_insights:
-            source_refs = []
-            for idx in insight.get("result_indices", []):
-                if 0 <= idx < len(search_results["results"]):
-                    result = search_results["results"][idx]
-                    source_refs.append(
-                        {
-                            "chunk_id": result["chunk_id"],
-                            "document_uri": result["document_uri"],
-                            "document_title": result["document_title"],
-                            "chunk_position": result["chunk_position"],
-                        }
-                    )
-
-            new_insights.append(
-                {
-                    "summary": insight["summary"],
-                    "confidence": insight.get("confidence", 0.7),
-                    "source_refs": source_refs,
-                }
-            )
-
-        ctx.deps.state.insights.extend(new_insights)
-        plan[question_id]["status"] = "done"
-        ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
-
-        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")
-
-        ctx.deps.state.phase = "evaluating"
-        ctx.deps.state.status = "Evaluating research confidence..."
-
-        confidences = [i.get("confidence", 0.5) for i in insights]
-        overall_confidence = sum(confidences) / len(confidences) if confidences else 0
-
-        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)
-
-        try:
-            evaluation = json.loads(response)
-            overall_confidence = evaluation.get("confidence", overall_confidence)
-        except json.JSONDecodeError:
-            pass
-
-        ctx.deps.state.confidence = overall_confidence
-        ctx.deps.state.status = f"Confidence: {overall_confidence:.0%}"
-
-        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")
-
-        ctx.deps.state.phase = "synthesizing"
-        ctx.deps.state.status = "Generating final report..."
-
-        insights_summary = []
-        for i in insights:
-            source_titles = [ref["document_title"] for ref in i.get("source_refs", [])]
-            unique_sources = list(dict.fromkeys(source_titles))
-            insights_summary.append(
-                f"- {i['summary']} (sources: {', '.join(unique_sources[:2])})"
-            )
-
-        report_prompt = f"""Generate a comprehensive research report answering: "{ctx.deps.state.question}"
-
-Based on these insights:
-{chr(10).join(insights_summary)}
-
-Create a structured report with:
-- Executive Summary (2-3 sentences)
-- Main Findings (bullet points)
-- Conclusions
-- Sources (list the document titles mentioned above)
-
-Return JSON with format:
-{{
-    "title": "...",
-    "summary": "...",
-    "findings": ["finding1", "finding2", ...],
-    "conclusions": ["conclusion1", ...],
-    "sources": ["source1", "source2", ...]
-}}"""
-
-        response = await ctx.deps.client.ask(report_prompt)
-
-        try:
-            report = json.loads(response)
-        except json.JSONDecodeError:
-            report = {
-                "title": ctx.deps.state.question,
-                "summary": response[:300],
-                "findings": [i["summary"] for i in insights],
-                "conclusions": ["See findings above"],
-                "sources": [],
-            }
-
-        citations = [
-            {
-                "document_uri": doc_uri,
-                "document_title": doc_info["title"],
-                "chunk_ids": doc_info["chunks_referenced"],
-            }
-            for doc_uri, doc_info in ctx.deps.state.document_registry.items()
-        ]
-        report["citations"] = citations
-
-        ctx.deps.state.final_report = report
-        ctx.deps.state.phase = "done"
-        ctx.deps.state.status = "Research complete"
-
-        return _as_state_snapshot(ctx)
-
-    @agent.tool
-    async def get_full_document(
-        ctx: RunContext[ResearchDeps],
-        document_uri: str,
-    ) -> StateSnapshotEvent:
-        """Retrieve and display the full content of a document by its URI."""
-        ctx.deps.state.status = f"Retrieving document: {document_uri}"
-        document = await ctx.deps.client.get_document_by_uri(document_uri)
-
-        if document is None:
-            ctx.deps.state.status = f"Document not found: {document_uri}"
-            ctx.deps.state.current_document = {
-                "uri": document_uri,
-                "title": "Not Found",
-                "content": f"Document with URI '{document_uri}' was not found.",
-                "total_chunks": 0,
-            }
-        else:
-            all_chunks = await ctx.deps.client.search(
-                query="", limit=1000, search_type="fts"
-            )
-            chunks_for_doc = [
-                c for c, _ in all_chunks if c.document_uri == document_uri
-            ]
-
-            ctx.deps.state.current_document = {
-                "uri": document.uri or document_uri,
-                "title": document.title or "Untitled",
-                "content": document.content,
-                "total_chunks": len(chunks_for_doc),
-                "metadata": document.metadata,
-            }
-            ctx.deps.state.status = f"Retrieved: {document.title or document_uri}"
-
-        return _as_state_snapshot(ctx)
-
-    return agent

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

@@ -1,87 +1,171 @@
 import logging
 import os
-from contextlib import asynccontextmanager
 from pathlib import Path
 
-from agent import ResearchDeps, ResearchState, create_agent
+from agent import AgentDeps, agent
+from anyio import create_memory_object_stream, create_task_group
+from anyio.streams.memory import MemoryObjectSendStream
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import JSONResponse
+from starlette.requests import Request
-from starlette.routing import Mount, Route
+from starlette.responses import JSONResponse, StreamingResponse
+from starlette.routing import Route
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
+from haiku.rag.config import load_yaml_config
+from haiku.rag.config.models import AppConfig
+from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
+from haiku.rag.graph.research.dependencies import ResearchContext
+from haiku.rag.graph.research.models import ResearchReport
+from haiku.rag.graph.research.state import ResearchState
 
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
 logger = logging.getLogger(__name__)
 
-client: HaikuRAG | None = None
+# Load config from mounted haiku.rag.yaml
-ag_ui_app = None
+config_path = Path("/app/haiku.rag.yaml")
+if config_path.exists():
+    yaml_data = load_yaml_config(config_path)
+    Config = AppConfig.model_validate(yaml_data)
+else:
+    # Fallback to default config
+    Config = AppConfig()
 
 
-@asynccontextmanager
+# Get DB path from environment
-async def lifespan(app):
+db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
-    global client
+db_path = Path(db_path_str)
-    db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
-    db_path = Path(db_path_str)
 
-    if not db_path.exists():
+if not db_path.exists():
-        logger.error(f"Database not found at {db_path}")
+    logger.error(f"Database not found at {db_path}")
-        logger.error("Run: haiku-rag add <path-to-documents>")
+    logger.error("Run: haiku-rag add <path-to-documents>")
-        raise RuntimeError(f"Database not found: {db_path}")
+    raise RuntimeError(f"Database not found: {db_path}")
 
-    logger.info(f"Initializing HaikuRAG client with database: {db_path}")
+logger.info(f"Initializing research assistant with database: {db_path}")
-    client = HaikuRAG(db_path)
+logger.info(
-    logger.info("Research assistant backend ready")
+    f"Research Provider: {Config.research.provider}, Model: {Config.research.model}"
-    logger.info(f"QA Provider: {Config.qa.provider}, Model: {Config.qa.model}")
+)
 
-    yield
+# Store client reference for proper lifecycle management
-
+_client_cache: dict[str, HaikuRAG] = {}
-    if client:
-        logger.info("Closing HaikuRAG client")
-        client.close()
 
 
-agent = create_agent()
+def get_client(effective_db_path: Path) -> HaikuRAG:
+    """Get or create cached client."""
+    path_key = str(effective_db_path)
+    if path_key not in _client_cache:
+        _client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=Config)
+    return _client_cache[path_key]
 
 
-async def health(request):
+async def stream_research_agent(request: Request) -> StreamingResponse:
-    db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
+    """Agent streaming endpoint with research graph integration."""
+    body = await request.json()
+    input_data = RunAgentInput(**body)
+
+    user_message = ""
+    if input_data.messages:
+        user_message = input_data.messages[-1].get("content", "")
+
+    send_stream, receive_stream = create_memory_object_stream[str]()
+
+    async def run_agent_with_streaming(
+        send_stream: MemoryObjectSendStream[str],
+    ) -> None:
+        """Execute agent and forward emitter events to memory stream."""
+        async with send_stream:
+            try:
+                # Create shared emitter
+                emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter(
+                    thread_id=input_data.thread_id,
+                    run_id=input_data.run_id,
+                    use_deltas=False,
+                )
+
+                # Get client
+                effective_db_path = input_data.config.get("db_path") or db_path
+                if isinstance(effective_db_path, str):
+                    effective_db_path = Path(effective_db_path)
+                client = get_client(effective_db_path)
+
+                # Create agent dependencies with shared emitter
+                agent_deps = AgentDeps(client=client, agui_emitter=emitter)
+
+                # Start run with empty initial state
+                emitter.start_run(
+                    initial_state=ResearchState.from_config(
+                        context=ResearchContext(original_question=""),
+                        config=Config,
+                    )
+                )
+
+                # Forward emitter events to stream
+                async def forward_events():
+                    async for event in emitter:
+                        # Filter out ACTIVITY_SNAPSHOT - not supported by CopilotKit
+                        if event.get("type") == "ACTIVITY_SNAPSHOT":
+                            continue
+                        await send_stream.send(format_sse_event(event))
+
+                # Run agent and event forwarding concurrently
+                async with create_task_group() as tg:
+                    tg.start_soon(forward_events)
+
+                    result = await agent.run(user_message, deps=agent_deps)
+                    emitter.log(result.output)
+                    await emitter.close()
+
+            except Exception as e:
+                logger.exception("Error executing agent")
+                try:
+                    await send_stream.send(
+                        format_sse_event({"type": "error", "error": str(e)})
+                    )
+                except Exception:
+                    pass
+
+    async def event_generator():
+        """Generate SSE events from memory stream."""
+        async with create_task_group() as tg:
+            tg.start_soon(run_agent_with_streaming, send_stream)
+            async with receive_stream:
+                async for event_str in receive_stream:
+                    yield event_str
+
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
+async def health_check(_: Request) -> JSONResponse:
+    """Health check endpoint with configuration info."""
     return JSONResponse(
         {
             "status": "healthy",
             "agent_model": str(agent.model),
-            "qa_provider": Config.qa.provider,
+            "research_provider": Config.research.provider,
-            "qa_model": Config.qa.model,
+            "research_model": Config.research.model,
-            "ollama_base_url": Config.providers.ollama.base_url,
+            "db_path": str(db_path),
-            "db_path": db_path_str,
+            "db_exists": db_path.exists(),
-            "db_exists": Path(db_path_str).exists(),
         }
     )
 
 
-def get_ag_ui_app():
+# Create Starlette app
-    global ag_ui_app
-    if ag_ui_app is None and client is not None:
-        research_deps = ResearchDeps(client=client, state=ResearchState())
-        logger.info("Creating AG-UI app")
-        ag_ui_app = agent.to_ag_ui(deps=research_deps)
-    return ag_ui_app
-
-
-async def agent_endpoint(scope, receive, send):
-    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)
-
-
 app = Starlette(
     routes=[
-        Route("/health", health),
+        Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
-        Mount("/agent", agent_endpoint),
+        Route("/health", health_check, methods=["GET"]),
     ],
     middleware=[
         Middleware(
@@ -92,17 +176,11 @@ app = Starlette(
             allow_headers=["*"],
         )
     ],
-    lifespan=lifespan,
 )
 
 if __name__ == "__main__":
     import uvicorn
 
-    print("Starting haiku.rag research assistant backend...")
-    print(f"Agent model: {agent.model}")
-    print(f"QA provider: {Config.qa.provider}")
-    print(f"QA model: {Config.qa.model}")
-
     uvicorn.run(
         "main:app",
         host="0.0.0.0",

examples/ag-ui-research/backend/pyproject.toml

@@ -9,7 +9,7 @@ dependencies = [
     "uvicorn[standard]>=0.34.2",
     "pydantic-ai-slim[ag-ui,openai]>=1.1.0",
     "python-dotenv>=1.0.1",
-    "haiku-rag>=0.12.1",
+    "haiku-rag-slim @ file:///Users/ggozad/dev/open-source/haiku.rag-agui/haiku_rag_slim",
 ]
 
 [dependency-groups]
@@ -18,6 +18,9 @@ dev = [
     "ruff>=0.13.0",
 ]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.hatch.build.targets.wheel]
 packages = ["."]
 

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

@@ -10,11 +10,16 @@ services:
       # API keys (set these in your shell or .env file)
       - OPENAI_API_KEY=${OPENAI_API_KEY}
       - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
+      # Ollama connection (use value from .env)
+      - OLLAMA_BASE_URL=${OLLAMA_BASE_URL}
+      # Prevent Python bytecode caching for development
+      - PYTHONDONTWRITEBYTECODE=1
     volumes:
-      - ./backend:/app
-      - /app/.venv
       - ${DB_PATH}:/app/data/haiku.rag.lancedb
       - ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
+      - ./backend/main.py:/app/main.py
+      - ./backend/agent.py:/app/agent.py
+      - ../../haiku_rag_slim/haiku:/app/.venv/lib/python3.13/site-packages/haiku
     networks:
       - ag-ui-network
     extra_hosts:

examples/ag-ui-research/frontend/app/api/copilotkit/route.ts

@@ -13,7 +13,7 @@ const runtime = new CopilotRuntime({
 	agents: {
 		// "research_agent" maps to the agent name used in useCoAgent() on the frontend
 		research_agent: new HttpAgent({
-			url: `${process.env.BACKEND_URL || "http://backend:8000"}/agent`,
+			url: `${process.env.BACKEND_URL || "http://backend:8000"}/v1/research/stream`,
 		}),
 	},
 });

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

@@ -1,208 +1,94 @@
 "use client";
 
-import {
+import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
-	CopilotKit,
-	useCoAgent,
-	useCoAgentStateRender,
-	useCopilotAction,
-} from "@copilotkit/react-core";
 import { CopilotChat } from "@copilotkit/react-ui";
 import "@copilotkit/react-ui/styles.css";
 import StateDisplay from "./StateDisplay";
 
-interface SourceRef {
+interface InsightRecord {
-	chunk_id: string;
+	id: string;
-	document_uri: string;
+	summary: string;
-	document_title: string;
+	status: string;
-	chunk_position: number;
+	notes?: string;
+	supporting_sources: string[];
+	originating_questions: string[];
+}
+
+interface GapRecord {
+	id: string;
+	description: string;
+	severity: string;
+	blocking: boolean;
+	resolved: boolean;
+	notes?: string;
+	supporting_sources: string[];
+	resolved_by: string[];
+}
+
+interface SearchAnswer {
+	query: string;
+	answer: string;
+	confidence: number;
+	context: string[];
+	sources: string[];
+}
+
+interface ResearchContext {
+	original_question: string;
+	sub_questions: string[];
+	qa_responses: SearchAnswer[];
+	insights: InsightRecord[];
+	gaps: GapRecord[];
+}
+
+interface EvaluationResult {
+	confidence: number;
+	reasoning: string;
+	should_continue: boolean;
+	gaps_identified: string[];
+	follow_up_questions: string[];
+}
+
+interface ResearchReport {
+	question: string;
+	summary: string;
+	findings: string[];
+	conclusions: string[];
+	insights_used: string[];
+	methodology: string;
 }
 
 interface ResearchState {
-	question: string;
+	context: ResearchContext;
-	phase: string;
+	iterations: number;
-	status: string;
+	max_iterations: number;
-	plan: Array<{
+	confidence_threshold: number;
-		id: number;
+	max_concurrency: number;
-		question: string;
+	last_eval: EvaluationResult | null;
-		status: string;
+	last_analysis: {
-		search_results?: {
+		insights_extracted: InsightRecord[];
-			type: string;
+		gaps_identified: GapRecord[];
-			results: Array<{
-				chunk: string;
-				chunk_id: string;
-				document_uri: string;
-				document_title: string;
-				chunk_position: number;
-				full_chunk_content: string;
-				score: number;
-				expanded: boolean;
-			}>;
-		};
-	}>;
-	current_question_index: number;
-	insights: Array<{
-		summary: string;
-		confidence: number;
-		source_refs: SourceRef[];
-	}>;
-	document_registry: Record<
-		string,
-		{
-			title: string;
-			chunks_referenced: string[];
-		}
-	>;
-	current_document: {
-		uri: string;
-		title: string;
-		content: string;
-		total_chunks: number;
-		metadata?: Record<string, unknown>;
-	} | null;
-	confidence: number;
-	final_report: {
-		title: string;
-		summary: string;
-		findings: string[];
-		conclusions: string[];
-		sources: string[];
-		citations: Array<{
-			document_uri: string;
-			document_title: string;
-			chunk_ids: string[];
-		}>;
 	} | null;
+	result?: ResearchReport;
 }
 
 function AgentContent() {
 	const { state } = useCoAgent<ResearchState>({
 		name: "research_agent",
 		initialState: {
-			question: "",
+			context: {
-			phase: "idle",
+				original_question: "",
-			status: "",
+				sub_questions: [],
-			plan: [],
+				qa_responses: [],
-			current_question_index: 0,
+				insights: [],
-			insights: [],
+				gaps: [],
-			document_registry: {},
+			},
-			current_document: null,
+			iterations: 0,
-			confidence: 0.0,
+			max_iterations: 3,
-			final_report: null,
+			confidence_threshold: 0.8,
-		},
+			max_concurrency: 1,
-	});
+			last_eval: null,
-
+			last_analysis: null,
-	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 }) => (
-			<div
-				style={{
-					padding: "1.5rem",
-					background: "white",
-					borderRadius: "8px",
-					border: "2px solid #4299e1",
-					marginBottom: "1rem",
-					boxShadow: "0 2px 8px rgba(0,0,0,0.1)",
-				}}
-			>
-				<h3
-					style={{
-						fontSize: "1.25rem",
-						fontWeight: "bold",
-						marginBottom: "1rem",
-						color: "#2d3748",
-					}}
-				>
-					Research Plan Approval
-				</h3>
-				<p
-					style={{
-						fontSize: "0.875rem",
-						color: "#4a5568",
-						marginBottom: "1rem",
-					}}
-				>
-					Please review the research plan in the right pane.
-				</p>
-
-				<div
-					style={{
-						display: "flex",
-						gap: "1rem",
-					}}
-					className={status !== "executing" ? "hidden" : ""}
-				>
-					<button
-						type="button"
-						onClick={() => respond?.("REVISE")}
-						disabled={status !== "executing"}
-						style={{
-							flex: 1,
-							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,
-						}}
-					>
-						Revise Plan
-					</button>
-					<button
-						type="button"
-						onClick={() => respond?.("APPROVED")}
-						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>
-		),
-	});
-
-	useCoAgentStateRender<ResearchState>({
-		name: "research_agent",
-		render: ({ state: newState }) => {
-			const phaseMessages: Record<string, string> = {
-				planning: "Planning research...",
-				searching: "Searching...",
-				analyzing: "Extracting insights...",
-				evaluating: `Evaluating confidence: ${(newState.confidence * 100).toFixed(0)}%`,
-				synthesizing: "Generating final report...",
-				done: "Research complete!",
-			};
-			const phaseMessage =
-				phaseMessages[newState.phase] || newState.status || "Ready";
-
-			return (
-				<div
-					style={{
-						padding: "1rem",
-						background: "#e6f7ff",
-						borderRadius: "4px",
-						marginBottom: "0.5rem",
-						border: "1px solid #91d5ff",
-					}}
-				>
-					<strong>Research Update:</strong> {phaseMessage}
-				</div>
-			);
 		},
 	});
 

examples/ag-ui-research/haiku.rag.yaml.example

@@ -1,20 +1,23 @@
 # haiku.rag configuration for ag-ui-research example
 # Copy to haiku.rag.yaml and customize
 
-qa:
+research:
   provider: ollama
   model: gpt-oss:latest
+  max_iterations: 3
+  confidence_threshold: 0.8
+  max_concurrency: 1
 
 providers:
   ollama:
     base_url: http://host.docker.internal:11434
 
 # For OpenAI:
-# qa:
+# research:
 #   provider: openai
 #   model: gpt-4o-mini
 
 # For Anthropic:
-# qa:
+# research:
 #   provider: anthropic
 #   model: claude-3-5-haiku-20241022

haiku_rag_slim/haiku/rag/graph/common/nodes.py

@@ -156,8 +156,11 @@ def create_search_node[AgentDepsT: GraphAgentDeps](
         deps: GraphDeps = ctx.deps  # type: ignore[assignment]
         sub_q = ctx.inputs
 
+        # Create unique step name from question text
+        step_name = f"search: {sub_q}"
+
         if deps.agui_emitter and with_step_wrapper:
-            deps.agui_emitter.start_step("search_one")
+            deps.agui_emitter.start_step(step_name)
 
         try:
             # Create semaphore if not already provided