verifiers 0.1.13.dev1__tar.gz → 0.1.13.dev3__tar.gz

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: verifiers
-Version: 0.1.13.dev1
+Version: 0.1.13.dev3
 Summary: Verifiers: Environments for LLM Reinforcement Learning
 Project-URL: Homepage, https://github.com/primeintellect-ai/verifiers
 Project-URL: Documentation, https://github.com/primeintellect-ai/verifiers

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/tests/test_composable_env.py

@@ -429,6 +429,83 @@ async def test_composable_env_no_upload_when_no_dirs(tmp_path, monkeypatch):
     assert env.upload_file.await_count == 0
 
 
+@pytest.mark.asyncio
+async def test_composable_env_uploads_harness_dirs(tmp_path):
+    taskset = MockSandboxTaskSet(dataset=_make_dataset(), name="test")
+    harness_dir = tmp_path / "agent-src"
+    harness_dir.mkdir()
+    (harness_dir / "marker.txt").write_text("agent\n")
+
+    env = ComposableEnv(
+        taskset=taskset,
+        harness=Harness(
+            run_command="true",
+            install_script="install-agent",
+            get_upload_dirs=lambda: {"agent_src": harness_dir},
+            upload_dir_mapping={"agent_src": "/tmp/agent-src"},
+        ),
+    )
+    env.sandbox_client = SimpleNamespace(
+        execute_command=AsyncMock(
+            return_value=SimpleNamespace(stdout="", stderr="", exit_code=0)
+        ),
+        teardown=lambda: None,
+    )
+    env.taskset.setup = AsyncMock()
+    env.upload_content = AsyncMock()
+    env.upload_file = AsyncMock()
+
+    await env.post_sandbox_setup({"sandbox_id": "sbx", "info": {"id": 0}})
+
+    env.upload_file.assert_awaited_once()
+    upload_call = env.upload_file.await_args
+    assert upload_call.args[0] == "sbx"
+    assert upload_call.args[1] == "/tmp/_upload_tmp_agent-src.tar.gz"
+
+    extract_call = env.sandbox_client.execute_command.await_args_list[1]
+    assert extract_call == call(
+        "sbx",
+        "mkdir -p /tmp && tar -xzf /tmp/_upload_tmp_agent-src.tar.gz -C / && rm -f /tmp/_upload_tmp_agent-src.tar.gz",
+        timeout=60,
+    )
+
+
+@pytest.mark.asyncio
+async def test_composable_env_rejects_duplicate_task_and_harness_upload_names(
+    tmp_path, monkeypatch
+):
+    mod, _ = _make_temp_taskset_package(tmp_path, monkeypatch, with_skills=True)
+    monkeypatch.setattr(MockSandboxTaskSetWithSkills, "__module__", mod.__name__)
+    taskset = MockSandboxTaskSetWithSkills(dataset=_make_dataset(), name="test")
+    harness_dir = tmp_path / "skills"
+    harness_dir.mkdir()
+
+    env = ComposableEnv(
+        taskset=taskset,
+        harness=Harness(
+            run_command="true",
+            install_script="install-agent",
+            get_upload_dirs=lambda: {"skills": harness_dir},
+            skills_path="/task/skills",
+        ),
+    )
+    env.sandbox_client = SimpleNamespace(
+        execute_command=AsyncMock(
+            return_value=SimpleNamespace(stdout="", stderr="", exit_code=0)
+        ),
+        teardown=lambda: None,
+    )
+    env.taskset.setup = AsyncMock()
+    env.upload_content = AsyncMock()
+    env.upload_file = AsyncMock()
+
+    with pytest.raises(
+        ValueError,
+        match="Upload directory names must be unique across task and harness",
+    ):
+        await env.post_sandbox_setup({"sandbox_id": "sbx", "info": {"id": 0}})
+
+
 # ── discover_sibling_dir ─────────────────────────────────────────────────
 
 

verifiers-0.1.13.dev3/tests/test_context_token_metrics.py

