Add token usage logging and completion output

main.py

@@ -61,35 +61,67 @@ def run_tournament(
     process_log = []
     hist_fig = None
     top_picks_str = ""
+    prompt_tokens = 0
+    completion_tokens = 0
+
+    def add_usage(usage):
+        nonlocal prompt_tokens, completion_tokens
+        if not usage:
+            return
+        pt = getattr(usage, "prompt_tokens", None)
+        if pt is None and isinstance(usage, dict):
+            pt = usage.get("prompt_tokens")
+        ct = getattr(usage, "completion_tokens", None)
+        if ct is None and isinstance(usage, dict):
+            ct = usage.get("completion_tokens")
+        if pt:
+            prompt_tokens += pt
+        if ct:
+            completion_tokens += ct
+
+    def usage_str():
+        return (
+            f"Prompt tokens: {prompt_tokens}\n"
+            f"Completion tokens: {completion_tokens}\n"
+            f"Total tokens: {prompt_tokens + completion_tokens}"
+        )
     def log(msg):
         process_log.append(msg)
         tqdm.write(msg)
-        yield "\n".join(process_log), hist_fig, top_picks_str
+        yield "\n".join(process_log), hist_fig, top_picks_str, usage_str()
     yield from log("Generating players …")
-    all_players = generate_players(
+    all_players, usage = generate_players(
         instruction,
         n_gen,
         model=generate_model,
         api_base=api_base,
         api_key=api_token,
+        return_usage=True,
     )
+    add_usage(usage)
     yield from log(f"{len(all_players)} players generated")
+    for i, p in enumerate(all_players, 1):
+        disp = p.replace("\n", " ")
+        if len(disp) > 100:
+            disp = disp[:100] + "…"
+        yield from log(f"Completion {i}: {disp}")
     def criteria_block():
         return "\n".join(f"{i + 1}) {c}" for i, c in enumerate(criteria_list))
 
     if enable_score_filter:
         def score(player):
-            data = _clean_json(
-                prompt_score(
-                    instruction,
-                    criteria_list,
-                    criteria_block(),
-                    player,
-                    model=score_model,
-                    api_base=api_base,
-                    api_key=api_token,
-                )
+            text, usage = prompt_score(
+                instruction,
+                criteria_list,
+                criteria_block(),
+                player,
+                model=score_model,
+                api_base=api_base,
+                api_key=api_token,
+                return_usage=True,
             )
+            add_usage(usage)
+            data = _clean_json(text)
             if "scores" in data and isinstance(data["scores"], list):
                 vals = data["scores"]
                 return sum(vals) / len(vals) if vals else 0.0
@@ -113,17 +145,18 @@ def run_tournament(
         top_players = all_players
     if enable_pairwise_filter:
         def play(a, b):
-            winner_label = _clean_json(
-                prompt_pairwise(
-                    instruction,
-                    criteria_block(),
-                    a,
-                    b,
-                    model=pairwise_model,
-                    api_base=api_base,
-                    api_key=api_token,
-                )
-            ).get("winner", "A")
+            text, usage = prompt_pairwise(
+                instruction,
+                criteria_block(),
+                a,
+                b,
+                model=pairwise_model,
+                api_base=api_base,
+                api_key=api_token,
+                return_usage=True,
+            )
+            add_usage(usage)
+            winner_label = _clean_json(text).get("winner", "A")
             return a if winner_label == "A" else b
 
         def tournament_round(pairs, executor):
@@ -178,7 +211,7 @@ def run_tournament(
     else:
         top_k = top_players[:num_top_picks]
     top_picks_str = "\n\n\n=====================================================\n\n\n".join(top_k)
-    yield "\n".join(process_log + ["Done"]), hist_fig, top_picks_str
+    yield "\n".join(process_log + ["Done"]), hist_fig, top_picks_str, usage_str()
 
 demo = gr.Interface(
     fn=run_tournament,
@@ -201,6 +234,7 @@ demo = gr.Interface(
         gr.Textbox(lines=10, label="Process"),
         gr.Plot(label="Score Distribution"),
         gr.Textbox(lines=50, label="Top picks"),
+        gr.Textbox(lines=5, label="Token Usage"),
     ],
     description="Generate multiple completions and use score and pairwise filters to find the best answers.",
 )

tests/test_main.py

@@ -82,10 +82,10 @@ def test_run_tournament_full_loop():
          patch('main.tqdm', new=dummy_tqdm), \
          patch('main.plt.figure', return_value='fig'), \
          patch('main.plt.hist'):
-        mock_gen.return_value = ['p1', 'p2', 'p3', 'p4']
+        mock_gen.return_value = (['p1', 'p2', 'p3', 'p4'], {'prompt_tokens':1,'completion_tokens':1})
         scores = {'p1':3, 'p2':2, 'p3':1, 'p4':0}
-        mock_score.side_effect = lambda instr, cl, block, player, **kw: json.dumps({'score': scores[player]})
-        mock_pair.side_effect = lambda instr, block, a, b, **kw: json.dumps({'winner': 'A'})
+        mock_score.side_effect = lambda instr, cl, block, player, **kw: (json.dumps({'score': scores[player]}), {'prompt_tokens':1,'completion_tokens':1})
+        mock_pair.side_effect = lambda instr, block, a, b, **kw: (json.dumps({'winner': 'A'}), {'prompt_tokens':1,'completion_tokens':1})
 
         results = list(main.run_tournament(
             api_base='b',
@@ -103,10 +103,11 @@ def test_run_tournament_full_loop():
             enable_pairwise_filter=True,
         ))
 
-    process_log, hist_fig, top_picks = results[-1]
+    process_log, hist_fig, top_picks, usage = results[-1]
     assert 'Done' in process_log
     assert hist_fig == 'fig'
     assert top_picks.strip() in {'p1', 'p2'}
-    mock_gen.assert_called_once_with('instr', 4, model='gm', api_base='b', api_key='k')
+    mock_gen.assert_called_once_with('instr', 4, model='gm', api_base='b', api_key='k', return_usage=True)
+    assert 'Prompt tokens' in usage
     assert mock_score.call_count == 4
     assert mock_pair.called

tournament_utils.py

@@ -18,15 +18,23 @@ def generate_players(
     *,
     api_base: str | None = None,
     api_key: str | None = None,
-):
-    """Request `n` completions for the instruction using the given model."""
+    return_usage: bool = False,
+) -> list[str] | tuple[list[str], object]:
+    """Request ``n`` completions for the instruction using the given model.
+
+    When ``return_usage`` is ``True`` the ``usage`` object from the completion
+    response is also returned.
+    """
     response = completion(
         model=model,
         messages=[{"role": "user", "content": instruction}],
         n=n,
         **_completion_kwargs(api_base, api_key),
     )
-    return [c.message.content.strip() for c in response.choices]
+    players = [c.message.content.strip() for c in response.choices]
+    if return_usage:
+        return players, getattr(response, "usage", None)
+    return players
 
 
 def prompt_score(
@@ -38,7 +46,8 @@ def prompt_score(
     *,
     api_base: str | None = None,
     api_key: str | None = None,
-) -> str:
+    return_usage: bool = False,
+) -> str | tuple[str, object]:
     """Return a JSON score string evaluating `player` on the criteria."""
     example_scores = ", ".join(["1-10"] * len(criteria_list)) or "1-10"
     prompt = f"""Evaluate the output below on the following criteria:
@@ -56,7 +65,10 @@ Output:
         messages=[{"role": "system", "content": prompt}],
         **_completion_kwargs(api_base, api_key),
     )
-    return response.choices[0].message.content.strip()
+    text = response.choices[0].message.content.strip()
+    if return_usage:
+        return text, getattr(response, "usage", None)
+    return text
 
 
 def prompt_pairwise(
@@ -68,7 +80,8 @@ def prompt_pairwise(
     *,
     api_base: str | None = None,
     api_key: str | None = None,
-) -> str:
+    return_usage: bool = False,
+) -> str | tuple[str, object]:
     """Return which player wins in JSON using the given criteria."""
     prompt = f"""Compare the two players below using:
 {criteria_block}
@@ -86,4 +99,7 @@ Players:
         messages=[{"role": "system", "content": prompt}],
         **_completion_kwargs(api_base, api_key),
     )
-    return response.choices[0].message.content.strip()
+    text = response.choices[0].message.content.strip()
+    if return_usage:
+        return text, getattr(response, "usage", None)
+    return text