haiku.rag 0.10.1__py3-none-any.whl → 0.11.0__py3-none-any.whl

haiku/rag/research/dependencies.py

@@ -1,8 +1,16 @@
+from collections.abc import Iterable
+
 from pydantic import BaseModel, Field
 from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.research.models import SearchAnswer
+from haiku.rag.research.models import (
+    GapRecord,
+    InsightAnalysis,
+    InsightRecord,
+    SearchAnswer,
+)
+from haiku.rag.research.stream import ResearchStream
 
 
 class ResearchContext(BaseModel):
@@ -15,10 +23,10 @@ class ResearchContext(BaseModel):
     qa_responses: list[SearchAnswer] = Field(
         default_factory=list, description="Structured QA pairs used during research"
     )
-    insights: list[str] = Field(
+    insights: list[InsightRecord] = Field(
         default_factory=list, description="Key insights discovered"
     )
-    gaps: list[str] = Field(
+    gaps: list[GapRecord] = Field(
         default_factory=list, description="Identified information gaps"
     )
 
@@ -26,15 +34,147 @@ class ResearchContext(BaseModel):
         """Add a structured QA response (minimal context already included)."""
         self.qa_responses.append(qa)
 
-    def add_insight(self, insight: str) -> None:
-        """Add a key insight."""
-        if insight not in self.insights:
-            self.insights.append(insight)
+    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 = next(
+                (ins for ins in self.insights if ins.id == candidate.id), None
+            )
+            if not existing:
+                existing = next(
+                    (ins for ins in self.insights if ins.summary == candidate.summary),
+                    None,
+                )
+
+            if existing:
+                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:
+                candidate = candidate.model_copy(deep=True)
+                if candidate.id is None:  # pragma: no cover - defensive
+                    raise ValueError(
+                        "InsightRecord.id must be populated after validation"
+                    )
+                candidate_id: str = candidate.id
+                candidate.id = self._allocate_insight_id(candidate_id)
+                self.insights.append(candidate)
+                merged.append(candidate)
+
+        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 = next((gap for gap in self.gaps if gap.id == candidate.id), None)
+            if not existing:
+                existing = next(
+                    (
+                        gap
+                        for gap in self.gaps
+                        if gap.description == candidate.description
+                    ),
+                    None,
+                )
+
+            if existing:
+                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:
+                candidate = candidate.model_copy(deep=True)
+                if candidate.id is None:  # pragma: no cover - defensive
+                    raise ValueError("GapRecord.id must be populated after validation")
+                candidate_id: str = candidate.id
+                candidate.id = self._allocate_gap_id(candidate_id)
+                self.gaps.append(candidate)
+                merged.append(candidate)
+
+        return merged
+
+    def mark_gap_resolved(
+        self, identifier: str, resolved_by: Iterable[str] | None = None
+    ) -> GapRecord | None:
+        """Mark a gap as resolved by identifier (id or description)."""
+
+        gap = self._find_gap(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 add_gap(self, gap: str) -> None:
-        """Identify an information gap."""
-        if gap not in self.gaps:
-            self.gaps.append(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 ins.id is not None]
+                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)
+
+    def _allocate_insight_id(self, candidate_id: str) -> str:
+        taken: set[str] = set()
+        for ins in self.insights:
+            if ins.id is not None:
+                taken.add(ins.id)
+        return _allocate_sequential_id(candidate_id, taken)
+
+    def _allocate_gap_id(self, candidate_id: str) -> str:
+        taken: set[str] = set()
+        for gap in self.gaps:
+            if gap.id is not None:
+                taken.add(gap.id)
+        return _allocate_sequential_id(candidate_id, taken)
+
+    def _find_gap(self, identifier: str) -> GapRecord | None:
+        normalized = identifier.lower().strip()
+        for gap in self.gaps:
+            if gap.id is not None and gap.id == normalized:
+                return gap
+            if gap.description.lower().strip() == normalized:
+                return gap
+        return None
 
 
 class ResearchDependencies(BaseModel):
@@ -45,3 +185,31 @@ class ResearchDependencies(BaseModel):
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
     console: Console | None = None
+    stream: ResearchStream | None = Field(
+        default=None, description="Optional research event stream"
+    )
+
+
+def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
+    """Merge two iterables preserving order while removing duplicates."""
+
+    merged = list(existing)
+    seen = {item for item in existing if item}
+    for item in incoming:
+        if item and item not in seen:
+            merged.append(item)
+            seen.add(item)
+    return merged
+
+
+def _allocate_sequential_id(candidate: str, taken: set[str]) -> str:
+    slug = candidate
+    if slug not in taken:
+        return slug
+    base = slug
+    counter = 2
+    while True:
+        slug = f"{base}-{counter}"
+        if slug not in taken:
+            return slug
+        counter += 1

haiku/rag/research/graph.py

@@ -1,7 +1,7 @@
 from pydantic_graph import Graph
 
 from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.nodes.evaluate import EvaluateNode
+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
@@ -10,7 +10,8 @@ from haiku.rag.research.state import ResearchDeps, ResearchState
 __all__ = [
     "PlanNode",
     "SearchDispatchNode",
-    "EvaluateNode",
+    "AnalyzeInsightsNode",
+    "DecisionNode",
     "SynthesizeNode",
     "ResearchState",
     "ResearchDeps",
@@ -23,7 +24,8 @@ def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]
         nodes=[
             PlanNode,
             SearchDispatchNode,
-            EvaluateNode,
+            AnalyzeInsightsNode,
+            DecisionNode,
             SynthesizeNode,
         ]
     )

haiku/rag/research/models.py

@@ -1,4 +1,134 @@
-from pydantic import BaseModel, Field
+import re
+from enum import Enum
+
+from pydantic import BaseModel, Field, model_validator
+
+_SLUG_RE = re.compile(r"[^a-z0-9]+")
+
+
+def _make_slug(text: str, prefix: str) -> str:
+    """Generate a lowercase slug with the given prefix as fallback."""
+
+    base = _SLUG_RE.sub("-", text.lower()).strip("-")
+    if not base:
+        base = prefix
+    # Trim overly long slugs but keep enough entropy for readability
+    return base[:48]
+
+
+class InsightStatus(str, Enum):
+    OPEN = "open"
+    VALIDATED = "validated"
+    TENTATIVE = "tentative"
+
+
+class GapSeverity(str, Enum):
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+
+
+class InsightRecord(BaseModel):
+    """Structured insight with provenance and lifecycle metadata."""
+
+    id: str | None = Field(
+        default=None,
+        description="Stable slug identifier for the insight (auto-generated if omitted)",
+    )
+    summary: str = Field(description="Concise description of the insight")
+    status: InsightStatus = Field(
+        default=InsightStatus.OPEN,
+        description="Lifecycle status for the insight",
+    )
+    supporting_sources: list[str] = Field(
+        default_factory=list,
+        description="Source identifiers backing the insight",
+    )
+    originating_questions: list[str] = Field(
+        default_factory=list,
+        description="Research sub-questions that produced this insight",
+    )
+    notes: str | None = Field(
+        default=None,
+        description="Optional elaboration or caveats for the insight",
+    )
+
+    @model_validator(mode="after")
+    def _set_defaults(self) -> "InsightRecord":
+        if not self.id:
+            self.id = _make_slug(self.summary, "insight")
+        self.id = self.id.lower()
+        self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
+        self.originating_questions = list(dict.fromkeys(self.originating_questions))
+        return self
+
+
+class GapRecord(BaseModel):
+    """Structured representation of an identified research gap."""
+
+    id: str | None = Field(
+        default=None,
+        description="Stable slug identifier for the gap (auto-generated if omitted)",
+    )
+    description: str = Field(description="Concrete statement of what is missing")
+    severity: GapSeverity = Field(
+        default=GapSeverity.MEDIUM,
+        description="Severity of the gap for answering the main question",
+    )
+    blocking: bool = Field(
+        default=True,
+        description="Whether this gap blocks a confident answer",
+    )
+    resolved: bool = Field(
+        default=False,
+        description="Flag indicating if the gap has been resolved",
+    )
+    resolved_by: list[str] = Field(
+        default_factory=list,
+        description="Insight IDs or notes explaining how the gap was closed",
+    )
+    supporting_sources: list[str] = Field(
+        default_factory=list,
+        description="Sources confirming the gap status (e.g., evidence of absence)",
+    )
+    notes: str | None = Field(
+        default=None,
+        description="Optional clarification about the gap or follow-up actions",
+    )
+
+    @model_validator(mode="after")
+    def _set_defaults(self) -> "GapRecord":
+        if not self.id:
+            self.id = _make_slug(self.description, "gap")
+        self.id = self.id.lower()
+        self.resolved_by = list(dict.fromkeys(self.resolved_by))
+        self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
+        return self
+
+
+class InsightAnalysis(BaseModel):
+    """Output of the insight aggregation agent."""
+
+    highlights: list[InsightRecord] = Field(
+        default_factory=list,
+        description="New or updated insights discovered this iteration",
+    )
+    gap_assessments: list[GapRecord] = Field(
+        default_factory=list,
+        description="New or updated gap records based on current evidence",
+    )
+    resolved_gaps: list[str] = Field(
+        default_factory=list,
+        description="Gap identifiers or descriptions considered resolved",
+    )
+    new_questions: list[str] = Field(
+        default_factory=list,
+        max_length=3,
+        description="Up to three follow-up sub-questions to pursue next",
+    )
+    commentary: str = Field(
+        description="Short narrative summary of the incremental findings",
+    )
 
 
 class ResearchPlan(BaseModel):
@@ -37,6 +167,9 @@ class EvaluationResult(BaseModel):
         max_length=3,
         default=[],
     )
+    gaps: list[str] = Field(
+        description="Concrete information gaps that remain", default_factory=list
+    )
     confidence_score: float = Field(
         description="Confidence level in the completeness of research (0-1)",
         ge=0.0,

haiku/rag/research/nodes/analysis.py

@@ -0,0 +1,181 @@
+from dataclasses import dataclass
+
+from pydantic_ai import Agent
+from pydantic_graph import BaseNode, GraphRunContext
+
+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
+
+
+@dataclass
+class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
+    ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
+        state = ctx.state
+        deps = ctx.deps
+
+        log(
+            deps,
+            state,
+            "\n[bold cyan]🧭 Synthesizing new insights and gap status...[/bold cyan]",
+        )
+
+        agent = Agent(
+            model=get_model(self.provider, self.model),
+            output_type=InsightAnalysis,
+            instructions=INSIGHT_AGENT_PROMPT,
+            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,
+            console=deps.console,
+            stream=deps.stream,
+        )
+        result = await agent.run(prompt, deps=agent_deps)
+        analysis: InsightAnalysis = result.output
+
+        state.context.integrate_analysis(analysis)
+        state.last_analysis = analysis
+
+        if analysis.commentary:
+            log(deps, state, f"   Summary: {analysis.commentary}")
+        if analysis.highlights:
+            log(deps, state, "   [bold]Updated insights:[/bold]")
+            for insight in analysis.highlights:
+                label = insight.status.value
+                log(
+                    deps,
+                    state,
+                    f"   • ({label}) {insight.summary}",
+                )
+        if analysis.gap_assessments:
+            log(deps, state, "   [bold yellow]Gap updates:[/bold yellow]")
+            for gap in analysis.gap_assessments:
+                status = "resolved" if gap.resolved else "open"
+                severity = gap.severity.value
+                log(
+                    deps,
+                    state,
+                    f"   • ({severity}/{status}) {gap.description}",
+                )
+        if analysis.resolved_gaps:
+            log(deps, state, "   [green]Resolved gaps:[/green]")
+            for resolved in analysis.resolved_gaps:
+                log(deps, state, f"   • {resolved}")
+        if analysis.new_questions:
+            log(deps, state, "   [cyan]Proposed follow-ups:[/cyan]")
+            for question in analysis.new_questions:
+                log(deps, state, f"   • {question}")
+
+        return DecisionNode(self.provider, self.model)
+
+
+@dataclass
+class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
+    ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
+        state = ctx.state
+        deps = ctx.deps
+
+        log(
+            deps,
+            state,
+            "\n[bold cyan]📊 Evaluating research sufficiency...[/bold cyan]",
+        )
+
+        agent = Agent(
+            model=get_model(self.provider, self.model),
+            output_type=EvaluationResult,
+            instructions=DECISION_AGENT_PROMPT,
+            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,
+            console=deps.console,
+            stream=deps.stream,
+        )
+        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)
+
+        if output.key_insights:
+            log(deps, state, "   [bold]Key insights:[/bold]")
+            for insight in output.key_insights:
+                log(deps, state, f"   • {insight}")
+
+        if output.gaps:
+            log(deps, state, "   [bold yellow]Remaining gaps:[/bold yellow]")
+            for gap in output.gaps:
+                log(deps, state, f"   • {gap}")
+
+        log(
+            deps,
+            state,
+            f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
+        )
+        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
+
+        if (
+            output.is_sufficient
+            and output.confidence_score >= state.confidence_threshold
+        ) or state.iterations >= state.max_iterations:
+            log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
+            return SynthesizeNode(self.provider, self.model)
+
+        return SearchDispatchNode(self.provider, self.model)

haiku/rag/research/nodes/plan.py

@@ -22,7 +22,7 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         state = ctx.state
         deps = ctx.deps
 
-        log(deps.console, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
+        log(deps, state, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
 
         plan_agent = Agent(
             model=get_model(self.provider, self.model),
@@ -45,19 +45,26 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
 
         prompt = (
             "Plan a focused research approach for the main question.\n\n"
-            f"Main question: {state.question}"
+            f"Main question: {state.context.original_question}"
         )
 
         agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
+            client=deps.client,
+            context=state.context,
+            console=deps.console,
+            stream=deps.stream,
         )
         plan_result = await plan_agent.run(prompt, deps=agent_deps)
-        state.sub_questions = list(plan_result.output.sub_questions)
+        state.context.sub_questions = list(plan_result.output.sub_questions)
 
-        log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]")
-        log(deps.console, f"   [bold]Main Question:[/bold] {state.question}")
-        log(deps.console, "   [bold]Sub-questions:[/bold]")
-        for i, sq in enumerate(state.sub_questions, 1):
-            log(deps.console, f"      {i}. {sq}")
+        log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
+        log(
+            deps,
+            state,
+            f"   [bold]Main Question:[/bold] {state.context.original_question}",
+        )
+        log(deps, state, "   [bold]Sub-questions:[/bold]")
+        for i, sq in enumerate(state.context.sub_questions, 1):
+            log(deps, state, f"      {i}. {sq}")
 
         return SearchDispatchNode(self.provider, self.model)

haiku/rag/research/nodes/search.py

@@ -24,20 +24,21 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
     ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
         state = ctx.state
         deps = ctx.deps
-        if not state.sub_questions:
-            from haiku.rag.research.nodes.evaluate import EvaluateNode
+        if not state.context.sub_questions:
+            from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode
 
-            return EvaluateNode(self.provider, self.model)
+            return AnalyzeInsightsNode(self.provider, self.model)
 
         # Take up to max_concurrency questions and answer them concurrently
         take = max(1, state.max_concurrency)
         batch: list[str] = []
-        while state.sub_questions and len(batch) < take:
-            batch.append(state.sub_questions.pop(0))
+        while state.context.sub_questions and len(batch) < take:
+            batch.append(state.context.sub_questions.pop(0))
 
         async def answer_one(sub_q: str) -> SearchAnswer | None:
             log(
-                deps.console,
+                deps,
+                state,
                 f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
             )
             agent = Agent(
@@ -71,12 +72,15 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
                 return format_as_xml(entries, root_tag="snippets")
 
             agent_deps = ResearchDependencies(
-                client=deps.client, context=state.context, console=deps.console
+                client=deps.client,
+                context=state.context,
+                console=deps.console,
+                stream=deps.stream,
             )
             try:
                 result = await agent.run(sub_q, deps=agent_deps)
             except Exception as e:
-                log(deps.console, f"[red]Search failed:[/red] {e}")
+                log(deps, state, f"[red]Search failed:[/red] {e}")
                 return None
 
             return result.output
@@ -86,8 +90,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             if ans is None:
                 continue
             state.context.add_qa_response(ans)
-            if deps.console:
-                preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
-                log(deps.console, f"   [green]✓[/green] {preview}")
+            preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
+            log(deps, state, f"   [green]✓[/green] {preview}")
 
         return SearchDispatchNode(self.provider, self.model)