@@ -0,0 +1,200 @@
+"""Tests for per-turn context token metrics.
+
+Tests the trajectory-based context token computation
+(final_input_tokens, final_output_tokens) which assumes a linear rollout
+using the last trajectory step.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from verifiers.utils.usage_utils import compute_context_token_metrics
+
+
+# =========================================================================
+# Helpers
+# =========================================================================
+
+SYS = {"role": "system", "content": "You are helpful"}
+USER = {"role": "user", "content": "hi"}
+
+
+def _make_response(prompt_tokens: int, completion_tokens: int) -> MagicMock:
+    response = MagicMock()
+    response.usage = MagicMock(
+        prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
+    )
+    return response
+
+
+def _asst(i: int) -> dict:
+    return {"role": "assistant", "content": f"response {i}"}
+
+
+# =========================================================================
+# compute_context_token_metrics
+# =========================================================================
+
+
+class TestContextMetrics:
+    def test_empty_trajectory(self):
+        metrics = compute_context_token_metrics([])
+        assert metrics["final_output_tokens"] == 0
+        assert metrics["final_input_tokens"] == 0
+
+    def test_single_turn(self):
+        trajectory = [
+            {
+                "prompt": [SYS, USER],
+                "completion": [_asst(0)],
+                "response": _make_response(100, 20),
+            },
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        assert metrics["final_output_tokens"] == 20
+        assert metrics["final_input_tokens"] == 100
+
+    def test_multi_turn(self):
+        trajectory = [
+            {
+                "response": _make_response(100, 20),
+            },
+            {
+                "response": _make_response(150, 25),
+            },
+            {
+                "response": _make_response(200, 30),
+            },
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        # Last step total = 200 + 30 = 230
+        # Sum of completion tokens = 20 + 25 + 30 = 75
+        assert metrics["final_output_tokens"] == 75
+        assert metrics["final_input_tokens"] == 230 - 75
+
+    def test_invariant_total_equals_last_step(self):
+        trajectory = [
+            {"response": _make_response(100, 20)},
+            {"response": _make_response(150, 25)},
+            {"response": _make_response(200, 30)},
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        total = metrics["final_output_tokens"] + metrics["final_input_tokens"]
+        # Total should equal last step's prompt_tokens + completion_tokens
+        assert total == 200 + 30
+
+    def test_no_response_on_any_step(self):
+        trajectory = [{"response": None}]
+        metrics = compute_context_token_metrics(trajectory)
+        assert metrics["final_output_tokens"] == 0
+        assert metrics["final_input_tokens"] == 0
+
+    def test_last_step_used_not_largest(self):
+        """Even if an earlier step has a larger context, we use the last step."""
+        trajectory = [
+            {"response": _make_response(500, 100)},  # larger context
+            {"response": _make_response(100, 20)},  # last step, smaller
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        # Last step total = 120, sum completions = 100 + 20 = 120
+        assert metrics["final_output_tokens"] == 120
+        assert metrics["final_input_tokens"] == 0  # clamped to 0
+
+    def test_skips_none_responses_for_last_step(self):
+        """Last step with response=None is skipped; uses previous step."""
+        trajectory = [
+            {"response": _make_response(100, 20)},
+            {"response": _make_response(200, 30)},
+            {"response": None},
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        # Last step with response is step 1: total = 230
+        # Sum completions from all steps with responses: 20 + 30 = 50
+        assert metrics["final_output_tokens"] == 50
+        assert metrics["final_input_tokens"] == 230 - 50
+
+    def test_skips_responses_without_usage(self):
+        """Responses with no .usage attribute are skipped entirely."""
+        no_usage = MagicMock()
+        no_usage.usage = None
+        trajectory = [
+            {"response": _make_response(100, 20)},
+            {"response": _make_response(200, 30)},
+            {"response": no_usage},  # last step, but no usage
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        # Should use step 1 (last with usage): total = 230
+        assert metrics["final_output_tokens"] == 50
+        assert metrics["final_input_tokens"] == 230 - 50
+
+    def test_all_responses_lack_usage(self):
+        """If no response has usage data, return zeros."""
+        no_usage = MagicMock()
+        no_usage.usage = None
+        trajectory = [
+            {"response": no_usage},
+            {"response": no_usage},
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        assert metrics["final_output_tokens"] == 0
+        assert metrics["final_input_tokens"] == 0
+
+    def test_final_input_tokens_clamped_to_zero(self):
+        """If sum of completions exceeds last step total, input is clamped to 0."""
+        trajectory = [
+            {"response": _make_response(10, 500)},  # huge completion
+            {"response": _make_response(50, 10)},
+        ]
+        metrics = compute_context_token_metrics(trajectory)
+        # Last step total = 60, sum completions = 510
+        assert metrics["final_output_tokens"] == 510
+        assert metrics["final_input_tokens"] == 0
+
+
+# =========================================================================
+# Metric classes
+# =========================================================================
+
+
+class TestContextTokenMetricClasses:
+    def test_input_tokens_metric(self):
+        from verifiers.utils.metric_utils import InputTokensMetric
+
+        m = InputTokensMetric()
+        m.add_output({"token_usage": {"input_tokens": 100.0}})
+        m.add_output({"token_usage": {"input_tokens": 200.0}})
+        assert m.compute() == pytest.approx(150.0)
+
+    def test_output_tokens_metric(self):
+        from verifiers.utils.metric_utils import OutputTokensMetric
+
+        m = OutputTokensMetric()
+        m.add_output({"token_usage": {"output_tokens": 40.0}})
+        m.add_output({"token_usage": {"output_tokens": 60.0}})
+        assert m.compute() == pytest.approx(50.0)
+
+    def test_final_input_tokens_metric(self):
+        from verifiers.utils.metric_utils import FinalInputTokensMetric
+
+        m = FinalInputTokensMetric()
+        m.add_output({"token_usage": {"final_input_tokens": 50.0}})
+        m.add_output({"token_usage": {"final_input_tokens": 100.0}})
+        assert m.compute() == pytest.approx(75.0)
+
+    def test_final_output_tokens_metric(self):
+        from verifiers.utils.metric_utils import FinalOutputTokensMetric
+
+        m = FinalOutputTokensMetric()
+        m.add_output({"token_usage": {"final_output_tokens": 150.0}})
+        m.add_output({"token_usage": {"final_output_tokens": 250.0}})
+        assert m.compute() == pytest.approx(200.0)
+
+    def test_skips_outputs_without_token_usage(self):
+        from verifiers.utils.metric_utils import FinalInputTokensMetric
+
+        m = FinalInputTokensMetric()
+        m.add_output({})
+        m.add_output({"token_usage": {}})
+        assert m.count == 0
+        assert m.compute() == 0.0

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/tests/test_environment_extra.py

@@ -237,7 +237,9 @@ async def test_state_to_output_uses_state_usage_not_trajectory(
     state["reward"] = 0.0
 
     output = state_to_output(state, state_columns=[])
-    assert output["token_usage"] == {"input_tokens": 5.0, "output_tokens": 4.0}
+    usage = output["token_usage"]
+    assert usage["input_tokens"] == 5.0
+    assert usage["output_tokens"] == 4.0
 
 
 @pytest.mark.asyncio

verifiers-0.1.13.dev3/tests/test_interception_utils.py

@@ -0,0 +1,133 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+from verifiers.errors import InfraError
+from verifiers.types import (
+    Response,
+    ResponseMessage,
+    TextContentPart,
+    ToolCall,
+    Usage,
+)
+from verifiers.utils import interception_utils
+from verifiers.utils.interception_utils import (
+    InterceptionServer,
+    StreamInterrupted,
+    create_empty_completion,
+    serialize_intercept_response,
+)
+
+
+def test_serialize_intercept_response_from_vf_response_uses_chat_completion_shape():
+    response = Response(
+        id="resp_1",
+        created=123,
+        model="test-model",
+        usage=Usage(
+            prompt_tokens=10,
+            reasoning_tokens=0,
+            completion_tokens=5,
+            total_tokens=15,
+        ),
+        message=ResponseMessage(
+            content=[TextContentPart(text="hello "), {"type": "text", "text": "world"}],
+            reasoning_content=None,
+            tool_calls=[
+                ToolCall(id="call_1", name="echo", arguments='{"x": 1}'),
+            ],
+            finish_reason="tool_calls",
+            is_truncated=False,
+            tokens=None,
+        ),
+    )
+
+    payload = serialize_intercept_response(response)
+
+    assert payload["id"] == "resp_1"
+    assert payload["object"] == "chat.completion"
+    assert payload["model"] == "test-model"
+    assert payload["choices"][0]["message"]["role"] == "assistant"
+    assert payload["choices"][0]["message"]["content"] == "hello world"
+    assert payload["choices"][0]["message"]["tool_calls"] == [
+        {
+            "id": "call_1",
+            "type": "function",
+            "function": {"name": "echo", "arguments": '{"x": 1}'},
+        }
+    ]
+    assert payload["choices"][0]["finish_reason"] == "tool_calls"
+    assert payload["usage"]["prompt_tokens"] == 10
+    assert payload["usage"]["completion_tokens"] == 5
+    assert payload["usage"]["total_tokens"] == 15
+
+
+def test_serialize_intercept_response_passthrough_native_chat_completion():
+    native = create_empty_completion("native-model")
+    payload = serialize_intercept_response(native)
+
+    assert payload["object"] == "chat.completion"
+    assert payload["model"] == "native-model"
+    assert len(payload["choices"]) == 1
+
+
+def test_set_rollout_error_attaches_stream_interrupted_to_state():
+    server = InterceptionServer(port=0)
+    state: dict = {}
+    server.register_rollout("r1", state=state)
+
+    err = StreamInterrupted("tunnel died")
+    server._set_rollout_error("r1", err)
+
+    assert state["error"] is err
+    assert isinstance(state["error"], InfraError)
+
+
+def test_set_rollout_error_does_not_clobber_existing_error():
+    # First error wins — later write failures must not hide the original cause.
+    server = InterceptionServer(port=0)
+    original = InfraError("original")
+    state: dict = {"error": original}
+    server.register_rollout("r1", state=state)
+
+    server._set_rollout_error("r1", StreamInterrupted("later"))
+
+    assert state["error"] is original
+
+
+async def test_streaming_write_failure_surfaces_to_state(monkeypatch):
+    """The real failure path: a mid-SSE transport close on the client side
+    raises out of ``response.write(...)``. The except branch must funnel
+    that into ``state["error"]`` so the rollout halts via ``has_error``."""
+    server = InterceptionServer(port=0)
+    state: dict = {}
+    server.register_rollout("r1", state=state)
+
+    # Mock StreamResponse whose second write raises (first write succeeds
+    # to prove we're in the streaming loop, not failing at prepare()).
+    writes: list[bytes] = []
+
+    async def fake_write(data: bytes) -> None:
+        writes.append(data)
+        if len(writes) >= 2:
+            raise ConnectionResetError("client closed transport")
+
+    fake_response = MagicMock()
+    fake_response.prepare = AsyncMock()
+    fake_response.write = AsyncMock(side_effect=fake_write)
+    fake_response.write_eof = AsyncMock()
+    monkeypatch.setattr(
+        interception_utils.web, "StreamResponse", lambda **_: fake_response
+    )
+
+    chunk_queue: asyncio.Queue = asyncio.Queue()
+    await chunk_queue.put({"choices": [{"delta": {"content": "hi"}}]})
+    await chunk_queue.put({"choices": [{"delta": {"content": " there"}}]})
+    intercept = {
+        "chunk_queue": chunk_queue,
+        "response_future": asyncio.Future(),
+    }
+
+    await server._handle_streaming_response(MagicMock(), "r1", intercept)
+
+    assert isinstance(state["error"], StreamInterrupted)
+    assert "ConnectionResetError" in str(state["error"])

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/tests/test_rlm_composable_env.py

@@ -6,6 +6,8 @@ fields and that the install script is generated correctly.
 
 import importlib
 import json
+from pathlib import Path
+import subprocess
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, call
 
@@ -18,9 +20,11 @@ from verifiers.envs.experimental.composable import (
     SandboxSpec,
     SandboxTaskSet,
 )
+from verifiers.envs.experimental.composable.harnesses import rlm as rlm_module
 from verifiers.envs.experimental.composable.harnesses.rlm import (
     build_install_script,
     rlm_harness,
+    resolve_local_checkout,
 )
 
 
@@ -86,28 +90,160 @@ def _make_temp_taskset_package(tmp_path, monkeypatch, *, with_skills: bool):
     return mod
 
 
+def _make_git_checkout(target: Path) -> Path:
+    checkout = target
+    checkout.mkdir()
+    (checkout / "install.sh").write_text("#!/usr/bin/env bash\n")
+    (checkout / "pyproject.toml").write_text("[project]\nname='rlm'\nversion='0.0.0'\n")
+    subprocess.run(["git", "init", "-b", "main"], cwd=checkout, check=True)
+    subprocess.run(
+        ["git", "add", "install.sh", "pyproject.toml"], cwd=checkout, check=True
+    )
+    subprocess.run(
+        [
+            "git",
+            "-c",
+            "user.name=Codex",
+            "-c",
+            "user.email=codex@example.com",
+            "commit",
+            "-m",
+            "init",
+        ],
+        cwd=checkout,
+        check=True,
+    )
+    return checkout
+
+
 # ── RLM harness ──────────────────────────────────────────────────────────
 
 
-def test_rlm_harness_install_script_downloads_repo_install_sh():
+def test_rlm_harness_install_script_requires_uploaded_checkout():
     script = build_install_script()
-    assert "git clone --depth 1 --branch main" in script
-    assert "github.com/PrimeIntellect-ai/rlm.git" in script
-    assert "bash /tmp/rlm-checkout/install.sh" in script
+    assert 'test -f "$RLM_CHECKOUT_PATH/install.sh"' in script
+    assert "git clone" not in script
+    assert 'bash "$RLM_CHECKOUT_PATH/install.sh"' in script
 
 
-def test_rlm_harness_sets_metrics_fields():
-    harness = rlm_harness()
+def test_rlm_harness_sets_metrics_fields(tmp_path):
+    harness = rlm_harness(local_checkout=_make_git_checkout(tmp_path / "rlm"))
     assert harness.metrics_path == "{workdir}/.rlm/sessions/*/meta.json"
     assert harness.metrics_key == "metrics"
     assert harness.metrics_prefix == "rlm_"
 
 
-def test_rlm_harness_sets_skills_path():
-    harness = rlm_harness()
+def test_rlm_harness_sets_skills_path(tmp_path):
+    harness = rlm_harness(local_checkout=_make_git_checkout(tmp_path / "rlm"))
     assert harness.skills_path == "/task/rlm-skills"
 
 
+def test_resolve_local_checkout_validates_explicit_path(tmp_path):
+    checkout = _make_git_checkout(tmp_path / "rlm")
+
+    resolved = resolve_local_checkout(checkout)
+
+    assert resolved == checkout.resolve()
+
+
+def test_rlm_harness_uploads_explicit_local_checkout(tmp_path):
+    checkout = _make_git_checkout(tmp_path / "rlm")
+
+    harness = rlm_harness(local_checkout=checkout)
+
+    assert harness.get_upload_dirs is not None
+    assert harness.get_upload_dirs() == {"rlm_checkout": checkout.resolve()}
+    assert harness.upload_dir_mapping == {"rlm_checkout": "/tmp/rlm-checkout"}
+
+
+def test_resolve_local_checkout_materializes_host_cache(tmp_path):
+    source_checkout = _make_git_checkout(tmp_path / "rlm-source")
+    checkout_dir = tmp_path / "checkout-root" / "rlm"
+
+    resolved = resolve_local_checkout(
+        local_checkout=checkout_dir,
+        rlm_repo_url=str(source_checkout),
+        rlm_branch="main",
+    )
+
+    assert resolved == checkout_dir.resolve()
+    assert (checkout_dir / ".git").is_dir()
+    assert (checkout_dir / "install.sh").is_file()
+    assert (checkout_dir / "pyproject.toml").is_file()
+
+
+def test_rlm_harness_uses_default_host_cache_when_local_checkout_unspecified(
+    tmp_path, monkeypatch
+):
+    source_checkout = _make_git_checkout(tmp_path / "rlm-source")
+    monkeypatch.setattr(
+        rlm_module,
+        "DEFAULT_RLM_LOCAL_CHECKOUT_CACHE_ROOT",
+        tmp_path / "cache-root",
+    )
+
+    harness = rlm_harness(
+        rlm_repo_url=str(source_checkout),
+        rlm_branch="main",
+    )
+
+    assert harness.get_upload_dirs is not None
+    upload_checkout = harness.get_upload_dirs()["rlm_checkout"]
+    assert isinstance(upload_checkout, Path)
+    assert upload_checkout.is_dir()
+    assert upload_checkout.name.startswith("rlm-source-main-")
+    assert harness.upload_dir_mapping == {"rlm_checkout": "/tmp/rlm-checkout"}
+
+
+def test_rlm_harness_always_uploads_checkout(tmp_path, monkeypatch):
+    source_checkout = _make_git_checkout(tmp_path / "rlm-source")
+    monkeypatch.setattr(
+        rlm_module,
+        "DEFAULT_RLM_LOCAL_CHECKOUT_CACHE_ROOT",
+        tmp_path / "cache-root",
+    )
+
+    harness = rlm_harness(
+        rlm_repo_url=str(source_checkout),
+        rlm_branch="main",
+    )
+
+    assert harness.get_upload_dirs is not None
+    assert harness.upload_dir_mapping is not None
+
+
+def test_resolve_local_checkout_redacts_gh_token_on_clone_failure(
+    tmp_path, monkeypatch
+):
+    failing_checkout = tmp_path / "checkout-root" / "rlm"
+    token = "super/secret token"
+    quoted_token = "super%2Fsecret%20token"
+
+    def _raise_clone_error(*args, **kwargs):
+        raise subprocess.CalledProcessError(
+            128,
+            args[0],
+            stderr=(
+                "fatal: could not read from "
+                f"https://{quoted_token}@github.com/PrimeIntellect-ai/rlm.git"
+            ),
+        )
+
+    monkeypatch.setattr(rlm_module.subprocess, "run", _raise_clone_error)
+
+    with pytest.raises(RuntimeError) as exc_info:
+        resolve_local_checkout(
+            local_checkout=failing_checkout,
+            rlm_repo_url="github.com/PrimeIntellect-ai/rlm.git",
+            rlm_branch="main",
+            gh_token=token,
+        )
+
+    message = str(exc_info.value)
+    assert token not in message
+    assert "<redacted>" in message
+
+
 # ── install_env ──────────────────────────────────────────────────────────
 
 
@@ -201,7 +337,7 @@ async def test_rlm_uploads_skills_before_install(tmp_path, monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_rlm_collects_logs_and_metrics():
+async def test_rlm_collects_logs_and_metrics(tmp_path):
     taskset = MockSandboxTaskSet(dataset=_make_dataset(), name="test")
     metrics = {
         "turns": 3,
@@ -209,7 +345,7 @@ async def test_rlm_collects_logs_and_metrics():
         "prompt_tokens": 100,
         "completion_tokens": 25,
     }
-    harness = rlm_harness()
+    harness = rlm_harness(local_checkout=_make_git_checkout(tmp_path / "rlm"))
     env = ComposableEnv(
         taskset=taskset,
         harness=Harness(

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/verifiers/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.13.dev1"
+__version__ = "0.1.13.dev3"
 
 import importlib
 import os

{verifiers-0.1.13.dev1 → verifiers-0.1.13.dev3}/verifiers/envs/environment.py

@@ -483,14 +483,12 @@ class Environment(ABC):
         usage = state.get("usage")
         if isinstance(usage, Mapping):
             try:
-                input_tokens = float(usage.get("input_tokens", 0.0))
-                output_tokens = float(usage.get("output_tokens", 0.0))
+                return {
+                    "input_tokens": float(usage.get("input_tokens", 0.0)),
+                    "output_tokens": float(usage.get("output_tokens", 0.0)),
+                }
             except (TypeError, ValueError):
                 return None
-            return {
-                "input_tokens": input_tokens,
-                "output_tokens": output_tokens,
-            }
         return None
 
     async def get_model_response(