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

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)

haiku/rag/research/nodes/synthesize.py

@@ -24,7 +24,8 @@ class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         deps = ctx.deps
 
         log(
-            deps.console,
+            deps,
+            state,
             "\n[bold cyan]📝 Generating final research report...[/bold cyan]",
         )
 
@@ -43,9 +44,12 @@ class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             "Create a detailed report that synthesizes all findings into a coherent response."
         )
         agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
+            client=deps.client,
+            context=state.context,
+            console=deps.console,
+            stream=deps.stream,
         )
         result = await agent.run(prompt, deps=agent_deps)
 
-        log(deps.console, "[bold green]✅ Research complete![/bold green]")
+        log(deps, state, "[bold green]✅ Research complete![/bold green]")
         return End(result.output)

haiku/rag/research/prompts.py

@@ -44,38 +44,77 @@ Answering rules:
 - Prefer concise phrasing; avoid copying long passages.
 - When evidence is partial, state the limits explicitly in the answer."""
 
-EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for
-the research workflow.
+INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
+research workflow.
 
 Inputs available:
-- Original research question
-- Question–answer pairs produced by search
-- Raw search results and source metadata
-- Previously identified insights
-
-ANALYSIS:
-1. Extract the most important, non‑obvious insights from the collected evidence.
-2. Identify patterns, agreements, and disagreements across sources.
-3. Note material uncertainties and assumptions.
-
-EVALUATION:
-1. Decide if we have sufficient information to answer the original question.
-2. Provide a confidence_score in [0,1] considering:
-   - Coverage of the main question’s aspects
-   - Quality, consistency, and diversity of sources
-   - Depth and specificity of evidence
-3. List concrete gaps that still need investigation.
-4. Propose up to 3 new sub_questions that would close the highest‑value gaps.
+- Original research question and sub-questions
+- Question–answer pairs with supporting snippets and sources
+- Existing insights and gaps (with status metadata)
+
+Tasks:
+1. Extract new or refined insights that advance understanding of the question.
+2. Update gap status, creating new gap entries when necessary and marking
+   resolved ones explicitly.
+3. Suggest up to 3 high-impact follow-up sub_questions that would close the
+   most important remaining gaps.
+
+Output format (map directly to fields):
+- highlights: list of insights with fields {summary, status, supporting_sources,
+  originating_questions, notes}. Use status one of {validated, open, tentative}.
+- gap_assessments: list of gaps with fields {description, severity, blocking,
+  resolved, resolved_by, supporting_sources, notes}. Severity must be one of
+  {low, medium, high}. resolved_by may reference related insight summaries if no
+  stable identifier yet.
+- resolved_gaps: list of identifiers or descriptions for gaps now closed.
+- new_questions: up to 3 standalone, specific sub-questions (no duplicates with
+  existing ones).
+- commentary: 1–3 sentences summarizing what changed this round.
+
+Guidance:
+- Be concise and avoid repeating previously recorded information unless it
+  changed materially.
+- Tie supporting_sources to the evidence used; omit if unavailable.
+- Only propose new sub_questions that directly address remaining gaps.
+- When marking a gap as resolved, ensure the rationale is clear via
+  resolved_by or notes."""
+
+DECISION_AGENT_PROMPT = """You are the research governor responsible for making
+stop/go decisions.
+
+Inputs available:
+- Original research question and current plan
+- Full insight ledger with status metadata
+- Up-to-date gap tracker, including resolved indicators
+- Latest insight analysis summary (highlights, gap changes, new questions)
+- Previous evaluation decision (if any)
+
+Tasks:
+1. Determine whether the collected evidence now answers the original question.
+2. Provide a confidence_score in [0,1] that reflects coverage, evidence quality,
+   and agreement across sources.
+3. List the highest-priority gaps that still block a confident answer. Reference
+   existing gap descriptions rather than inventing new ones.
+4. Optionally propose up to 3 new sub_questions only if they are not already in
+   the current backlog.
 
 Strictness:
-- Only mark research as sufficient when all major aspects are addressed with
-  consistent, reliable evidence and no critical gaps remain.
-
-New sub_questions must:
-- Be genuinely new (not answered or duplicative; check qa_responses).
-- Be standalone and specific (entities, scope, timeframe/region if relevant).
-- Be actionable and scoped to the knowledge base (narrow if necessary).
-- Be ordered by expected impact (most valuable first)."""
+- Only mark research as sufficient when every critical aspect of the main
+  question is addressed with reliable, corroborated evidence.
+- Treat unresolved high-severity or blocking gaps as a hard stop.
+
+Output fields must line up with EvaluationResult:
+- key_insights: concise bullet-ready statements of the most decision-relevant
+  insights (cite status if helpful).
+- new_questions: follow-up sub-questions (max 3) meeting the specificity rules.
+- gaps: list remaining blockers; reuse wording from the tracked gaps when
+  possible to aid downstream reconciliation.
+- confidence_score: numeric in [0,1].
+- is_sufficient: true only when no blocking gaps remain.
+- reasoning: short narrative tying the decision to evidence coverage.
+
+Remember: prefer maintaining continuity with the structured context over
+introducing new terminology."""
 
 SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
 research report.

haiku/rag/research/state.py

@@ -1,25 +1,32 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research.dependencies import ResearchContext
-from haiku.rag.research.models import EvaluationResult
+from haiku.rag.research.models import EvaluationResult, InsightAnalysis
+from haiku.rag.research.stream import ResearchStream
 
 
 @dataclass
 class ResearchDeps:
     client: HaikuRAG
     console: Console | None = None
+    stream: ResearchStream | None = None
+
+    def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
+        if self.console:
+            self.console.print(message)
+        if self.stream:
+            self.stream.log(message, state)
 
 
 @dataclass
 class ResearchState:
-    question: str
     context: ResearchContext
-    sub_questions: list[str] = field(default_factory=list)
     iterations: int = 0
     max_iterations: int = 3
     max_concurrency: int = 1
     confidence_threshold: float = 0.8
     last_eval: EvaluationResult | None = None
+    last_analysis: InsightAnalysis | None = None

haiku/rag/research/stream.py

@@ -0,0 +1,177 @@
+import asyncio
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+from haiku.rag.research.models import ResearchReport
+
+if TYPE_CHECKING:  # pragma: no cover
+    from haiku.rag.research.state import ResearchState
+
+
+@dataclass(slots=True)
+class ResearchStateSnapshot:
+    question: str
+    sub_questions: list[str]
+    iterations: int
+    max_iterations: int
+    max_concurrency: int
+    confidence_threshold: float
+    pending_sub_questions: int
+    answered_questions: int
+    insights: list[str]
+    gaps: list[str]
+    last_confidence: float | None
+    last_sufficient: bool | None
+
+    @classmethod
+    def from_state(cls, state: "ResearchState") -> "ResearchStateSnapshot":
+        context = state.context
+        last_confidence: float | None = None
+        last_sufficient: bool | None = None
+        if state.last_eval:
+            last_confidence = state.last_eval.confidence_score
+            last_sufficient = state.last_eval.is_sufficient
+
+        return cls(
+            question=context.original_question,
+            sub_questions=list(context.sub_questions),
+            iterations=state.iterations,
+            max_iterations=state.max_iterations,
+            max_concurrency=state.max_concurrency,
+            confidence_threshold=state.confidence_threshold,
+            pending_sub_questions=len(context.sub_questions),
+            answered_questions=len(context.qa_responses),
+            insights=[
+                f"{insight.status.value}:{insight.summary}"
+                for insight in context.insights
+            ],
+            gaps=[
+                f"{gap.severity.value}/{'resolved' if gap.resolved else 'open'}:{gap.description}"
+                for gap in context.gaps
+            ],
+            last_confidence=last_confidence,
+            last_sufficient=last_sufficient,
+        )
+
+
+@dataclass(slots=True)
+class ResearchStreamEvent:
+    type: Literal["log", "report", "error"]
+    message: str | None = None
+    state: ResearchStateSnapshot | None = None
+    report: ResearchReport | None = None
+    error: str | None = None
+
+
+class ResearchStream:
+    """Queue-backed stream for research graph events."""
+
+    def __init__(self) -> None:
+        self._queue: asyncio.Queue[ResearchStreamEvent | None] = asyncio.Queue()
+        self._closed = False
+
+    def _snapshot(self, state: "ResearchState | None") -> ResearchStateSnapshot | None:
+        if state is None:
+            return None
+        return ResearchStateSnapshot.from_state(state)
+
+    def log(self, message: str, state: "ResearchState | None" = None) -> None:
+        if self._closed:
+            return
+        event = ResearchStreamEvent(
+            type="log", message=message, state=self._snapshot(state)
+        )
+        self._queue.put_nowait(event)
+
+    def report(self, report: ResearchReport, state: "ResearchState") -> None:
+        if self._closed:
+            return
+        event = ResearchStreamEvent(
+            type="report",
+            report=report,
+            state=self._snapshot(state),
+        )
+        self._queue.put_nowait(event)
+
+    def error(self, error: Exception, state: "ResearchState | None" = None) -> None:
+        if self._closed:
+            return
+        event = ResearchStreamEvent(
+            type="error",
+            message=str(error),
+            error=str(error),
+            state=self._snapshot(state),
+        )
+        self._queue.put_nowait(event)
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await self._queue.put(None)
+
+    def __aiter__(self) -> AsyncIterator[ResearchStreamEvent]:
+        return self._iter_events()
+
+    async def _iter_events(self) -> AsyncIterator[ResearchStreamEvent]:
+        while True:
+            event = await self._queue.get()
+            if event is None:
+                break
+            yield event
+
+
+async def stream_research_graph(
+    graph,
+    start,
+    state: "ResearchState",
+    deps,
+) -> AsyncIterator[ResearchStreamEvent]:
+    """Run the research graph and yield streaming events as they occur."""
+
+    from contextlib import suppress
+
+    from haiku.rag.research.state import ResearchDeps  # Local import to avoid cycle
+
+    if not isinstance(deps, ResearchDeps):
+        raise TypeError("deps must be an instance of ResearchDeps")
+
+    stream = ResearchStream()
+    deps.stream = stream
+
+    async def _execute() -> None:
+        try:
+            report = None
+            try:
+                result = await graph.run(start, state=state, deps=deps)
+                report = result.output
+            except Exception:
+                from pydantic_graph import End
+
+                async with graph.iter(start, state=state, deps=deps) as run:
+                    node = run.next_node
+                    while not isinstance(node, End):
+                        node = await run.next(node)
+                    if run.result:
+                        report = run.result.output
+
+            if report is None:
+                raise RuntimeError("Graph did not produce a report")
+
+            stream.report(report, state)
+        except Exception as exc:  # pragma: no cover - defensive path
+            stream.error(exc, state)
+        finally:
+            await stream.close()
+
+    runner = asyncio.create_task(_execute())
+
+    try:
+        async for event in stream:
+            yield event
+    finally:
+        if not runner.done():
+            runner.cancel()
+        with suppress(asyncio.CancelledError):
+            await runner