fix: anthropic output_tokens reset by zero usage chunks

[andrewwhitecdw]

This PR addresses the following issue in src/polar/gateway/transform/anthropic.py: anthropic output_tokens reset by zero usage chunks.

Changes

• src/polar/gateway/transform/anthropic.py: anthropic output_tokens reset by zero usage chunks.

Details

--- a/src/polar/gateway/transform/anthropic.py
+++ b/src/polar/gateway/transform/anthropic.py
@@ -1,3 +1,5 @@
-        usage = chunk.get("usage", {})
-        if usage:
-            self.output_tokens = usage.get("completion_tokens", self.output_tokens)
+        usage = chunk.get("usage", {})
+        if usage:
+            completion_tokens = usage.get("completion_tokens")
+            if isinstance(completion_tokens, int):
+                self.output_tokens = max(self.output_tokens, completion_tokens)

Tests

• tests/unit/gateway/transform/test_anthropic_output_tokens.py

--- /dev/null
+++ b/tests/unit/gateway/transform/test_anthropic_output_tokens.py
@@ -0,0 +1,26 @@
+from polar.gateway.transform.anthropic import AnthropicStreamState, AnthropicTransformer
+
+
+def test_stream_state_output_tokens_not_reset_by_zero_usage_chunk():
+    state = AnthropicStreamState(
+        "claude-3", AnthropicTransformer.FINISH_TO_STOP_REASON
+    )
+    state.process_chunk(
+        {"usage": {"completion_tokens": 5}, "choices": [{"delta": {"content": "hello"}}]}
+    )
+    assert state.output_tokens == 5
+    state.process_chunk({"usage": {"completion_tokens": 0}, "choices": []})
+    assert state.output_tokens == 5
+    final = state.finalize()
+    message_delta = next(e for e in final if e["type"] == "message_delta")
+    assert message_delta["usage"]["output_tokens"] == 5
+
+
+def test_stream_state_output_tokens_updates_with_larger_usage_chunk():
+    state = AnthropicStreamState(
+        "claude-3", AnthropicTransformer.FINISH_TO_STOP_REASON
+    )
+    state.process_chunk(
+        {"usage": {"completion_tokens": 3}, "choices": [{"delta": {"content": "hi"}}]}
+    )
+    assert state.output_tokens == 3
+    state.process_chunk(
+        {"usage": {"completion_tokens": 10}, "choices": [{"delta": {"content": " world"}}]}
+    )
+    assert state.output_tokens == 10
+    final = state.finalize()
+    message_delta = next(e for e in final if e["type"] == "message_delta")
+    assert message_delta["usage"]["output_tokens"] == 10