From fbb43c06c73d5bc1cdef9fdab1289fe366806af7 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Fri, 19 Sep 2025 17:25:04 +0300
Subject: [PATCH] Cleanup console logging

---
 src/haiku/rag/app.py                     |  6 ++++++
 src/haiku/rag/research/nodes/evaluate.py | 23 +++++++++++------------
 src/haiku/rag/research/nodes/plan.py     | 11 +++++------
 src/haiku/rag/research/nodes/search.py   | 11 +++++------
 4 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py
index 73e46392..d0948ea4 100644
--- a/src/haiku/rag/app.py
+++ b/src/haiku/rag/app.py
@@ -145,6 +145,12 @@ class HaikuRAGApp:
                 self.console.print(report.executive_summary)
                 self.console.print()
 
+                # Confidence (from last evaluation)
+                if state.last_eval:
+                    conf = state.last_eval.confidence_score  # type: ignore[attr-defined]
+                    self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
+                    self.console.print()
+
                 # Main Findings
                 if report.main_findings:
                     self.console.print("[bold cyan]Main Findings:[/bold cyan]")
diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py
index 79cb80d5..7270d0a0 100644
--- a/src/haiku/rag/research/nodes/evaluate.py
+++ b/src/haiku/rag/research/nodes/evaluate.py
@@ -57,16 +57,16 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         state.last_eval = output
         state.iterations += 1
 
-        if deps.console:
-            if output.key_insights:
-                deps.console.print("   [bold]Key insights:[/bold]")
-                for ins in output.key_insights:
-                    deps.console.print(f"   • {ins}")
-            deps.console.print(
-                f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
-            )
-            status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
-            deps.console.print(f"   Sufficient: {status}")
+        if output.key_insights:
+            log(deps.console, "   [bold]Key insights:[/bold]")
+            for ins in output.key_insights:
+                log(deps.console, f"   • {ins}")
+        log(
+            deps.console,
+            f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
+        )
+        status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
+        log(deps.console, f"   Sufficient: {status}")
 
         from haiku.rag.research.nodes.search import SearchDispatchNode
 
@@ -74,8 +74,7 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             output.is_sufficient
             and output.confidence_score >= state.confidence_threshold
         ) or state.iterations >= state.max_iterations:
-            if deps.console:
-                deps.console.print("\n[bold green]✅ Stopping research.[/bold green]")
+            log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
             return SynthesizeNode(self.provider, self.model)
 
         return SearchDispatchNode(self.provider, self.model)
diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py
index a412639f..653c12c5 100644
--- a/src/haiku/rag/research/nodes/plan.py
+++ b/src/haiku/rag/research/nodes/plan.py
@@ -54,11 +54,10 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         plan_result = await plan_agent.run(prompt, deps=agent_deps)
         state.sub_questions = list(plan_result.output.sub_questions)
 
-        if deps.console:
-            deps.console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
-            deps.console.print(f"   [bold]Main Question:[/bold] {state.question}")
-            deps.console.print("   [bold]Sub-questions:[/bold]")
-            for i, sq in enumerate(state.sub_questions, 1):
-                deps.console.print(f"      {i}. {sq}")
+        log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]")
+        log(deps.console, f"   [bold]Main Question:[/bold] {state.question}")
+        log(deps.console, "   [bold]Sub-questions:[/bold]")
+        for i, sq in enumerate(state.sub_questions, 1):
+            log(deps.console, f"      {i}. {sq}")
 
         return SearchDispatchNode(self.provider, self.model)
diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py
index 4640b3c1..0943338a 100644
--- a/src/haiku/rag/research/nodes/search.py
+++ b/src/haiku/rag/research/nodes/search.py
@@ -36,11 +36,10 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             batch.append(state.sub_questions.pop(0))
 
         async def answer_one(sub_q: str) -> SearchAnswer | None:
-            if deps.console:
-                deps.console.print(
-                    f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}"
-                )
-
+            log(
+                deps.console,
+                f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
+            )
             agent = Agent(
                 model=get_model(self.provider, self.model),
                 output_type=ToolOutput(SearchAnswer, max_retries=3),
@@ -87,6 +86,6 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             state.context.add_qa_response(ans)
             if deps.console:
                 preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
-                deps.console.log(f"   [green]✓[/green] {preview}")
+                log(deps.console, f"   [green]✓[/green] {preview}")
 
         return SearchDispatchNode(self.provider, self.model)