Use StateSnapshot updates

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:
+    def _render_run_started(self, _event: AGUIEvent) -> None:
-        """Render run start event.
+        """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.
+        """Render error event."""
-
-        Args:
-            event: RunError 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.
+        """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]")
+            self.console.print(
+                f"\n[bold cyan][STEP_STARTED][/bold cyan] {display_name}"
+            )
 
-    def _render_step_finished(self, event: AGUIEvent) -> None:
+    def _render_step_finished(self, _event: AGUIEvent) -> None:
-        """Render step finish event.
+        """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.
+        """Render complete text message."""
-
-        Args:
-            event: TextMessageChunk event
-        """
         delta = event.get("delta", "")
-        # The delta contains the text content to display
+        self.console.print(f"[magenta][TEXT_MESSAGE][/magenta] {delta}")
-        self.console.print(delta)
 
     def _render_text_content(self, event: AGUIEvent) -> None:
-        """Render streaming text content delta.
+        """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.
+        """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]")
+            self.console.print(f"[yellow][ACTIVITY][/yellow] {content}")
 
-    def _apply_state_delta(self, event: AGUIEvent) -> None:
+    def _render_state_snapshot(self, new_state: dict | None) -> None:
-        """Apply state delta to current state.
+        """Render state snapshot showing only what changed."""
+        if not new_state:
+            return
 
-        Args:
+        old_state = self._state or {}
-            event: StateDelta event
+        diff = self._compute_diff(old_state, new_state)
-        """
+
-        # Currently not applying deltas - could implement state patching later
+        if not diff:
-        # For now, the text messages contain all the information we need to display
+            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."""

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
+            # State now contains the plan - emit state update and narrate
-            log(deps, state, f"Main Question: {state.context.original_question}")
+            if deps.agui_emitter:
-            log(deps, state, "Sub-questions:")
+                deps.agui_emitter.update_state(state)
-            for i, sq in enumerate(state.context.sub_questions, 1):
+                count = len(state.context.sub_questions)
-                log(deps, state, f"  {i}. {sq}")
+                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:
+            # State updated with insights/gaps - emit state update and narrate
-                log(deps, state, f"Summary: {analysis.commentary}")
+            if deps.agui_emitter:
-            if analysis.highlights:
+                deps.agui_emitter.update_state(state)
-                log(deps, state, "Updated insights:")
+                highlights = len(analysis.highlights) if analysis.highlights else 0
-                for insight in analysis.highlights:
+                gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
-                    label = insight.status.value
+                resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
-                    log(deps, state, f"  • ({label}) {insight.summary}")
+                parts = []
-            if analysis.gap_assessments:
+                if highlights:
-                log(deps, state, "Gap updates:")
+                    parts.append(f"{highlights} insights")
-                for gap in analysis.gap_assessments:
+                if gaps:
-                    status = "resolved" if gap.resolved else "open"
+                    parts.append(f"{gaps} gaps")
-                    severity = gap.severity.value
+                if resolved:
-                    log(deps, state, f"  • ({severity}/{status}) {gap.description}")
+                    parts.append(f"{resolved} resolved")
-            if analysis.resolved_gaps:
+                summary = ", ".join(parts) if parts else "No updates"
-                log(deps, state, "Resolved gaps:")
+                deps.agui_emitter.update_activity("analyzing", f"Analysis: {summary}")
-                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()
@@ -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:
+            # State updated with evaluation - emit state update and narrate
-                log(deps, state, "Key insights:")
+            if deps.agui_emitter:
-                for insight in output.key_insights:
+                deps.agui_emitter.update_state(state)
-                    log(deps, state, f"  • {insight}")
+                sufficient = "Yes" if output.is_sufficient else "No"
-
+                deps.agui_emitter.update_activity(
-            if output.gaps:
+                    "evaluating",
-                log(deps, state, "Remaining gaps:")
+                    f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
-                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}")
 
             should_continue = (
                 not output.is_sufficient