Go back and separate common patters in graph_common

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

@@ -8,7 +8,7 @@ from pydantic_ai.ag_ui import StateDeps
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
-from haiku.rag.graph.common import get_model
+from haiku.rag.graph_common import get_model
 
 
 class ResearchState(BaseModel):

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.qa.deep.common import get_model
+from haiku.rag.graph_common import get_model
 
 from .context import load_message_history, save_message_history
 from .models import AgentDependencies, SearchResult

haiku_rag_slim/haiku/rag/graph_common/__init__.py

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

haiku_rag_slim/haiku/rag/graph_common/models.py

@@ -0,0 +1,42 @@
+"""Common models used across different graph implementations."""
+
+from pydantic import BaseModel, Field, field_validator
+
+
+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 SearchAnswer(BaseModel):
+    """Answer from a search operation with sources."""
+
+    query: str = Field(..., description="The question that was answered")
+    answer: str = Field(..., description="The comprehensive answer to the question")
+    context: list[str] = Field(
+        default_factory=list,
+        description="Relevant snippets that directly support the answer",
+    )
+    sources: list[str] = Field(
+        default_factory=list,
+        description="Source URIs or titles that contributed to this answer",
+    )
+    confidence: float = Field(
+        default=1.0,
+        description="Confidence score for this answer (0-1)",
+        ge=0.0,
+        le=1.0,
+    )

haiku_rag_slim/haiku/rag/graph_common/prompts.py

@@ -0,0 +1,46 @@
+"""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.
+- 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/graph_common/utils.py

@@ -0,0 +1,64 @@
+"""Common utilities for all graph implementations."""
+
+from typing import Any, Protocol
+
+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
+
+
+class HasEmitLog(Protocol):
+    """Protocol for objects that can emit log messages."""
+
+    def emit_log(self, message: str, state: Any = None) -> None: ...
+
+
+def get_model(provider: str, model: str) -> OpenAIChatModel | str:
+    """
+    Get a model instance for the specified provider and model name.
+
+    Args:
+        provider: The model provider ("ollama", "vllm", or other)
+        model: The model name
+
+    Returns:
+        A configured model instance
+
+    Raises:
+        ValueError: If the provider is unknown
+    """
+    if provider == "ollama":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
+        )
+    elif provider == "vllm":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OpenAIProvider(
+                base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1",
+                api_key="none",
+            ),
+        )
+    elif provider in ("openai", "anthropic", "gemini", "groq", "bedrock"):
+        # These providers use string format
+        return f"{provider}:{model}"
+    else:
+        raise ValueError(
+            f"Unknown model provider: {provider}. "
+            f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock"
+        )
+
+
+def log(deps: HasEmitLog, state: Any, message: str) -> None:
+    """
+    Emit a log message through the dependencies.
+
+    Args:
+        deps: Dependencies object with emit_log method
+        state: Current state (passed to emit_log)
+        message: The message to log
+    """
+    deps.emit_log(message, state)

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

@@ -1,41 +0,0 @@
-from typing import Any, Protocol
-
-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.qa.deep.models import SearchAnswer
-
-
-class HasEmitLog(Protocol):
-    def emit_log(self, message: str, state: Any = None) -> None: ...
-
-
-def get_model(provider: str, model: str) -> Any:
-    if provider == "ollama":
-        return OpenAIChatModel(
-            model_name=model,
-            provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
-        )
-    elif provider == "vllm":
-        return OpenAIChatModel(
-            model_name=model,
-            provider=OpenAIProvider(
-                base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1",
-                api_key="none",
-            ),
-        )
-    else:
-        return f"{provider}:{model}"
-
-
-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.qa.deep.models import SearchAnswer
+from haiku.rag.graph_common.models import SearchAnswer
 
 
 class DeepQAContext(BaseModel):

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

@@ -3,27 +3,25 @@ from typing import Any
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.output import ToolOutput
-from pydantic_graph.beta import GraphBuilder, StepContext
+from pydantic_graph.beta import Graph, GraphBuilder, StepContext
+from pydantic_graph.beta.join import reduce_list_append
 
-from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log
+from haiku.rag.graph_common import get_model, log
+from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer
+from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.qa.deep.dependencies import DeepQADependencies
-from haiku.rag.qa.deep.models import (
-    DeepQAAnswer,
-    DeepQAEvaluation,
-    ResearchPlan,
-    SearchAnswer,
-)
+from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
 from haiku.rag.qa.deep.prompts import (
     DECISION_PROMPT,
-    PLAN_PROMPT,
-    SEARCH_AGENT_PROMPT,
     SYNTHESIS_PROMPT,
     SYNTHESIS_PROMPT_WITH_CITATIONS,
 )
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 
-def build_deep_qa_graph(provider: str, model: str):
+def build_deep_qa_graph(
+    provider: str, model: str
+) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
     g = GraphBuilder(
         state_type=DeepQAState,
         deps_type=DeepQADeps,
@@ -82,7 +80,7 @@ def build_deep_qa_graph(provider: str, model: str):
     @g.step
     async def search_one(
         ctx: StepContext[DeepQAState, DeepQADeps, str],
-    ) -> SearchAnswer | None:
+    ) -> SearchAnswer:
         state = ctx.state
         deps = ctx.deps
         sub_q = ctx.inputs
@@ -130,17 +128,22 @@ def build_deep_qa_graph(provider: str, model: str):
         )
         try:
             result = await agent.run(sub_q, deps=agent_deps)
+            answer = result.output
+            if answer:
+                state.context.add_qa_response(answer)
+                preview = answer.answer[:150] + (
+                    "…" if len(answer.answer) > 150 else ""
+                )
+                log(deps, state, f"   [green]✓[/green] {preview}")
+            return answer
         except Exception as e:
             log(deps, state, f"[red]Search failed:[/red] {e}")
-            return None
-
-        answer = result.output
-        if answer:
-            state.context.add_qa_response(answer)
-            preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "")
-            log(deps, state, f"   [green]✓[/green] {preview}")
-
-        return answer
+            failure_answer = SearchAnswer(
+                query=sub_q,
+                answer=f"Search failed after retries: {str(e)}",
+                confidence=0.0,
+            )
+            return failure_answer
 
     @g.step
     async def get_batch(
@@ -291,8 +294,8 @@ def build_deep_qa_graph(provider: str, model: str):
 
     # Build the graph structure
     collect_answers = g.join(
-        collect_answers_reducer,
-        initial_factory=lambda: [],
+        reduce_list_append,
+        initial_factory=list[SearchAnswer],
     )
 
     g.add(

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

@@ -1,29 +1,6 @@
 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,48 +1,4 @@
-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."""
+"""Deep QA specific prompts."""
 
 SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers.
 

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

