Remove graph common and move its content to deep QA, research.

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
       - id: debug-statements
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.4
+    rev: v0.14.3
     hooks:
       # Run the linter.
       - id: ruff
@@ -17,6 +17,6 @@ repos:
       - id: ruff-format
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.399
+    rev: v1.1.407
     hooks:
       - id: pyright

haiku_rag_slim/haiku/rag/a2a/__init__.py

@@ -6,7 +6,7 @@ import logfire
 from pydantic_ai import Agent, RunContext
 
 from haiku.rag.config import Config
-from haiku.rag.graph.common import get_model
+from haiku.rag.qa.deep.common import get_model
 
 from .context import load_message_history, save_message_history
 from .models import AgentDependencies, SearchResult
@@ -138,7 +138,11 @@ def create_a2a_app(
     if security_schemes or security:
         # Monkey-patch the agent card endpoint to include security
         async def _agent_card_endpoint_with_security(request):
-            from fasta2a.schema import AgentCapabilities, AgentCard, agent_card_ta
+            from fasta2a.schema import (  # type: ignore
+                AgentCapabilities,
+                AgentCard,
+                agent_card_ta,
+            )
             from starlette.responses import Response
 
             if app._agent_card_json_schema is None:

haiku_rag_slim/haiku/rag/graph/models.py

@@ -1,24 +0,0 @@
-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,
-    )

haiku_rag_slim/haiku/rag/graph/prompts.py

@@ -1,45 +0,0 @@
-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."""

haiku_rag_slim/haiku/rag/qa/deep/common.py

@@ -5,6 +5,7 @@ from pydantic_ai.providers.ollama import OllamaProvider
 from pydantic_ai.providers.openai import OpenAIProvider
 
 from haiku.rag.config import Config
+from haiku.rag.qa.deep.models import SearchAnswer
 
 
 class HasEmitLog(Protocol):
@@ -31,3 +32,10 @@ def get_model(provider: str, model: str) -> Any:
 
 def log(deps: HasEmitLog, state: Any, message: str) -> None:
     deps.emit_log(message, state)
+
+
+def collect_answers_reducer(
+    acc: list[SearchAnswer], item: SearchAnswer | None
+) -> list[SearchAnswer]:
+    """Reducer function to collect search answers, filtering out None values."""
+    return acc + [item] if item else acc

haiku_rag_slim/haiku/rag/qa/deep/dependencies.py

@@ -2,7 +2,7 @@ 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.qa.deep.models import SearchAnswer
 
 
 class DeepQAContext(BaseModel):

haiku_rag_slim/haiku/rag/qa/deep/graph.py

@@ -5,13 +5,18 @@ from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.output import ToolOutput
 from pydantic_graph.beta import GraphBuilder, StepContext
 
-from haiku.rag.graph.common import get_model, log
+from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log
-from haiku.rag.graph.models import ResearchPlan, SearchAnswer
-from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.qa.deep.dependencies import DeepQADependencies
-from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
+from haiku.rag.qa.deep.models import (
+    DeepQAAnswer,
+    DeepQAEvaluation,
+    ResearchPlan,
+    SearchAnswer,
+)
 from haiku.rag.qa.deep.prompts import (
     DECISION_PROMPT,
+    PLAN_PROMPT,
+    SEARCH_AGENT_PROMPT,
     SYNTHESIS_PROMPT,
     SYNTHESIS_PROMPT_WITH_CITATIONS,
 )
@@ -52,7 +57,7 @@ def build_deep_qa_graph(provider: str, model: str):
             return "\n\n".join(chunk.content for chunk, _ in expanded)
 
         prompt = (
-            "Plan a focused approach for answering the main question.\n\n"
+            "Plan a focused approach for the main question.\n\n"
             f"Main question: {state.context.original_question}"
         )
 
@@ -62,9 +67,7 @@ def build_deep_qa_graph(provider: str, model: str):
             console=deps.console,
         )
         plan_result = await plan_agent.run(prompt, deps=agent_deps)
-        state.context.sub_questions = list(plan_result.output.sub_questions)[
+        state.context.sub_questions = list(plan_result.output.sub_questions)
-            : state.max_sub_questions
-        ]
 
         log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
         log(
@@ -139,9 +142,21 @@ def build_deep_qa_graph(provider: str, model: str):
 
         return answer
 
+    @g.step
+    async def get_batch(
+        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
+    ) -> list[str] | None:
+        """Get next batch of questions from state."""
+        state = ctx.state
+        take = max(1, state.max_concurrency)
+        batch: list[str] = []
+        while state.context.sub_questions and len(batch) < take:
+            batch.append(state.context.sub_questions.pop(0))
+        return batch if batch else None
+
     @g.step
     async def decide(
-        ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer | None]],
+        ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
     ) -> bool:
         state = ctx.state
         deps = ctx.deps
@@ -222,18 +237,6 @@ def build_deep_qa_graph(provider: str, model: str):
 
         return should_continue
 
-    @g.step
-    async def get_batch(
-        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
-    ) -> list[str] | None:
-        """Get next batch of questions from state."""
-        state = ctx.state
-        take = max(1, state.max_concurrency)
-        batch: list[str] = []
-        while state.context.sub_questions and len(batch) < take:
-            batch.append(state.context.sub_questions.pop(0))
-        return batch if batch else None
-
     @g.step
     async def synthesize(
         ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
@@ -287,13 +290,8 @@ def build_deep_qa_graph(provider: str, model: str):
         return result.output
 
     # Build the graph structure
-    def collect_reducer(
-        acc: list[SearchAnswer | None], item: SearchAnswer | None
-    ) -> list[SearchAnswer | None]:
-        return acc + [item] if item else acc
-
     collect_answers = g.join(
-        collect_reducer,
+        collect_answers_reducer,
         initial_factory=lambda: [],
     )
 

haiku_rag_slim/haiku/rag/qa/deep/models.py

@@ -1,6 +1,29 @@
 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,
+    )
+
+
 class DeepQAEvaluation(BaseModel):
     is_sufficient: bool = Field(
         description="Whether we have sufficient information to answer the question"

haiku_rag_slim/haiku/rag/qa/deep/prompts.py

@@ -1,3 +1,49 @@
+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."""
+
 SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers.
 
 Task:

haiku_rag_slim/haiku/rag/research/__init__.py

@@ -1,3 +1,3 @@
-from haiku.rag.graph.models import SearchAnswer
+from haiku.rag.qa.deep.models import SearchAnswer
 from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
 from haiku.rag.research.models import EvaluationResult, ResearchReport

haiku_rag_slim/haiku/rag/research/dependencies.py

@@ -4,7 +4,7 @@ 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.qa.deep.models import SearchAnswer
 from haiku.rag.research.models import (
     GapRecord,
     InsightAnalysis,

haiku_rag_slim/haiku/rag/research/graph.py

@@ -5,9 +5,9 @@ from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.output import ToolOutput
 from pydantic_graph.beta import GraphBuilder, StepContext
 
-from haiku.rag.graph.common import get_model, log
+from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log
-from haiku.rag.graph.models import ResearchPlan, SearchAnswer
+from haiku.rag.qa.deep.models import ResearchPlan, SearchAnswer
-from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
+from haiku.rag.qa.deep.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.research.common import (
     format_analysis_for_prompt,
     format_context_for_prompt,
@@ -60,7 +60,7 @@ def build_research_graph(provider: str, model: str):
             return "\n\n".join(chunk.content for chunk, _ in expanded)
 
         prompt = (
-            "Plan a focused research approach for the main question.\n\n"
+            "Plan a focused approach for the main question.\n\n"
             f"Main question: {state.context.original_question}"
         )
 
@@ -73,7 +73,7 @@ def build_research_graph(provider: str, model: str):
         plan_result = await plan_agent.run(prompt, deps=agent_deps)
         state.context.sub_questions = list(plan_result.output.sub_questions)
 
-        log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
+        log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
         log(
             deps,
             state,
@@ -147,9 +147,21 @@ def build_research_graph(provider: str, model: str):
 
         return answer
 
+    @g.step
+    async def get_batch(
+        ctx: StepContext[ResearchState, ResearchDeps, None | bool],
+    ) -> list[str] | None:
+        """Get next batch of questions from state."""
+        state = ctx.state
+        take = max(1, state.max_concurrency)
+        batch: list[str] = []
+        while state.context.sub_questions and len(batch) < take:
+            batch.append(state.context.sub_questions.pop(0))
+        return batch if batch else None
+
     @g.step
     async def analyze_insights(
-        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer | None]],
+        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
     ) -> None:
         state = ctx.state
         deps = ctx.deps
@@ -297,18 +309,6 @@ def build_research_graph(provider: str, model: str):
 
         return should_continue
 
-    @g.step
-    async def get_batch(
-        ctx: StepContext[ResearchState, ResearchDeps, None | bool],
-    ) -> list[str] | None:
-        """Get next batch of questions from state."""
-        state = ctx.state
-        take = max(1, state.max_concurrency)
-        batch: list[str] = []
-        while state.context.sub_questions and len(batch) < take:
-            batch.append(state.context.sub_questions.pop(0))
-        return batch if batch else None
-
     @g.step
     async def synthesize(
         ctx: StepContext[ResearchState, ResearchDeps, None | bool],
@@ -348,13 +348,8 @@ def build_research_graph(provider: str, model: str):
         return result.output
 
     # Build the graph structure
-    def collect_reducer(
-        acc: list[SearchAnswer | None], item: SearchAnswer | None
-    ) -> list[SearchAnswer | None]:
-        return acc + [item] if item else acc
-
     collect_answers = g.join(
-        collect_reducer,
+        collect_answers_reducer,
         initial_factory=lambda: [],
     )
 

tests/test_app.py

@@ -383,11 +383,9 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
     from haiku.rag.qa.deep.models import DeepQAAnswer
 
     mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
-    mock_result = MagicMock()
-    mock_result.output = mock_output
 
     mock_graph = AsyncMock()
-    mock_graph.run.return_value = mock_result
+    mock_graph.run.return_value = mock_output
 
     mock_client = AsyncMock()
     mock_client.__aenter__.return_value = mock_client
@@ -415,11 +413,9 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
     mock_output = DeepQAAnswer(
         answer="Deep QA answer with citations [test.md]", sources=["test.md"]
     )
-    mock_result = MagicMock()
-    mock_result.output = mock_output
 
     mock_graph = AsyncMock()
-    mock_graph.run.return_value = mock_result
+    mock_graph.run.return_value = mock_output
 
     mock_client = AsyncMock()
     mock_client.__aenter__.return_value = mock_client
@@ -445,11 +441,9 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
     from haiku.rag.qa.deep.models import DeepQAAnswer
 
     mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
-    mock_result = MagicMock()
-    mock_result.output = mock_output
 
     mock_graph = AsyncMock()
-    mock_graph.run.return_value = mock_result
+    mock_graph.run.return_value = mock_output
 
     mock_client = AsyncMock()
     mock_client.__aenter__.return_value = mock_client

tests/test_deep_qa.py

@@ -2,9 +2,9 @@ import pytest
 from pydantic_ai.models.test import TestModel
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.qa.deep.dependencies import DeepQAContext
 from haiku.rag.qa.deep.graph import build_deep_qa_graph
+from haiku.rag.qa.deep.models import SearchAnswer
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 
@@ -16,7 +16,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
     def test_model_factory(provider, model):
         return TestModel()
 
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
 
     graph = build_deep_qa_graph(provider="test", model="test")
@@ -50,7 +50,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
     def test_model_factory(provider, model):
         return TestModel()
 
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
 
     graph = build_deep_qa_graph(provider="test", model="test")

tests/test_mcp.py

@@ -257,7 +257,7 @@ async def test_mcp_ask_question_deep():
 
             mock_graph = AsyncMock()
             mock_result = AsyncMock()
-            mock_result.output.answer = "Deep answer"
+            mock_result.answer = "Deep answer"
             mock_graph.run = AsyncMock(return_value=mock_result)
             mock_graph_builder.return_value = mock_graph
 
@@ -299,9 +299,7 @@ async def test_mcp_research_question():
             mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
 
             mock_graph = AsyncMock()
-            mock_result = AsyncMock()
+            mock_graph.run = AsyncMock(return_value=mock_report)
-            mock_result.output = mock_report
-            mock_graph.run = AsyncMock(return_value=mock_result)
             mock_graph_builder.return_value = mock_graph
 
             tools = await mcp.get_tools()

tests/test_research_graph_integration.py

@@ -17,7 +17,7 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
     def test_model_factory(provider, model):
         return TestModel()
 
-    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
 
     graph = build_research_graph(provider="test", model="test")