haiku.rag-slim 0.16.0__py3-none-any.whl

haiku/rag/graph/deep_qa/prompts.py

@@ -0,0 +1,59 @@
+"""Deep QA specific prompts."""
+
+SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers.
+
+Task:
+- Combine the gathered information from sub-questions into a single comprehensive answer
+- Answer the original question directly and completely
+- Base your answer strictly on the provided evidence
+- Be clear, accurate, and well-structured
+
+Output format:
+- answer: The complete answer to the original question (2-4 paragraphs)
+- sources: List of document titles/URIs used (extract from the sub-answers)
+
+Guidelines:
+- Start directly with the answer - no preamble like "Based on the research..."
+- Use a clear, professional tone
+- Organize information logically
+- If evidence is incomplete, state limitations clearly
+- Do not include any claims not supported by the gathered information"""
+
+SYNTHESIS_PROMPT_WITH_CITATIONS = """You are an expert at synthesizing information into clear, concise answers with proper citations.
+
+Task:
+- Combine the gathered information from sub-questions into a single comprehensive answer
+- Answer the original question directly and completely
+- Base your answer strictly on the provided evidence
+- Include inline citations using [Source Title] format
+
+Output format:
+- answer: The complete answer with inline citations (2-4 paragraphs)
+- sources: List of document titles/URIs used (extract from the sub-answers)
+
+Guidelines:
+- Start directly with the answer - no preamble like "Based on the research..."
+- Add citations after each claim: [Source Title]
+- Use a clear, professional tone
+- Organize information logically
+- If evidence is incomplete, state limitations clearly
+- Do not include any claims not supported by the gathered information"""
+
+DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question.
+
+Task:
+- Review the original question and all gathered sub-question answers
+- Determine if we have enough information to provide a comprehensive answer
+- If insufficient, suggest specific new sub-questions to fill the gaps
+
+Output format:
+- is_sufficient: Boolean indicating if we can answer the question comprehensively
+- reasoning: Clear explanation of your assessment
+- new_questions: List of specific follow-up questions needed (empty if sufficient)
+
+Guidelines:
+- Be strict but reasonable in your assessment
+- Focus on whether core aspects of the question are addressed
+- New questions should be specific and distinct from what's been asked
+- Limit new questions to 2-3 maximum
+- Consider whether additional searches would meaningfully improve the answer"""

haiku/rag/graph/deep_qa/state.py