@@ -21,5 +21,5 @@ class DeepQAState:
     context: DeepQAContext
     max_sub_questions: int = 3
     max_iterations: int = 2
-    max_concurrency: int = 3
+    max_concurrency: int = 1
     iterations: int = 0

haiku_rag_slim/haiku/rag/research/__init__.py

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

haiku_rag_slim/haiku/rag/research/graph.py

@@ -3,11 +3,12 @@ from typing import Any
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.output import ToolOutput
-from pydantic_graph.beta import GraphBuilder, StepContext
+from pydantic_graph.beta import Graph, GraphBuilder, StepContext
+from pydantic_graph.beta.join import reduce_list_append
 
-from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log
-from haiku.rag.qa.deep.models import ResearchPlan, SearchAnswer
-from haiku.rag.qa.deep.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
+from haiku.rag.graph_common import get_model, log
+from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer
+from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.research.common import (
     format_analysis_for_prompt,
     format_context_for_prompt,
@@ -26,7 +27,9 @@ from haiku.rag.research.prompts import (
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
-def build_research_graph(provider: str, model: str):
+def build_research_graph(
+    provider: str, model: str
+) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
     g = GraphBuilder(
         state_type=ResearchState,
         deps_type=ResearchDeps,
@@ -86,7 +89,7 @@ def build_research_graph(provider: str, model: str):
     @g.step
     async def search_one(
         ctx: StepContext[ResearchState, ResearchDeps, str],
-    ) -> SearchAnswer | None:
+    ) -> SearchAnswer:
         state = ctx.state
         deps = ctx.deps
         sub_q = ctx.inputs
@@ -135,17 +138,22 @@ def build_research_graph(provider: str, model: str):
         )
         try:
             result = await agent.run(sub_q, deps=agent_deps)
+            answer = result.output
+            if answer:
+                state.context.add_qa_response(answer)
+                preview = answer.answer[:150] + (
+                    "…" if len(answer.answer) > 150 else ""
+                )
+                log(deps, state, f"   [green]✓[/green] {preview}")
+            return answer
         except Exception as e:
             log(deps, state, f"[red]Search failed:[/red] {e}")
-            return None
-
-        answer = result.output
-        if answer:
-            state.context.add_qa_response(answer)
-            preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "")
-            log(deps, state, f"   [green]✓[/green] {preview}")
-
-        return answer
+            failure_answer = SearchAnswer(
+                query=sub_q,
+                answer=f"Search failed after retries: {str(e)}",
+                confidence=0.0,
+            )
+            return failure_answer
 
     @g.step
     async def get_batch(
@@ -349,8 +357,8 @@ def build_research_graph(provider: str, model: str):
 
     # Build the graph structure
     collect_answers = g.join(
-        collect_answers_reducer,
-        initial_factory=lambda: [],
+        reduce_list_append,
+        initial_factory=list[SearchAnswer],
     )
 
     g.add(

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_common.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.qa.deep.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph_common.utils.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.qa.deep.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph_common.utils.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_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.qa.deep.common.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
     monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
 
     graph = build_research_graph(provider="test", model="test")