diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index a2a690c5..3edce247 100644 --- a/src/haiku/rag/app.py +++ b/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]") diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index e48953e4..b034748a 100644 --- a/src/haiku/rag/research/__init__.py +++ b/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", ] diff --git a/src/haiku/rag/research/common.py b/src/haiku/rag/research/common.py index 4c821d3c..ca877a31 100644 --- a/src/haiku/rag/research/common.py +++ b/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: diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index c075f852..729264e6 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/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" + ) diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py index 7270d0a0..23181136 100644 --- a/src/haiku/rag/research/nodes/evaluate.py +++ b/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) diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py index 653c12c5..0f3726cb 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/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) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py index ae863d37..fddc1b9e 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/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) diff --git a/src/haiku/rag/research/nodes/synthesize.py b/src/haiku/rag/research/nodes/synthesize.py index e2ec6be4..4f8eee13 100644 --- a/src/haiku/rag/research/nodes/synthesize.py +++ b/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) diff --git a/src/haiku/rag/research/state.py b/src/haiku/rag/research/state.py index c153b4ca..ad0920b2 100644 --- a/src/haiku/rag/research/state.py +++ b/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 diff --git a/src/haiku/rag/research/stream.py b/src/haiku/rag/research/stream.py new file mode 100644 index 00000000..0df6b8fb --- /dev/null +++ b/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 diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index ea719155..66c411f5 100644 --- a/tests/test_research_graph_integration.py +++ b/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)