Use StateSnapshot updates
This commit is contained in:
Yiorgis Gozadinos
parent
663a7fac9a
commit
2c1f79614a
2 changed files with 89 additions and 101 deletions
|
|
@ -3,7 +3,6 @@
|
|||
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
|
||||
|
|
@ -23,7 +22,7 @@ class AGUIConsoleRenderer:
|
|||
console: Optional Rich console instance (creates new one if not provided)
|
||||
"""
|
||||
self.console = console or Console()
|
||||
self._state: BaseModel | None = None
|
||||
self._state: dict | None = None
|
||||
|
||||
async def render(self, events: AsyncIterator[AGUIEvent]) -> Any | None:
|
||||
"""Process events and render to console, return final result.
|
||||
|
|
@ -59,7 +58,9 @@ class AGUIConsoleRenderer:
|
|||
elif event_type == "TEXT_MESSAGE_END":
|
||||
pass # End of streaming message, no output needed
|
||||
elif event_type == "STATE_SNAPSHOT":
|
||||
self._state = event.get("snapshot")
|
||||
new_state = event.get("snapshot")
|
||||
self._render_state_snapshot(new_state)
|
||||
self._state = new_state
|
||||
elif event_type == "STATE_DELTA":
|
||||
self._apply_state_delta(event)
|
||||
elif event_type == "ACTIVITY_SNAPSHOT":
|
||||
|
|
@ -69,84 +70,72 @@ class AGUIConsoleRenderer:
|
|||
|
||||
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_started(self, _event: AGUIEvent) -> None:
|
||||
"""Render run start event."""
|
||||
|
||||
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
|
||||
"""
|
||||
"""Render error event."""
|
||||
message = event.get("message", "Unknown error")
|
||||
self.console.print(f"[bold red]❌ Error:[/bold red] {message}")
|
||||
self.console.print(f"[bold red][RUN_ERROR][/bold red] {message}")
|
||||
|
||||
def _render_step_started(self, event: AGUIEvent) -> None:
|
||||
"""Render step start event.
|
||||
|
||||
Args:
|
||||
event: StepStarted event
|
||||
"""
|
||||
"""Render step start 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]")
|
||||
self.console.print(
|
||||
f"\n[bold cyan][STEP_STARTED][/bold cyan] {display_name}"
|
||||
)
|
||||
|
||||
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_step_finished(self, _event: AGUIEvent) -> None:
|
||||
"""Render step finish event."""
|
||||
|
||||
def _render_text_message(self, event: AGUIEvent) -> None:
|
||||
"""Render complete text message.
|
||||
|
||||
Args:
|
||||
event: TextMessageChunk event
|
||||
"""
|
||||
"""Render complete text message."""
|
||||
delta = event.get("delta", "")
|
||||
# The delta contains the text content to display
|
||||
self.console.print(delta)
|
||||
self.console.print(f"[magenta][TEXT_MESSAGE][/magenta] {delta}")
|
||||
|
||||
def _render_text_content(self, event: AGUIEvent) -> None:
|
||||
"""Render streaming text content delta.
|
||||
|
||||
Args:
|
||||
event: TextMessageContent event
|
||||
"""
|
||||
"""Render streaming text content delta."""
|
||||
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
|
||||
"""
|
||||
"""Render activity update."""
|
||||
content = event.get("content", "")
|
||||
|
||||
# Render activity content with emphasis
|
||||
if content:
|
||||
self.console.print(f"[dim]{content}[/dim]")
|
||||
self.console.print(f"[yellow][ACTIVITY][/yellow] {content}")
|
||||
|
||||
def _apply_state_delta(self, event: AGUIEvent) -> None:
|
||||
"""Apply state delta to current state.
|
||||
def _render_state_snapshot(self, new_state: dict | None) -> None:
|
||||
"""Render state snapshot showing only what changed."""
|
||||
if not new_state:
|
||||
return
|
||||
|
||||
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
|
||||
old_state = self._state or {}
|
||||
diff = self._compute_diff(old_state, new_state)
|
||||
|
||||
if not diff:
|
||||
return
|
||||
|
||||
self.console.print("[blue][STATE_SNAPSHOT][/blue]")
|
||||
self.console.print(diff, style="dim")
|
||||
|
||||
def _compute_diff(self, old: dict, new: dict) -> dict:
|
||||
"""Compute difference between old and new state."""
|
||||
diff = {}
|
||||
for key, new_value in new.items():
|
||||
old_value = old.get(key)
|
||||
if old_value != new_value:
|
||||
if isinstance(new_value, dict) and isinstance(old_value, dict):
|
||||
nested_diff = self._compute_diff(old_value, new_value)
|
||||
if nested_diff:
|
||||
diff[key] = nested_diff
|
||||
else:
|
||||
diff[key] = new_value
|
||||
return diff
|
||||
|
||||
def _apply_state_delta(self, _event: AGUIEvent) -> None:
|
||||
"""Apply state delta to current state."""
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from pydantic_graph.beta.join import reduce_list_append
|
|||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.graph_common import get_model, log
|
||||
from haiku.rag.graph_common import get_model
|
||||
from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer
|
||||
from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
|
||||
from haiku.rag.research.common import (
|
||||
|
|
@ -90,11 +90,13 @@ def build_research_graph(
|
|||
plan_result = await plan_agent.run(prompt, deps=agent_deps)
|
||||
state.context.sub_questions = list(plan_result.output.sub_questions)
|
||||
|
||||
# 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}")
|
||||
# State now contains the plan - emit state update and narrate
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_state(state)
|
||||
count = len(state.context.sub_questions)
|
||||
deps.agui_emitter.update_activity(
|
||||
"planning", f"Created plan with {count} sub-questions"
|
||||
)
|
||||
finally:
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.finish_step()
|
||||
|
|
@ -171,10 +173,18 @@ def build_research_graph(
|
|||
answer = result.output
|
||||
if answer:
|
||||
state.context.add_qa_response(answer)
|
||||
log(deps, state, f"Answer: {answer.answer}")
|
||||
# State updated with new answer - emit state update and narrate
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_state(state)
|
||||
deps.agui_emitter.update_activity(
|
||||
"searching",
|
||||
f"Found answer with {answer.confidence:.0%} confidence",
|
||||
)
|
||||
return answer
|
||||
except Exception as e:
|
||||
log(deps, state, f"Search failed: {e}")
|
||||
# Narrate the error
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_activity("searching", f"Search failed: {e}")
|
||||
failure_answer = SearchAnswer(
|
||||
query=sub_q,
|
||||
answer=f"Search failed after retries: {str(e)}",
|
||||
|
|
@ -236,27 +246,21 @@ def build_research_graph(
|
|||
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, "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}")
|
||||
# State updated with insights/gaps - emit state update and narrate
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_state(state)
|
||||
highlights = len(analysis.highlights) if analysis.highlights else 0
|
||||
gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
|
||||
resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
|
||||
parts = []
|
||||
if highlights:
|
||||
parts.append(f"{highlights} insights")
|
||||
if gaps:
|
||||
parts.append(f"{gaps} gaps")
|
||||
if resolved:
|
||||
parts.append(f"{resolved} resolved")
|
||||
summary = ", ".join(parts) if parts else "No updates"
|
||||
deps.agui_emitter.update_activity("analyzing", f"Analysis: {summary}")
|
||||
finally:
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.finish_step()
|
||||
|
|
@ -314,19 +318,14 @@ def build_research_graph(
|
|||
if new_q not in state.context.sub_questions:
|
||||
state.context.sub_questions.append(new_q)
|
||||
|
||||
if output.key_insights:
|
||||
log(deps, state, "Key insights:")
|
||||
for insight in output.key_insights:
|
||||
log(deps, state, f" • {insight}")
|
||||
|
||||
if output.gaps:
|
||||
log(deps, state, "Remaining gaps:")
|
||||
for gap in output.gaps:
|
||||
log(deps, state, f" • {gap}")
|
||||
|
||||
log(deps, state, f"Confidence: {output.confidence_score:.1%}")
|
||||
status = "Yes" if output.is_sufficient else "No"
|
||||
log(deps, state, f"Sufficient: {status}")
|
||||
# State updated with evaluation - emit state update and narrate
|
||||
if deps.agui_emitter:
|
||||
deps.agui_emitter.update_state(state)
|
||||
sufficient = "Yes" if output.is_sufficient else "No"
|
||||
deps.agui_emitter.update_activity(
|
||||
"evaluating",
|
||||
f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
|
||||
)
|
||||
|
||||
should_continue = (
|
||||
not output.is_sufficient
|
||||
|
|
|
|||
Loading…
Reference in a new issue