diff --git a/haiku_rag_slim/haiku/rag/agui/__init__.py b/haiku_rag_slim/haiku/rag/agui/__init__.py
index fe38fc94..470d1b74 100644
--- a/haiku_rag_slim/haiku/rag/agui/__init__.py
+++ b/haiku_rag_slim/haiku/rag/agui/__init__.py
@@ -1,5 +1,6 @@
"""Generic AG-UI protocol support for haiku.rag graphs."""
+from haiku.rag.agui.cli_renderer import AGUIConsoleRenderer
from haiku.rag.agui.emitter import AGUIEmitter
from haiku.rag.agui.events import (
AGUIEvent,
@@ -18,8 +19,10 @@ from haiku.rag.agui.events import (
emit_text_message_start,
)
from haiku.rag.agui.state import compute_state_delta
+from haiku.rag.agui.stream import stream_graph
__all__ = [
+ "AGUIConsoleRenderer",
"AGUIEmitter",
"AGUIEvent",
"compute_state_delta",
@@ -36,4 +39,5 @@ __all__ = [
"emit_text_message_content",
"emit_text_message_end",
"emit_text_message_start",
+ "stream_graph",
]
diff --git a/haiku_rag_slim/haiku/rag/agui/cli_renderer.py b/haiku_rag_slim/haiku/rag/agui/cli_renderer.py
new file mode 100644
index 00000000..0d146b61
--- /dev/null
+++ b/haiku_rag_slim/haiku/rag/agui/cli_renderer.py
@@ -0,0 +1,152 @@
+"""Generic CLI renderer for AG-UI events with Rich console output."""
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+from pydantic import BaseModel
+from rich.console import Console
+
+from haiku.rag.agui.events import AGUIEvent
+
+
+class AGUIConsoleRenderer:
+ """Renders AG-UI events to Rich console with formatted output.
+
+ Generic renderer that processes AG-UI protocol events and renders them
+ with Rich formatting. Works with any graph that emits AG-UI events.
+ """
+
+ def __init__(self, console: Console | None = None):
+ """Initialize the renderer.
+
+ Args:
+ console: Optional Rich console instance (creates new one if not provided)
+ """
+ self.console = console or Console()
+ self._state: BaseModel | None = None
+
+ async def render(self, events: AsyncIterator[AGUIEvent]) -> Any | None:
+ """Process events and render to console, return final result.
+
+ Args:
+ events: Async iterator of AG-UI events
+
+ Returns:
+ The final result from RunFinished event, or None
+ """
+ result = None
+
+ async for event in events:
+ event_type = event.get("type")
+
+ if event_type == "RUN_STARTED":
+ self._render_run_started(event)
+ elif event_type == "RUN_FINISHED":
+ result = event.get("result")
+ self._render_run_finished()
+ elif event_type == "RUN_ERROR":
+ self._render_error(event)
+ elif event_type == "STEP_STARTED":
+ self._render_step_started(event)
+ elif event_type == "STEP_FINISHED":
+ self._render_step_finished(event)
+ elif event_type == "TEXT_MESSAGE_CHUNK":
+ self._render_text_message(event)
+ elif event_type == "TEXT_MESSAGE_START":
+ pass # Start of streaming message, no output needed
+ elif event_type == "TEXT_MESSAGE_CONTENT":
+ self._render_text_content(event)
+ elif event_type == "TEXT_MESSAGE_END":
+ pass # End of streaming message, no output needed
+ elif event_type == "STATE_SNAPSHOT":
+ self._state = event.get("snapshot")
+ elif event_type == "STATE_DELTA":
+ self._apply_state_delta(event)
+ elif event_type == "ACTIVITY_SNAPSHOT":
+ self._render_activity(event)
+ elif event_type == "ACTIVITY_DELTA":
+ pass # Activity deltas don't need separate rendering
+
+ return result
+
+ def _render_run_started(self, event: AGUIEvent) -> None:
+ """Render run start event.
+
+ Args:
+ event: RunStarted event
+ """
+ # Currently silent - could render run metadata later
+
+ def _render_run_finished(self) -> None:
+ """Render run completion."""
+ # Currently silent - the result is rendered separately
+
+ def _render_error(self, event: AGUIEvent) -> None:
+ """Render error event.
+
+ Args:
+ event: RunError event
+ """
+ message = event.get("message", "Unknown error")
+ self.console.print(f"[bold red]❌ Error:[/bold red] {message}")
+
+ def _render_step_started(self, event: AGUIEvent) -> None:
+ """Render step start event.
+
+ Args:
+ event: StepStarted event
+ """
+ step_name = event.get("stepName", "")
+ if step_name:
+ # Format step name for display
+ display_name = step_name.replace("_", " ").title()
+ self.console.print(f"\n[bold cyan]{display_name}[/bold cyan]")
+
+ def _render_step_finished(self, event: AGUIEvent) -> None:
+ """Render step finish event.
+
+ Args:
+ event: StepFinished event
+ """
+ # Step completion is implicit from the next step or activity
+
+ def _render_text_message(self, event: AGUIEvent) -> None:
+ """Render complete text message.
+
+ Args:
+ event: TextMessageChunk event
+ """
+ delta = event.get("delta", "")
+ # The delta contains the text content to display
+ self.console.print(delta)
+
+ def _render_text_content(self, event: AGUIEvent) -> None:
+ """Render streaming text content delta.
+
+ Args:
+ event: TextMessageContent event
+ """
+ delta = event.get("delta", "")
+ # Print delta without newline for streaming effect
+ self.console.print(delta, end="")
+
+ def _render_activity(self, event: AGUIEvent) -> None:
+ """Render activity update.
+
+ Args:
+ event: ActivitySnapshot event
+ """
+ content = event.get("content", "")
+
+ # Render activity content with emphasis
+ if content:
+ self.console.print(f"[dim]{content}[/dim]")
+
+ def _apply_state_delta(self, event: AGUIEvent) -> None:
+ """Apply state delta to current state.
+
+ Args:
+ event: StateDelta event
+ """
+ # Currently not applying deltas - could implement state patching later
+ # For now, the text messages contain all the information we need to display
diff --git a/haiku_rag_slim/haiku/rag/agui/emitter.py b/haiku_rag_slim/haiku/rag/agui/emitter.py
index 4f9e0794..b634f771 100644
--- a/haiku_rag_slim/haiku/rag/agui/emitter.py
+++ b/haiku_rag_slim/haiku/rag/agui/emitter.py
@@ -13,7 +13,6 @@ from haiku.rag.agui.events import (
emit_run_error,
emit_run_finished,
emit_run_started,
- emit_state_delta,
emit_state_snapshot,
emit_step_finished,
emit_step_started,
@@ -60,18 +59,19 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
"""Get the run ID for this emitter."""
return self._run_id
- def start_run(self, input_data: str, initial_state: StateT) -> None:
+ def start_run(self, initial_state: StateT) -> None:
"""Emit RunStarted and initial StateSnapshot.
Args:
- input_data: The input that started the run
initial_state: The initial state of the graph
"""
- # If thread_id wasn't provided, generate from input hash
+ # If thread_id wasn't provided, generate from state hash
if not self._thread_id or self._thread_id == str(uuid4()):
- self._thread_id = self._generate_thread_id(input_data)
+ state_json = initial_state.model_dump_json()
+ self._thread_id = self._generate_thread_id(state_json)
- self._emit(emit_run_started(self._thread_id, self._run_id, input_data))
+ # RunStarted (state snapshot follows immediately with full state)
+ self._emit(emit_run_started(self._thread_id, self._run_id))
self._emit(emit_state_snapshot(initial_state))
self._last_state = initial_state
@@ -100,21 +100,13 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
self._emit(emit_text_message(message, role))
def update_state(self, new_state: StateT) -> None:
- """Emit StateDelta for state change, or StateSnapshot if no previous state.
+ """Emit StateSnapshot for state change.
Args:
new_state: The updated state
"""
- if self._last_state:
- # Emit delta if we have a previous state
- delta_event = emit_state_delta(self._last_state, new_state)
- # Only emit if there are actual changes
- if delta_event.get("operations"):
- self._emit(delta_event)
- else:
- # Emit snapshot if this is the first state update
- self._emit(emit_state_snapshot(new_state))
-
+ # Always emit full snapshot (not delta) for complete state visibility
+ self._emit(emit_state_snapshot(new_state))
self._last_state = new_state
def update_activity(
@@ -125,9 +117,11 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
Args:
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
- message_id: Optional message ID to associate activity with
+ message_id: Optional message ID to associate activity with (auto-generated if None)
"""
- self._emit(emit_activity(activity_type, content, message_id))
+ if message_id is None:
+ message_id = str(uuid4())
+ self._emit(emit_activity(message_id, activity_type, content))
def finish_run(self, result: ResultT) -> None:
"""Emit RunFinished event.
diff --git a/haiku_rag_slim/haiku/rag/agui/events.py b/haiku_rag_slim/haiku/rag/agui/events.py
index fca360b1..d439d005 100644
--- a/haiku_rag_slim/haiku/rag/agui/events.py
+++ b/haiku_rag_slim/haiku/rag/agui/events.py
@@ -11,23 +11,27 @@ from haiku.rag.agui.state import compute_state_delta
AGUIEvent = dict[str, Any]
-def emit_run_started(thread_id: str, run_id: str, input_data: str) -> dict[str, Any]:
+def emit_run_started(
+ thread_id: str, run_id: str, input_data: str | None = None
+) -> dict[str, Any]:
"""Create a RunStarted event.
Args:
thread_id: Unique identifier for the conversation thread
run_id: Unique identifier for this run
- input_data: The input that started the run
+ input_data: Optional input that started the run
Returns:
RunStarted event dict
"""
- return {
+ event: dict[str, Any] = {
"type": "RUN_STARTED",
"threadId": thread_id,
"runId": run_id,
- "input": input_data,
}
+ if input_data:
+ event["input"] = input_data
+ return event
def emit_run_finished(thread_id: str, run_id: str, result: Any) -> dict[str, Any]:
@@ -119,7 +123,7 @@ def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]:
"type": "TEXT_MESSAGE_CHUNK",
"messageId": message_id,
"role": role,
- "content": content,
+ "delta": content,
}
@@ -197,46 +201,47 @@ def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, An
Returns:
StateDelta event dict
"""
- operations = compute_state_delta(old_state, new_state)
+ delta = compute_state_delta(old_state, new_state)
return {
"type": "STATE_DELTA",
- "operations": operations,
+ "delta": delta,
}
def emit_activity(
+ message_id: str,
activity_type: str,
content: str,
- message_id: str | None = None,
) -> dict[str, Any]:
"""Create an ActivitySnapshot event.
Args:
+ message_id: Message ID to associate activity with (required)
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
- message_id: Optional message ID to associate activity with
Returns:
ActivitySnapshot event dict
"""
- event: dict[str, Any] = {
+ return {
"type": "ACTIVITY_SNAPSHOT",
+ "messageId": message_id,
"activityType": activity_type,
"content": content,
}
- if message_id:
- event["messageId"] = message_id
- return event
def emit_activity_delta(
- message_id: str, operations: list[dict[str, Any]]
+ message_id: str,
+ activity_type: str,
+ patch: list[dict[str, Any]],
) -> dict[str, Any]:
"""Create an ActivityDelta event with JSON Patch operations.
Args:
message_id: Message ID of the activity being updated
- operations: JSON Patch operations to apply
+ activity_type: Type of activity being updated
+ patch: JSON Patch operations to apply
Returns:
ActivityDelta event dict
@@ -244,5 +249,6 @@ def emit_activity_delta(
return {
"type": "ACTIVITY_DELTA",
"messageId": message_id,
- "operations": operations,
+ "activityType": activity_type,
+ "patch": patch,
}
diff --git a/haiku_rag_slim/haiku/rag/agui/stream.py b/haiku_rag_slim/haiku/rag/agui/stream.py
new file mode 100644
index 00000000..e91b9b64
--- /dev/null
+++ b/haiku_rag_slim/haiku/rag/agui/stream.py
@@ -0,0 +1,80 @@
+"""Generic graph streaming with AG-UI events."""
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import suppress
+from typing import Any, Protocol
+
+from pydantic import BaseModel
+
+from haiku.rag.agui.emitter import AGUIEmitter
+from haiku.rag.agui.events import AGUIEvent
+
+
+class GraphDeps(Protocol):
+ """Protocol for graph dependencies that support AG-UI emission."""
+
+ agui_emitter: AGUIEmitter[Any, Any] | None
+
+
+async def stream_graph(
+ graph: Any,
+ state: BaseModel,
+ deps: GraphDeps,
+) -> AsyncIterator[AGUIEvent]:
+ """Run a graph and yield AG-UI events as they occur.
+
+ This is a generic streaming function that works with any pydantic-graph
+ that follows the AG-UI pattern:
+ - State must be a Pydantic BaseModel
+ - Deps must have an optional agui_emitter attribute
+ - Graph must be a pydantic-graph Graph instance
+
+ Args:
+ graph: The pydantic-graph Graph to execute
+ state: Initial state (Pydantic BaseModel)
+ deps: Graph dependencies with agui_emitter support
+
+ Yields:
+ AG-UI event dictionaries
+
+ Raises:
+ TypeError: If deps doesn't support agui_emitter
+ RuntimeError: If graph doesn't produce a result
+ """
+ if not hasattr(deps, "agui_emitter"):
+ raise TypeError("deps must have an 'agui_emitter' attribute")
+
+ # Create AG-UI emitter
+ emitter: AGUIEmitter[Any, Any] = AGUIEmitter()
+ deps.agui_emitter = emitter
+
+ async def _execute() -> None:
+ try:
+ # Start the run with initial state
+ emitter.start_run(initial_state=state)
+
+ # Execute the graph
+ result = await graph.run(state=state, deps=deps)
+
+ if result is None:
+ raise RuntimeError("Graph did not produce a result")
+
+ # Finish the run with the result
+ emitter.finish_run(result)
+ except Exception as exc:
+ # Emit error event
+ emitter.error(exc)
+ finally:
+ await emitter.close()
+
+ runner = asyncio.create_task(_execute())
+
+ try:
+ async for event in emitter:
+ yield event
+ finally:
+ if not runner.done():
+ runner.cancel()
+ with suppress(asyncio.CancelledError):
+ await runner
diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py
index 704d9610..6b1fe382 100644
--- a/haiku_rag_slim/haiku/rag/app.py
+++ b/haiku_rag_slim/haiku/rag/app.py
@@ -8,6 +8,7 @@ from rich.console import Console
from rich.markdown import Markdown
from rich.progress import Progress
+from haiku.rag.agui import AGUIConsoleRenderer, stream_graph
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, Config
from haiku.rag.mcp import create_mcp_server
@@ -15,7 +16,6 @@ from haiku.rag.monitor import FileWatcher
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
-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
@@ -243,46 +243,36 @@ class HaikuRAGApp:
except Exception as e:
self.console.print(f"[red]Error: {e}[/red]")
- async def research(
- self,
- question: str,
- verbose: bool = False,
- ):
+ async def research(self, question: str):
"""Run research via the pydantic-graph pipeline.
Args:
question: The research question
- verbose: Show verbose output
"""
async with HaikuRAG(db_path=self.db_path, config=self.config) as client:
try:
- if verbose:
- self.console.print("[bold cyan]Starting research[/bold cyan]")
- self.console.print(f"[bold blue]Question:[/bold blue] {question}")
- self.console.print()
+ self.console.print("[bold cyan]Starting research[/bold cyan]")
+ self.console.print(f"[bold blue]Question:[/bold blue] {question}")
+ self.console.print()
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=self.config)
- deps = ResearchDeps(
- client=client, console=self.console if verbose else None
- )
+ deps = ResearchDeps(client=client)
- report = None
- async for event in stream_research_graph(graph, 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
+ # Use AG-UI renderer to process events
+ renderer = AGUIConsoleRenderer(self.console)
+ report_dict = await renderer.render(stream_graph(graph, state, deps))
- if report is None:
+ if report_dict is None:
self.console.print("[red]Research did not produce a report.[/red]")
return
+ # Convert dict to ResearchReport model
+ from haiku.rag.research.models import ResearchReport
+
+ report = ResearchReport.model_validate(report_dict)
+
# Display the report
self.console.print("[bold green]Research Report[/bold green]")
self.console.rule()
diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py
index 054078f7..e92dbe01 100644
--- a/haiku_rag_slim/haiku/rag/cli.py
+++ b/haiku_rag_slim/haiku/rag/cli.py
@@ -299,14 +299,9 @@ def research(
"--db",
help="Path to the LanceDB database file",
),
- verbose: bool = typer.Option(
- False,
- "--verbose",
- help="Show verbose progress output",
- ),
):
app = create_app(db)
- asyncio.run(app.research(question=question, verbose=verbose))
+ asyncio.run(app.research(question=question))
@cli.command("settings", help="Display current configuration settings")
diff --git a/haiku_rag_slim/haiku/rag/research/dependencies.py b/haiku_rag_slim/haiku/rag/research/dependencies.py
index 09376b53..39bf641c 100644
--- a/haiku_rag_slim/haiku/rag/research/dependencies.py
+++ b/haiku_rag_slim/haiku/rag/research/dependencies.py
@@ -1,7 +1,6 @@
from collections.abc import Iterable
from pydantic import BaseModel, Field
-from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.graph_common.models import SearchAnswer
@@ -10,7 +9,6 @@ from haiku.rag.research.models import (
InsightAnalysis,
InsightRecord,
)
-from haiku.rag.research.stream import ResearchStream
class ResearchContext(BaseModel):
@@ -184,10 +182,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"
- )
def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
diff --git a/haiku_rag_slim/haiku/rag/research/graph.py b/haiku_rag_slim/haiku/rag/research/graph.py
index a756476e..62a3f9a5 100644
--- a/haiku_rag_slim/haiku/rag/research/graph.py
+++ b/haiku_rag_slim/haiku/rag/research/graph.py
@@ -53,50 +53,51 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
- log(deps, state, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
+ if deps.agui_emitter:
+ deps.agui_emitter.start_step("plan")
+ deps.agui_emitter.update_activity("planning", "Creating research plan")
- plan_agent = Agent(
- model=get_model(provider, model),
- output_type=ResearchPlan,
- instructions=(
- PLAN_PROMPT
- + "\n\nUse the gather_context tool once on the main question before planning."
- ),
- retries=3,
- deps_type=ResearchDependencies,
- )
+ try:
+ plan_agent = Agent(
+ model=get_model(provider, model),
+ output_type=ResearchPlan,
+ instructions=(
+ PLAN_PROMPT
+ + "\n\nUse the gather_context tool once on the main question before planning."
+ ),
+ retries=3,
+ output_retries=3,
+ deps_type=ResearchDependencies,
+ )
- @plan_agent.tool
- async def gather_context(
- ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
- ) -> str:
- results = await ctx2.deps.client.search(query, limit=limit)
- expanded = await ctx2.deps.client.expand_context(results)
- return "\n\n".join(chunk.content for chunk, _ in expanded)
+ @plan_agent.tool
+ async def gather_context(
+ ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
+ ) -> str:
+ results = await ctx2.deps.client.search(query, limit=limit)
+ expanded = await ctx2.deps.client.expand_context(results)
+ return "\n\n".join(chunk.content for chunk, _ in expanded)
- prompt = (
- "Plan a focused approach for the main question.\n\n"
- f"Main question: {state.context.original_question}"
- )
+ prompt = (
+ "Plan a focused approach for the main question.\n\n"
+ f"Main question: {state.context.original_question}"
+ )
- agent_deps = ResearchDependencies(
- client=deps.client,
- context=state.context,
- console=deps.console,
- stream=deps.stream,
- )
- plan_result = await plan_agent.run(prompt, deps=agent_deps)
- state.context.sub_questions = list(plan_result.output.sub_questions)
+ agent_deps = ResearchDependencies(
+ client=deps.client,
+ context=state.context,
+ )
+ plan_result = await plan_agent.run(prompt, deps=agent_deps)
+ state.context.sub_questions = list(plan_result.output.sub_questions)
- log(deps, state, "\n[bold green]✅ 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}")
+ # Log the plan results
+ log(deps, state, f"Main Question: {state.context.original_question}")
+ log(deps, state, "Sub-questions:")
+ for i, sq in enumerate(state.context.sub_questions, 1):
+ log(deps, state, f" {i}. {sq}")
+ finally:
+ if deps.agui_emitter:
+ deps.agui_emitter.finish_step()
@g.step
async def search_one(
@@ -106,26 +107,30 @@ def build_research_graph(
deps = ctx.deps
sub_q = ctx.inputs
- # Create semaphore if not already provided
- if deps.semaphore is None:
- import asyncio
+ if deps.agui_emitter:
+ deps.agui_emitter.start_step("search_one")
- deps.semaphore = asyncio.Semaphore(state.max_concurrency)
+ try:
+ # Create semaphore if not already provided
+ if deps.semaphore is None:
+ import asyncio
- # Use semaphore to control concurrency
- async with deps.semaphore:
- return await _do_search(state, deps, sub_q)
+ deps.semaphore = asyncio.Semaphore(state.max_concurrency)
+
+ # Use semaphore to control concurrency
+ async with deps.semaphore:
+ return await _do_search(state, deps, sub_q)
+ finally:
+ if deps.agui_emitter:
+ deps.agui_emitter.finish_step()
async def _do_search(
state: ResearchState,
deps: ResearchDeps,
sub_q: str,
) -> SearchAnswer:
- log(
- deps,
- state,
- f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
- )
+ if deps.agui_emitter:
+ deps.agui_emitter.update_activity("searching", f"Searching: {sub_q}")
agent = Agent(
model=get_model(provider, model),
@@ -160,21 +165,16 @@ def build_research_graph(
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
- console=deps.console,
- stream=deps.stream,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
answer = result.output
if answer:
state.context.add_qa_response(answer)
- preview = answer.answer[:150] + (
- "…" if len(answer.answer) > 150 else ""
- )
- log(deps, state, f" [green]✓[/green] {preview}")
+ log(deps, state, f"Answer: {answer.answer}")
return answer
except Exception as e:
- log(deps, state, f"[red]Search failed:[/red] {e}")
+ log(deps, state, f"Search failed: {e}")
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
@@ -204,148 +204,139 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
- log(
- deps,
- state,
- "\n[bold cyan]🧭 Synthesizing new insights and gap status...[/bold cyan]",
- )
+ if deps.agui_emitter:
+ deps.agui_emitter.start_step("analyze_insights")
+ deps.agui_emitter.update_activity(
+ "analyzing", "Synthesizing insights and gaps"
+ )
- agent = Agent(
- model=get_model(provider, model),
- output_type=InsightAnalysis,
- instructions=INSIGHT_AGENT_PROMPT,
- retries=3,
- deps_type=ResearchDependencies,
- )
+ try:
+ agent = Agent(
+ model=get_model(provider, model),
+ output_type=InsightAnalysis,
+ instructions=INSIGHT_AGENT_PROMPT,
+ retries=3,
+ output_retries=3,
+ deps_type=ResearchDependencies,
+ )
- context_xml = format_context_for_prompt(state.context)
- prompt = (
- "Review the latest research context and update the shared ledger of insights, gaps,"
- " and follow-up questions.\n\n"
- f"{context_xml}"
- )
- agent_deps = ResearchDependencies(
- client=deps.client,
- context=state.context,
- console=deps.console,
- stream=deps.stream,
- )
- result = await agent.run(prompt, deps=agent_deps)
- analysis: InsightAnalysis = result.output
+ context_xml = format_context_for_prompt(state.context)
+ prompt = (
+ "Review the latest research context and update the shared ledger of insights, gaps,"
+ " and follow-up questions.\n\n"
+ f"{context_xml}"
+ )
+ agent_deps = ResearchDependencies(
+ client=deps.client,
+ context=state.context,
+ )
+ result = await agent.run(prompt, deps=agent_deps)
+ analysis: InsightAnalysis = result.output
- state.context.integrate_analysis(analysis)
- state.last_analysis = analysis
+ state.context.integrate_analysis(analysis)
+ state.last_analysis = analysis
- if analysis.commentary:
- log(deps, state, f" Summary: {analysis.commentary}")
- if analysis.highlights:
- log(deps, state, " [bold]Updated insights:[/bold]")
- for insight in analysis.highlights:
- label = insight.status.value
- log(
- deps,
- state,
- f" • ({label}) {insight.summary}",
- )
- if analysis.gap_assessments:
- log(deps, state, " [bold yellow]Gap updates:[/bold yellow]")
- for gap in analysis.gap_assessments:
- status = "resolved" if gap.resolved else "open"
- severity = gap.severity.value
- log(
- deps,
- state,
- f" • ({severity}/{status}) {gap.description}",
- )
- if analysis.resolved_gaps:
- log(deps, state, " [green]Resolved gaps:[/green]")
- for resolved in analysis.resolved_gaps:
- log(deps, state, f" • {resolved}")
- if analysis.new_questions:
- log(deps, state, " [cyan]Proposed follow-ups:[/cyan]")
- for question in analysis.new_questions:
- log(deps, state, f" • {question}")
+ if analysis.commentary:
+ log(deps, state, f"Summary: {analysis.commentary}")
+ if analysis.highlights:
+ log(deps, state, "Updated insights:")
+ for insight in analysis.highlights:
+ label = insight.status.value
+ log(deps, state, f" • ({label}) {insight.summary}")
+ if analysis.gap_assessments:
+ log(deps, state, "Gap updates:")
+ for gap in analysis.gap_assessments:
+ status = "resolved" if gap.resolved else "open"
+ severity = gap.severity.value
+ log(deps, state, f" • ({severity}/{status}) {gap.description}")
+ if analysis.resolved_gaps:
+ log(deps, state, "Resolved gaps:")
+ for resolved in analysis.resolved_gaps:
+ log(deps, state, f" • {resolved}")
+ if analysis.new_questions:
+ log(deps, state, "Proposed follow-ups:")
+ for question in analysis.new_questions:
+ log(deps, state, f" • {question}")
+ finally:
+ if deps.agui_emitter:
+ deps.agui_emitter.finish_step()
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
state = ctx.state
deps = ctx.deps
- log(
- deps,
- state,
- "\n[bold cyan]📊 Evaluating research sufficiency...[/bold cyan]",
- )
-
- agent = Agent(
- model=get_model(provider, model),
- output_type=EvaluationResult,
- instructions=DECISION_AGENT_PROMPT,
- retries=3,
- deps_type=ResearchDependencies,
- )
-
- context_xml = format_context_for_prompt(state.context)
- analysis_xml = format_analysis_for_prompt(state.last_analysis)
- prompt_parts = [
- "Assess whether the research now answers the original question with adequate confidence.",
- context_xml,
- analysis_xml,
- ]
- if state.last_eval is not None:
- prev = state.last_eval
- prompt_parts.append(
- ""
- f"{prev.confidence_score:.2f}"
- f"{str(prev.is_sufficient).lower()}"
- f"{prev.reasoning}"
- ""
+ if deps.agui_emitter:
+ deps.agui_emitter.start_step("decide")
+ deps.agui_emitter.update_activity(
+ "evaluating", "Evaluating research sufficiency"
)
- prompt = "\n\n".join(part for part in prompt_parts if part)
- agent_deps = ResearchDependencies(
- client=deps.client,
- context=state.context,
- console=deps.console,
- stream=deps.stream,
- )
- decision_result = await agent.run(prompt, deps=agent_deps)
- output = decision_result.output
+ try:
+ agent = Agent(
+ model=get_model(provider, model),
+ output_type=EvaluationResult,
+ instructions=DECISION_AGENT_PROMPT,
+ retries=3,
+ output_retries=3,
+ deps_type=ResearchDependencies,
+ )
- state.last_eval = output
- state.iterations += 1
+ context_xml = format_context_for_prompt(state.context)
+ analysis_xml = format_analysis_for_prompt(state.last_analysis)
+ prompt_parts = [
+ "Assess whether the research now answers the original question with adequate confidence.",
+ context_xml,
+ analysis_xml,
+ ]
+ if state.last_eval is not None:
+ prev = state.last_eval
+ prompt_parts.append(
+ ""
+ f"{prev.confidence_score:.2f}"
+ f"{str(prev.is_sufficient).lower()}"
+ f"{prev.reasoning}"
+ ""
+ )
+ prompt = "\n\n".join(part for part in prompt_parts if part)
- for new_q in output.new_questions:
- if new_q not in state.context.sub_questions:
- state.context.sub_questions.append(new_q)
+ agent_deps = ResearchDependencies(
+ client=deps.client,
+ context=state.context,
+ )
+ decision_result = await agent.run(prompt, deps=agent_deps)
+ output = decision_result.output
- if output.key_insights:
- log(deps, state, " [bold]Key insights:[/bold]")
- for insight in output.key_insights:
- log(deps, state, f" • {insight}")
+ state.last_eval = output
+ state.iterations += 1
- if output.gaps:
- log(deps, state, " [bold yellow]Remaining gaps:[/bold yellow]")
- for gap in output.gaps:
- log(deps, state, f" • {gap}")
+ for new_q in output.new_questions:
+ if new_q not in state.context.sub_questions:
+ state.context.sub_questions.append(new_q)
- log(
- deps,
- state,
- f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
- )
- status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
- log(deps, state, f" Sufficient: {status}")
+ if output.key_insights:
+ log(deps, state, "Key insights:")
+ for insight in output.key_insights:
+ log(deps, state, f" • {insight}")
- should_continue = (
- not output.is_sufficient
- or output.confidence_score < state.confidence_threshold
- ) and state.iterations < state.max_iterations
+ if output.gaps:
+ log(deps, state, "Remaining gaps:")
+ for gap in output.gaps:
+ log(deps, state, f" • {gap}")
- if not should_continue:
- log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
+ log(deps, state, f"Confidence: {output.confidence_score:.1%}")
+ status = "Yes" if output.is_sufficient else "No"
+ log(deps, state, f"Sufficient: {status}")
- return should_continue
+ should_continue = (
+ not output.is_sufficient
+ or output.confidence_score < state.confidence_threshold
+ ) and state.iterations < state.max_iterations
+
+ return should_continue
+ finally:
+ if deps.agui_emitter:
+ deps.agui_emitter.finish_step()
@g.step
async def synthesize(
@@ -354,36 +345,37 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
- log(
- deps,
- state,
- "\n[bold cyan]📝 Generating final research report...[/bold cyan]",
- )
+ if deps.agui_emitter:
+ deps.agui_emitter.start_step("synthesize")
+ deps.agui_emitter.update_activity(
+ "synthesizing", "Generating final research report"
+ )
- agent = Agent(
- model=get_model(provider, model),
- output_type=ResearchReport,
- instructions=SYNTHESIS_AGENT_PROMPT,
- retries=3,
- deps_type=ResearchDependencies,
- )
+ try:
+ agent = Agent(
+ model=get_model(provider, model),
+ output_type=ResearchReport,
+ instructions=SYNTHESIS_AGENT_PROMPT,
+ retries=3,
+ output_retries=3,
+ deps_type=ResearchDependencies,
+ )
- context_xml = format_context_for_prompt(state.context)
- prompt = (
- "Generate a comprehensive research report based on all gathered information.\n\n"
- f"{context_xml}\n\n"
- "Create a detailed report that synthesizes all findings into a coherent response."
- )
- agent_deps = ResearchDependencies(
- client=deps.client,
- context=state.context,
- console=deps.console,
- stream=deps.stream,
- )
- result = await agent.run(prompt, deps=agent_deps)
-
- log(deps, state, "[bold green]✅ Research complete![/bold green]")
- return result.output
+ context_xml = format_context_for_prompt(state.context)
+ prompt = (
+ "Generate a comprehensive research report based on all gathered information.\n\n"
+ f"{context_xml}\n\n"
+ "Create a detailed report that synthesizes all findings into a coherent response."
+ )
+ agent_deps = ResearchDependencies(
+ client=deps.client,
+ context=state.context,
+ )
+ result = await agent.run(prompt, deps=agent_deps)
+ return result.output
+ finally:
+ if deps.agui_emitter:
+ deps.agui_emitter.finish_step()
# Build the graph structure
collect_answers = g.join(
diff --git a/haiku_rag_slim/haiku/rag/research/state.py b/haiku_rag_slim/haiku/rag/research/state.py
index d9396055..72e341f5 100644
--- a/haiku_rag_slim/haiku/rag/research/state.py
+++ b/haiku_rag_slim/haiku/rag/research/state.py
@@ -3,29 +3,35 @@ from dataclasses import dataclass
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
-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, InsightAnalysis
-from haiku.rag.research.stream import ResearchStream
+from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport
if TYPE_CHECKING:
+ from haiku.rag.agui.emitter import AGUIEmitter
from haiku.rag.config.models import AppConfig
@dataclass
class ResearchDeps:
+ """Dependencies for research graph execution."""
+
client: HaikuRAG
- console: Console | None = None
- stream: ResearchStream | None = None
+ agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
semaphore: asyncio.Semaphore | 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)
+ """Emit a log message through AG-UI events.
+
+ Args:
+ message: The message to log
+ state: Optional state to include in state update
+ """
+ if self.agui_emitter:
+ self.agui_emitter.log(message)
+ if state:
+ self.agui_emitter.update_state(state)
class ResearchState(BaseModel):
diff --git a/haiku_rag_slim/haiku/rag/research/stream.py b/haiku_rag_slim/haiku/rag/research/stream.py
deleted file mode 100644
index 5a2b1950..00000000
--- a/haiku_rag_slim/haiku/rag/research/stream.py
+++ /dev/null
@@ -1,162 +0,0 @@
-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
- 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,
- confidence_threshold=state.confidence_threshold,
- pending_sub_questions=len(context.sub_questions),
- answered_questions=len(context.qa_responses),
- insights=[
- f"{insight.status.value}:{insight.summary}"
- for insight in context.insights
- ],
- gaps=[
- f"{gap.severity.value}/{'resolved' if gap.resolved else 'open'}:{gap.description}"
- for gap in 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,
- 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
-
- 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 = await graph.run(state=state, deps=deps)
-
- if report is None:
- raise RuntimeError("Graph did not produce a report")
-
- stream.report(report, state)
- except Exception as exc:
- 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