@@ -0,0 +1,56 @@
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
+from haiku.rag.graph.deep_qa.models import DeepQAAnswer
+
+if TYPE_CHECKING:
+    from haiku.rag.config.models import AppConfig
+    from haiku.rag.graph.agui.emitter import AGUIEmitter
+
+
+@dataclass
+class DeepQADeps:
+    client: HaikuRAG
+    agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None
+    semaphore: asyncio.Semaphore | None = None
+
+
+class DeepQAState(BaseModel):
+    """Deep QA state for multi-agent question answering."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    context: DeepQAContext = Field(description="Shared QA context")
+    max_sub_questions: int = Field(
+        default=3, description="Maximum number of sub-questions"
+    )
+    max_iterations: int = Field(
+        default=2, description="Maximum number of QA iterations"
+    )
+    max_concurrency: int = Field(
+        default=1, description="Maximum parallel sub-question searches"
+    )
+    iterations: int = Field(default=0, description="Current iteration number")
+
+    @classmethod
+    def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":
+        """Create a DeepQAState from an AppConfig.
+
+        Args:
+            context: The DeepQAContext containing the question and settings
+            config: The AppConfig object (uses config.qa for state parameters)
+
+        Returns:
+            A configured DeepQAState instance
+        """
+        return cls(
+            context=context,
+            max_sub_questions=config.qa.max_sub_questions,
+            max_iterations=config.qa.max_iterations,
+            max_concurrency=config.qa.max_concurrency,
+        )

haiku/rag/graph/research/__init__.py

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

haiku/rag/graph/research/common.py

@@ -0,0 +1,87 @@
+from pydantic_ai import format_as_xml
+
+from haiku.rag.graph.research.dependencies import ResearchContext
+from haiku.rag.graph.research.models import InsightAnalysis
+
+
+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,
+                "context_snippets": qa.context,
+                "sources": qa.sources,  # pyright: ignore[reportAttributeAccessIssue]
+            }
+            for qa in context.qa_responses
+        ],
+        "insights": [
+            {
+                "id": insight.id,
+                "summary": insight.summary,
+                "status": insight.status.value,
+                "supporting_sources": insight.supporting_sources,
+                "originating_questions": insight.originating_questions,
+                "notes": insight.notes,
+            }
+            for insight in context.insights
+        ],
+        "gaps": [
+            {
+                "id": gap.id,
+                "description": gap.description,
+                "severity": gap.severity.value,
+                "blocking": gap.blocking,
+                "resolved": gap.resolved,
+                "resolved_by": gap.resolved_by,
+                "supporting_sources": gap.supporting_sources,
+                "notes": gap.notes,
+            }
+            for gap in context.gaps
+        ],
+    }
+    return format_as_xml(context_data, root_tag="research_context")
+
+
+def format_analysis_for_prompt(
+    analysis: InsightAnalysis | None,
+) -> str:
+    """Format the latest insight analysis as XML for prompts."""
+
+    if analysis is None:
+        return "<latest_analysis />"
+
+    data = {
+        "commentary": analysis.commentary,
+        "highlights": [
+            {
+                "id": insight.id,
+                "summary": insight.summary,
+                "status": insight.status.value,
+                "supporting_sources": insight.supporting_sources,
+                "originating_questions": insight.originating_questions,
+                "notes": insight.notes,
+            }
+            for insight in analysis.highlights
+        ],
+        "gap_assessments": [
+            {
+                "id": gap.id,
+                "description": gap.description,
+                "severity": gap.severity.value,
+                "blocking": gap.blocking,
+                "resolved": gap.resolved,
+                "resolved_by": gap.resolved_by,
+                "supporting_sources": gap.supporting_sources,
+                "notes": gap.notes,
+            }
+            for gap in analysis.gap_assessments
+        ],
+        "resolved_gaps": analysis.resolved_gaps,
+        "new_questions": analysis.new_questions,
+    }
+    return format_as_xml(data, root_tag="latest_analysis")

haiku/rag/graph/research/dependencies.py

@@ -0,0 +1,151 @@
+from collections.abc import Iterable
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.common.models import SearchAnswer
+from haiku.rag.graph.research.models import (
+    GapRecord,
+    InsightAnalysis,
+    InsightRecord,
+)
+
+
+class ResearchContext(BaseModel):
+    """Context shared across research agents."""
+
+    original_question: str = Field(description="The original research question")
+    sub_questions: list[str] = Field(
+        default_factory=list, description="Decomposed sub-questions"
+    )
+    qa_responses: list[SearchAnswer] = Field(
+        default_factory=list, description="Structured QA pairs used during research"
+    )
+    insights: list[InsightRecord] = Field(
+        default_factory=list, description="Key insights discovered"
+    )
+    gaps: list[GapRecord] = Field(
+        default_factory=list, description="Identified information gaps"
+    )
+
+    # Private dict indexes for O(1) lookups
+    _insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict)
+    _gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict)
+
+    def model_post_init(self, __context: object) -> None:
+        """Build indexes after initialization."""
+        self._insights_by_id = {ins.id: ins for ins in self.insights}
+        self._gaps_by_id = {gap.id: gap for gap in self.gaps}
+
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        """Add a structured QA response (minimal context already included)."""
+        self.qa_responses.append(qa)
+
+    def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
+        """Merge one or more insights into the shared context with deduplication."""
+        merged: list[InsightRecord] = []
+
+        for record in records:
+            candidate = InsightRecord.model_validate(record)
+            existing = self._insights_by_id.get(candidate.id)
+
+            if existing:
+                # Update existing insight
+                existing.summary = candidate.summary
+                existing.status = candidate.status
+                if candidate.notes:
+                    existing.notes = candidate.notes
+                existing.supporting_sources = _merge_unique(
+                    existing.supporting_sources, candidate.supporting_sources
+                )
+                existing.originating_questions = _merge_unique(
+                    existing.originating_questions, candidate.originating_questions
+                )
+                merged.append(existing)
+            else:
+                # Add new insight
+                new_insight = candidate.model_copy(deep=True)
+                self.insights.append(new_insight)
+                self._insights_by_id[new_insight.id] = new_insight
+                merged.append(new_insight)
+
+        return merged
+
+    def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]:
+        """Merge one or more gap records into the shared context with deduplication."""
+        merged: list[GapRecord] = []
+
+        for record in records:
+            candidate = GapRecord.model_validate(record)
+            existing = self._gaps_by_id.get(candidate.id)
+
+            if existing:
+                # Update existing gap
+                existing.description = candidate.description
+                existing.severity = candidate.severity
+                existing.blocking = candidate.blocking
+                existing.resolved = candidate.resolved
+                if candidate.notes:
+                    existing.notes = candidate.notes
+                existing.supporting_sources = _merge_unique(
+                    existing.supporting_sources, candidate.supporting_sources
+                )
+                existing.resolved_by = _merge_unique(
+                    existing.resolved_by, candidate.resolved_by
+                )
+                merged.append(existing)
+            else:
+                # Add new gap
+                new_gap = candidate.model_copy(deep=True)
+                self.gaps.append(new_gap)
+                self._gaps_by_id[new_gap.id] = new_gap
+                merged.append(new_gap)
+
+        return merged
+
+    def mark_gap_resolved(
+        self, identifier: str, resolved_by: Iterable[str] | None = None
+    ) -> GapRecord | None:
+        """Mark a gap as resolved by identifier."""
+        gap = self._gaps_by_id.get(identifier)
+        if gap is None:
+            return None
+
+        gap.resolved = True
+        gap.blocking = False
+        if resolved_by:
+            gap.resolved_by = _merge_unique(gap.resolved_by, list(resolved_by))
+        return gap
+
+    def integrate_analysis(self, analysis: InsightAnalysis) -> None:
+        """Apply an analysis result to the shared context."""
+        merged_insights: list[InsightRecord] = []
+        if analysis.highlights:
+            merged_insights = self.upsert_insights(analysis.highlights)
+            analysis.highlights = merged_insights
+        if analysis.gap_assessments:
+            merged_gaps = self.upsert_gaps(analysis.gap_assessments)
+            analysis.gap_assessments = merged_gaps
+        if analysis.resolved_gaps:
+            resolved_by_list = (
+                [ins.id for ins in merged_insights] if merged_insights else None
+            )
+            for resolved in analysis.resolved_gaps:
+                self.mark_gap_resolved(resolved, resolved_by=resolved_by_list)
+        for question in analysis.new_questions:
+            if question not in self.sub_questions:
+                self.sub_questions.append(question)
+
+
+class ResearchDependencies(BaseModel):
+    """Dependencies for research agents with multi-agent context."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    context: ResearchContext = Field(description="Shared research context")
+
+
+def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
+    """Merge two iterables preserving order while removing duplicates."""
+    return [k for k in dict.fromkeys([*existing, *incoming]) if k]

