Stream research

src/haiku/rag/app.py

@@ -18,6 +18,7 @@ from haiku.rag.research.graph import (
     ResearchState,
     build_research_graph,
 )
+from haiku.rag.research.stream import stream_research_graph
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 
@@ -236,22 +237,20 @@ class HaikuRAGApp:
                     provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
                     model=Config.RESEARCH_MODEL or Config.QA_MODEL,
                 )
-                # Prefer graph.run; fall back to iter if unavailable
                 report = None
-                try:
-                    result = await graph.run(start, state=state, deps=deps)
-                    report = result.output
-                except Exception:
-                    from pydantic_graph import End
+                async for event in stream_research_graph(graph, start, state, deps):
+                    if event.type == "report":
+                        report = event.report
+                        break
+                    if event.type == "error":
+                        self.console.print(
+                            f"[red]Error during research: {event.message}[/red]"
+                        )
+                        return
 
-                    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")
+                    self.console.print("[red]Research did not produce a report.[/red]")
+                    return
 
                 # Display the report
                 self.console.print("[bold green]Research Report[/bold green]")

src/haiku/rag/research/__init__.py

@@ -6,6 +6,11 @@ from haiku.rag.research.graph import (
     build_research_graph,
 )
 from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
+from haiku.rag.research.stream import (
+    ResearchStateSnapshot,
+    ResearchStreamEvent,
+    stream_research_graph,
+)
 
 __all__ = [
     "ResearchDependencies",
@@ -17,4 +22,7 @@ __all__ = [
     "ResearchState",
     "PlanNode",
     "build_research_graph",
+    "stream_research_graph",
+    "ResearchStreamEvent",
+    "ResearchStateSnapshot",
 ]

src/haiku/rag/research/common.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic_ai import format_as_xml
 from pydantic_ai.models.openai import OpenAIChatModel
@@ -8,6 +8,9 @@ from pydantic_ai.providers.openai import OpenAIProvider
 from haiku.rag.config import Config
 from haiku.rag.research.dependencies import ResearchContext
 
+if TYPE_CHECKING:  # pragma: no cover
+    from haiku.rag.research.state import ResearchDeps, ResearchState
+
 
 def get_model(provider: str, model: str) -> Any:
     if provider == "ollama":
@@ -27,9 +30,8 @@ def get_model(provider: str, model: str) -> Any:
         return f"{provider}:{model}"
 
 
-def log(console, msg: str) -> None:
-    if console:
-        console.print(msg)
+def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
+    deps.emit_log(msg, state)
 
 
 def format_context_for_prompt(context: ResearchContext) -> str:

src/haiku/rag/research/dependencies.py

@@ -3,6 +3,7 @@ from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research.models import SearchAnswer
+from haiku.rag.research.stream import ResearchStream
 
 
 class ResearchContext(BaseModel):
@@ -45,3 +46,6 @@ 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"
+    )

src/haiku/rag/research/nodes/evaluate.py

@@ -25,7 +25,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         deps = ctx.deps
 
         log(
-            deps.console,
+            deps,
+            state,
             "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
         )
 
@@ -43,7 +44,10 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             f"{context_xml}"
         )
         agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
+            client=deps.client,
+            context=state.context,
+            console=deps.console,
+            stream=deps.stream,
         )
         eval_result = await agent.run(prompt, deps=agent_deps)
         output = eval_result.output
@@ -58,15 +62,16 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         state.iterations += 1
 
         if output.key_insights:
-            log(deps.console, "   [bold]Key insights:[/bold]")
+            log(deps, state, "   [bold]Key insights:[/bold]")
             for ins in output.key_insights:
-                log(deps.console, f"   • {ins}")
+                log(deps, state, f"   • {ins}")
         log(
-            deps.console,
+            deps,
+            state,
             f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
         )
         status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
-        log(deps.console, f"   Sufficient: {status}")
+        log(deps, state, f"   Sufficient: {status}")
 
         from haiku.rag.research.nodes.search import SearchDispatchNode
 
@@ -74,7 +79,7 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             output.is_sufficient
             and output.confidence_score >= state.confidence_threshold
         ) or state.iterations >= state.max_iterations:
-            log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
+            log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
             return SynthesizeNode(self.provider, self.model)
 
         return SearchDispatchNode(self.provider, self.model)

src/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),
@@ -49,15 +49,18 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         )
 
         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)
 
-        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]")
+        log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
+        log(deps, state, f"   [bold]Main Question:[/bold] {state.question}")
+        log(deps, state, "   [bold]Sub-questions:[/bold]")
         for i, sq in enumerate(state.sub_questions, 1):
-            log(deps.console, f"      {i}. {sq}")
+            log(deps, state, f"      {i}. {sq}")
 
         return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/research/nodes/search.py

@@ -37,7 +37,8 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
 
         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)

src/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)

src/haiku/rag/research/state.py

@@ -5,12 +5,20 @@ 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.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

src/haiku/rag/research/stream.py

@@ -0,0 +1,171 @@
+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=state.question,
+            sub_questions=list(state.sub_questions),
+            iterations=state.iterations,
+            max_iterations=state.max_iterations,
+            max_concurrency=state.max_concurrency,
+            confidence_threshold=state.confidence_threshold,
+            pending_sub_questions=len(state.sub_questions),
+            answered_questions=len(context.qa_responses),
+            insights=list(context.insights),
+            gaps=list(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

tests/test_research_graph_integration.py

@@ -13,6 +13,7 @@ from haiku.rag.research.graph import (
     build_research_graph,
 )
 from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
+from haiku.rag.research.stream import stream_research_graph
 
 
 @pytest.mark.asyncio
@@ -35,6 +36,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
             "Describe haiku.rag in one sentence",
             "List core components of haiku.rag",
         ]
+        ctx.deps.emit_log("planning", ctx.state)
         return SearchDispatchNode(self.provider, self.model)
 
     async def fake_search_dispatch_run(self, ctx) -> Any:
@@ -45,6 +47,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
             ctx.state.context.add_qa_response(
                 SearchAnswer(query=q, answer="A", context=["x"], sources=["s"])  # pyright: ignore[reportCallIssue]
             )
+            ctx.deps.emit_log(f"answered:{q}", ctx.state)
         return EvaluateNode(self.provider, self.model)
 
     async def fake_evaluate_run(self, ctx) -> Any:
@@ -56,6 +59,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
             reasoning="done",
         )
         ctx.state.iterations += 1
+        ctx.deps.emit_log("evaluated", ctx.state)
         return SynthesizeNode(self.provider, self.model)
 
     async def fake_synthesize_run(self, ctx) -> Any:
@@ -81,9 +85,16 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
 
     start = PlanNode(provider="test", model="test")
 
-    result = await graph.run(start, state=state, deps=deps)
-    report = result.output
+    collected = []
+    async for event in stream_research_graph(graph, start, state, deps):
+        collected.append(event)
+        if event.type == "report":
+            report = event.report
+            break
+    else:  # pragma: no cover - defensive guard
+        report = None
 
     assert isinstance(report, ResearchReport)
     assert report.title == "Haiku RAG"
     assert len(state.context.qa_responses) == 2
+    assert any(evt.type == "log" for evt in collected)