Merge pull request #76 from ggozad/feat/streaming-research
Streaming support for research
This commit is contained in:
commit
5caca8229f
16 changed files with 373 additions and 78 deletions
39
README.md
39
README.md
|
|
@ -64,11 +64,12 @@ haiku-rag serve
|
||||||
```python
|
```python
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.research import (
|
from haiku.rag.research import (
|
||||||
|
PlanNode,
|
||||||
ResearchContext,
|
ResearchContext,
|
||||||
ResearchDeps,
|
ResearchDeps,
|
||||||
ResearchState,
|
ResearchState,
|
||||||
build_research_graph,
|
build_research_graph,
|
||||||
PlanNode,
|
stream_research_graph,
|
||||||
)
|
)
|
||||||
|
|
||||||
async with HaikuRAG("database.lancedb") as client:
|
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)
|
# Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
|
||||||
graph = build_research_graph()
|
graph = build_research_graph()
|
||||||
state = ResearchState(
|
|
||||||
question = (
|
question = (
|
||||||
"What are the main drivers and trends of global temperature "
|
"What are the main drivers and trends of global temperature "
|
||||||
"anomalies since 1990?"
|
"anomalies since 1990?"
|
||||||
),
|
)
|
||||||
context=ResearchContext(original_question="…"),
|
state = ResearchState(
|
||||||
|
context=ResearchContext(original_question=question),
|
||||||
max_iterations=2,
|
max_iterations=2,
|
||||||
confidence_threshold=0.8,
|
confidence_threshold=0.8,
|
||||||
max_concurrency=3,
|
max_concurrency=2,
|
||||||
)
|
)
|
||||||
deps = ResearchDeps(client=client)
|
deps = ResearchDeps(client=client)
|
||||||
start = PlanNode(provider=None, model=None)
|
|
||||||
result = await graph.run(start, state=state, deps=deps)
|
# Blocking run (final result only)
|
||||||
report = result.output
|
result = await graph.run(
|
||||||
print(report.title)
|
PlanNode(provider="openai", model="gpt-4o-mini"),
|
||||||
print(report.executive_summary)
|
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
|
## MCP Server
|
||||||
|
|
|
||||||
|
|
@ -76,30 +76,75 @@ haiku-rag research "How does haiku.rag organize and query documents?" \
|
||||||
--verbose
|
--verbose
|
||||||
```
|
```
|
||||||
|
|
||||||
Python usage:
|
Python usage (blocking result):
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.research import (
|
from haiku.rag.research import (
|
||||||
|
PlanNode,
|
||||||
ResearchContext,
|
ResearchContext,
|
||||||
ResearchDeps,
|
ResearchDeps,
|
||||||
ResearchState,
|
ResearchState,
|
||||||
build_research_graph,
|
build_research_graph,
|
||||||
PlanNode,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
async with HaikuRAG(path_to_db) as client:
|
async with HaikuRAG(path_to_db) as client:
|
||||||
graph = build_research_graph()
|
graph = build_research_graph()
|
||||||
|
question = "What are the main drivers and trends of global temperature anomalies since 1990?"
|
||||||
state = ResearchState(
|
state = ResearchState(
|
||||||
question="What are the main drivers and trends of global temperature anomalies since 1990?",
|
context=ResearchContext(original_question=question),
|
||||||
context=ResearchContext(original_question=... ),
|
|
||||||
max_iterations=2,
|
max_iterations=2,
|
||||||
confidence_threshold=0.8,
|
confidence_threshold=0.8,
|
||||||
max_concurrency=3,
|
max_concurrency=2,
|
||||||
)
|
)
|
||||||
deps = ResearchDeps(client=client)
|
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
|
report = result.output
|
||||||
print(report.title)
|
print(report.title)
|
||||||
print(report.executive_summary)
|
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)
|
||||||
|
```
|
||||||
|
|
|
||||||
|
|
@ -113,6 +113,8 @@ Flags:
|
||||||
- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
|
- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
|
||||||
- `--verbose`: show planning, searching previews, evaluation summary, and stop reason
|
- `--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
|
## Server
|
||||||
|
|
||||||
Start the MCP server:
|
Start the MCP server:
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,7 @@ from haiku.rag.research.graph import (
|
||||||
ResearchState,
|
ResearchState,
|
||||||
build_research_graph,
|
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.chunk import Chunk
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
|
|
@ -221,9 +222,9 @@ class HaikuRAGApp:
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
||||||
graph = build_research_graph()
|
graph = build_research_graph()
|
||||||
|
context = ResearchContext(original_question=question)
|
||||||
state = ResearchState(
|
state = ResearchState(
|
||||||
question=question,
|
context=context,
|
||||||
context=ResearchContext(original_question=question),
|
|
||||||
max_iterations=max_iterations,
|
max_iterations=max_iterations,
|
||||||
confidence_threshold=confidence_threshold,
|
confidence_threshold=confidence_threshold,
|
||||||
max_concurrency=max_concurrency,
|
max_concurrency=max_concurrency,
|
||||||
|
|
@ -236,22 +237,20 @@ class HaikuRAGApp:
|
||||||
provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
|
provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
|
||||||
model=Config.RESEARCH_MODEL or Config.QA_MODEL,
|
model=Config.RESEARCH_MODEL or Config.QA_MODEL,
|
||||||
)
|
)
|
||||||
# Prefer graph.run; fall back to iter if unavailable
|
|
||||||
report = None
|
report = None
|
||||||
try:
|
async for event in stream_research_graph(graph, start, state, deps):
|
||||||
result = await graph.run(start, state=state, deps=deps)
|
if event.type == "report":
|
||||||
report = result.output
|
report = event.report
|
||||||
except Exception:
|
break
|
||||||
from pydantic_graph import End
|
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:
|
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
|
# Display the report
|
||||||
self.console.print("[bold green]Research Report[/bold green]")
|
self.console.print("[bold green]Research Report[/bold green]")
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,11 @@ from haiku.rag.research.graph import (
|
||||||
build_research_graph,
|
build_research_graph,
|
||||||
)
|
)
|
||||||
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
||||||
|
from haiku.rag.research.stream import (
|
||||||
|
ResearchStateSnapshot,
|
||||||
|
ResearchStreamEvent,
|
||||||
|
stream_research_graph,
|
||||||
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"ResearchDependencies",
|
"ResearchDependencies",
|
||||||
|
|
@ -17,4 +22,7 @@ __all__ = [
|
||||||
"ResearchState",
|
"ResearchState",
|
||||||
"PlanNode",
|
"PlanNode",
|
||||||
"build_research_graph",
|
"build_research_graph",
|
||||||
|
"stream_research_graph",
|
||||||
|
"ResearchStreamEvent",
|
||||||
|
"ResearchStateSnapshot",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from pydantic_ai import format_as_xml
|
from pydantic_ai import format_as_xml
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
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.config import Config
|
||||||
from haiku.rag.research.dependencies import ResearchContext
|
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:
|
def get_model(provider: str, model: str) -> Any:
|
||||||
if provider == "ollama":
|
if provider == "ollama":
|
||||||
|
|
@ -27,9 +30,8 @@ def get_model(provider: str, model: str) -> Any:
|
||||||
return f"{provider}:{model}"
|
return f"{provider}:{model}"
|
||||||
|
|
||||||
|
|
||||||
def log(console, msg: str) -> None:
|
def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
|
||||||
if console:
|
deps.emit_log(msg, state)
|
||||||
console.print(msg)
|
|
||||||
|
|
||||||
|
|
||||||
def format_context_for_prompt(context: ResearchContext) -> str:
|
def format_context_for_prompt(context: ResearchContext) -> str:
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ from rich.console import Console
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.research.models import SearchAnswer
|
from haiku.rag.research.models import SearchAnswer
|
||||||
|
from haiku.rag.research.stream import ResearchStream
|
||||||
|
|
||||||
|
|
||||||
class ResearchContext(BaseModel):
|
class ResearchContext(BaseModel):
|
||||||
|
|
@ -45,3 +46,6 @@ class ResearchDependencies(BaseModel):
|
||||||
client: HaikuRAG = Field(description="RAG client for document operations")
|
client: HaikuRAG = Field(description="RAG client for document operations")
|
||||||
context: ResearchContext = Field(description="Shared research context")
|
context: ResearchContext = Field(description="Shared research context")
|
||||||
console: Console | None = None
|
console: Console | None = None
|
||||||
|
stream: ResearchStream | None = Field(
|
||||||
|
default=None, description="Optional research event stream"
|
||||||
|
)
|
||||||
|
|
|
||||||
|
|
@ -37,6 +37,9 @@ class EvaluationResult(BaseModel):
|
||||||
max_length=3,
|
max_length=3,
|
||||||
default=[],
|
default=[],
|
||||||
)
|
)
|
||||||
|
gaps: list[str] = Field(
|
||||||
|
description="Concrete information gaps that remain", default_factory=list
|
||||||
|
)
|
||||||
confidence_score: float = Field(
|
confidence_score: float = Field(
|
||||||
description="Confidence level in the completeness of research (0-1)",
|
description="Confidence level in the completeness of research (0-1)",
|
||||||
ge=0.0,
|
ge=0.0,
|
||||||
|
|
|
||||||
|
|
@ -25,7 +25,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
deps = ctx.deps
|
deps = ctx.deps
|
||||||
|
|
||||||
log(
|
log(
|
||||||
deps.console,
|
deps,
|
||||||
|
state,
|
||||||
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
|
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -43,7 +44,10 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
f"{context_xml}"
|
f"{context_xml}"
|
||||||
)
|
)
|
||||||
agent_deps = ResearchDependencies(
|
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)
|
eval_result = await agent.run(prompt, deps=agent_deps)
|
||||||
output = eval_result.output
|
output = eval_result.output
|
||||||
|
|
@ -51,22 +55,29 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
for insight in output.key_insights:
|
for insight in output.key_insights:
|
||||||
state.context.add_insight(insight)
|
state.context.add_insight(insight)
|
||||||
for new_q in output.new_questions:
|
for new_q in output.new_questions:
|
||||||
if new_q not in state.sub_questions:
|
if new_q not in state.context.sub_questions:
|
||||||
state.sub_questions.append(new_q)
|
state.context.sub_questions.append(new_q)
|
||||||
|
for gap in output.gaps:
|
||||||
|
state.context.add_gap(gap)
|
||||||
|
|
||||||
state.last_eval = output
|
state.last_eval = output
|
||||||
state.iterations += 1
|
state.iterations += 1
|
||||||
|
|
||||||
if output.key_insights:
|
if output.key_insights:
|
||||||
log(deps.console, " [bold]Key insights:[/bold]")
|
log(deps, state, " [bold]Key insights:[/bold]")
|
||||||
for ins in output.key_insights:
|
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(
|
log(
|
||||||
deps.console,
|
deps,
|
||||||
|
state,
|
||||||
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
|
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
|
||||||
)
|
)
|
||||||
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
|
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
|
from haiku.rag.research.nodes.search import SearchDispatchNode
|
||||||
|
|
||||||
|
|
@ -74,7 +85,7 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
output.is_sufficient
|
output.is_sufficient
|
||||||
and output.confidence_score >= state.confidence_threshold
|
and output.confidence_score >= state.confidence_threshold
|
||||||
) or state.iterations >= state.max_iterations:
|
) 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 SynthesizeNode(self.provider, self.model)
|
||||||
|
|
||||||
return SearchDispatchNode(self.provider, self.model)
|
return SearchDispatchNode(self.provider, self.model)
|
||||||
|
|
|
||||||
|
|
@ -22,7 +22,7 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
state = ctx.state
|
state = ctx.state
|
||||||
deps = ctx.deps
|
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(
|
plan_agent = Agent(
|
||||||
model=get_model(self.provider, self.model),
|
model=get_model(self.provider, self.model),
|
||||||
|
|
@ -45,19 +45,26 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
|
|
||||||
prompt = (
|
prompt = (
|
||||||
"Plan a focused research approach for the main question.\n\n"
|
"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(
|
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)
|
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, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
|
||||||
log(deps.console, f" [bold]Main Question:[/bold] {state.question}")
|
log(
|
||||||
log(deps.console, " [bold]Sub-questions:[/bold]")
|
deps,
|
||||||
for i, sq in enumerate(state.sub_questions, 1):
|
state,
|
||||||
log(deps.console, f" {i}. {sq}")
|
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)
|
return SearchDispatchNode(self.provider, self.model)
|
||||||
|
|
|
||||||
|
|
@ -24,7 +24,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
|
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
|
||||||
state = ctx.state
|
state = ctx.state
|
||||||
deps = ctx.deps
|
deps = ctx.deps
|
||||||
if not state.sub_questions:
|
if not state.context.sub_questions:
|
||||||
from haiku.rag.research.nodes.evaluate import EvaluateNode
|
from haiku.rag.research.nodes.evaluate import EvaluateNode
|
||||||
|
|
||||||
return EvaluateNode(self.provider, self.model)
|
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 up to max_concurrency questions and answer them concurrently
|
||||||
take = max(1, state.max_concurrency)
|
take = max(1, state.max_concurrency)
|
||||||
batch: list[str] = []
|
batch: list[str] = []
|
||||||
while state.sub_questions and len(batch) < take:
|
while state.context.sub_questions and len(batch) < take:
|
||||||
batch.append(state.sub_questions.pop(0))
|
batch.append(state.context.sub_questions.pop(0))
|
||||||
|
|
||||||
async def answer_one(sub_q: str) -> SearchAnswer | None:
|
async def answer_one(sub_q: str) -> SearchAnswer | None:
|
||||||
log(
|
log(
|
||||||
deps.console,
|
deps,
|
||||||
|
state,
|
||||||
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
|
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
|
||||||
)
|
)
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
|
|
@ -71,12 +72,15 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
return format_as_xml(entries, root_tag="snippets")
|
return format_as_xml(entries, root_tag="snippets")
|
||||||
|
|
||||||
agent_deps = ResearchDependencies(
|
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:
|
try:
|
||||||
result = await agent.run(sub_q, deps=agent_deps)
|
result = await agent.run(sub_q, deps=agent_deps)
|
||||||
except Exception as e:
|
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 None
|
||||||
|
|
||||||
return result.output
|
return result.output
|
||||||
|
|
@ -86,8 +90,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
if ans is None:
|
if ans is None:
|
||||||
continue
|
continue
|
||||||
state.context.add_qa_response(ans)
|
state.context.add_qa_response(ans)
|
||||||
if deps.console:
|
|
||||||
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
|
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
|
||||||
log(deps.console, f" [green]✓[/green] {preview}")
|
log(deps, state, f" [green]✓[/green] {preview}")
|
||||||
|
|
||||||
return SearchDispatchNode(self.provider, self.model)
|
return SearchDispatchNode(self.provider, self.model)
|
||||||
|
|
|
||||||
|
|
@ -24,7 +24,8 @@ class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
|
||||||
deps = ctx.deps
|
deps = ctx.deps
|
||||||
|
|
||||||
log(
|
log(
|
||||||
deps.console,
|
deps,
|
||||||
|
state,
|
||||||
"\n[bold cyan]📝 Generating final research report...[/bold cyan]",
|
"\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."
|
"Create a detailed report that synthesizes all findings into a coherent response."
|
||||||
)
|
)
|
||||||
agent_deps = ResearchDependencies(
|
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)
|
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)
|
return End(result.output)
|
||||||
|
|
|
||||||
|
|
@ -1,23 +1,29 @@
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass
|
||||||
|
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.research.dependencies import ResearchContext
|
from haiku.rag.research.dependencies import ResearchContext
|
||||||
from haiku.rag.research.models import EvaluationResult
|
from haiku.rag.research.models import EvaluationResult
|
||||||
|
from haiku.rag.research.stream import ResearchStream
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ResearchDeps:
|
class ResearchDeps:
|
||||||
client: HaikuRAG
|
client: HaikuRAG
|
||||||
console: Console | None = None
|
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
|
@dataclass
|
||||||
class ResearchState:
|
class ResearchState:
|
||||||
question: str
|
|
||||||
context: ResearchContext
|
context: ResearchContext
|
||||||
sub_questions: list[str] = field(default_factory=list)
|
|
||||||
iterations: int = 0
|
iterations: int = 0
|
||||||
max_iterations: int = 3
|
max_iterations: int = 3
|
||||||
max_concurrency: int = 1
|
max_concurrency: int = 1
|
||||||
|
|
|
||||||
171
src/haiku/rag/research/stream.py
Normal file
171
src/haiku/rag/research/stream.py
Normal 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
|
||||||
|
|
@ -9,7 +9,6 @@ def test_build_graph_and_state():
|
||||||
assert graph is not None
|
assert graph is not None
|
||||||
|
|
||||||
state = ResearchState(
|
state = ResearchState(
|
||||||
question="What are the key features of haiku.rag?",
|
|
||||||
context=ResearchContext(
|
context=ResearchContext(
|
||||||
original_question="What are the key features of haiku.rag?"
|
original_question="What are the key features of haiku.rag?"
|
||||||
),
|
),
|
||||||
|
|
@ -17,7 +16,7 @@ def test_build_graph_and_state():
|
||||||
confidence_threshold=0.8,
|
confidence_threshold=0.8,
|
||||||
)
|
)
|
||||||
assert state.iterations == 0
|
assert state.iterations == 0
|
||||||
assert state.sub_questions == []
|
assert state.context.sub_questions == []
|
||||||
|
|
||||||
|
|
||||||
def test_async_loop_available():
|
def test_async_loop_available():
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ from haiku.rag.research.graph import (
|
||||||
build_research_graph,
|
build_research_graph,
|
||||||
)
|
)
|
||||||
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
||||||
|
from haiku.rag.research.stream import stream_research_graph
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -20,7 +21,6 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
|
||||||
graph = build_research_graph()
|
graph = build_research_graph()
|
||||||
|
|
||||||
state = ResearchState(
|
state = ResearchState(
|
||||||
question="What is haiku.rag?",
|
|
||||||
context=ResearchContext(original_question="What is haiku.rag?"),
|
context=ResearchContext(original_question="What is haiku.rag?"),
|
||||||
max_iterations=1,
|
max_iterations=1,
|
||||||
confidence_threshold=0.5,
|
confidence_threshold=0.5,
|
||||||
|
|
@ -31,31 +31,36 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
|
||||||
) # client unused in patched nodes
|
) # client unused in patched nodes
|
||||||
|
|
||||||
async def fake_plan_run(self, ctx) -> Any:
|
async def fake_plan_run(self, ctx) -> Any:
|
||||||
ctx.state.sub_questions = [
|
ctx.state.context.sub_questions = [
|
||||||
"Describe haiku.rag in one sentence",
|
"Describe haiku.rag in one sentence",
|
||||||
"List core components of haiku.rag",
|
"List core components of haiku.rag",
|
||||||
]
|
]
|
||||||
|
ctx.deps.emit_log("planning", ctx.state)
|
||||||
return SearchDispatchNode(self.provider, self.model)
|
return SearchDispatchNode(self.provider, self.model)
|
||||||
|
|
||||||
async def fake_search_dispatch_run(self, ctx) -> Any:
|
async def fake_search_dispatch_run(self, ctx) -> Any:
|
||||||
# Answer all pending questions deterministically, then move to evaluation
|
# Answer all pending questions deterministically, then move to evaluation
|
||||||
while ctx.state.sub_questions:
|
while ctx.state.context.sub_questions:
|
||||||
q = ctx.state.sub_questions.pop(0)
|
q = ctx.state.context.sub_questions.pop(0)
|
||||||
# pydantic BaseModel kwargs not fully typed for pyright
|
# pydantic BaseModel kwargs not fully typed for pyright
|
||||||
ctx.state.context.add_qa_response(
|
ctx.state.context.add_qa_response(
|
||||||
SearchAnswer(query=q, answer="A", context=["x"], sources=["s"]) # pyright: ignore[reportCallIssue]
|
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)
|
return EvaluateNode(self.provider, self.model)
|
||||||
|
|
||||||
async def fake_evaluate_run(self, ctx) -> Any:
|
async def fake_evaluate_run(self, ctx) -> Any:
|
||||||
ctx.state.last_eval = EvaluationResult(
|
ctx.state.last_eval = EvaluationResult(
|
||||||
key_insights=["ok"],
|
key_insights=["ok"],
|
||||||
new_questions=[],
|
new_questions=[],
|
||||||
|
gaps=["gap"],
|
||||||
confidence_score=1.0,
|
confidence_score=1.0,
|
||||||
is_sufficient=True,
|
is_sufficient=True,
|
||||||
reasoning="done",
|
reasoning="done",
|
||||||
)
|
)
|
||||||
ctx.state.iterations += 1
|
ctx.state.iterations += 1
|
||||||
|
ctx.state.context.add_gap("gap")
|
||||||
|
ctx.deps.emit_log("evaluated", ctx.state)
|
||||||
return SynthesizeNode(self.provider, self.model)
|
return SynthesizeNode(self.provider, self.model)
|
||||||
|
|
||||||
async def fake_synthesize_run(self, ctx) -> Any:
|
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")
|
start = PlanNode(provider="test", model="test")
|
||||||
|
|
||||||
result = await graph.run(start, state=state, deps=deps)
|
collected = []
|
||||||
report = result.output
|
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 isinstance(report, ResearchReport)
|
||||||
assert report.title == "Haiku RAG"
|
assert report.title == "Haiku RAG"
|
||||||
assert len(state.context.qa_responses) == 2
|
assert len(state.context.qa_responses) == 2
|
||||||
|
assert any(evt.type == "log" for evt in collected)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue