Flatten graph structure, simplify

haiku_rag_slim/haiku/rag/client.py

@@ -27,7 +27,7 @@ from haiku.rag.store.repositories.settings import SettingsRepository
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
 
-    from haiku.rag.graph.common.models import Citation
+    from haiku.rag.graph.research.models import Citation
 
 logger = logging.getLogger(__name__)
 

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

@@ -1,5 +0,0 @@
-"""Common utilities for graph implementations."""
-
-from haiku.rag.utils import get_model
-
-__all__ = ["get_model"]

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

@@ -1,107 +0,0 @@
-"""Common models used across different graph implementations."""
-
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel, Field, field_validator
-
-if TYPE_CHECKING:
-    from haiku.rag.store.models import SearchResult
-
-
-class ResearchPlan(BaseModel):
-    """A structured research plan with sub-questions to explore."""
-
-    sub_questions: list[str] = Field(
-        ...,
-        description="Specific questions to research, phrased as complete questions",
-    )
-
-    @field_validator("sub_questions")
-    @classmethod
-    def validate_sub_questions(cls, v: list[str]) -> list[str]:
-        if len(v) < 1:
-            raise ValueError("Must have at least 1 sub-question")
-        if len(v) > 12:
-            raise ValueError("Cannot have more than 12 sub-questions")
-        return v
-
-
-class Citation(BaseModel):
-    """Resolved citation with full metadata for display/visual grounding."""
-
-    document_id: str
-    chunk_id: str
-    document_uri: str
-    document_title: str | None = None
-    page_numbers: list[int] = Field(default_factory=list)
-    headings: list[str] | None = None
-    content: str
-
-
-class RawSearchAnswer(BaseModel):
-    """Answer to a search query with chunk references."""
-
-    query: str = Field(..., description="The question that was answered")
-    answer: str = Field(..., description="The answer to the question")
-    cited_chunks: list[str] = Field(
-        default_factory=list,
-        description="IDs of chunks used to form the answer",
-    )
-    confidence: float = Field(
-        default=1.0,
-        description="Confidence score for this answer (0-1)",
-        ge=0.0,
-        le=1.0,
-    )
-
-
-class SearchAnswer(RawSearchAnswer):
-    """Answer to a search query with resolved citations."""
-
-    citations: list[Citation] = Field(
-        default_factory=list,
-        description="Resolved citations with full metadata",
-    )
-
-    @classmethod
-    def from_raw(
-        cls,
-        raw: RawSearchAnswer,
-        search_results: "list[SearchResult]",
-    ) -> "SearchAnswer":
-        """Create SearchAnswer from RawSearchAnswer with resolved citations."""
-        citations = resolve_citations(raw.cited_chunks, search_results)
-        return cls(
-            query=raw.query,
-            answer=raw.answer,
-            cited_chunks=raw.cited_chunks,
-            confidence=raw.confidence,
-            citations=citations,
-        )
-
-
-def resolve_citations(
-    cited_chunk_ids: list[str],
-    search_results: "list[SearchResult]",
-) -> list[Citation]:
-    """Resolve chunk IDs to full Citation objects with metadata."""
-    # Build lookup by chunk_id
-    by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
-
-    citations = []
-    for chunk_id in cited_chunk_ids:
-        r = by_id.get(chunk_id)
-        if not r:
-            continue
-        citations.append(
-            Citation(
-                document_id=r.document_id or "",
-                chunk_id=chunk_id,
-                document_uri=r.document_uri or "",
-                document_title=r.document_title,
-                page_numbers=r.page_numbers,
-                headings=r.headings,
-                content=r.content,
-            )
-        )
-    return citations

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

