Consolidate format_stats into single function

Merge format_stats and format_stats_from_trace — one function
handles both Trace objects and dicts.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

app.py

@@ -156,31 +156,27 @@ def save_uploaded_files(file_paths: list[str]) -> Path:
     return workspace
 
 
-def format_stats(trace: object) -> str:
-    cached = trace.cached_tokens
-    cache_str = f" ({cached} cached)" if cached else ""
-    model_name = trace.model.split("/")[-1] if trace.model else ""
-    parts = [
-        model_name,
-        f"{trace.prompt_tokens + trace.completion_tokens:,} tokens{cache_str}",
-        f"{len(trace.tool_calls)} tool calls",
-        f"{trace.wall_time_s:.1f}s",
-    ]
-    if trace.cost:
-        parts.append(f"${trace.cost:.4f}")
-    return " · ".join(parts)
-
+def format_stats(trace) -> str:
+    """Format trace stats for display. Accepts a Trace object or dict."""
+    if isinstance(trace, dict):
+        cached = trace.get("cached_tokens", 0)
+        model = trace.get("model", "")
+        prompt = trace.get("prompt_tokens", 0)
+        completion = trace.get("completion_tokens", 0)
+        tool_calls = trace.get("tool_calls", [])
+        wall = trace.get("wall_time_s", 0)
+        cost = trace.get("cost")
+    else:
+        cached = trace.cached_tokens
+        model = trace.model
+        prompt = trace.prompt_tokens
+        completion = trace.completion_tokens
+        tool_calls = trace.tool_calls
+        wall = trace.wall_time_s
+        cost = trace.cost
 
-def format_stats_from_trace(trace: dict) -> str:
-    cached = trace.get("cached_tokens", 0)
     cache_str = f" ({cached} cached)" if cached else ""
-    model = trace.get("model", "")
     model_name = model.split("/")[-1] if model else ""
-    prompt = trace.get("prompt_tokens", 0)
-    completion = trace.get("completion_tokens", 0)
-    tool_calls = trace.get("tool_calls", [])
-    wall = trace.get("wall_time_s", 0)
-    cost = trace.get("cost")
     parts = [
         model_name,
         f"{prompt + completion:,} tokens{cache_str}",
@@ -396,7 +392,7 @@ def build_app() -> gr.Blocks:
                     "question": data.get("question", ""),
                     "answer": clean_answer,
                     "sources": sources,
-                    "stats": format_stats_from_trace(trace),
+                    "stats": format_stats(trace),
                     "source_tag": data.get("source", ""),
                     "trace_html": trace_html,
                     "filename": safe_name,