From 89368c899f65b5842775a8612ebf5da0602daf6a Mon Sep 17 00:00:00 2001
From: Mika Senghaas <mail@mikasenghaas.de>
Date: Tue, 9 Jun 2026 18:18:45 +0000
Subject: [PATCH] feat(v1): parse vLLM token ids + logprobs in the chat client

The openai_chat_completions client now best-effort parses the prompt and
completion token ids and sampling logprobs that vLLM returns (return_token_ids
+ logprobs) into Response.tokens, so MITO training (no renderer) can train on
real on-policy tokens instead of re-tokenizing the messages downstream.

Sampling args still pass straight through; tokens stay None when the provider
returns neither token ids nor logprobs (e.g. eval, or non-vLLM providers).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 verifiers/v1/clients/openai.py | 29 ++++++++++++++++++++++++++---
 verifiers/v1/types.py          |  7 ++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/verifiers/v1/clients/openai.py b/verifiers/v1/clients/openai.py
index b9b5b338da..92af79c815 100644
--- a/verifiers/v1/clients/openai.py
+++ b/verifiers/v1/clients/openai.py
@@ -1,9 +1,11 @@
 """OpenAI-compatible chat-completions client.
 
 Distilled from v1's 545-line client: message<->wire translation, tool schemas,
-best-effort reasoning_content. Token-id/logprob/routed-experts/audio handling is
-dropped (training-only). This is the one place raw provider dicts cross into our
-typed `Response`.
+best-effort reasoning_content. Sampling args pass straight through; when the
+response carries vLLM's token ids + sampling logprobs (the caller asked for
+`logprobs` and `return_token_ids`), we parse them into the response's `tokens`
+so MITO training needs no renderer. Routed-experts/audio handling stays dropped.
+This is the one place raw provider dicts cross into our typed `Response`.
 """
 
 from openai import AsyncOpenAI, OpenAIError
@@ -19,6 +21,7 @@
     SamplingConfig,
     Tool,
     ToolCall,
+    TurnTokens,
     Usage,
 )
 
@@ -58,6 +61,25 @@ def tool_to_wire(tool: Tool) -> dict:
     return {"type": "function", "function": function}
 
 
+def tokens_from_wire(completion, choice) -> TurnTokens | None:
+    """Parse vLLM's token ids + sampling logprobs into `TurnTokens`, for training.
+
+    vLLM surfaces the completion ids on the choice (`return_token_ids`), the prompt
+    ids on the completion, and the sampled logprobs as one `logprobs.content` entry
+    per generated token (`logprobs=True`). All are absent on providers that don't
+    return them, so this is best-effort: no completion ids means no `tokens`.
+    """
+    completion_ids = getattr(choice, "token_ids", None)
+    if not completion_ids:
+        return None
+    content = choice.logprobs.content if choice.logprobs else None
+    return TurnTokens(
+        prompt_ids=list(getattr(completion, "prompt_token_ids", None) or []),
+        completion_ids=list(completion_ids),
+        completion_logprobs=[lp.logprob for lp in content] if content else [],
+    )
+
+
 def response_from_wire(completion) -> Response:
     choice = completion.choices[0]
     message = choice.message
@@ -87,6 +109,7 @@ def response_from_wire(completion) -> Response:
         ),
         finish_reason=finish,
         usage=usage,
+        tokens=tokens_from_wire(completion, choice),
     )
 
 
diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py
index 2e705d1754..64321e1574 100644
--- a/verifiers/v1/types.py
+++ b/verifiers/v1/types.py
@@ -97,8 +97,9 @@ def total_tokens(self) -> int:
 
 
 class TurnTokens(StrictBaseModel):
-    """Token ids + sampling logprobs for one response, for training. Populated only
-    by the renderer client (which tokenizes client-side); None otherwise."""
+    """Token ids + sampling logprobs for one response, for training. Populated by the
+    renderer client (client-side tokenization) or the chat client (parsed from vLLM's
+    token ids); None when the provider returns neither."""
 
     prompt_ids: list[int] = Field(default_factory=list)
     completion_ids: list[int] = Field(default_factory=list)
@@ -115,7 +116,7 @@ class Response(StrictBaseModel):
     finish_reason: FinishReason
     usage: Usage | None = None
     tokens: TurnTokens | None = None
-    """Client-side token ids + logprobs (renderer client only)."""
+    """Token ids + logprobs for training (renderer client, or chat client via vLLM)."""
 
 
 # --- sampling -----------------------------------------------------------------