@@ -1,315 +0,0 @@
-"""Common node implementations for graph workflows."""
-
-import asyncio
-from collections.abc import Awaitable, Callable
-from typing import Any, Protocol
-
-from pydantic_ai import Agent, RunContext
-from pydantic_ai.output import ToolOutput
-from pydantic_graph.beta import StepContext
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
-from haiku.rag.config.models import AppConfig, ModelConfig
-from haiku.rag.graph.agui.emitter import AGUIEmitter
-from haiku.rag.graph.common import get_model
-from haiku.rag.graph.common.models import RawSearchAnswer, ResearchPlan, SearchAnswer
-from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
-from haiku.rag.store.models import SearchResult
-
-
-class GraphContext(Protocol):
-    """Protocol for graph context objects."""
-
-    original_question: str
-    sub_questions: list[str]
-
-    def add_qa_response(self, qa: SearchAnswer) -> None:
-        """Add a QA response to context."""
-        ...
-
-
-class GraphState(Protocol):
-    """Protocol for graph state objects."""
-
-    context: GraphContext
-    max_concurrency: int
-    search_filter: str | None
-
-
-class GraphDeps(Protocol):
-    """Protocol for graph dependencies."""
-
-    client: HaikuRAG
-    agui_emitter: AGUIEmitter[Any, Any] | None
-    semaphore: asyncio.Semaphore | None
-
-
-class GraphAgentDeps(Protocol):
-    """Protocol for agent dependencies."""
-
-    client: HaikuRAG
-    context: GraphContext
-    search_results: list[SearchResult]
-
-
-def create_plan_node[AgentDepsT: GraphAgentDeps](
-    model_config: ModelConfig,
-    deps_type: type[AgentDepsT],
-    activity_message: str = "Creating plan",
-    output_retries: int | None = None,
-    config: AppConfig = Config,
-) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
-    """Create a plan node for any graph.
-
-    Args:
-        model_config: ModelConfig with provider, model, and settings
-        deps_type: Type of dependencies for the agent (e.g., ResearchDependencies)
-        activity_message: Message to show during planning activity
-        output_retries: Number of output retries for the agent (optional)
-        config: AppConfig object (defaults to global Config)
-
-    Returns:
-        Async function that can be used as a graph step
-    """
-
-    async def plan(ctx: StepContext[Any, Any, None], /) -> None:
-        state: GraphState = ctx.state  # type: ignore[assignment]
-        deps: GraphDeps = ctx.deps  # type: ignore[assignment]
-
-        if deps.agui_emitter:
-            deps.agui_emitter.start_step("plan")
-            deps.agui_emitter.update_activity(
-                "planning", {"stepName": "plan", "message": activity_message}
-            )
-
-        try:
-            # Build agent configuration
-            agent_config = {
-                "model": get_model(model_config, config),
-                "output_type": ResearchPlan,
-                "instructions": (
-                    PLAN_PROMPT
-                    + "\n\nUse the gather_context tool once on the main question before planning."
-                ),
-                "retries": 3,
-                "deps_type": deps_type,
-            }
-            if output_retries is not None:
-                agent_config["output_retries"] = output_retries
-
-            plan_agent = Agent(**agent_config)
-
-            # Capture search filter for use in tool
-            search_filter = state.search_filter
-
-            @plan_agent.tool
-            async def gather_context(
-                ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
-            ) -> str:
-                results = await ctx2.deps.client.search(
-                    query, limit=limit, filter=search_filter
-                )
-                results = await ctx2.deps.client.expand_context(results)
-                return "\n\n".join(r.content for r in results)
-
-            # Tool is registered via decorator above
-            _ = gather_context
-
-            prompt = (
-                "Plan a focused approach for the main question.\n\n"
-                f"Main question: {state.context.original_question}"
-            )
-
-            # Create agent dependencies
-            agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]
-            plan_result = await plan_agent.run(prompt, deps=agent_deps)
-            state.context.sub_questions = list(plan_result.output.sub_questions)
-
-            # State now contains the plan - emit state update and narrate
-            if deps.agui_emitter:
-                deps.agui_emitter.update_state(state)
-                count = len(state.context.sub_questions)
-                deps.agui_emitter.update_activity(
-                    "planning",
-                    {
-                        "stepName": "plan",
-                        "message": f"Created plan with {count} sub-questions",
-                        "sub_questions": list(state.context.sub_questions),
-                    },
-                )
-        finally:
-            if deps.agui_emitter:
-                deps.agui_emitter.finish_step()
-
-    return plan
-
-
-def create_search_node[AgentDepsT: GraphAgentDeps](
-    model_config: ModelConfig,
-    deps_type: type[AgentDepsT],
-    with_step_wrapper: bool = True,
-    success_message_format: str = "Answered: {sub_q}",
-    handle_exceptions: bool = False,
-    config: AppConfig = Config,
-) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
-    """Create a search_one node for any graph.
-
-    Args:
-        model_config: ModelConfig with provider, model, and settings
-        deps_type: Type of dependencies for the agent
-        with_step_wrapper: Whether to wrap with agui_emitter start/finish step
-        success_message_format: Format string for success activity message
-        handle_exceptions: Whether to handle exceptions with fallback answer
-        config: AppConfig object (defaults to global Config)
-
-    Returns:
-        Async function that can be used as a graph step
-    """
-
-    async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
-        state: GraphState = ctx.state  # type: ignore[assignment]
-        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(step_name)
-
-        try:
-            # Create semaphore if not already provided
-            if deps.semaphore is None:
-                deps.semaphore = asyncio.Semaphore(state.max_concurrency)
-
-            # Use semaphore to control concurrency
-            async with deps.semaphore:
-                return await _do_search(
-                    state,
-                    deps,
-                    sub_q,
-                    model_config,
-                    deps_type,
-                    success_message_format,
-                    handle_exceptions,
-                    config,
-                )
-        finally:
-            if deps.agui_emitter and with_step_wrapper:
-                deps.agui_emitter.finish_step()
-
-    return search_one
-
-
-async def _do_search[AgentDepsT: GraphAgentDeps](
-    state: GraphState,
-    deps: GraphDeps,
-    sub_q: str,
-    model_config: ModelConfig,
-    deps_type: type[AgentDepsT],
-    success_message_format: str,
-    handle_exceptions: bool,
-    config: AppConfig,
-) -> SearchAnswer:
-    """Internal search implementation."""
-    if deps.agui_emitter:
-        deps.agui_emitter.update_activity(
-            "searching",
-            {
-                "stepName": "search_one",
-                "message": f"Searching: {sub_q}",
-                "query": sub_q,
-            },
-        )
-
-    agent = Agent(
-        model=get_model(model_config, config),
-        output_type=ToolOutput(RawSearchAnswer, max_retries=3),
-        instructions=SEARCH_AGENT_PROMPT,
-        retries=3,
-        deps_type=deps_type,
-    )
-
-    # Capture search filter for use in tool
-    search_filter = state.search_filter
-
-    @agent.tool
-    async def search_and_answer(
-        ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None
-    ) -> str:
-        """Search the knowledge base for relevant documents.
-
-        Returns results with chunk IDs and relevance scores.
-        Reference results by their chunk_id in cited_chunks.
-        """
-        results = await ctx2.deps.client.search(
-            query, limit=limit, filter=search_filter
-        )
-        results = await ctx2.deps.client.expand_context(results)
-        # Store results for citation resolution
-        ctx2.deps.search_results = results
-
-        # Format with metadata for agent context
-        parts = [r.format_for_agent() for r in results]
-
-        if not parts:
-            return f"No relevant information found in the knowledge base for: {query}"
-
-        return "\n\n".join(parts)
-
-    # Tool is registered via decorator above
-    _ = search_and_answer
-
-    agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]
-
-    try:
-        result = await agent.run(sub_q, deps=agent_deps)
-        raw_answer = result.output
-        if raw_answer:
-            # Convert RawSearchAnswer to SearchAnswer with resolved citations
-            answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
-            state.context.add_qa_response(answer)
-            # State updated with new answer - emit state update and narrate
-            if deps.agui_emitter:
-                deps.agui_emitter.update_state(state)
-                # Format the success message
-                if "{confidence" in success_message_format:
-                    message = success_message_format.format(
-                        sub_q=sub_q, confidence=answer.confidence
-                    )
-                else:
-                    message = success_message_format.format(sub_q=sub_q)
-                deps.agui_emitter.update_activity(
-                    "searching",
-                    {
-                        "stepName": "search_one",
-                        "message": message,
-                        "query": sub_q,
-                        "confidence": answer.confidence,
-                    },
-                )
-            return answer
-        # Return empty SearchAnswer if no result
-        return SearchAnswer(query=sub_q, answer="", confidence=0.0)
-    except Exception as e:
-        if handle_exceptions:
-            # Narrate the error
-            if deps.agui_emitter:
-                deps.agui_emitter.update_activity(
-                    "searching",
-                    {
-                        "stepName": "search_one",
-                        "message": f"Search failed: {e}",
-                        "query": sub_q,
-                        "error": str(e),
-                    },
-                )
-            failure_answer = SearchAnswer(
-                query=sub_q,
-                answer=f"Search failed after retries: {str(e)}",
-                confidence=0.0,
-            )
-            return failure_answer
-        else:
-            raise

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

