Shuffle reusable nodes under graph/

src/haiku/rag/graph/__init__.py

@@ -0,0 +1 @@
+from haiku.rag.graph.models import ResearchPlan, SearchAnswer

src/haiku/rag/graph/base.py

@@ -0,0 +1,31 @@
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.models import SearchAnswer
+
+
+@runtime_checkable
+class GraphContext(Protocol):
+    """Protocol for graph context objects."""
+
+    original_question: str
+    sub_questions: list[str]
+    qa_responses: list[SearchAnswer]
+
+    def add_qa_response(self, qa: SearchAnswer) -> None: ...
+
+
+class BaseGraphDeps(BaseModel):
+    """Base dependencies for graph nodes."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    console: Console | None = None
+
+    def emit_log(self, message: str) -> None:
+        if self.console:
+            self.console.print(message)

src/haiku/rag/graph/common.py

@@ -0,0 +1,32 @@
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from haiku.rag.config import Config
+
+if TYPE_CHECKING:  # pragma: no cover
+    from haiku.rag.research.state import ResearchDeps, ResearchState
+
+
+def get_model(provider: str, model: str) -> Any:
+    if provider == "ollama":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+    elif provider == "vllm":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OpenAIProvider(
+                base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
+                api_key="none",
+            ),
+        )
+    else:
+        return f"{provider}:{model}"
+
+
+def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
+    deps.emit_log(msg, state)

src/haiku/rag/graph/models.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel, Field
+
+
+class ResearchPlan(BaseModel):
+    main_question: str
+    sub_questions: list[str]
+
+
+class SearchAnswer(BaseModel):
+    query: str = Field(description="The search query that was performed")
+    answer: str = Field(description="The answer generated based on the context")
+    context: list[str] = Field(
+        description=(
+            "Only the minimal set of relevant snippets (verbatim) that directly "
+            "support the answer"
+        )
+    )
+    sources: list[str] = Field(
+        description=(
+            "Document titles (if available) or URIs corresponding to the"
+            " snippets actually used in the answer (one per snippet; omit if none)"
+        ),
+        default_factory=list,
+    )

src/haiku/rag/graph/nodes/analysis.py

@@ -3,15 +3,13 @@ from dataclasses import dataclass
 from pydantic_ai import Agent
 from pydantic_graph import BaseNode, GraphRunContext
 
+from haiku.rag.graph.common import get_model, log
 from haiku.rag.research.common import (
     format_analysis_for_prompt,
     format_context_for_prompt,
-    get_model,
-    log,
 )
 from haiku.rag.research.dependencies import ResearchDependencies
 from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport
-from haiku.rag.research.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.prompts import DECISION_AGENT_PROMPT, INSIGHT_AGENT_PROMPT
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
@@ -89,6 +87,8 @@ class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport])
             for question in analysis.new_questions:
                 log(deps, state, f"   • {question}")
 
+        from haiku.rag.graph.nodes.analysis import DecisionNode
+
         return DecisionNode(self.provider, self.model)
 
 
@@ -169,7 +169,8 @@ class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
         log(deps, state, f"   Sufficient: {status}")
 
-        from haiku.rag.research.nodes.search import SearchDispatchNode
+        from haiku.rag.graph.nodes.search import SearchDispatchNode
+        from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 
         if (
             output.is_sufficient

src/haiku/rag/graph/nodes/plan.py

@@ -3,11 +3,11 @@ from dataclasses import dataclass
 from pydantic_ai import Agent, RunContext
 from pydantic_graph import BaseNode, GraphRunContext
 
-from haiku.rag.research.common import get_model, log
+from haiku.rag.graph.common import get_model, log
+from haiku.rag.graph.models import ResearchPlan
+from haiku.rag.graph.prompts import PLAN_PROMPT
 from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.research.models import ResearchPlan, ResearchReport
-from haiku.rag.research.nodes.search import SearchDispatchNode
-from haiku.rag.research.prompts import PLAN_PROMPT
+from haiku.rag.research.models import ResearchReport
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
@@ -67,4 +67,6 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         for i, sq in enumerate(state.context.sub_questions, 1):
             log(deps, state, f"      {i}. {sq}")
 
+        from haiku.rag.graph.nodes.search import SearchDispatchNode
+
         return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/graph/nodes/search.py

@@ -7,10 +7,11 @@ from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.output import ToolOutput
 from pydantic_graph import BaseNode, GraphRunContext
 
-from haiku.rag.research.common import get_model, log
+from haiku.rag.graph.common import get_model, log
+from haiku.rag.graph.models import SearchAnswer
+from haiku.rag.graph.prompts import SEARCH_AGENT_PROMPT
 from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.research.models import ResearchReport, SearchAnswer
-from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
+from haiku.rag.research.models import ResearchReport
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
@@ -25,7 +26,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         state = ctx.state
         deps = ctx.deps
         if not state.context.sub_questions:
-            from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode
+            from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode
 
             return AnalyzeInsightsNode(self.provider, self.model)
 

src/haiku/rag/graph/nodes/synthesize.py

@@ -3,10 +3,9 @@ from dataclasses import dataclass
 from pydantic_ai import Agent
 from pydantic_graph import BaseNode, End, GraphRunContext
 
-from haiku.rag.research.common import format_context_for_prompt, get_model, log
-from haiku.rag.research.dependencies import (
-    ResearchDependencies,
-)
+from haiku.rag.graph.common import get_model, log
+from haiku.rag.research.common import format_context_for_prompt
+from haiku.rag.research.dependencies import ResearchDependencies
 from haiku.rag.research.models import ResearchReport
 from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
 from haiku.rag.research.state import ResearchDeps, ResearchState

src/haiku/rag/graph/prompts.py

@@ -0,0 +1,45 @@
+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.
+- 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)."""
+
+SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist.
+
+Tasks:
+1. Search the knowledge base for relevant evidence.
+2. Analyze retrieved snippets.
+3. Provide an answer strictly grounded in that evidence.
+
+Tool usage:
+- Always call search_and_answer before drafting any answer.
+- The tool returns snippets with verbatim `text`, a relevance `score`, and the
+  originating document identifier (document title if available, otherwise URI).
+- You may call the tool multiple times to refine or broaden context, but do not
+  exceed 3 total calls. Favor precision over volume.
+- Use scores to prioritize evidence, but include only the minimal subset of
+  snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
+- Set SearchAnswer.sources to the corresponding document identifiers for the
+  snippets you used (title if available, otherwise URI; one per snippet; same
+  order as context). Context must be text‑only.
+- If no relevant information is found, clearly say so and return an empty
+  context list and sources list.
+
+Answering rules:
+- Be direct and specific; avoid meta commentary about the process.
+- Do not include any claims not supported by the provided snippets.
+- Prefer concise phrasing; avoid copying long passages.
+- When evidence is partial, state the limits explicitly in the answer."""

src/haiku/rag/research/__init__.py

@@ -1,28 +1,3 @@
+from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.research.graph import (
-    PlanNode,
-    ResearchDeps,
-    ResearchState,
-    build_research_graph,
-)
-from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
-from haiku.rag.research.stream import (
-    ResearchStateSnapshot,
-    ResearchStreamEvent,
-    stream_research_graph,
-)
-
-__all__ = [
-    "ResearchDependencies",
-    "ResearchContext",
-    "SearchAnswer",
-    "EvaluationResult",
-    "ResearchReport",
-    "ResearchDeps",
-    "ResearchState",
-    "PlanNode",
-    "build_research_graph",
-    "stream_research_graph",
-    "ResearchStreamEvent",
-    "ResearchStateSnapshot",
-]
+from haiku.rag.research.models import EvaluationResult, ResearchReport

src/haiku/rag/research/common.py

@@ -1,39 +1,8 @@
-from typing import TYPE_CHECKING, Any
-
 from pydantic_ai import format_as_xml
-from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.providers.ollama import OllamaProvider
-from pydantic_ai.providers.openai import OpenAIProvider
 
-from haiku.rag.config import Config
 from haiku.rag.research.dependencies import ResearchContext
 from haiku.rag.research.models import InsightAnalysis
 
-if TYPE_CHECKING:  # pragma: no cover
-    from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-def get_model(provider: str, model: str) -> Any:
-    if provider == "ollama":
-        return OpenAIChatModel(
-            model_name=model,
-            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
-        )
-    elif provider == "vllm":
-        return OpenAIChatModel(
-            model_name=model,
-            provider=OpenAIProvider(
-                base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
-                api_key="none",
-            ),
-        )
-    else:
-        return f"{provider}:{model}"
-
-
-def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
-    deps.emit_log(msg, state)
-
 
 def format_context_for_prompt(context: ResearchContext) -> str:
     """Format the research context as XML for inclusion in prompts."""

src/haiku/rag/research/dependencies.py

@@ -4,11 +4,11 @@ from pydantic import BaseModel, Field
 from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.research.models import (
     GapRecord,
     InsightAnalysis,
     InsightRecord,
-    SearchAnswer,
 )
 from haiku.rag.research.stream import ResearchStream
 

src/haiku/rag/research/graph.py

@@ -1,23 +1,12 @@
 from pydantic_graph import Graph
 
+from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode
+from haiku.rag.graph.nodes.plan import PlanNode
+from haiku.rag.graph.nodes.search import SearchDispatchNode
+from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode, DecisionNode
-from haiku.rag.research.nodes.plan import PlanNode
-from haiku.rag.research.nodes.search import SearchDispatchNode
-from haiku.rag.research.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
-__all__ = [
-    "PlanNode",
-    "SearchDispatchNode",
-    "AnalyzeInsightsNode",
-    "DecisionNode",
-    "SynthesizeNode",
-    "ResearchState",
-    "ResearchDeps",
-    "build_research_graph",
-]
-
 
 def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
     return Graph(

src/haiku/rag/research/models.py

@@ -131,31 +131,6 @@ class InsightAnalysis(BaseModel):
     )
 
 
-class ResearchPlan(BaseModel):
-    main_question: str
-    sub_questions: list[str]
-
-
-class SearchAnswer(BaseModel):
-    """Structured output for the SearchSpecialist agent."""
-
-    query: str = Field(description="The search query that was performed")
-    answer: str = Field(description="The answer generated based on the context")
-    context: list[str] = Field(
-        description=(
-            "Only the minimal set of relevant snippets (verbatim) that directly "
-            "support the answer"
-        )
-    )
-    sources: list[str] = Field(
-        description=(
-            "Document titles (if available) or URIs corresponding to the"
-            " snippets actually used in the answer (one per snippet; omit if none)"
-        ),
-        default_factory=list,
-    )
-
-
 class EvaluationResult(BaseModel):
     """Result of analysis and evaluation."""
 

src/haiku/rag/research/prompts.py

@@ -1,49 +1,3 @@
-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.
-- 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)."""
-
-SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist.
-
-Tasks:
-1. Search the knowledge base for relevant evidence.
-2. Analyze retrieved snippets.
-3. Provide an answer strictly grounded in that evidence.
-
-Tool usage:
-- Always call search_and_answer before drafting any answer.
-- The tool returns snippets with verbatim `text`, a relevance `score`, and the
-  originating document identifier (document title if available, otherwise URI).
-- You may call the tool multiple times to refine or broaden context, but do not
-  exceed 3 total calls. Favor precision over volume.
-- Use scores to prioritize evidence, but include only the minimal subset of
-  snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
-- Set SearchAnswer.sources to the corresponding document identifiers for the
-  snippets you used (title if available, otherwise URI; one per snippet; same
-  order as context). Context must be text‑only.
-- If no relevant information is found, clearly say so and return an empty
-  context list and sources list.
-
-Answering rules:
-- Be direct and specific; avoid meta commentary about the process.
-- Do not include any claims not supported by the provided snippets.
-- Prefer concise phrasing; avoid copying long passages.
-- When evidence is partial, state the limits explicitly in the answer."""
-
 INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
 research workflow.
 

tests/test_research_graph_integration.py

@@ -2,15 +2,15 @@ from typing import Any, cast
 
 import pytest
 
+from haiku.rag.graph.models import SearchAnswer
+from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode
+from haiku.rag.graph.nodes.plan import PlanNode
+from haiku.rag.graph.nodes.search import SearchDispatchNode
+from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.dependencies import ResearchContext
 from haiku.rag.research.graph import (
-    AnalyzeInsightsNode,
-    DecisionNode,
-    PlanNode,
     ResearchDeps,
     ResearchState,
-    SearchDispatchNode,
-    SynthesizeNode,
     build_research_graph,
 )
 from haiku.rag.research.models import (
@@ -21,7 +21,6 @@ from haiku.rag.research.models import (
     InsightRecord,
     InsightStatus,
     ResearchReport,
-    SearchAnswer,
 )
 from haiku.rag.research.stream import stream_research_graph