From 5c3355dd0eef8fc72944b501e052d88086987885 Mon Sep 17 00:00:00 2001
From: Manu Nicholas Jacob <manunicholasjacob@gmail.com>
Date: Wed, 5 Aug 2026 12:10:17 -0500
Subject: [PATCH] fix(analyze): handle non-string prompts in count_tokens

count_tokens called len() directly on r["prompt"], which is a dict for reports
written by generators that store structured prompts, so the token count was the
number of dict keys rather than the prompt length. The outputs path already had
a local _to_text helper for exactly this; hoist it to module scope and use it
for the prompt as well.

Adds a regression test that fails on the old code path.

Signed-off-by: Manu Nicholas Jacob <manunicholasjacob@gmail.com>
---
 garak/analyze/count_tokens.py      | 26 +++++++--------
 tests/analyze/test_count_tokens.py | 52 ++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 13 deletions(-)
 create mode 100644 tests/analyze/test_count_tokens.py

diff --git a/garak/analyze/count_tokens.py b/garak/analyze/count_tokens.py
index 454910751..f3bd85929 100644
--- a/garak/analyze/count_tokens.py
+++ b/garak/analyze/count_tokens.py
@@ -19,6 +19,18 @@
 import garak
 
 
+def _to_text(value) -> str:
+    """Extract text from report values while tolerating legacy shapes."""
+    if isinstance(value, str):
+        return value
+    if isinstance(value, dict):
+        for key in ("text", "content", "response"):
+            nested = value.get(key)
+            if isinstance(nested, (str, dict)):
+                return _to_text(nested)
+    return str(value)
+
+
 def count_tokens(report_path: str) -> None:
     calls = 0
     input_length = 0
@@ -35,22 +47,10 @@ def count_tokens(report_path: str) -> None:
                 generations = r["run.generations"]
                 continue
             if "status" in r and r["status"] == 2:
-                input_length += len(r["prompt"]) * generations
+                input_length += len(_to_text(r["prompt"])) * generations
                 calls += generations
                 outputs = r.get("outputs", [])
                 if isinstance(outputs, list):
-
-                    def _to_text(o):
-                        if isinstance(o, str):
-                            return o
-                        if isinstance(o, dict):
-                            # common keys seen in generator outputs
-                            for k in ("text", "content", "response"):
-                                v = o.get(k)
-                                if isinstance(v, str):
-                                    return v
-                        return str(o)
-
                     output_text = "".join(_to_text(o) for o in outputs)
                 else:
                     output_text = str(outputs)
diff --git a/tests/analyze/test_count_tokens.py b/tests/analyze/test_count_tokens.py
new file mode 100644
index 000000000..40c15df62
--- /dev/null
+++ b/tests/analyze/test_count_tokens.py
@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from garak.analyze.count_tokens import count_tokens
+
+
+def _write_report(tmp_path, prompt, generations):
+    report_path = tmp_path / "count_tokens.report.jsonl"
+    entries = [
+        {"run.generations": generations},
+        {
+            "status": 2,
+            "prompt": prompt,
+            "outputs": [],
+        },
+    ]
+    report_path.write_text(
+        "".join(json.dumps(entry) + "\n" for entry in entries),
+        encoding="utf-8",
+    )
+    return report_path
+
+
+def test_count_tokens_counts_nested_prompt_text(tmp_path, capsys):
+    prompt_text = "known prompt text"
+    generations = 3
+    report_path = _write_report(
+        tmp_path,
+        {
+            "role": "user",
+            "content": {"text": prompt_text},
+        },
+        generations,
+    )
+
+    count_tokens(str(report_path))
+
+    output = capsys.readouterr().out
+    assert f"Calls: {generations}" in output
+    assert f"Input chars: {len(prompt_text) * generations}" in output
+
+
+def test_count_tokens_accepts_legacy_string_prompt(tmp_path, capsys):
+    prompt_text = "legacy prompt"
+    report_path = _write_report(tmp_path, prompt_text, generations=2)
+
+    count_tokens(str(report_path))
+
+    output = capsys.readouterr().out
+    assert f"Input chars: {len(prompt_text) * 2}" in output