@@ -1,67 +0,0 @@
-"""Common prompts used across different graph implementations."""
-
-PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
-
-Responsibilities:
-1. Understand and decompose the main question
-2. Propose a minimal, high-leverage plan
-3. Coordinate specialized agents to gather evidence
-4. Iterate based on gaps and new findings
-
-Plan requirements:
-- Produce at most 3 sub_questions that together cover the main question.
-- sub_questions must be a list of plain strings, where each string is a complete
-  question. Do NOT use objects with nested fields like {question, details}.
-- Each sub_question must be a standalone, self-contained query that can run
-  without extra context. Include concrete entities, scope, timeframe, and any
-  qualifiers. Avoid ambiguous pronouns (it/they/this/that).
-- Prioritize the highest-value aspects first; avoid redundancy and overlap.
-- Prefer questions that are likely answerable from the current knowledge base;
-  if coverage is uncertain, make scopes narrower and specific.
-- Order sub_questions by execution priority (most valuable first).
-
-Use the gather_context tool once on the main question before planning."""
-
-SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
-
-Process:
-1. Call search_and_answer with relevant keywords from the question.
-2. Review the results and their relevance scores.
-3. If needed, perform follow-up searches with different keywords (max 3 total).
-4. Provide a concise answer based strictly on the retrieved content.
-
-The search tool returns results like:
-[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
-Source: "Document Title" > Section > Subsection
-Type: paragraph
-Content:
-The actual text content here...
-
-[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
-Source: "Another Document"
-Type: table
-Content:
-| Column 1 | Column 2 |
-...
-
-Each result includes:
-- chunk_id in brackets and relevance score
-- Source: document title and section hierarchy (when available)
-- Type: content type like paragraph, table, code, list_item (when available)
-- Content: the actual text
-
-Output format:
-- query: Echo the question you are answering
-- answer: Your concise answer based on the retrieved content
-- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
-- confidence: A score from 0.0 to 1.0 indicating answer confidence
-
-IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
-
-Guidelines:
-- Base answers strictly on retrieved content - do not use external knowledge.
-- Use the Source and Type metadata to understand context.
-- If multiple results are relevant, synthesize them coherently.
-- If information is insufficient, say so clearly.
-- Be concise and direct; avoid meta commentary about the process.
-- Higher scores indicate more relevant results."""