haiku/rag/graph/research/graph.py

@@ -0,0 +1,295 @@
+from pydantic_ai import Agent
+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.common.models import SearchAnswer
+from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
+from haiku.rag.graph.research.common import (
+    format_analysis_for_prompt,
+    format_context_for_prompt,
+)
+from haiku.rag.graph.research.dependencies import ResearchDependencies
+from haiku.rag.graph.research.models import (
+    EvaluationResult,
+    InsightAnalysis,
+    ResearchReport,
+)
+from haiku.rag.graph.research.prompts import (
+    DECISION_AGENT_PROMPT,
+    INSIGHT_AGENT_PROMPT,
+    SYNTHESIS_AGENT_PROMPT,
+)
+from haiku.rag.graph.research.state import ResearchDeps, ResearchState
+
+
+def build_research_graph(
+    config: AppConfig = Config,
+) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
+    """Build the Research graph.
+
+    Args:
+        config: AppConfig object (uses config.research for provider, model, and graph parameters)
+
+    Returns:
+        Configured Research graph
+    """
+    provider = config.research.provider
+    model = config.research.model
+    g = GraphBuilder(
+        state_type=ResearchState,
+        deps_type=ResearchDeps,
+        output_type=ResearchReport,
+    )
+
+    # Create and register the plan node using the factory
+    plan = g.step(
+        create_plan_node(
+            provider=provider,
+            model=model,
+            deps_type=ResearchDependencies,  # type: ignore[arg-type]
+            activity_message="Creating research plan",
+            output_retries=3,
+        )
+    )  # type: ignore[arg-type]
+
+    # Create and register the search_one node using the factory
+    search_one = g.step(
+        create_search_node(
+            provider=provider,
+            model=model,
+            deps_type=ResearchDependencies,  # type: ignore[arg-type]
+            with_step_wrapper=True,
+            success_message_format="Found answer with {confidence:.0%} confidence",
+            handle_exceptions=True,
+        )
+    )  # type: ignore[arg-type]
+
+    @g.step
+    async def get_batch(
+        ctx: StepContext[ResearchState, ResearchDeps, None | bool],
+    ) -> list[str] | None:
+        """Get all remaining questions for this iteration."""
+        state = ctx.state
+
+        if not state.context.sub_questions:
+            return None
+
+        # Take ALL remaining questions and process them in parallel
+        batch = list(state.context.sub_questions)
+        state.context.sub_questions.clear()
+        return batch
+
+    @g.step
+    async def analyze_insights(
+        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
+    ) -> None:
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("analyze_insights")
+            deps.agui_emitter.update_activity(
+                "analyzing", "Synthesizing insights and gaps"
+            )
+
+        try:
+            agent = Agent(
+                model=get_model(provider, model),
+                output_type=InsightAnalysis,
+                instructions=INSIGHT_AGENT_PROMPT,
+                retries=3,
+                output_retries=3,
+                deps_type=ResearchDependencies,
+            )
+
+            context_xml = format_context_for_prompt(state.context)
+            prompt = (
+                "Review the latest research context and update the shared ledger of insights, gaps,"
+                " and follow-up questions.\n\n"
+                f"{context_xml}"
+            )
+            agent_deps = ResearchDependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+            analysis: InsightAnalysis = result.output
+
+            state.context.integrate_analysis(analysis)
+            state.last_analysis = analysis
+
+            # State updated with insights/gaps - emit state update and narrate
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                highlights = len(analysis.highlights) if analysis.highlights else 0
+                gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
+                resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
+                parts = []
+                if highlights:
+                    parts.append(f"{highlights} insights")
+                if gaps:
+                    parts.append(f"{gaps} gaps")
+                if resolved:
+                    parts.append(f"{resolved} resolved")
+                summary = ", ".join(parts) if parts else "No updates"
+                deps.agui_emitter.update_activity("analyzing", f"Analysis: {summary}")
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    @g.step
+    async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("decide")
+            deps.agui_emitter.update_activity(
+                "evaluating", "Evaluating research sufficiency"
+            )
+
+        try:
+            agent = Agent(
+                model=get_model(provider, model),
+                output_type=EvaluationResult,
+                instructions=DECISION_AGENT_PROMPT,
+                retries=3,
+                output_retries=3,
+                deps_type=ResearchDependencies,
+            )
+
+            context_xml = format_context_for_prompt(state.context)
+            analysis_xml = format_analysis_for_prompt(state.last_analysis)
+            prompt_parts = [
+                "Assess whether the research now answers the original question with adequate confidence.",
+                context_xml,
+                analysis_xml,
+            ]
+            if state.last_eval is not None:
+                prev = state.last_eval
+                prompt_parts.append(
+                    "<previous_evaluation>"
+                    f"<confidence>{prev.confidence_score:.2f}</confidence>"
+                    f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
+                    f"<reasoning>{prev.reasoning}</reasoning>"
+                    "</previous_evaluation>"
+                )
+            prompt = "\n\n".join(part for part in prompt_parts if part)
+
+            agent_deps = ResearchDependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            decision_result = await agent.run(prompt, deps=agent_deps)
+            output = decision_result.output
+
+            state.last_eval = output
+            state.iterations += 1
+
+            for new_q in output.new_questions:
+                if new_q not in state.context.sub_questions:
+                    state.context.sub_questions.append(new_q)
+
+            # State updated with evaluation - emit state update and narrate
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                sufficient = "Yes" if output.is_sufficient else "No"
+                deps.agui_emitter.update_activity(
+                    "evaluating",
+                    f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
+                )
+
+            should_continue = (
+                not output.is_sufficient
+                or output.confidence_score < state.confidence_threshold
+            ) and state.iterations < state.max_iterations
+
+            return should_continue
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    @g.step
+    async def synthesize(
+        ctx: StepContext[ResearchState, ResearchDeps, None | bool],
+    ) -> ResearchReport:
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("synthesize")
+            deps.agui_emitter.update_activity(
+                "synthesizing", "Generating final research report"
+            )
+
+        try:
+            agent = Agent(
+                model=get_model(provider, model),
+                output_type=ResearchReport,
+                instructions=SYNTHESIS_AGENT_PROMPT,
+                retries=3,
+                output_retries=3,
+                deps_type=ResearchDependencies,
+            )
+
+            context_xml = format_context_for_prompt(state.context)
+            prompt = (
+                "Generate a comprehensive research report based on all gathered information.\n\n"
+                f"{context_xml}\n\n"
+                "Create a detailed report that synthesizes all findings into a coherent response."
+            )
+            agent_deps = ResearchDependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+            return result.output
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    # Build the graph structure
+    collect_answers = g.join(
+        reduce_list_append,
+        initial_factory=list[SearchAnswer],
+    )
+
+    g.add(
+        g.edge_from(g.start_node).to(plan),
+        g.edge_from(plan).to(get_batch),
+    )
+
+    # Branch based on whether we have questions
+    g.add(
+        g.edge_from(get_batch).to(
+            g.decision()
+            .branch(g.match(list).label("Has questions").map().to(search_one))
+            .branch(g.match(type(None)).label("No questions").to(synthesize))
+        ),
+        g.edge_from(search_one).to(collect_answers),
+        g.edge_from(collect_answers).to(analyze_insights),
+        g.edge_from(analyze_insights).to(decide),
+    )
+
+    # Branch based on decision
+    g.add(
+        g.edge_from(decide).to(
+            g.decision()
+            .branch(
+                g.match(bool, matches=lambda x: x)
+                .label("Continue research")
+                .to(get_batch)
+            )
+            .branch(
+                g.match(bool, matches=lambda x: not x)
+                .label("Done researching")
+                .to(synthesize)
+            )
+        ),
+        g.edge_from(synthesize).to(g.end_node),
+    )
+
+    return g.build()