From 086eb56d3445c1fd87db52e487f8c5fbd07afa0b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Sep 2025 11:02:06 +0300 Subject: [PATCH 1/4] Stream research --- src/haiku/rag/app.py | 25 ++- src/haiku/rag/research/__init__.py | 8 + src/haiku/rag/research/common.py | 10 +- src/haiku/rag/research/dependencies.py | 4 + src/haiku/rag/research/nodes/evaluate.py | 19 ++- src/haiku/rag/research/nodes/plan.py | 15 +- src/haiku/rag/research/nodes/search.py | 15 +- src/haiku/rag/research/nodes/synthesize.py | 10 +- src/haiku/rag/research/state.py | 8 + src/haiku/rag/research/stream.py | 171 +++++++++++++++++++++ tests/test_research_graph_integration.py | 15 +- 11 files changed, 259 insertions(+), 41 deletions(-) create mode 100644 src/haiku/rag/research/stream.py 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) From cf9d6a116f880c13c9bd968a36eb99e79a891d9d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Sep 2025 13:43:07 +0300 Subject: [PATCH 2/4] Do not duplicate ResearchContext in ResearchState --- src/haiku/rag/app.py | 4 ++-- src/haiku/rag/research/nodes/evaluate.py | 4 ++-- src/haiku/rag/research/nodes/plan.py | 12 ++++++++---- src/haiku/rag/research/nodes/search.py | 6 +++--- src/haiku/rag/research/state.py | 4 +--- src/haiku/rag/research/stream.py | 6 +++--- tests/test_research_graph.py | 3 +-- tests/test_research_graph_integration.py | 7 +++---- 8 files changed, 23 insertions(+), 23 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 3edce247..0a0302f9 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -222,9 +222,9 @@ class HaikuRAGApp: self.console.print() graph = build_research_graph() + context = ResearchContext(original_question=question) state = ResearchState( - question=question, - context=ResearchContext(original_question=question), + context=context, max_iterations=max_iterations, confidence_threshold=confidence_threshold, max_concurrency=max_concurrency, diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py index 23181136..46634b52 100644 --- a/src/haiku/rag/research/nodes/evaluate.py +++ b/src/haiku/rag/research/nodes/evaluate.py @@ -55,8 +55,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): for insight in output.key_insights: state.context.add_insight(insight) for new_q in output.new_questions: - if new_q not in state.sub_questions: - state.sub_questions.append(new_q) + if new_q not in state.context.sub_questions: + state.context.sub_questions.append(new_q) state.last_eval = output state.iterations += 1 diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py index 0f3726cb..63612a55 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/research/nodes/plan.py @@ -45,7 +45,7 @@ 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( @@ -55,12 +55,16 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): 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, state, "\n[bold green]βœ… Research Plan Created:[/bold green]") - log(deps, state, f" [bold]Main Question:[/bold] {state.question}") + 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.sub_questions, 1): + for i, sq in enumerate(state.context.sub_questions, 1): 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 fddc1b9e..c7e471a9 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/research/nodes/search.py @@ -24,7 +24,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]: state = ctx.state deps = ctx.deps - if not state.sub_questions: + if not state.context.sub_questions: from haiku.rag.research.nodes.evaluate import EvaluateNode return EvaluateNode(self.provider, self.model) @@ -32,8 +32,8 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): # 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( diff --git a/src/haiku/rag/research/state.py b/src/haiku/rag/research/state.py index ad0920b2..238accf6 100644 --- a/src/haiku/rag/research/state.py +++ b/src/haiku/rag/research/state.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, field +from dataclasses import dataclass from rich.console import Console @@ -23,9 +23,7 @@ class ResearchDeps: @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 diff --git a/src/haiku/rag/research/stream.py b/src/haiku/rag/research/stream.py index 0df6b8fb..9ecf2d80 100644 --- a/src/haiku/rag/research/stream.py +++ b/src/haiku/rag/research/stream.py @@ -34,13 +34,13 @@ class ResearchStateSnapshot: last_sufficient = state.last_eval.is_sufficient return cls( - question=state.question, - sub_questions=list(state.sub_questions), + 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(state.sub_questions), + pending_sub_questions=len(context.sub_questions), answered_questions=len(context.qa_responses), insights=list(context.insights), gaps=list(context.gaps), diff --git a/tests/test_research_graph.py b/tests/test_research_graph.py index 8b50a6f3..b986eeba 100644 --- a/tests/test_research_graph.py +++ b/tests/test_research_graph.py @@ -9,7 +9,6 @@ def test_build_graph_and_state(): assert graph is not None state = ResearchState( - question="What are the key features of haiku.rag?", context=ResearchContext( original_question="What are the key features of haiku.rag?" ), @@ -17,7 +16,7 @@ def test_build_graph_and_state(): confidence_threshold=0.8, ) assert state.iterations == 0 - assert state.sub_questions == [] + assert state.context.sub_questions == [] def test_async_loop_available(): diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 66c411f5..64863f9f 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -21,7 +21,6 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch): graph = build_research_graph() state = ResearchState( - question="What is haiku.rag?", context=ResearchContext(original_question="What is haiku.rag?"), max_iterations=1, confidence_threshold=0.5, @@ -32,7 +31,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch): ) # client unused in patched nodes async def fake_plan_run(self, ctx) -> Any: - ctx.state.sub_questions = [ + ctx.state.context.sub_questions = [ "Describe haiku.rag in one sentence", "List core components of haiku.rag", ] @@ -41,8 +40,8 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch): async def fake_search_dispatch_run(self, ctx) -> Any: # Answer all pending questions deterministically, then move to evaluation - while ctx.state.sub_questions: - q = ctx.state.sub_questions.pop(0) + while ctx.state.context.sub_questions: + q = ctx.state.context.sub_questions.pop(0) # pydantic BaseModel kwargs not fully typed for pyright ctx.state.context.add_qa_response( SearchAnswer(query=q, answer="A", context=["x"], sources=["s"]) # pyright: ignore[reportCallIssue] From 57bac178ddbcc6bb57f05171d9a6baab94939de4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Sep 2025 13:43:49 +0300 Subject: [PATCH 3/4] Update docs --- README.md | 43 ++++++++++++++++++++++++++----------- docs/agents.md | 57 ++++++++++++++++++++++++++++++++++++++++++++------ docs/cli.md | 2 ++ 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 8368d909..23c1c63e 100644 --- a/README.md +++ b/README.md @@ -64,11 +64,12 @@ haiku-rag serve ```python from haiku.rag.client import HaikuRAG from haiku.rag.research import ( + PlanNode, ResearchContext, ResearchDeps, ResearchState, build_research_graph, - PlanNode, + stream_research_graph, ) async with HaikuRAG("database.lancedb") as client: @@ -90,22 +91,40 @@ async with HaikuRAG("database.lancedb") as client: # Multi‑agent research pipeline (Plan β†’ Search β†’ Evaluate β†’ Synthesize) graph = build_research_graph() + question = ( + "What are the main drivers and trends of global temperature " + "anomalies since 1990?" + ) state = ResearchState( - question=( - "What are the main drivers and trends of global temperature " - "anomalies since 1990?" - ), - context=ResearchContext(original_question="…"), + context=ResearchContext(original_question=question), max_iterations=2, confidence_threshold=0.8, - max_concurrency=3, + max_concurrency=2, ) deps = ResearchDeps(client=client) - start = PlanNode(provider=None, model=None) - result = await graph.run(start, state=state, deps=deps) - report = result.output - print(report.title) - print(report.executive_summary) + + # Blocking run (final result only) + result = await graph.run( + PlanNode(provider="openai", model="gpt-4o-mini"), + state=state, + deps=deps, + ) + print(result.output.title) + + # Streaming progress (log/report/error events) + async for event in stream_research_graph( + graph, + PlanNode(provider="openai", model="gpt-4o-mini"), + state, + deps, + ): + if event.type == "log": + iteration = event.state.iterations if event.state else state.iterations + print(f"[{iteration}] {event.message}") + elif event.type == "report": + print("\nResearch complete!\n") + print(event.report.title) + print(event.report.executive_summary) ``` ## MCP Server diff --git a/docs/agents.md b/docs/agents.md index 59e314cc..5e6374ec 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -76,30 +76,75 @@ haiku-rag research "How does haiku.rag organize and query documents?" \ --verbose ``` -Python usage: +Python usage (blocking result): ```python from haiku.rag.client import HaikuRAG from haiku.rag.research import ( + PlanNode, ResearchContext, ResearchDeps, ResearchState, build_research_graph, - PlanNode, ) async with HaikuRAG(path_to_db) as client: graph = build_research_graph() + question = "What are the main drivers and trends of global temperature anomalies since 1990?" state = ResearchState( - question="What are the main drivers and trends of global temperature anomalies since 1990?", - context=ResearchContext(original_question=... ), + context=ResearchContext(original_question=question), max_iterations=2, confidence_threshold=0.8, - max_concurrency=3, + max_concurrency=2, ) deps = ResearchDeps(client=client) - result = await graph.run(PlanNode(provider=None, model=None), state=state, deps=deps) + + result = await graph.run( + PlanNode(provider="openai", model="gpt-4o-mini"), + state=state, + deps=deps, + ) + report = result.output print(report.title) print(report.executive_summary) ``` + +Python usage (streamed events): + +```python +from haiku.rag.client import HaikuRAG +from haiku.rag.research import ( + PlanNode, + ResearchContext, + ResearchDeps, + ResearchState, + build_research_graph, + stream_research_graph, +) + +async with HaikuRAG(path_to_db) as client: + graph = build_research_graph() + question = "What are the main drivers and trends of global temperature anomalies since 1990?" + state = ResearchState( + context=ResearchContext(original_question=question), + max_iterations=2, + confidence_threshold=0.8, + max_concurrency=2, + ) + deps = ResearchDeps(client=client) + + async for event in stream_research_graph( + graph, + PlanNode(provider="openai", model="gpt-4o-mini"), + state, + deps, + ): + if event.type == "log": + iteration = event.state.iterations if event.state else state.iterations + print(f"[{iteration}] {event.message}") + elif event.type == "report": + print("\nResearch complete!\n") + print(event.report.title) + print(event.report.executive_summary) +``` diff --git a/docs/cli.md b/docs/cli.md index 4afc96aa..751d40e7 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -113,6 +113,8 @@ Flags: - `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3) - `--verbose`: show planning, searching previews, evaluation summary, and stop reason +When `--verbose` is set the CLI also consumes the internal research stream, printing every `log` event as agents progress through planning, search, evaluation, and synthesis. If you build your own integration, call `stream_research_graph` to access the same `log`, `report`, and `error` events and render them however you like while the graph is running. + ## Server Start the MCP server: From 84e94f825788fb56f68a81060d883ba7c7fcf351 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Sep 2025 14:00:31 +0300 Subject: [PATCH 4/4] Wire gaps that were forgotten --- src/haiku/rag/research/models.py | 3 +++ src/haiku/rag/research/nodes/evaluate.py | 6 ++++++ tests/test_research_graph_integration.py | 2 ++ 3 files changed, 11 insertions(+) diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py index 30cb97d1..ab16601d 100644 --- a/src/haiku/rag/research/models.py +++ b/src/haiku/rag/research/models.py @@ -37,6 +37,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, diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py index 46634b52..6289934e 100644 --- a/src/haiku/rag/research/nodes/evaluate.py +++ b/src/haiku/rag/research/nodes/evaluate.py @@ -57,6 +57,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): for new_q in output.new_questions: if new_q not in state.context.sub_questions: state.context.sub_questions.append(new_q) + for gap in output.gaps: + state.context.add_gap(gap) state.last_eval = output state.iterations += 1 @@ -65,6 +67,10 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): log(deps, state, " [bold]Key insights:[/bold]") for ins in output.key_insights: log(deps, state, f" β€’ {ins}") + 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, diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 64863f9f..d869af7d 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -53,11 +53,13 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch): ctx.state.last_eval = EvaluationResult( key_insights=["ok"], new_questions=[], + gaps=["gap"], confidence_score=1.0, is_sufficient=True, reasoning="done", ) ctx.state.iterations += 1 + ctx.state.context.add_gap("gap") ctx.deps.emit_log("evaluated", ctx.state) return SynthesizeNode(self.provider, self.model)