danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #76 from ggozad/feat/streaming-research

Browse source 
Streaming support for research
This commit is contained in:
Yiorgis Gozadinos 2025-09-24 16:25:09 +03:00 committed by GitHub
commit 5caca8229f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 373 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
README.md
View file

@ -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

57
docs/agents.md
View file

@ -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)
```

2
docs/cli.md
View file

@ -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:

29
src/haiku/rag/app.py
View file

@ -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
@ -221,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,
@ -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]")

8
src/haiku/rag/research/__init__.py
View file

@ -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",
]

10
src/haiku/rag/research/common.py
View file

@ -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:

4
src/haiku/rag/research/dependencies.py
View file

@ -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"
)

3
src/haiku/rag/research/models.py
View file

@ -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,

29
src/haiku/rag/research/nodes/evaluate.py
View file

@ -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
@ -51,22 +55,29 @@ 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)
for gap in output.gaps:
state.context.add_gap(gap)
state.last_eval = output
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}")
if output.gaps:
log(deps, state, " [bold yellow]Remaining gaps:[/bold yellow]")
for gap in output.gaps:
log(deps, state, f" • {gap}")
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 +85,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)

25
src/haiku/rag/research/nodes/plan.py
View file

@ -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)

21
src/haiku/rag/research/nodes/search.py
View file

@ -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,12 +32,13 @@ 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(
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)

10
src/haiku/rag/research/nodes/synthesize.py
View file

@ -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)

12
src/haiku/rag/research/state.py
View file

@ -1,23 +1,29 @@
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.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

171
src/haiku/rag/research/stream.py Normal file
View file

@ -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=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=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

3
tests/test_research_graph.py
View file

@ -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():

24
tests/test_research_graph_integration.py
View file

@ -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
@ -20,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,
@ -31,31 +31,36 @@ 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",
]
ctx.deps.emit_log("planning", ctx.state)
return SearchDispatchNode(self.provider, self.model)
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]
)
ctx.deps.emit_log(f"answered:{q}", ctx.state)
return EvaluateNode(self.provider, self.model)
async def fake_evaluate_run(self, ctx) -> Any:
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)
async def fake_synthesize_run(self, ctx) -> Any:
@ -81,9 +86,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)