verified rollouts

Author: benjibc

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -139,16 +139,22 @@ async def process_row(row: EvaluationRow) -> EvaluationRow:
                     tool_calls=converted_tool_calls,
                 )
             ]
-            row.execution_metadata.usage = (
-                CompletionUsage(  # Note: LiteLLM sets usage dynamically via setattr(), not as a typed field
-                    prompt_tokens=response.usage.prompt_tokens,  # pyright: ignore[reportAttributeAccessIssue]
-                    completion_tokens=response.usage.completion_tokens,  # pyright: ignore[reportAttributeAccessIssue]
-                    total_tokens=response.usage.total_tokens,  # pyright: ignore[reportAttributeAccessIssue]
-                )
-            )
-
             row.messages = messages
 
+            usage = getattr(response, "usage", None)
+            if usage is not None:
+                prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
+                completion_tokens = getattr(usage, "completion_tokens", 0) or 0
+                total_tokens = getattr(usage, "total_tokens", None)
+                if total_tokens is None:
+                    total_tokens = prompt_tokens + completion_tokens
+
+                row.execution_metadata.usage = CompletionUsage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=total_tokens,
+                )
+
             row.execution_metadata.duration_seconds = time.perf_counter() - start_time
 
             default_logger.log(row)

tests/pytest/test_single_turn_rollout_processor.py

@@ -116,3 +116,50 @@ async def fake_acompletion(**kwargs):
     assert [m["role"] for m in sent_msgs] == ["user", "assistant"]
     assert [m.role for m in out.messages] == ["user", "assistant", "assistant"]
     assert out.messages[-1].content == "Hello again"
+
+
+@pytest.mark.asyncio
+async def test_single_turn_handles_missing_usage_block(monkeypatch):
+    row = EvaluationRow(messages=[Message(role="user", content="Describe the picture")])
+
+    import eval_protocol.pytest.default_single_turn_rollout_process as mod
+
+    class StubChoices:
+        pass
+
+    class StubModelResponse:
+        def __init__(self, text: str):
+            self.choices = [StubChoices()]
+            self.choices[0].message = SimpleNamespace(content=text, tool_calls=None)
+            self.usage = None
+
+    async def fake_acompletion(**kwargs):
+        return StubModelResponse(text="It looks like creme brulee")
+
+    class StubLogger:
+        def __init__(self):
+            self.logged = []
+
+        def log(self, row):
+            self.logged.append(row)
+
+        def read(self, rollout_id=None):
+            return list(self.logged)
+
+    stub_logger = StubLogger()
+
+    monkeypatch.setattr(mod, "ModelResponse", StubModelResponse, raising=True)
+    monkeypatch.setattr(mod, "Choices", StubChoices, raising=True)
+    monkeypatch.setattr(mod, "acompletion", fake_acompletion, raising=True)
+    monkeypatch.setattr(mod, "default_logger", stub_logger, raising=False)
+
+    processor = SingleTurnRolloutProcessor()
+    config = _DummyConfig()
+
+    tasks = processor([row], config)
+    out = await tasks[0]
+
+    assert [m.role for m in out.messages] == ["user", "assistant"]
+    assert out.messages[-1].content == "It looks like creme brulee"
+    # Usage should remain unset when the provider omits it
+    assert out.execution_metadata.usage is None