feat: parse vLLM token ids + logprobs in the chat client

[mikasenghaas]

Summary

• The openai_chat_completions client now best-effort parses the prompt/completion token ids and sampling logprobs that vLLM returns (return_token_ids + logprobs) into Response.tokens.
• This lets MITO training (no renderer) train on real on-policy tokens instead of re-tokenizing messages downstream — the chat client now carries TurnTokens just like the renderer client does.
• New tokens_from_wire(completion, choice): reads choice.token_ids (completion ids), completion.prompt_token_ids (prompt ids), and one logprobs.content[i].logprob per generated token.
• Pure pass-through on the request side: sampling args are forwarded unchanged (the caller, e.g. prime-rl, opts in by sending logprobs + return_token_ids). tokens stays None when the provider returns neither (eval, non-vLLM providers).
• Doc fixes on TurnTokens / Response.tokens to reflect that the chat client can populate them too.

Base is the feat/nano-as-v1 vf branch.

Verification

Confirmed end-to-end against a running vLLM server:

• Dummy chat request with prime-rl's training sampling args (temperature, top_p=1.0, logprobs=True) + return_token_ids=True: prompt_token_ids (== usage.prompt_tokens), choice.token_ids (== usage.completion_tokens, ends in <|im_end|>), and one logprob per completion token — all aligned and accessible via the OpenAI SDK objects.
• 20-step reverse-text-v1 RL run in MITO mode with the orchestrator-side backfill disabled, so tokens could only come from this parser: Trainable 128/128 (100%) every step, 0 "No trainable samples" warnings, reward 0.16 → 0.78, 0% errors. A saved rollout's response.tokens decoded back to the correct prompt and a real reversed-text answer.

Note

Parse vLLM token ids and logprobs into Response.tokens in the OpenAI chat client

Adds a tokens_from_wire helper in openai.py that extracts choice.token_ids, completion.prompt_token_ids, and choice.logprobs from a vLLM completion response and constructs a TurnTokens object. response_from_wire now calls this helper to populate Response.tokens when the provider returns token ids; previously tokens was always None in this client. Returns None when token ids are absent, so non-vLLM providers are unaffected.

^{Macroscope summarized 89368c8.}

---

Note

Low Risk
Additive, best-effort parsing on the response path only; providers without token ids behave as before with tokens=None.

Overview
The OpenAI chat-completions client now best-effort fills Response.tokens from vLLM when the caller enables logprobs and return_token_ids, so MITO training can use on-policy token ids and sampling logprobs without a renderer client.

A new tokens_from_wire helper reads choice.token_ids, completion.prompt_token_ids, and per-token logprobs from choice.logprobs.content; response_from_wire attaches the result to every response. If completion token ids are missing (eval or non-vLLM providers), tokens stays None—request forwarding is unchanged.

Docstrings on TurnTokens and Response.tokens now state the chat client can populate them via vLLM, not only the renderer path.

^{Reviewed by Cursor Bugbot for commit 89368c8. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 89368c8. Configure here.}

[cursor]

Partial tokens without prompt ids

Medium Severity

tokens_from_wire builds TurnTokens whenever choice.token_ids is present, but missing completion.prompt_token_ids becomes an empty prompt_ids list instead of skipping tokens. Downstream MITO training can treat the rollout as on-policy while the prompt side is empty.

 

^{Reviewed by Cursor Bugbot for commit 89368c8. Configure here.}

[macroscopeapp]

Approvability

Verdict: Needs human review

Unresolved review comment identifies a potential bug where missing prompt_token_ids leads to empty prompt IDs being passed to training, which could affect MITO training correctness. This warrants human review to determine if the behavior is intentional or needs a fix.

^{You can customize Macroscope's approvability policy. Learn more.}