haiku_rag_slim/haiku/rag/graph/research/__init__.py

@@ -1,3 +1,6 @@
-from haiku.rag.graph.common.models import SearchAnswer
 from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
+from haiku.rag.graph.research.models import (
+    EvaluationResult,
+    ResearchReport,
+    SearchAnswer,
+)

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

@@ -1,29 +0,0 @@
-from pydantic_ai import format_as_xml
-
-from haiku.rag.graph.research.dependencies import ResearchContext
-
-
-def format_context_for_prompt(context: ResearchContext) -> str:
-    """Format the research context as XML for inclusion in prompts."""
-    context_data = {
-        "original_question": context.original_question,
-        "unanswered_questions": context.sub_questions,
-        "qa_responses": [
-            {
-                "question": qa.query,
-                "answer": qa.answer,
-                "confidence": qa.confidence,
-                "sources": [
-                    {
-                        "document_uri": c.document_uri,
-                        "document_title": c.document_title,
-                        "page_numbers": c.page_numbers,
-                        "headings": c.headings,
-                    }
-                    for c in qa.citations
-                ],
-            }
-            for qa in context.qa_responses
-        ],
-    }
-    return format_as_xml(context_data, root_tag="research_context")

haiku_rag_slim/haiku/rag/graph/research/dependencies.py

@@ -1,9 +1,13 @@
+from typing import TYPE_CHECKING, Any
+
 from pydantic import BaseModel, Field
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.common.models import SearchAnswer
 from haiku.rag.store.models import SearchResult
 
+if TYPE_CHECKING:
+    from haiku.rag.graph.research.models import SearchAnswer
+
 
 class ResearchContext(BaseModel):
     """Context shared across research agents."""
@@ -12,11 +16,11 @@ class ResearchContext(BaseModel):
     sub_questions: list[str] = Field(
         default_factory=list, description="Decomposed sub-questions"
     )
-    qa_responses: list[SearchAnswer] = Field(
+    qa_responses: list[Any] = Field(
         default_factory=list, description="Structured QA pairs used during research"
     )
 
-    def add_qa_response(self, qa: SearchAnswer) -> None:
+    def add_qa_response(self, qa: "SearchAnswer") -> None:
         """Add a structured QA response."""
         self.qa_responses.append(qa)
 

haiku_rag_slim/haiku/rag/graph/research/graph.py

@@ -1,20 +1,54 @@
-from pydantic_ai import Agent
+import asyncio
+
+from pydantic_ai import Agent, RunContext, format_as_xml
+from pydantic_ai.output import ToolOutput
 from pydantic_graph.beta import Graph, GraphBuilder, StepContext
 from pydantic_graph.beta.join import reduce_list_append
 
 from haiku.rag.config import Config
 from haiku.rag.config.models import AppConfig
-from haiku.rag.graph.common import get_model
+from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.graph.common.models import SearchAnswer
+from haiku.rag.graph.research.models import (
-from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
+    EvaluationResult,
-from haiku.rag.graph.research.common import format_context_for_prompt
+    RawSearchAnswer,
-from haiku.rag.graph.research.dependencies import ResearchDependencies
+    ResearchPlan,
-from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
+    ResearchReport,
+    SearchAnswer,
+)
 from haiku.rag.graph.research.prompts import (
-    DECISION_AGENT_PROMPT,
+    DECISION_PROMPT,
-    SYNTHESIS_AGENT_PROMPT,
+    PLAN_PROMPT,
+    SEARCH_PROMPT,
+    SYNTHESIS_PROMPT,
 )
 from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+from haiku.rag.utils import get_model
+
+
+def format_context_for_prompt(context: ResearchContext) -> str:
+    """Format the research context as XML for inclusion in prompts."""
+    context_data = {
+        "original_question": context.original_question,
+        "unanswered_questions": context.sub_questions,
+        "qa_responses": [
+            {
+                "question": qa.query,
+                "answer": qa.answer,
+                "confidence": qa.confidence,
+                "sources": [
+                    {
+                        "document_uri": c.document_uri,
+                        "document_title": c.document_title,
+                        "page_numbers": c.page_numbers,
+                        "headings": c.headings,
+                    }
+                    for c in qa.citations
+                ],
+            }
+            for qa in context.qa_responses
+        ],
+    }
+    return format_as_xml(context_data, root_tag="research_context")
 
 
 def build_research_graph(
@@ -37,26 +71,172 @@ def build_research_graph(
         output_type=ResearchReport,
     )
 
-    plan = g.step(
+    @g.step
-        create_plan_node(
+    async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
-            model_config=model_config,
+        """Create research plan with sub-questions."""
-            deps_type=ResearchDependencies,  # type: ignore[arg-type]
+        state = ctx.state
-            activity_message="Creating research plan",
+        deps = ctx.deps
-            output_retries=3,
-            config=config,
-        )
-    )  # type: ignore[arg-type]
 
-    search_one = g.step(
+        if deps.agui_emitter:
-        create_search_node(
+            deps.agui_emitter.start_step("plan")
-            model_config=model_config,
+            deps.agui_emitter.update_activity(
-            deps_type=ResearchDependencies,  # type: ignore[arg-type]
+                "planning", {"stepName": "plan", "message": "Creating research plan"}
-            with_step_wrapper=True,
+            )
-            success_message_format="Found answer with {confidence:.0%} confidence",
+
-            handle_exceptions=True,
+        try:
-            config=config,
+            plan_agent = Agent(
-        )
+                model=get_model(model_config, config),
-    )  # type: ignore[arg-type]
+                output_type=ResearchPlan,
+                instructions=(
+                    PLAN_PROMPT
+                    + "\n\nUse the gather_context tool once on the main question before planning."
+                ),
+                retries=3,
+                output_retries=3,
+                deps_type=ResearchDependencies,
+            )
+
+            search_filter = state.search_filter
+
+            @plan_agent.tool
+            async def gather_context(
+                ctx2: RunContext[ResearchDependencies],
+                query: str,
+                limit: int | None = None,
+            ) -> str:
+                results = await ctx2.deps.client.search(
+                    query, limit=limit, filter=search_filter
+                )
+                results = await ctx2.deps.client.expand_context(results)
+                return "\n\n".join(r.content for r in results)
+
+            _ = gather_context
+
+            prompt = (
+                "Plan a focused approach for the main question.\n\n"
+                f"Main question: {state.context.original_question}"
+            )
+
+            agent_deps = ResearchDependencies(client=deps.client, context=state.context)
+            plan_result = await plan_agent.run(prompt, deps=agent_deps)
+            state.context.sub_questions = list(plan_result.output.sub_questions)
+
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                count = len(state.context.sub_questions)
+                deps.agui_emitter.update_activity(
+                    "planning",
+                    {
+                        "stepName": "plan",
+                        "message": f"Created plan with {count} sub-questions",
+                        "sub_questions": list(state.context.sub_questions),
+                    },
+                )
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    @g.step
+    async def search_one(
+        ctx: StepContext[ResearchState, ResearchDeps, str],
+    ) -> SearchAnswer:
+        """Answer a single sub-question using the knowledge base."""
+        state = ctx.state
+        deps = ctx.deps
+        sub_q = ctx.inputs
+        step_name = f"search: {sub_q}"
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step(step_name)
+
+        try:
+            if deps.semaphore is None:
+                deps.semaphore = asyncio.Semaphore(state.max_concurrency)
+
+            async with deps.semaphore:
+                if deps.agui_emitter:
+                    deps.agui_emitter.update_activity(
+                        "searching",
+                        {
+                            "stepName": "search_one",
+                            "message": f"Searching: {sub_q}",
+                            "query": sub_q,
+                        },
+                    )
+
+                agent = Agent(
+                    model=get_model(model_config, config),
+                    output_type=ToolOutput(RawSearchAnswer, max_retries=3),
+                    instructions=SEARCH_PROMPT,
+                    retries=3,
+                    deps_type=ResearchDependencies,
+                )
+
+                search_filter = state.search_filter
+
+                @agent.tool
+                async def search_and_answer(
+                    ctx2: RunContext[ResearchDependencies],
+                    query: str,
+                    limit: int | None = None,
+                ) -> str:
+                    """Search the knowledge base for relevant documents."""
+                    results = await ctx2.deps.client.search(
+                        query, limit=limit, filter=search_filter
+                    )
+                    results = await ctx2.deps.client.expand_context(results)
+                    ctx2.deps.search_results = results
+                    parts = [r.format_for_agent() for r in results]
+                    if not parts:
+                        return f"No relevant information found for: {query}"
+                    return "\n\n".join(parts)
+
+                _ = search_and_answer
+
+                agent_deps = ResearchDependencies(
+                    client=deps.client, context=state.context
+                )
+
+                try:
+                    result = await agent.run(sub_q, deps=agent_deps)
+                    raw_answer = result.output
+                    if raw_answer:
+                        answer = SearchAnswer.from_raw(
+                            raw_answer, agent_deps.search_results
+                        )
+                        state.context.add_qa_response(answer)
+                        if deps.agui_emitter:
+                            deps.agui_emitter.update_state(state)
+                            deps.agui_emitter.update_activity(
+                                "searching",
+                                {
+                                    "stepName": "search_one",
+                                    "message": f"Found answer with {answer.confidence:.0%} confidence",
+                                    "query": sub_q,
+                                    "confidence": answer.confidence,
+                                },
+                            )
+                        return answer
+                    return SearchAnswer(query=sub_q, answer="", confidence=0.0)
+                except Exception as e:
+                    if deps.agui_emitter:
+                        deps.agui_emitter.update_activity(
+                            "searching",
+                            {
+                                "stepName": "search_one",
+                                "message": f"Search failed: {e}",
+                                "query": sub_q,
+                                "error": str(e),
+                            },
+                        )
+                    return SearchAnswer(
+                        query=sub_q,
+                        answer=f"Search failed: {str(e)}",
+                        confidence=0.0,
+                    )
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
 
     @g.step
     async def get_batch(
@@ -76,6 +256,7 @@ def build_research_graph(
     async def decide(
         ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
     ) -> bool:
+        """Evaluate research sufficiency and decide whether to continue."""
         state = ctx.state
         deps = ctx.deps
 
@@ -89,7 +270,7 @@ def build_research_graph(
             agent = Agent(
                 model=get_model(model_config, config),
                 output_type=EvaluationResult,
-                instructions=DECISION_AGENT_PROMPT,
+                instructions=DECISION_PROMPT,
                 retries=3,
                 output_retries=3,
                 deps_type=ResearchDependencies,
@@ -152,6 +333,7 @@ def build_research_graph(
     async def synthesize(
         ctx: StepContext[ResearchState, ResearchDeps, None | bool],
     ) -> ResearchReport:
+        """Generate final research report."""
         state = ctx.state
         deps = ctx.deps
 
@@ -165,7 +347,7 @@ def build_research_graph(
             agent = Agent(
                 model=get_model(model_config, config),
                 output_type=ResearchReport,
-                instructions=SYNTHESIS_AGENT_PROMPT,
+                instructions=SYNTHESIS_PROMPT,
                 retries=3,
                 output_retries=3,
                 deps_type=ResearchDependencies,
@@ -201,7 +383,6 @@ def build_research_graph(
     else:
         g.add(g.edge_from(g.start_node).to(get_batch))
 
-    # Branch based on whether we have questions
     g.add(
         g.edge_from(get_batch).to(
             g.decision()
@@ -209,12 +390,9 @@ def build_research_graph(
             .branch(g.match(type(None)).label("No questions").to(synthesize))
         ),
         g.edge_from(search_one).to(collect_answers),
-        g.edge_from(collect_answers).to(
+        g.edge_from(collect_answers).to(decide),
-            decide
-        ),  # Direct: collect → decide (no analyze_insights)
     )
 
-    # Branch based on decision
     g.add(
         g.edge_from(decide).to(
             g.decision()

haiku_rag_slim/haiku/rag/graph/research/models.py

@@ -1,4 +1,107 @@
-from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field, field_validator
+
+if TYPE_CHECKING:
+    from haiku.rag.store.models import SearchResult
+
+
+class ResearchPlan(BaseModel):
+    """A structured research plan with sub-questions to explore."""
+
+    sub_questions: list[str] = Field(
+        ...,
+        description="Specific questions to research, phrased as complete questions",
+    )
+
+    @field_validator("sub_questions")
+    @classmethod
+    def validate_sub_questions(cls, v: list[str]) -> list[str]:
+        if len(v) < 1:
+            raise ValueError("Must have at least 1 sub-question")
+        if len(v) > 12:
+            raise ValueError("Cannot have more than 12 sub-questions")
+        return v
+
+
+class Citation(BaseModel):
+    """Resolved citation with full metadata for display/visual grounding."""
+
+    document_id: str
+    chunk_id: str
+    document_uri: str
+    document_title: str | None = None
+    page_numbers: list[int] = Field(default_factory=list)
+    headings: list[str] | None = None
+    content: str
+
+
+class RawSearchAnswer(BaseModel):
+    """Answer to a search query with chunk references."""
+
+    query: str = Field(..., description="The question that was answered")
+    answer: str = Field(..., description="The answer to the question")
+    cited_chunks: list[str] = Field(
+        default_factory=list,
+        description="IDs of chunks used to form the answer",
+    )
+    confidence: float = Field(
+        default=1.0,
+        description="Confidence score for this answer (0-1)",
+        ge=0.0,
+        le=1.0,
+    )
+
+
+class SearchAnswer(RawSearchAnswer):
+    """Answer to a search query with resolved citations."""
+
+    citations: list[Citation] = Field(
+        default_factory=list,
+        description="Resolved citations with full metadata",
+    )
+
+    @classmethod
+    def from_raw(
+        cls,
+        raw: RawSearchAnswer,
+        search_results: "list[SearchResult]",
+    ) -> "SearchAnswer":
+        """Create SearchAnswer from RawSearchAnswer with resolved citations."""
+        citations = resolve_citations(raw.cited_chunks, search_results)
+        return cls(
+            query=raw.query,
+            answer=raw.answer,
+            cited_chunks=raw.cited_chunks,
+            confidence=raw.confidence,
+            citations=citations,
+        )
+
+
+def resolve_citations(
+    cited_chunk_ids: list[str],
+    search_results: "list[SearchResult]",
+) -> list[Citation]:
+    """Resolve chunk IDs to full Citation objects with metadata."""
+    by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
+
+    citations = []
+    for chunk_id in cited_chunk_ids:
+        r = by_id.get(chunk_id)
+        if not r:
+            continue
+        citations.append(
+            Citation(
+                document_id=r.document_id or "",
+                chunk_id=chunk_id,
+                document_uri=r.document_uri or "",
+                document_title=r.document_title,
+                page_numbers=r.page_numbers,
+                headings=r.headings,
+                content=r.content,
+            )
+        )
+    return citations
 
 
 class EvaluationResult(BaseModel):

haiku_rag_slim/haiku/rag/graph/research/prompts.py

@@ -1,4 +1,70 @@
-DECISION_AGENT_PROMPT = """You are the research evaluator responsible for assessing
+PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
+
+Responsibilities:
+1. Understand and decompose the main question
+2. Propose a minimal, high-leverage plan
+3. Coordinate specialized agents to gather evidence
+4. Iterate based on gaps and new findings
+
+Plan requirements:
+- Produce at most 3 sub_questions that together cover the main question.
+- sub_questions must be a list of plain strings, where each string is a complete
+  question. Do NOT use objects with nested fields like {question, details}.
+- Each sub_question must be a standalone, self-contained query that can run
+  without extra context. Include concrete entities, scope, timeframe, and any
+  qualifiers. Avoid ambiguous pronouns (it/they/this/that).
+- Prioritize the highest-value aspects first; avoid redundancy and overlap.
+- Prefer questions that are likely answerable from the current knowledge base;
+  if coverage is uncertain, make scopes narrower and specific.
+- Order sub_questions by execution priority (most valuable first).
+
+Use the gather_context tool once on the main question before planning."""
+
+SEARCH_PROMPT = """You are a search and question-answering specialist.
+
+Process:
+1. Call search_and_answer with relevant keywords from the question.
+2. Review the results and their relevance scores.
+3. If needed, perform follow-up searches with different keywords (max 3 total).
+4. Provide a concise answer based strictly on the retrieved content.
+
+The search tool returns results like:
+[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
+Source: "Document Title" > Section > Subsection
+Type: paragraph
+Content:
+The actual text content here...
+
+[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
+Source: "Another Document"
+Type: table
+Content:
+| Column 1 | Column 2 |
+...
+
+Each result includes:
+- chunk_id in brackets and relevance score
+- Source: document title and section hierarchy (when available)
+- Type: content type like paragraph, table, code, list_item (when available)
+- Content: the actual text
+
+Output format:
+- query: Echo the question you are answering
+- answer: Your concise answer based on the retrieved content
+- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
+- confidence: A score from 0.0 to 1.0 indicating answer confidence
+
+IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
+
+Guidelines:
+- Base answers strictly on retrieved content - do not use external knowledge.
+- Use the Source and Type metadata to understand context.
+- If multiple results are relevant, synthesize them coherently.
+- If information is insufficient, say so clearly.
+- Be concise and direct; avoid meta commentary about the process.
+- Higher scores indicate more relevant results."""
+
+DECISION_PROMPT = """You are the research evaluator responsible for assessing
 whether gathered evidence sufficiently answers the research question.
 
 Inputs available:
@@ -19,7 +85,7 @@ Output fields:
 
 Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
 
-SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
+SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
 research report that directly answers the original question.
 
 Goals:
@@ -47,15 +113,3 @@ Style:
 - Be professional, objective, and specific.
 - NEVER use meta-commentary like "This report covers..." or "The findings show...".
   Instead, state the actual information directly."""
-
-PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
-
-Task:
-- Call gather_context once on the main question to obtain relevant text from
-  the knowledge base (KB).
-- Read that context and produce a short natural-language summary of what the
-  KB appears to contain relative to the question.
-
-Rules:
-- Base the summary strictly on the provided text; do not invent.
-- Output only the summary as plain text (one short paragraph)."""

haiku_rag_slim/haiku/rag/qa/agent.py

@@ -5,10 +5,10 @@ from pydantic_ai.output import ToolOutput
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.config.models import ModelConfig
-from haiku.rag.graph.common import get_model
+from haiku.rag.graph.research.models import Citation, RawSearchAnswer, resolve_citations
-from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations
 from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.store.models import SearchResult
+from haiku.rag.utils import get_model
 
 
 class Dependencies(BaseModel):

haiku_rag_slim/haiku/rag/utils.py

@@ -8,7 +8,7 @@ from packaging.version import Version, parse
 if TYPE_CHECKING:
     from rich.console import RenderableType
 
-    from haiku.rag.graph.common.models import Citation
+    from haiku.rag.graph.research.models import Citation
 
 
 def apply_common_settings(

tests/graph/test_research_graph.py

@@ -18,8 +18,6 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
 
     # Patch all locations where get_model is imported
     monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
 
     graph = build_research_graph()

tests/graph/test_search_filter.py

@@ -67,8 +67,6 @@ async def test_research_graph_uses_search_filter(monkeypatch, client_with_docs):
         return TestModel()
 
     monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
 
     graph = build_research_graph()
@@ -114,8 +112,6 @@ async def test_search_filter_none_searches_all(monkeypatch, client_with_docs):
         return TestModel()
 
     monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
 
     graph = build_research_graph()

tests/test_app.py

@@ -310,7 +310,7 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
 @pytest.mark.asyncio
 async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
     """Test asking a question with citations."""
-    from haiku.rag.graph.common.models import Citation
+    from haiku.rag.graph.research.models import Citation
 
     mock_answer = "Test answer with citations"
     mock_citations = [