From 2c1f79614a73b605f609703d70628a4d122d3e50 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Nov 2025 11:10:51 +0200 Subject: [PATCH] Use StateSnapshot updates --- haiku_rag_slim/haiku/rag/agui/cli_renderer.py | 107 ++++++++---------- haiku_rag_slim/haiku/rag/research/graph.py | 83 +++++++------- 2 files changed, 89 insertions(+), 101 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agui/cli_renderer.py b/haiku_rag_slim/haiku/rag/agui/cli_renderer.py index 0d146b61..c09e62b6 100644 --- a/haiku_rag_slim/haiku/rag/agui/cli_renderer.py +++ b/haiku_rag_slim/haiku/rag/agui/cli_renderer.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/research/graph.py b/haiku_rag_slim/haiku/rag/research/graph.py index 62a3f9a5..b2019e00 100644 --- a/haiku_rag_slim/haiku/rag/research/graph.py +++ b/haiku_rag_slim/haiku/rag/research/graph.py @@ -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