diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index bdd1f313bc..d735f0fca9 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -16,7 +16,7 @@ import pytest from grid import agent_seats, runtime_config -from verifiers.v1.clients import EvalClientConfig, ModelContext +from verifiers.v1.clients import ClientConfig, ModelContext from verifiers.v1.configs.cli.eval import EvalConfig from verifiers.v1.configs.harness import HarnessConfig from verifiers.v1.episode import Episode @@ -201,7 +201,7 @@ async def live_ctx(): # Endpoint config only — each rollout builds and closes its own client. yield ModelContext( model=CI_MODEL, - client=EvalClientConfig(), + client=ClientConfig(), sampling=SamplingConfig(max_tokens=2048, reasoning_effort=CI_REASONING_EFFORT), ) @@ -283,7 +283,7 @@ async def answer(request: web.Request) -> web.StreamResponse: if request.path.endswith(suffix) ) raw = build(str(body.get("model", "scripted"))) - if not dialect.streaming(body): + if not body.get("stream"): return web.json_response(raw) response = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) await response.prepare(request) diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index a4d8a25779..7ef32ac677 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -539,20 +539,13 @@ def second_turn(trace, prompt_ids): ] -def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): +def test_expanded_prompt_tokens_and_attribution_are_reused(): trace = vf.Trace( agent=vf.AgentInfo(config=vf.AgentConfig()), task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")), ) user = vf.UserMessage(content="image") assistant = vf.AssistantMessage(content="looked") - assert graph.MessageNode(message=user, token_ids=[9]).logical_ids == [9] - assert ( - graph.MessageNode( - message=user, token_ids=[9], renderer_token_ids=[] - ).logical_ids - == [] - ) graph.prepare_turn(trace, [user]).commit( vf.Response( id="a", @@ -562,11 +555,9 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): finish_reason="stop", tokens=TurnTokens( prompt_ids=[1, 9, 9, 3], - renderer_prompt_ids=[1, 9, 3], completion_ids=[4], - message_spans=[(0, 2)], - is_content=[False, True, False], - mm_placeholders=[(1, 2)], + message_spans=[(0, 3)], + is_content=[False, True, True, False], mm_token_type_id_map={9: 1}, ), ) @@ -577,25 +568,10 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): assert trace.nodes[0].token_ids == [1, 9, 9] assert trace.nodes[0].is_content == [False, True, True] turn = graph.prepare_turn(trace, [user, assistant, vf.UserMessage(content="next")]) - assert turn.previous_renderer_token_ids() == ([1, 9, 3], [4]) - - with pytest.raises(ValueError, match="exactly extend"): - turn.commit( - vf.Response( - id="b", - created=0, - model="t", - message=vf.AssistantMessage(content="bad"), - finish_reason="stop", - tokens=TurnTokens( - prompt_ids=[1, 9, 8, 3, 4, 5], - renderer_prompt_ids=[1, 9, 3, 4, 5], - bridged=True, - completion_ids=[6], - message_spans=[None, None, (4, 5)], - ), - ) - ) + assert [trace.nodes[nid].token_ids for nid in turn.prefix_node_ids] == [ + [1, 9, 9], + [3, 4], + ] tool = vf.ToolMessage(content="result", tool_call_id="call_0") graph.prepare_turn(trace, [user, assistant, tool]).commit( @@ -607,12 +583,9 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): finish_reason="stop", tokens=TurnTokens( prompt_ids=[1, 9, 9, 3, 4, 5, 6, 7], - renderer_prompt_ids=[1, 9, 3, 4, 5, 6, 7], - bridged=True, completion_ids=[8], - message_spans=[None, None, (4, 6)], - is_content=[False, True, False, True, False, True, False], - mm_placeholders=[(1, 2)], + message_spans=[None, None, (5, 7)], + is_content=[False, True, True, False, True, False, True, False], mm_token_type_id_map={9: 1}, ), ) @@ -621,7 +594,6 @@ def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens(): tool_node = trace.nodes[-2] assert tool_node.message == tool assert tool_node.token_ids == [5, 6] - assert tool_node.renderer_token_ids == [5, 6] assert tool_node.is_content == [False, True] diff --git a/tests/unit/test_trace.py b/tests/unit/test_trace.py index bf1b35a90f..d152f6b4ee 100644 --- a/tests/unit/test_trace.py +++ b/tests/unit/test_trace.py @@ -243,12 +243,12 @@ def test_transcript(make_full_trace): assert "SECRET REASONING" not in transcript # reasoning is excluded -def test_sampling_mask_from_packed_matches_lists(): +def test_packed_training_records_match_native_arrays(): import base64 import numpy as np - from verifiers.v1.types import SamplingMask + from verifiers.v1.dialects.chat import response_from_wire rows = [[4, 7], [], [9]] ids, counts = np.array([4, 7, 9], np.int32), np.array([2, 0, 1], np.int32) @@ -260,9 +260,46 @@ def test_sampling_mask_from_packed_matches_lists(): } for name, a in (("ids", ids), ("counts", counts)) } - for mask in ( - SamplingMask.from_sampling_mask(rows), - SamplingMask.from_sampling_mask(packed), + logprobs = np.array([-0.5, -1.0, -1.5], np.float32) + for mask_payload, logprob_field in ( + ( + rows, + { + "logprobs": { + "content": [ + {"token": f"token_id:{token_id}", "logprob": float(logprob)} + for token_id, logprob in zip([4, 7, 9], logprobs, strict=True) + ] + } + }, + ), + ( + packed, + { + "completion_logprobs": { + "data": base64.b64encode(logprobs.tobytes()).decode(), + "shape": [3], + "dtype": "float32", + } + }, + ), ): + response = response_from_wire( + { + "prompt_token_ids": [1, 2], + "choices": [ + { + "message": {"role": "assistant", "content": "done"}, + "token_ids": [4, 7, 9], + "sampling_mask": mask_payload, + **logprob_field, + } + ], + } + ) + assert response.tokens is not None + np.testing.assert_array_equal(response.tokens.completion_logprobs, logprobs) + mask = response.tokens.sampling_mask + assert mask is not None np.testing.assert_array_equal(mask.ids, ids) np.testing.assert_array_equal(mask.counts, counts) diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 0d7ca105a7..42e8aa48d1 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -5,13 +5,9 @@ from verifiers.v1.acp import ACPConfig, ACPHarness, ACPTurn from verifiers.v1.agent import Agent, Agents, Interaction, Segment, make_agent from verifiers.v1.clients import ( - BaseClientConfig, Client, ClientConfig, - EvalClientConfig, ModelContext, - TrainClientConfig, - resolve_client, ) from verifiers.v1.configs.agent import AgentConfig from verifiers.v1.configs.cli.env import ( @@ -287,11 +283,7 @@ "TunnelError", # clients "Client", - "BaseClientConfig", "ClientConfig", - "EvalClientConfig", - "TrainClientConfig", - "resolve_client", # taskset / harness / runtime / environment "Taskset", "TaskConfig", diff --git a/verifiers/v1/agent.py b/verifiers/v1/agent.py index 0bd6d960cb..f472b7d042 100644 --- a/verifiers/v1/agent.py +++ b/verifiers/v1/agent.py @@ -19,7 +19,7 @@ from typing_extensions import TypeVar from verifiers.v1.clients import ( - EvalClientConfig, + ClientConfig, ModelContext, ) from verifiers.v1.configs.agent import ( @@ -565,7 +565,7 @@ def __init__( self.harness = load_harness(config.harness) self.ctx = ModelContext( model=config.model, - client=config.client or EvalClientConfig(), + client=config.client or ClientConfig(), sampling=config.sampling, ) self._closed = False diff --git a/verifiers/v1/clients/__init__.py b/verifiers/v1/clients/__init__.py index 0d2f8fcc19..290d9df05c 100644 --- a/verifiers/v1/clients/__init__.py +++ b/verifiers/v1/clients/__init__.py @@ -1,21 +1,4 @@ -from verifiers.v1.clients.client import Client, ModelContext, resolve_client -from verifiers.v1.clients.eval import EvalClient -from verifiers.v1.clients.train import TrainClient -from verifiers.v1.configs.client import ( - BaseClientConfig, - ClientConfig, - EvalClientConfig, - TrainClientConfig, -) +from verifiers.v1.clients.client import Client, ModelContext +from verifiers.v1.configs.client import ClientConfig -__all__ = [ - "BaseClientConfig", - "Client", - "ClientConfig", - "EvalClient", - "EvalClientConfig", - "ModelContext", - "TrainClient", - "TrainClientConfig", - "resolve_client", -] +__all__ = ["Client", "ClientConfig", "ModelContext"] diff --git a/verifiers/v1/clients/base.py b/verifiers/v1/clients/base.py index ae8ec5f6e8..abf938af9a 100644 --- a/verifiers/v1/clients/base.py +++ b/verifiers/v1/clients/base.py @@ -6,7 +6,7 @@ from openai import AsyncOpenAI from verifiers.v1.configs.client import ( - BaseClientConfig, + ClientConfig, resolve_api_key, resolve_headers, ) @@ -23,13 +23,13 @@ VERSION_SEGMENT = re.compile(r"v\d+") -def build_async_httpx(config: BaseClientConfig) -> httpx.AsyncClient: +def build_async_httpx(config: ClientConfig) -> httpx.AsyncClient: return httpx.AsyncClient( timeout=httpx.Timeout(**config.timeout.model_dump()), limits=DEFAULT_LIMITS ) -def build_async_openai(config: BaseClientConfig) -> AsyncOpenAI: +def build_async_openai(config: ClientConfig) -> AsyncOpenAI: http_client = build_async_httpx(config) # Pass it explicitly: SDK defaults can replace a transport timeout equal to HTTPX's. return AsyncOpenAI( diff --git a/verifiers/v1/clients/client.py b/verifiers/v1/clients/client.py index 24d17a888c..43876e3eec 100644 --- a/verifiers/v1/clients/client.py +++ b/verifiers/v1/clients/client.py @@ -1,35 +1,256 @@ -"""Client interfaces for model inference and relay.""" +"""SDK inference transport shared by evaluation and training.""" -from abc import ABC, abstractmethod +import asyncio from collections.abc import Mapping from dataclasses import dataclass, field -from verifiers.v1.configs.client import ( - BaseClientConfig, - ClientConfig, - TrainClientConfig, +import anthropic +import httpx2 as httpx +import openai + +from verifiers.v1.clients.base import ( + MAX_RETRIES, + SESSION_ID_HEADER, + build_async_httpx, + join_url, +) +from verifiers.v1.configs.client import ClientConfig, resolve_api_key, resolve_headers +from verifiers.v1.dialects import AnthropicDialect, ChatDialect, Dialect +from verifiers.v1.errors import RolloutError, model_error +from verifiers.v1.semantic import ACP_EXTENSION_HEADERS +from verifiers.v1.types import Response, Sampling + +# Local authentication, framing, and signatures must not reach the provider. +_BLOCKED_REQUEST_HEADERS = ( + frozenset( + { + # The harness uses this rollout secret to authenticate with the localhost server. + # The SDK supplies provider authorization after filtering. + "authorization", + "x-api-key", + "x-stainless-raw-response", + # HTTPX recalculates these for the provider URL, JSON bytes, and supported decoders. + "accept-encoding", + "content-encoding", + "content-length", + "content-type", + "host", + "transfer-encoding", + # These control only the localhost HTTP exchange. + "expect", + "keep-alive", + "proxy-authorization", + "proxy-connection", + "te", + "trailer", + "upgrade", + # Provider affinity must not compete with the rollout-wide session header below. + "session_id", + # The eval owns the model and sampling settings, so it changes those JSON fields before + # sending upstream. Hashes and signatures calculated from the intercepted body are stale. + "content-digest", + "content-md5", + "digest", + "repr-digest", + "signature", + "signature-input", + } + ) + | ACP_EXTENSION_HEADERS ) -from verifiers.v1.dialects import Dialect -from verifiers.v1.graph import PendingTurn -from verifiers.v1.types import Response, Sampling, SamplingConfig -class Client(ABC): - @abstractmethod +class Client: + """Preserve native payloads; project SDK responses into the trace.""" + + def __init__(self, config: ClientConfig) -> None: + self.base_url = config.base_url + self.headers = resolve_headers(config) + http_client = build_async_httpx(config) + options = { + "base_url": config.base_url, + "api_key": resolve_api_key(config), + "max_retries": MAX_RETRIES, + "timeout": http_client.timeout, + "http_client": http_client, + } + self.openai = openai.AsyncOpenAI(**options) + self.anthropic = anthropic.AsyncAnthropic(**options) + async def _complete( self, dialect: Dialect, body: dict, - sampling: SamplingConfig, session_id: str | None = None, - turn: PendingTurn | None = None, headers: Mapping[str, str] | None = None, ) -> tuple[Response, bytes | None]: - """Complete an effective native request, retaining provider SSE for atomic replay. + # Exact token records require JSON; the interception server can emit SSE after commit. + if body.get("return_token_ids") or body.get("return_training_metadata"): + body = {**body, "stream": False} + body.pop("stream_options", None) + client, reply = await self._request( + dialect, dialect.upstream_path, body, headers, session_id + ) + try: + events = None + if body.get("stream"): + # Parse during receipt, retaining decoded bytes for atomic delivery/replay. + # Bounded chunks keep cached or coalesced responses from starving other turns. + recorded = bytearray() + chunks = reply.iter_bytes(chunk_size=16_384) + + async def capture(): + async for chunk in chunks: + recorded.extend(chunk) + yield chunk + await asyncio.sleep(0) + + stream_type = ( + anthropic.AsyncStream + if isinstance(client, anthropic.AsyncAnthropic) + else openai.AsyncStream + ) + async with stream_type( + cast_to=dialect.event_type, + response=httpx.Response( + reply.status_code, + request=reply.http_response.request, + content=capture(), + ), + client=client, + ) as stream: + raw = await dialect.read_stream(stream) + # Chat's SDK stops at [DONE]; retain any trailing provider bytes too. + async for chunk in chunks: + recorded.extend(chunk) + await asyncio.sleep(0) + events = bytes(recorded) + if isinstance(dialect, ChatDialect) and not await asyncio.to_thread( + any, + ( + event.data == "[DONE]" + for event in anthropic.Stream.raw_events( + httpx.Response(200, content=events) + ) + ), + ): + raise ValueError("Chat stream ended without [DONE]") + else: + raw = await reply.json() + response = dialect.parse_response(raw) + if body.get("return_token_ids") or body.get("return_training_metadata"): + if response.tokens is None: + raise ValueError("the inference endpoint did not return token IDs") + if body.get("logprobs") and len( + response.tokens.completion_logprobs + ) != len(response.tokens.completion_ids): + raise ValueError( + "the inference endpoint did not return completion logprobs" + ) + if body.get("return_training_metadata") and ( + response.tokens is None or response.tokens.message_spans is None + ): + raise ValueError( + "the inference endpoint did not return training metadata" + ) + if ( + response.tokens is not None + and response.tokens.message_spans is not None + and len(response.tokens.message_spans) + != len(dialect.parse_request(body)[0].messages) + ): + raise ValueError( + "message spans must have one entry per request message" + ) + except ( + openai.OpenAIError, + anthropic.AnthropicError, + httpx.HTTPError, + ConnectionResetError, + ) as error: + raise model_error(error) from error + except RolloutError: + raise + except Exception as error: + # SDK accumulation can reject malformed event sequences. + raise model_error( + f"malformed upstream response: {error}", status_code=502 + ) from error + finally: + await reply.close() + # Training attribution belongs to the trace, not the harness response. + raw.pop("training_metadata", None) + # Training records stay in the trace; native token-ID requests keep their arrays. + if body.get("return_training_metadata"): + raw.pop("prompt_token_ids", None) + for choice in raw.get("choices", []): + for field in ( + "token_ids", + "completion_logprobs", + "sampling_mask", + "routed_experts", + ): + choice.pop(field, None) + response.raw = raw + return response, events + + def _headers( + self, + incoming: Mapping[str, str] | None, + session_id: str | None, + ) -> dict[str, str]: + """Retain feature headers; let the SDK supply provider auth and HTTP framing.""" + headers = httpx.Headers(incoming) + connection = headers.pop("connection", "") + for name in _BLOCKED_REQUEST_HEADERS | set( + map(str.strip, connection.lower().split(",")) + ): + headers.pop(name, None) + headers.update(self.headers) + # Configured headers cannot replace the provider key with a local rollout secret. + headers.pop("authorization", None) + headers.pop("x-api-key", None) + if session_id: + headers[SESSION_ID_HEADER] = session_id + return dict(headers) - Eval forwards the request; training renders its graph-resolved prompt to tokens. - The response always carries native JSON; generated responses have no provider SSE. - """ + async def _request( + self, + dialect: Dialect, + route: str, + body: dict, + headers: Mapping[str, str] | None, + session_id: str | None = None, + ): + client, response_type = ( + (self.anthropic, anthropic.AsyncAPIResponse) + if isinstance(dialect, AnthropicDialect) + else (self.openai, openai.AsyncAPIResponse) + ) + request_headers = self._headers(headers, session_id) + # SDK defaults are plain dicts: match their casing so forwarded headers replace + # them instead of producing duplicate User-Agent or provider feature headers. + default_names = {name.lower(): name for name in client.default_headers} + request_headers = { + default_names.get(name.lower(), name): value + for name, value in request_headers.items() + } + if body.get("stream"): + # Ask the SDK for its unread response wrapper instead of a typed SSE iterator. + request_headers["X-Stainless-Raw-Response"] = "stream" + try: + # Low-level SDK requests preserve extension fields without signature filtering. + # The raw wrapper retains bytes and supports the SDK's permissive typed parse. + reply = await client.post( + join_url(self.base_url, route), + body=body, + cast_to=response_type[dict], + options={"headers": request_headers}, + stream=bool(body.get("stream")), + ) + except (openai.OpenAIError, anthropic.AnthropicError) as error: + raise model_error(error) from error + return client, reply async def relay_aux( self, @@ -38,22 +259,15 @@ async def relay_aux( body: dict, headers: Mapping[str, str] | None = None, ) -> dict: - """Relay a non-model-turn side request (an `aux_route`, e.g. Anthropic's `count_tokens`) - as native JSON and return the provider JSON. Only the relay (eval) client supports it.""" - raise NotImplementedError(f"{type(self).__name__} does not relay aux routes") + _, reply = await self._request(dialect, route, body, headers) + try: + return await reply.json() + finally: + await reply.close() async def close(self) -> None: - pass - - -def resolve_client(config: BaseClientConfig) -> Client: - if isinstance(config, TrainClientConfig): - from verifiers.v1.clients.train import TrainClient - - return TrainClient(config) - from verifiers.v1.clients.eval import EvalClient - - return EvalClient(config) + await self.openai.close() + await self.anthropic.close() @dataclass(frozen=True) diff --git a/verifiers/v1/clients/eval.py b/verifiers/v1/clients/eval.py deleted file mode 100644 index f391cde401..0000000000 --- a/verifiers/v1/clients/eval.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Forward native evaluation requests through the provider SDKs.""" - -import asyncio -from collections.abc import Mapping - -import anthropic -import httpx2 as httpx -import openai - -from verifiers.v1.clients.base import ( - MAX_RETRIES, - SESSION_ID_HEADER, - build_async_httpx, - join_url, -) -from verifiers.v1.clients.client import Client -from verifiers.v1.configs.client import ( - BaseClientConfig, - resolve_api_key, - resolve_headers, -) -from verifiers.v1.dialects import AnthropicDialect, ChatDialect, Dialect -from verifiers.v1.errors import RolloutError, model_error -from verifiers.v1.graph import PendingTurn -from verifiers.v1.semantic import ACP_EXTENSION_HEADERS -from verifiers.v1.types import Response, SamplingConfig - -# Local authentication, framing, and signatures must not reach the provider. -_BLOCKED_REQUEST_HEADERS = ( - frozenset( - { - # The harness uses this rollout secret to authenticate with the localhost server. - # The SDK supplies provider authorization after filtering. - "authorization", - "x-api-key", - "x-stainless-raw-response", - # HTTPX recalculates these for the provider URL, JSON bytes, and supported decoders. - "accept-encoding", - "content-encoding", - "content-length", - "content-type", - "host", - "transfer-encoding", - # These control only the localhost HTTP exchange. - "expect", - "keep-alive", - "proxy-authorization", - "proxy-connection", - "te", - "trailer", - "upgrade", - # Provider affinity must not compete with the rollout-wide session header below. - "session_id", - # The eval owns the model and sampling settings, so it changes those JSON fields before - # sending upstream. Hashes and signatures calculated from the intercepted body are stale. - "content-digest", - "content-md5", - "digest", - "repr-digest", - "signature", - "signature-input", - } - ) - | ACP_EXTENSION_HEADERS -) - - -class EvalClient(Client): - """Preserve native payloads; project SDK responses into the trace.""" - - def __init__(self, config: BaseClientConfig) -> None: - self.base_url = config.base_url - self.headers = resolve_headers(config) - http_client = build_async_httpx(config) - options = { - "base_url": config.base_url, - "api_key": resolve_api_key(config), - "max_retries": MAX_RETRIES, - "timeout": http_client.timeout, - "http_client": http_client, - } - self.openai = openai.AsyncOpenAI(**options) - self.anthropic = anthropic.AsyncAnthropic(**options) - - async def _complete( - self, - dialect: Dialect, - body: dict, - sampling: SamplingConfig, - session_id: str | None = None, - turn: PendingTurn | None = None, - headers: Mapping[str, str] | None = None, - ) -> tuple[Response, bytes | None]: - client, reply = await self._request( - dialect, dialect.upstream_path, body, headers, session_id - ) - try: - events = None - if body.get("stream"): - # Parse during receipt, retaining decoded bytes for atomic delivery/replay. - # Bounded chunks keep cached or coalesced responses from starving other turns. - recorded = bytearray() - chunks = reply.iter_bytes(chunk_size=16_384) - - async def capture(): - async for chunk in chunks: - recorded.extend(chunk) - yield chunk - await asyncio.sleep(0) - - stream_type = ( - anthropic.AsyncStream - if isinstance(client, anthropic.AsyncAnthropic) - else openai.AsyncStream - ) - async with stream_type( - cast_to=dialect.event_type, - response=httpx.Response( - reply.status_code, - request=reply.http_response.request, - content=capture(), - ), - client=client, - ) as stream: - raw = await dialect.read_stream(stream) - # Chat's SDK stops at [DONE]; retain any trailing provider bytes too. - async for chunk in chunks: - recorded.extend(chunk) - await asyncio.sleep(0) - events = bytes(recorded) - if isinstance(dialect, ChatDialect) and not await asyncio.to_thread( - any, - ( - event.data == "[DONE]" - for event in anthropic.Stream.raw_events( - httpx.Response(200, content=events) - ) - ), - ): - raise ValueError("Chat stream ended without [DONE]") - else: - raw = await reply.json() - response = dialect.parse_response(raw) - except ( - openai.OpenAIError, - anthropic.AnthropicError, - httpx.HTTPError, - ConnectionResetError, - ) as error: - raise model_error(error) from error - except RolloutError: - raise - except Exception as error: - # SDK accumulation can reject malformed event sequences. - raise model_error( - f"malformed upstream response: {error}", status_code=502 - ) from error - finally: - await reply.close() - response.raw = raw - return response, events - - def _headers( - self, - incoming: Mapping[str, str] | None, - session_id: str | None, - ) -> dict[str, str]: - """Retain feature headers; let the SDK supply provider auth and HTTP framing.""" - headers = httpx.Headers(incoming) - connection = headers.pop("connection", "") - for name in _BLOCKED_REQUEST_HEADERS | set( - map(str.strip, connection.lower().split(",")) - ): - headers.pop(name, None) - headers.update(self.headers) - # Configured headers cannot replace the provider key with a local rollout secret. - headers.pop("authorization", None) - headers.pop("x-api-key", None) - if session_id: - headers[SESSION_ID_HEADER] = session_id - return dict(headers) - - async def _request( - self, - dialect: Dialect, - route: str, - body: dict, - headers: Mapping[str, str] | None, - session_id: str | None = None, - ): - client, response_type = ( - (self.anthropic, anthropic.AsyncAPIResponse) - if isinstance(dialect, AnthropicDialect) - else (self.openai, openai.AsyncAPIResponse) - ) - request_headers = self._headers(headers, session_id) - # SDK defaults are plain dicts: match their casing so forwarded headers replace - # them instead of producing duplicate User-Agent or provider feature headers. - default_names = {name.lower(): name for name in client.default_headers} - request_headers = { - default_names.get(name.lower(), name): value - for name, value in request_headers.items() - } - if body.get("stream"): - # Ask the SDK for its unread response wrapper instead of a typed SSE iterator. - request_headers["X-Stainless-Raw-Response"] = "stream" - try: - # Low-level SDK requests preserve extension fields without signature filtering. - # The raw wrapper retains bytes and supports the SDK's permissive typed parse. - reply = await client.post( - join_url(self.base_url, route), - body=body, - cast_to=response_type[dict], - options={"headers": request_headers}, - stream=bool(body.get("stream")), - ) - except (openai.OpenAIError, anthropic.AnthropicError) as error: - raise model_error(error) from error - return client, reply - - async def relay_aux( - self, - dialect: Dialect, - route: str, - body: dict, - headers: Mapping[str, str] | None = None, - ) -> dict: - _, reply = await self._request(dialect, route, body, headers) - try: - return await reply.json() - finally: - await reply.close() - - async def close(self) -> None: - await self.openai.close() - await self.anthropic.close() diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py deleted file mode 100644 index 12f98ebe61..0000000000 --- a/verifiers/v1/clients/train.py +++ /dev/null @@ -1,473 +0,0 @@ -"""Train client: renders prompts to token ids and calls a vLLM generate endpoint.""" - -import asyncio -import json -import logging -import threading -from collections.abc import AsyncIterator, Callable, Mapping -from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from typing import Any, ClassVar, TypeVar - -from openai import OpenAIError -from renderers import OverlongPromptError, RenderedTokens, Renderer, RendererConfig -from renderers.base import ToolCallParseStatus, is_multimodal - -from verifiers.v1.clients.base import SESSION_ID_HEADER, build_async_openai -from verifiers.v1.clients.client import Client -from verifiers.v1.configs.client import TrainClientConfig -from verifiers.v1.dialects import FINISH_REASONS, ChatDialect, Dialect -from verifiers.v1.dialects.chat import message_to_wire -from verifiers.v1.errors import ProviderError, model_error -from verifiers.v1.graph import PendingTurn -from verifiers.v1.types import ( - AssistantMessage, - FinishReason, - Response, - SamplingConfig, - SamplingMask, - Tool, - ToolCall, - TurnTokens, - Usage, -) - -logger = logging.getLogger(__name__) - -T = TypeVar("T") - - -def tool_to_wire(tool: Tool) -> dict: - if tool.type != "function" or tool.namespace: - raise NotImplementedError( - "The renderer client only supports unnamespaced function tools." - ) - function: dict = { - "name": tool.name, - "description": tool.description, - "parameters": tool.parameters, - } - if tool.strict is not None: - function["strict"] = tool.strict - return {"type": "function", "function": function} - - -def serialize_completion(response: Response, model: str) -> dict: - """A vf `Response` -> an OpenAI chat.completion dict the program's SDK expects. The renderer - sets this on `Response.raw` (it generates, so has no provider response to relay).""" - message: dict = {"role": "assistant", "content": response.message.content} - if response.message.reasoning_content is not None: - message["reasoning_content"] = response.message.reasoning_content - if response.message.tool_calls: - message["tool_calls"] = [ - { - "id": c.id, - "type": c.type, - c.type: { - "name": c.name, - **({"namespace": c.namespace} if c.namespace else {}), - "input" if c.type == "custom" else "arguments": c.arguments, - }, - } - for c in response.message.tool_calls - ] - usage: dict | None = None - if response.usage: - # Usage is validated earlier in the pipeline; building its wire dict directly saves time. - usage = { - "completion_tokens": response.usage.completion_tokens, - "prompt_tokens": response.usage.input_tokens, - "total_tokens": response.usage.total_tokens, - } - if response.usage.reasoning_tokens is not None: - usage["completion_tokens_details"] = { - "reasoning_tokens": response.usage.reasoning_tokens - } - if response.usage.cached_input_tokens is not None: - usage["prompt_tokens_details"] = { - "cached_tokens": response.usage.cached_input_tokens - } - return { - "id": response.id or "vf-intercept", - "object": "chat.completion", - "created": response.created, - "model": response.model or model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": response.finish_reason or "stop", - } - ], - "usage": usage, - } - - -def response_from_generate( - result: dict, - model: str, - bridged_turn: PendingTurn | None = None, - mm_token_type_id_map: dict[int, int] | None = None, -) -> Response: - """Parse a `renderers.client.generate` result dict into a typed `Response`, - mirroring the chat client's `response_from_wire` (plus the token encoding).""" - finish: FinishReason = ( - result["finish_reason"] - if result.get("finish_reason") in FINISH_REASONS - else None - ) - tool_calls = [ - ToolCall( - id=tc.id or f"call_{i}", - name=tc.name, - arguments=tc.arguments - if isinstance(tc.arguments, str) - else json.dumps(tc.arguments or {}), - ) - for i, tc in enumerate(result.get("tool_calls") or []) - if getattr(tc, "name", None) - # TODO: we need a better way for renderers to expose this - and getattr(tc, "status", None) != ToolCallParseStatus.UNKNOWN_TOOL - ] or None - prompt_ids = result.get("prompt_ids") or [] - completion_ids = result.get("completion_ids") or [] - # Per-message token spans (the renderer's attribution) let the trace graph store each - # message's tokens once; carried transiently on TurnTokens and consumed by turn.commit(). - attribution = result.get("prompt_attribution") - if attribution is None: - message_spans = None - elif bridged_turn is not None: - message_spans = bridged_turn.prompt_message_spans(attribution) - else: - message_spans = attribution.message_token_spans() - raw_mm_placeholders = result.get("mm_placeholders") - mm_placeholders = ( - sorted( - (placeholder["offset"], placeholder["length"]) - for ranges in raw_mm_placeholders.values() - for placeholder in ranges - ) - if raw_mm_placeholders is not None - else None - ) - return Response( - id=result.get("request_id", ""), - created=0, - model=model, - message=AssistantMessage( - content=result.get("content") or None, - reasoning_content=result.get("reasoning_content"), - tool_calls=tool_calls, - ), - finish_reason=finish, - # /inference/v1/generate returns exact token ids but no usage details, so the - # completion's reasoning-token subset is unknown. - usage=Usage( - prompt_tokens=len(prompt_ids), completion_tokens=len(completion_ids) - ), - # generate() returns owned, typed lists. Skip revalidation here to avoid copying - # million-token contexts synchronously on the event loop. - tokens=TurnTokens.model_construct( - prompt_ids=prompt_ids, - renderer_prompt_ids=result.get("renderer_prompt_ids"), - bridged=bridged_turn is not None, - completion_ids=completion_ids, - completion_logprobs=result.get("completion_logprobs") or [], - message_spans=message_spans, - is_content=attribution.is_content if attribution is not None else None, - mm_placeholders=mm_placeholders, - mm_token_type_id_map=mm_token_type_id_map, - routed_experts=result.get("routed_experts"), - sampling_mask=SamplingMask.from_sampling_mask(mask) - if (mask := result.get("sampling_mask")) - else None, - ), - ) - - -def _is_valid_incremental_tail(messages: list[dict[str, Any]]) -> bool: - """Renderer bridges may extend sampled assistant turns with tool calls and/or a new user.""" - if not messages: - return False - roles = [] - for message in messages: - role = message.get("role") - roles.append(role if isinstance(role, str) else None) - if roles[-1] == "user": - return all(role == "tool" for role in roles[:-1]) - return all(role == "tool" for role in roles) - - -def _has_multimodal_content(messages) -> bool: - for message in messages: - content = getattr(message, "content", None) - if not isinstance(content, list): - continue - if any(getattr(part, "type", None) == "image_url" for part in content): - return True - return False - - -@dataclass -class RendererSlot: - """One renderer and the rollouts currently holding it. Encoding mutates a fast - tokenizer's state, so `run` serializes encode-side work (render, bridge) on a thread; - decode-side work is a pure read and needs neither the lock nor the hop.""" - - renderer: Renderer - load: int = 0 - lock: threading.Lock = field(default_factory=threading.Lock) - - async def run(self, fn: Callable[[], T]) -> T: - def locked() -> T: - with self.lock: - return fn() - - return await asyncio.to_thread(locked) - - -class ElasticRendererPool: - """Process-shared renderers, multiplexed and auto-growing. A pool object is a cheap - per-client view: the renderers themselves are the shared state, keyed by what builds - them, so every client with the same build inputs works one list.""" - - _renderers: ClassVar[dict[tuple, list[RendererSlot]]] = {} - """The process's renderers, keyed by build inputs. Unlike the owned interception pool, - renderers have no owner to inject them (clients are built per rollout) and must outlive - loops and clients, so the shared state lives on the type — and needs no `start`/`stop`: - a renderer is pure memory.""" - - _locks: ClassVar[dict[tuple, tuple[asyncio.AbstractEventLoop, asyncio.Lock]]] = {} - """Per-key single-flight lock for `grow`, minted fresh on a loop change: an asyncio - lock binds to the loop that first awaits it, while renderers outlive loops.""" - - def __init__( - self, - renderer_model: str, - config: RendererConfig | None, - *, - chat_template_kwargs: Mapping[str, Any] | None = None, - multiplex: int, - ) -> None: - self.renderer_model = renderer_model - self.config = config - self.chat_template_kwargs = chat_template_kwargs - self.multiplex = multiplex - self.key = ( - renderer_model, - config.model_dump_json() if config is not None else None, - json.dumps(dict(chat_template_kwargs), sort_keys=True) - if chat_template_kwargs - else None, - ) - self.renderers = self._renderers.setdefault(self.key, []) - - def warm(self) -> None: - """Start building the first renderer if none exists, so the tokenizer loads while - the rollout provisions rather than in front of its first turn. A no-op off the - event loop — `acquire` builds on demand anyway.""" - if self.renderers: - return - try: - task = asyncio.get_running_loop().create_task(self.grow()) - except RuntimeError: - return - # A failed warm is not an event: acquire retries the build and surfaces the error. - task.add_done_callback(lambda t: t.cancelled() or t.exception()) - - async def grow(self) -> RendererSlot: - """A renderer with spare capacity — reuse one under `multiplex`, else load one - more tokenizer on a thread. Single-flight per key: concurrent cold acquires wait - for one build instead of stacking tokenizers.""" - loop = asyncio.get_running_loop() - bound = self._locks.get(self.key) - if bound is None or bound[0] is not loop: - bound = self._locks[self.key] = (loop, asyncio.Lock()) - async with bound[1]: - for slot in self.renderers: - if slot.load < self.multiplex: - return slot - from renderers import create_renderer - from renderers.base import load_tokenizer - - def build(): - return create_renderer( - load_tokenizer(self.renderer_model), - self.config, - chat_template_kwargs=self.chat_template_kwargs, - ) - - slot = RendererSlot(await asyncio.to_thread(build)) - self.renderers.append(slot) - logger.info( - "renderer pool: %d renderer(s), multiplex=%d", - len(self.renderers), - self.multiplex, - ) - return slot - - @asynccontextmanager - async def acquire(self) -> AsyncIterator[RendererSlot]: - """A renderer to render this turn with, growing the pool when every one already - carries `multiplex` rollouts. The slot is held for the turn, so `load` counts - rollouts in flight rather than renders in progress.""" - slot = await self.grow() - slot.load += 1 - try: - yield slot - finally: - slot.load -= 1 - - -class TrainClient(Client): - """Renders prompts to token ids and calls a vLLM `/inference/v1/generate` engine. - - Owned by the interception server and shared by the rollouts it multiplexes: they reuse - its engine connection pool, and each turn takes a slot on the shared - `ElasticRendererPool`.""" - - def __init__(self, config: TrainClientConfig) -> None: - self.config = config - self.client = build_async_openai(config) - # The per-request model is only known at call time; a config that pins the renderer - # model can warm now, which is every training run (prime-rl always pins it). - if config.renderer_model_name is not None: - ElasticRendererPool( - config.renderer_model_name, - config.renderer, - multiplex=config.multiplex, - ).warm() - - async def _complete( - self, - dialect: Dialect, - body: dict, - sampling: SamplingConfig, - session_id: str | None = None, - turn: PendingTurn | None = None, - headers: Mapping[str, str] | None = None, - ) -> tuple[Response, bytes | None]: - # The renderer tokenizes the typed prompt for training (it needs per-token ids + logprobs - # back), so it can't forward the raw request — it parses `body` via the dialect and renders - # it with a chat template. It leaves `Response.raw` unset; the interception server serializes - # its `Response` for the program instead of relaying provider bytes. - if not isinstance(dialect, ChatDialect): - # The renderer renders a chat template, so it's only validated for chat-completions - # input; other dialects' semantics (Responses reasoning items, Anthropic thinking) may - # not round-trip faithfully through chat-template tokenization. Refuse them explicitly. - raise NotImplementedError( - f"The renderer client only supports the chat-completions dialect, got " - f"{type(dialect).__name__}. Use the proxy client for this dialect, or add " - f"renderer support for it." - ) - if turn is not None: - prompt = turn.prompt - tools = turn.tools - else: - request, _ = dialect.parse_request(body) - prompt = request.messages - tools = request.tools - from renderers.client import generate - - wire_tools = [tool_to_wire(t) for t in tools] if tools else None - wire_messages = [message_to_wire(m) for m in prompt] - wire_tail = [message_to_wire(m) for m in turn.tail] if turn is not None else [] - prompt_ids: list[int] | None = None - prompt_attribution: RenderedTokens | None = None - model = body["model"] - sampling_params = sampling.model_dump(exclude_none=True) - chat_template_kwargs = sampling_params.pop("chat_template_kwargs", None) - cache_salt = sampling_params.pop("cache_salt", None) - pool = ElasticRendererPool( - self.config.renderer_model_name or model, - self.config.renderer, - chat_template_kwargs=chat_template_kwargs, - multiplex=self.config.multiplex, - ) - bridged_turn: PendingTurn | None = None - - async with pool.acquire() as slot: - renderer = slot.renderer - mm_token_type_id_map = ( - renderer.mm_token_type_id_map if is_multimodal(renderer) else None - ) - has_images = _has_multimodal_content(prompt) - process_multimodal = not has_images - if has_images and not getattr( - renderer, "supports_process_multimodal", False - ): - raise NotImplementedError( - f"{type(renderer).__name__} does not support process_multimodal=False" - ) - render_kwargs = {} if process_multimodal else {"process_multimodal": False} - # Only build the O(context) previous token stream for a bridgeable tail. - can_bridge = turn is not None and _is_valid_incremental_tail(wire_tail) - previous_ids = turn.previous_renderer_token_ids() if can_bridge else None - if previous_ids is not None: - previous_prompt_ids, previous_completion_ids = previous_ids - - def bridge(): - return renderer.bridge_to_next_turn( - previous_prompt_ids, - previous_completion_ids, - wire_tail, - tools=wire_tools, - **render_kwargs, - ) - - bridged = await slot.run(bridge) - if bridged is not None: - prompt_ids = bridged.token_ids - prompt_attribution = bridged - bridged_turn = turn - sampling_params["routed_experts_prompt_start"] = max( - turn.path_len - 1, 0 - ) - - # Render here (encode-side, so through the slot) rather than inside `generate`: - # handed prebuilt prompt_ids, generate's own renderer touches are decode-side - # and stop-id reads, safe on a bare renderer without lock or thread hop. - if prompt_ids is None: - rendered = await slot.run( - lambda: renderer.render( - wire_messages, - tools=wire_tools, - add_generation_prompt=True, - **render_kwargs, - ) - ) - prompt_ids = rendered.token_ids - prompt_attribution = rendered - - try: - result = await generate( - client=self.client, - renderer=renderer, - messages=wire_messages, - model=model, - prompt_ids=prompt_ids, - prompt_attribution=prompt_attribution, - tools=wire_tools, - sampling_params=sampling_params, - process_multimodal=process_multimodal, - cache_salt=cache_salt, - extra_headers={SESSION_ID_HEADER: session_id} - if session_id - else None, - ) - except OverlongPromptError as e: - # The renderer's pre-flight overflow never reached the provider: a - # deterministic 400, so the harness SDK never retries it. - raise ProviderError(str(e), status_code=400) from e - except OpenAIError as e: - raise model_error(e) from e - response = response_from_generate( - result, model, bridged_turn, mm_token_type_id_map - ) - # No provider response to relay (we generated), so serialize one for the program; the - # interception server hands `Response.raw` back regardless of client. - response.raw = serialize_completion(response, model) - return response, None - - async def close(self) -> None: - await self.client.close() diff --git a/verifiers/v1/configs/cli/eval.py b/verifiers/v1/configs/cli/eval.py index 944bd4b38f..3da6f8154a 100644 --- a/verifiers/v1/configs/cli/eval.py +++ b/verifiers/v1/configs/cli/eval.py @@ -6,7 +6,7 @@ from pydantic import AliasChoices, Field, PrivateAttr, SerializeAsAny, model_validator from pydantic_config import BaseConfig -from verifiers.v1.clients import ClientConfig, EvalClientConfig +from verifiers.v1.clients import ClientConfig from verifiers.v1.configs.cli.env import narrowed_env_annotation, resolve_env_field from verifiers.v1.configs.env import EnvConfig from verifiers.v1.configs.select import SelectCLIConfig @@ -83,7 +83,7 @@ class EvalConfig(BaseConfig): "deepseek/deepseek-v4-flash", validation_alias=AliasChoices("model", "m") ) """Model id.""" - client: ClientConfig = EvalClientConfig() + client: ClientConfig = ClientConfig() sampling: SamplingConfig = SamplingConfig() select: SelectCLIConfig = SelectCLIConfig() """Which of the taskset's tasks to evaluate, under `--select.*` (`-n` sets diff --git a/verifiers/v1/configs/client.py b/verifiers/v1/configs/client.py index 0ef124bac0..90d1771cb6 100644 --- a/verifiers/v1/configs/client.py +++ b/verifiers/v1/configs/client.py @@ -1,23 +1,11 @@ -"""Client configs: describe an OpenAI-compatible endpoint. - -A `BaseClientConfig` is an OpenAI-compatible endpoint (base_url + API-key env var -+ extra headers); `clients.resolve_client` turns one into a live `Client` — the -interception server builds one per distinct config and shares it across the rollouts -it multiplexes. The default Prime endpoint, API key, and team fall back to -the active Prime CLI config, so direct `uv run vf-eval` calls behave like `prime eval`. -Both the eval entrypoint (its model client) and in-env LLM calls (e.g. a judge reward) -build clients from these. `ClientConfig` is the CLI-selectable discriminated union -(eval | train). -""" +"""Endpoint configuration shared by evaluation, training, and in-env model calls.""" import os from itertools import chain -from typing import Annotated, Literal from urllib.parse import urlparse from pydantic import Field, model_validator from pydantic_config import BaseConfig -from renderers import RendererConfig from verifiers.v1.utils.prime import load_prime_config @@ -42,7 +30,7 @@ class ClientTimeoutConfig(BaseConfig): """Time to acquire a connection from the pool.""" -class BaseClientConfig(BaseConfig): +class ClientConfig(BaseConfig): """An OpenAI-compatible endpoint. The API key is read from an env var.""" base_url: str = DEFAULT_PRIME_INFERENCE_URL @@ -55,7 +43,7 @@ class BaseClientConfig(BaseConfig): timeout: ClientTimeoutConfig = ClientTimeoutConfig() @model_validator(mode="after") - def apply_prime_config(self) -> "BaseClientConfig": + def apply_prime_config(self) -> "ClientConfig": if self.api_key_var != "PRIME_API_KEY": return self prime_config = load_prime_config() @@ -72,40 +60,7 @@ def apply_prime_config(self) -> "BaseClientConfig": return self -class EvalClientConfig(BaseClientConfig): - """The default (eval): forward each request to a matching endpoint via `EvalClient`.""" - - type: Literal["eval"] = "eval" - - -class TrainClientConfig(BaseClientConfig): - """Training: a vLLM `/inference/v1/generate` endpoint with client-side tokenization (via - `TrainClient`), so responses carry token IDs and logprobs. Needs a running vLLM engine.""" - - type: Literal["train"] = "train" - renderer: RendererConfig | None = None - """The `renderers.RendererConfig` to use (the same shared type prime-rl configures). - `None` auto-resolves from the model — which falls back to the default renderer (no - tool support) for models not in the renderer map, so set it explicitly for - fine-tunes / tool-using envs.""" - renderer_model_name: str | None = None - """Model the tokenizer/renderer pool is built for. Pin to the base model so a LoRA - adapter name (served only for sampling) never drives tokenizer loading. Falls back to - the per-request model when None.""" - multiplex: int = Field(256, ge=1) - """Rollouts that share one renderer (~75-95 MB each): the pool warms one and grows on - demand, so N concurrent rollouts hold ~N/multiplex tokenizers. A renderer is only busy - for the render itself (ms against a multi-second turn), so one absorbs many rollouts; - lower this when rendering is the slow part (very long prompts, frequent bridge misses).""" - - -# Discriminated union for a CLI-selectable client (`--client.type eval|train`). -ClientConfig = Annotated[ - EvalClientConfig | TrainClientConfig, Field(discriminator="type") -] - - -def resolve_api_key(config: BaseClientConfig) -> str: +def resolve_api_key(config: ClientConfig) -> str: """The API key for `config`: its env var, falling back to the Prime CLI config for a `PRIME_API_KEY`-keyed pinference endpoint. `"EMPTY"` when unset.""" api_key = os.environ.get(config.api_key_var) @@ -119,7 +74,7 @@ def resolve_api_key(config: BaseClientConfig) -> str: return api_key or "EMPTY" -def resolve_headers(config: BaseClientConfig) -> dict[str, str]: +def resolve_headers(config: ClientConfig) -> dict[str, str]: """Resolve request headers for this destination without storing credentials in config.""" # Keep original spelling for SDKs that merge default headers case-sensitively. names: dict[str, str] = {} diff --git a/verifiers/v1/configs/judge.py b/verifiers/v1/configs/judge.py index c315876804..ed9c3b0cd5 100644 --- a/verifiers/v1/configs/judge.py +++ b/verifiers/v1/configs/judge.py @@ -7,11 +7,11 @@ from pydantic import BaseModel, FiniteFloat, SerializeAsAny -from verifiers.v1.clients import BaseClientConfig +from verifiers.v1.clients import ClientConfig from verifiers.v1.types import ID, SamplingConfig -class JudgeConfig(BaseClientConfig): +class JudgeConfig(ClientConfig): id: ID = "" """Plugin id; empty for a judge called directly by task code.""" name: str = "" diff --git a/verifiers/v1/dialects/base.py b/verifiers/v1/dialects/base.py index e61fb20c22..28f20b6cde 100644 --- a/verifiers/v1/dialects/base.py +++ b/verifiers/v1/dialects/base.py @@ -5,7 +5,7 @@ registered dialect's `routes` (see `dialects.DIALECTS`), so a request's format is resolved from the endpoint the program's SDK posts to — the harness declares nothing. -The eval client relays a request's native JSON with explicit eval settings applied, and +The client relays a request's native JSON with explicit sampling settings applied, and the dialect reads copies of the JSON into vf types. The provider SDK assembles streams. Task hooks edit the typed copy, and their edits flow back into the native request through the setters `parse_request` returns. What a @@ -81,7 +81,7 @@ class Dialect(ABC): aux_routes: ClassVar[tuple[str, ...]] = () """Side endpoints the SDK may call that aren't model turns (e.g. Anthropic's - `count_tokens`): relayed as native JSON by the eval client, never recorded on the trace.""" + `count_tokens`): relayed as native JSON by the client, never recorded on the trace.""" upstream_path: ClassVar[str] """The provider endpoint the proxy forwards to for this format (e.g. `/chat/completions`).""" diff --git a/verifiers/v1/dialects/chat.py b/verifiers/v1/dialects/chat.py index e59cb94f0c..6aaeb11278 100644 --- a/verifiers/v1/dialects/chat.py +++ b/verifiers/v1/dialects/chat.py @@ -11,6 +11,7 @@ from functools import partial from typing import Any +import numpy as np from openai.types import CompletionUsage from openai.types.chat import ChatCompletionChunk @@ -23,12 +24,15 @@ NativeContentPart, Request, Response, + SamplingMask, SystemMessage, Tool, ToolCall, ToolMessage, + TurnTokens, Usage, UserMessage, + _decode_packed, content_to_parts, ) from verifiers.v1.utils.chat import read_chat_completion @@ -135,9 +139,8 @@ def parse_tools(raw: list[dict] | None) -> list[Tool] | None: # --- vf -> chat wire ---------------------------------------------------------- -# `message_to_wire` (chat-only): used by the bash harness (a Messages prompt) and the train -# client (its generate request). The proxy preserves its parsed native JSON independently and -# does not use this serializer. +# `message_to_wire` serves typed harness prompts and judge calls. Intercepted requests keep +# their native JSON independently and do not use this serializer. def _content_to_wire(content): @@ -202,12 +205,71 @@ def _write_message(native: dict, message: Message) -> None: def response_from_wire(completion: dict) -> Response: - """An OpenAI chat.completion -> a vf `Response` (the one place raw provider objects cross - into our typed `Response`). No token ids: training tokens come from the renderer client.""" + """Parse the native completion and, when present, its exact inference tokens.""" choice = completion["choices"][0] finish_reason = choice.get("finish_reason") finish: FinishReason = finish_reason if finish_reason in FINISH_REASONS else None usage = completion.get("usage") + tokens = None + prompt_ids = completion.get("prompt_token_ids") + completion_ids = choice.get("token_ids") + if prompt_ids is not None and completion_ids is not None: + packed_logprobs = choice.get("completion_logprobs") + logprobs = (choice.get("logprobs") or {}).get("content") + if logprobs is not None: + if len(logprobs) != len(completion_ids): + raise ValueError("completion token IDs and logprobs must align") + for token_id, entry in zip(completion_ids, logprobs, strict=True): + token = entry["token"] + if token.startswith("token_id:") and token != f"token_id:{token_id}": + raise ValueError("logprobs refer to different completion token IDs") + completion_logprobs = ( + _decode_packed(packed_logprobs, np.float64) + if packed_logprobs is not None + else np.asarray([entry["logprob"] for entry in logprobs or []]) + ) + if (packed_logprobs is not None or logprobs is not None) and ( + len(completion_logprobs) != len(completion_ids) + or not ( + np.isfinite(completion_logprobs) & (completion_logprobs > -9999.0) + ).all() + ): + raise ValueError( + "completion token IDs require aligned finite sampling logprobs" + ) + metadata = completion.get("training_metadata") or {} + spans = metadata.get("message_spans") + previous_end = 0 + for span in spans or []: + if span is not None: + start, end = span + if not previous_end <= start <= end <= len(prompt_ids): + raise ValueError( + "message spans must be ordered within prompt token IDs" + ) + previous_end = end + content = metadata.get("is_content") + if content is not None and len(content) != len(prompt_ids): + raise ValueError("content attribution must align with prompt token IDs") + sampling_mask = choice.get("sampling_mask") + if sampling_mask is not None: + sampling_mask = SamplingMask.from_sampling_mask(sampling_mask) + if ( + len(sampling_mask.counts) != len(completion_ids) + or (sampling_mask.counts < 0).any() + or int(sampling_mask.counts.sum()) != len(sampling_mask.ids) + ): + raise ValueError("sampling masks must align with completion token IDs") + tokens = TurnTokens( + prompt_ids=prompt_ids, + completion_ids=completion_ids, + completion_logprobs=completion_logprobs.tolist(), + message_spans=spans, + is_content=content, + mm_token_type_id_map=metadata.get("mm_token_type_id_map"), + routed_experts=choice.get("routed_experts"), + sampling_mask=sampling_mask, + ) return Response( id=completion.get("id") or "", created=completion.get("created") or 0, @@ -216,6 +278,7 @@ def response_from_wire(completion: dict) -> Response: finish_reason=finish, # The SDK's lenient model carries provider extensions such as `cost`. usage=Usage.from_openai(CompletionUsage.construct(**usage)) if usage else None, + tokens=tokens, ) @@ -243,6 +306,9 @@ class ChatDialect(Dialect): "tool_choice", "parallel_tool_calls", "extra_body", + "return_token_ids", + "return_tokens_as_token_ids", + "return_training_metadata", } ) max_tokens_keys = ("max_tokens", "max_completion_tokens") @@ -278,11 +344,18 @@ def parse_response(self, response: dict) -> Response: return response_from_wire(response) def rewrite_response(self, raw: dict, text: str) -> None: + # Replacement text was not sampled, so none of the original token records apply. + raw.pop("prompt_token_ids", None) + raw.pop("training_metadata", None) for choice in raw.get("choices") or []: if isinstance(choice.get("message"), dict): choice["message"] = {"role": "assistant", "content": text} choice["finish_reason"] = "stop" choice.pop("logprobs", None) + choice.pop("completion_logprobs", None) + choice.pop("token_ids", None) + choice.pop("sampling_mask", None) + choice.pop("routed_experts", None) def stream_events(self, raw: dict) -> list[bytes]: choice = (raw.get("choices") or [{}])[0] @@ -297,6 +370,7 @@ def stream_events(self, raw: dict) -> list[bytes]: "object": "chat.completion.chunk", "choices": [ { + **{key: value for key, value in choice.items() if key != "message"}, "index": 0, "delta": message, "finish_reason": choice.get("finish_reason"), diff --git a/verifiers/v1/gepa/config.py b/verifiers/v1/gepa/config.py index 0aa5b2a993..1b69e6a2e1 100644 --- a/verifiers/v1/gepa/config.py +++ b/verifiers/v1/gepa/config.py @@ -13,7 +13,7 @@ from pydantic import AliasChoices, Field, SerializeAsAny, model_validator from pydantic_config import BaseConfig -from verifiers.v1.clients import EvalClientConfig +from verifiers.v1.clients import ClientConfig from verifiers.v1.configs.cli.env import narrowed_env_annotation, resolve_env_field from verifiers.v1.configs.cli.eval import RunConfig, default_run_name from verifiers.v1.configs.env import EnvConfig @@ -43,11 +43,11 @@ def _resolve_env(cls, data): "deepseek/deepseek-v4-flash", validation_alias=AliasChoices("model", "m") ) """Model id for rollouts under optimization.""" - client: EvalClientConfig = EvalClientConfig() + client: ClientConfig = ClientConfig() sampling: SamplingConfig = SamplingConfig() reflection_model: str | None = None """Teacher model that proposes new system prompts. None = reuse `model`.""" - reflection_client: EvalClientConfig | None = None + reflection_client: ClientConfig | None = None """Endpoint for `reflection_model`. None = reuse `client`.""" select: SelectCLIConfig = SelectCLIConfig() diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index fb73986e3f..4b154010af 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -10,7 +10,7 @@ and making a branch's training sample a cheap concat of node `token_ids`/`mask`/`logprobs` along its path. -Token attribution (renderer client): the renderer reports, per prompt, each message's token +Token attribution: the inference server reports, per prompt, each message's token span (`RenderedTokens.message_token_spans()`, carried on `TurnTokens.message_spans`). A new input message's node gets its span plus the leading template scaffold since the previous message; the trailing scaffold (the generation prompt) goes on the assistant node, prefixed @@ -22,6 +22,7 @@ import binascii import hashlib +import io import json import time from dataclasses import dataclass @@ -37,7 +38,6 @@ field_validator, ) from pydantic.json_schema import SkipJsonSchema -from renderers.base import RenderedTokens from verifiers.v1.semantic import ParentLink from verifiers.v1.types import ( @@ -116,10 +116,6 @@ class MessageNode(BaseModel): template scaffold + its own tokens — for an assistant, the generation-prompt scaffold followed by the sampled completion. Concatenated along a path, these reproduce the exact `prompt_ids + completion_ids` the model saw.""" - renderer_token_ids: list[int] | None = Field(default=None, exclude=True) - """Logical renderer tokens retained only while extending a live rollout. - None means they are identical to `token_ids`; an empty list is a real empty slice.""" - mask: list[bool] = Field(default_factory=list) """Per-token, parallel to `token_ids`: True for trainable, model-sampled tokens (only an assistant node's completion span); False for template scaffold and every input-message @@ -155,7 +151,7 @@ class MessageNode(BaseModel): routed_experts: SkipJsonSchema[np.ndarray | None] = None """This node's slice of the MoE expert-routing array — uint8 `[len(token_ids), layers, top_k]`, the expert ids inference selected for exactly this node's tokens. Attributed from - the turn's `generate` payload by `_attribute_routed_experts`; `Branch.routed_experts` + the turn's inference payload by `_attribute_routed_experts`; `Branch.routed_experts` concatenates these along the path into the trainer's router-replay input. Rides the wire as a raw-bytes `__nd__` dict; kept off disk by the dump-site `exclude` in prime-rl.""" sampling_mask: SkipJsonSchema[SamplingMask | None] = None @@ -167,14 +163,6 @@ class MessageNode(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - @property - def logical_ids(self) -> list[int]: - return ( - self.token_ids - if self.renderer_token_ids is None - else self.renderer_token_ids - ) - @field_serializer( "logprobs", "advantages", @@ -371,7 +359,6 @@ def _matching_node( parent: int | None, message: Message, token_ids: list[int] | None = None, - renderer_token_ids: list[int] | None = None, tools: list[Tool] | None = None, ) -> int | None: """Find an existing child, optionally requiring its exact token spans. @@ -382,24 +369,18 @@ def _matching_node( """ key = _node_key(parent, message, tools) indexed = _head_index(trace).get(key) - if ( - indexed is not None - and (token_ids is None or trace.nodes[indexed].token_ids == token_ids) - and ( - renderer_token_ids is None - or trace.nodes[indexed].logical_ids == renderer_token_ids - ) + if indexed is not None and ( + token_ids is None or trace.nodes[indexed].token_ids == token_ids ): return indexed - if token_ids is None and renderer_token_ids is None: + if token_ids is None: return indexed for node_id in range(len(trace.nodes) - 1, -1, -1): if node_id == indexed: continue node = trace.nodes[node_id] if ( - (token_ids is None or node.token_ids == token_ids) - and (renderer_token_ids is None or node.logical_ids == renderer_token_ids) + node.token_ids == token_ids and _node_key(node.parent, node.message, node.tools) == key ): return node_id @@ -413,9 +394,6 @@ def _matching_prefix_node( prompt_ids: list[int], start: int, stop: int, - renderer_prompt_ids: list[int], - renderer_start: int, - renderer_stop: int, tools: list[Tool] | None = None, ) -> int | None: """Find the longest content-equivalent child matching inside `[start, stop]`. @@ -447,11 +425,6 @@ def _matching_prefix_node( if start + len(trace.nodes[node_id].token_ids) <= stop and prompt_ids[start : start + len(trace.nodes[node_id].token_ids)] == trace.nodes[node_id].token_ids - and renderer_start + len(trace.nodes[node_id].logical_ids) <= renderer_stop - and renderer_prompt_ids[ - renderer_start : renderer_start + len(trace.nodes[node_id].logical_ids) - ] - == trace.nodes[node_id].logical_ids ] return max( matches, key=lambda node_id: len(trace.nodes[node_id].token_ids), default=None @@ -462,9 +435,9 @@ def _matching_prefix_node( class PendingTurn: """A resolved prompt waiting on model inference. - `prepare_turn` resolves the graph prefix used for renderer bridging before inference. `commit` - preserves that inference-producing prefix, extends it with any matching nodes committed while - inference was in flight, then adds only the remaining prompt tail and sampled response. + `prepare_turn` resolves the matching message prefix before inference. `commit` checks + its token identity against the server response, reuses exact prefixes, then adds the + remaining prompt tail and sampled response. """ trace: Trace @@ -481,55 +454,6 @@ def tail_start(self) -> int: def tail(self) -> list[Message]: return self.prompt[self.tail_start :] - def previous_renderer_token_ids(self) -> tuple[list[int], list[int]] | None: - """Return the logical renderer prompt and completion for a bridge anchor. - - The anchor must end at a sampled assistant node. That node stores generation-prompt - scaffold followed by sampled completion tokens, so split off the sampled suffix. - """ - if not self.prefix_node_ids: - return None - last = self.trace.nodes[self.prefix_node_ids[-1]] - if not last.sampled: - return None - num_sampled = sum(last.mask) - if not num_sampled: - return None - - renderer_prompt_ids: list[int] = [] - for nid in self.prefix_node_ids[:-1]: - node = self.trace.nodes[nid] - renderer_prompt_ids.extend(node.logical_ids) - last_ids = last.logical_ids - renderer_prompt_ids.extend(last_ids[:-num_sampled]) - completion_ids = last_ids[-num_sampled:] - if not renderer_prompt_ids or not completion_ids: - return None - return renderer_prompt_ids, completion_ids - - def prompt_message_spans( - self, tail_attribution: RenderedTokens - ) -> list[tuple[int, int] | None]: - """Convert bridge-tail attribution into full-prompt message spans.""" - # Reused bridge tokens are unattributed, so scan only the newly rendered tail. - tail_spans = RenderedTokens( - message_indices=tail_attribution.message_indices[self.renderer_path_len :], - message_roles=tail_attribution.message_roles, - ).message_token_spans() - # Tail spans are slice-relative; restore their full-prompt token offsets. - return [None] * self.tail_start + [ - None - if span is None - else (span[0] + self.renderer_path_len, span[1] + self.renderer_path_len) - for span in tail_spans - ] - - @property - def renderer_path_len(self) -> int: - return sum( - len(self.trace.nodes[nid].logical_ids) for nid in self.prefix_node_ids - ) - def commit(self, response: Response) -> int: """Add this turn to the graph; returns the committed assistant node's id.""" assistant_id = _commit_turn(self, response) @@ -661,25 +585,28 @@ def _attribute_routed_experts( prefix_node_ids: list[int], new_node_ids: list[int], path_len: int, - payload: Any, + arr: np.ndarray | None, + start: int | None, ) -> None: - """Attach each new node's slice of this turn's MoE expert-routing array. The `generate` - payload's array covers the turn's prompt+completion from `payload["start"]` (0 = from token + """Attach each new node's slice of this turn's MoE expert-routing array. A compact + payload's array covers the turn's prompt+completion from `start` (0 = from token 0); the nodes created this turn tile sequence positions `[path_len:]` in creation order, so we hand each node `arr[off : off+len(node.token_ids)]` and advance. Reused-prefix nodes keep the routing attributed when they were first created, except for the one position this turn corrects (see `_replace_placeholder_routing_row`). A node whose slice falls outside the array (a `start` past `path_len`, e.g. an unexpected prefix-cache delta) is left unset — the branch then reports no routing rather than misaligning.""" - if payload is None: + if arr is None: return - raw = binascii.a2b_base64(payload["data"]) - arr = np.frombuffer(raw, dtype=np.dtype(payload.get("dtype", "uint8"))).reshape( - payload["shape"] - ) - off = path_len - int(payload.get("start", 0) or 0) + new_tokens = sum(len(trace.nodes[nid].token_ids) for nid in new_node_ids) + if start is None: + # Native vLLM .npy omits the final sampled token and any skipped prompt prefix; + # its row count therefore locates the new suffix without an explicit start field. + off = arr.shape[0] + 1 - new_tokens + else: + off = path_len - start _replace_placeholder_routing_row(trace, prefix_node_ids, arr, off) - needed = off + sum(len(trace.nodes[nid].token_ids) for nid in new_node_ids) + needed = off + new_tokens for nid in new_node_ids: n = len(trace.nodes[nid].token_ids) end = off + n @@ -709,128 +636,64 @@ def _attribute_sampling_mask( node.sampling_mask = payload -def _project_prompt_attribution( - renderer_prompt_ids: list[int], - prompt_ids: list[int], - mm_token_type_id_map: dict[int, int], - mm_placeholders: list[tuple[int, int]] | None, - message_spans: list[tuple[int, int] | None] | None, - is_content: list[bool] | None, -) -> tuple[list[tuple[int, int] | None] | None, list[bool] | None]: - """Project logical attribution using vLLM's multimodal placeholder ranges.""" - if renderer_prompt_ids == prompt_ids: - return message_spans, is_content - if not mm_token_type_id_map or mm_placeholders is None: - raise ValueError( - "cannot align renderer and vLLM prompt tokens without multimodal placeholders" - ) - - offsets = [0] - prompt_offset = 0 - placeholder_index = 0 - for token_id in renderer_prompt_ids: - if token_id in mm_token_type_id_map: - if placeholder_index >= len(mm_placeholders): - raise ValueError( - "vLLM multimodal placeholders do not align with renderer prompt" - ) - offset, length = mm_placeholders[placeholder_index] - if offset != prompt_offset or length < 1: - raise ValueError( - "vLLM multimodal placeholders do not align with renderer prompt" - ) - prompt_offset += length - placeholder_index += 1 - else: - if ( - prompt_offset >= len(prompt_ids) - or prompt_ids[prompt_offset] != token_id - ): - raise ValueError( - "renderer prompt does not align with vLLM prompt tokens" - ) - prompt_offset += 1 - offsets.append(prompt_offset) - if prompt_offset != len(prompt_ids) or placeholder_index != len(mm_placeholders): - raise ValueError("renderer prompt does not align with vLLM prompt tokens") - - projected_spans = None - if message_spans is not None: - projected_spans = [] - for span in message_spans: - if span is None: - projected_spans.append(None) - continue - start, end = span - if not 0 <= start <= end <= len(renderer_prompt_ids): - raise ValueError("message span exceeds renderer prompt tokens") - projected_spans.append((offsets[start], offsets[end])) - - projected_is_content = is_content - if is_content: - if len(is_content) != len(renderer_prompt_ids): - raise ValueError( - "content attribution does not match renderer prompt tokens" - ) - projected_is_content = [] - for index, value in enumerate(is_content): - projected_is_content.extend([value] * (offsets[index + 1] - offsets[index])) - return projected_spans, projected_is_content - - def _commit_turn(turn: PendingTurn, response: Response) -> int: trace = turn.trace prompt = turn.prompt tokens = response.tokens + # Decode routing before publishing nodes: a malformed payload must leave the graph intact. + routed_experts = None + routing_start = None + if tokens is not None and tokens.routed_experts is not None: + payload = tokens.routed_experts + if isinstance(payload, str): + routed_experts = np.load( + io.BytesIO(binascii.a2b_base64(payload)), allow_pickle=False + ) + else: + raw = binascii.a2b_base64(payload["data"]) + routed_experts = np.frombuffer( + raw, dtype=np.dtype(payload.get("dtype", "uint8")) + ).reshape(payload["shape"]) + routing_start = int(payload.get("start", 0) or 0) + if ( + not isinstance(routed_experts, np.ndarray) + or routed_experts.ndim != 3 + or routed_experts.dtype.kind not in "iu" + ): + raise ValueError( + "routed experts must be an integer [tokens, layers, top_k] array" + ) # Constant per renderer, so re-stamping every turn is idempotent. if tokens is not None and tokens.mm_token_type_id_map: trace.mm_token_type_id_map = tokens.mm_token_type_id_map prompt_ids = tokens.prompt_ids if tokens else [] - renderer_prompt_ids = ( - tokens.renderer_prompt_ids - if tokens and tokens.renderer_prompt_ids is not None - else prompt_ids - ) - renderer_spans = tokens.message_spans if tokens else None - renderer_is_content = tokens.is_content if tokens else None + spans = tokens.message_spans if tokens else None + is_content = tokens.is_content if tokens else None idx = _head_index(trace) + # Token-based prefix reuse. `prepare_turn` matched the prefix by message hash (content); when + # this turn carries token ids, tighten that to token identity — the stored prefix must be an + # exact token prefix of what the model saw this turn (`prompt_ids`). Reuse whole nodes within + # the longest common token prefix and fork at the first divergence, so a retokenized prior + # (BPE drift, dropped ``, rewritten tool calls) branches off with this turn's real + # tokens instead of silently inheriting stale ones. Comparing the *concatenated* prefix (not + # per-message spans) is what makes this correct: a prior assistant's stored generation form + # and its re-rendered input form place the turn-close scaffold in different nodes but at the + # same position, so only a genuine content/token change shifts the common prefix. Requests + # without token ids keep the message-hash prefix. prefix = turn.prefix_node_ids path_len = turn.path_len - renderer_path_len = turn.renderer_path_len if tokens is not None and prefix: keep = 0 off = 0 - renderer_off = 0 for nid in prefix: node = trace.nodes[nid] - node_renderer_ids = node.logical_ids - if ( - prompt_ids[off : off + len(node.token_ids)] != node.token_ids - or renderer_prompt_ids[ - renderer_off : renderer_off + len(node_renderer_ids) - ] - != node_renderer_ids - ): + if prompt_ids[off : off + len(node.token_ids)] != node.token_ids: break off += len(node.token_ids) - renderer_off += len(node_renderer_ids) keep += 1 - if tokens.bridged and keep != len(prefix): - raise ValueError( - "vLLM prompt tokens do not exactly extend the stored rollout prefix" - ) prefix = prefix[:keep] path_len = off - renderer_path_len = renderer_off - spans, is_content = _project_prompt_attribution( - renderer_prompt_ids, - prompt_ids, - trace.mm_token_type_id_map, - tokens.mm_placeholders if tokens else None, - renderer_spans, - renderer_is_content, - ) has_is_content = is_content is not None and len(is_content) == len(prompt_ids) # A parallel request may have committed more of this prompt after `prepare_turn` resolved @@ -843,10 +706,6 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: i = len(prefix) span = spans[i] if spans and i < len(spans) else None end = span[1] if span else path_len - renderer_span = ( - renderer_spans[i] if renderer_spans and i < len(renderer_spans) else None - ) - renderer_end = renderer_span[1] if renderer_span else renderer_path_len parent = prefix[-1] if prefix else None if tokens is not None and span is None: next_start = next( @@ -857,16 +716,6 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: ), len(prompt_ids), ) - next_renderer_start = next( - ( - later_span[0] - for later_span in ( - renderer_spans[i + 1 :] if renderer_spans else [] - ) - if later_span is not None - ), - len(renderer_prompt_ids), - ) existing = _matching_prefix_node( trace, parent, @@ -874,35 +723,27 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: prompt_ids, path_len, next_start, - renderer_prompt_ids, - renderer_path_len, - next_renderer_start, tools=turn.tools, ) if existing is not None: end += len(trace.nodes[existing].token_ids) - renderer_end += len(trace.nodes[existing].logical_ids) else: node_tokens = prompt_ids[path_len:end] - renderer_node_tokens = renderer_prompt_ids[renderer_path_len:renderer_end] existing = _matching_node( trace, parent, prompt[i], node_tokens if tokens is not None else None, - renderer_node_tokens if tokens is not None else None, tools=turn.tools, ) if existing is None: break prefix.append(existing) path_len = end - renderer_path_len = renderer_end num_reused = len(prefix) parent = prefix[-1] if prefix else None cursor: int | None = None - renderer_cursor: int | None = None # Track new nodes separately so routed-expert attribution needs only node ids, not this path. new_node_ids: list[int] = [] for i, msg in enumerate(prompt[num_reused:], start=num_reused): @@ -911,20 +752,12 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: span = spans[i] if spans and i < len(spans) else None end = span[1] if span else start node_tokens = prompt_ids[start:end] - renderer_start = ( - renderer_path_len if renderer_cursor is None else renderer_cursor - ) - renderer_span = ( - renderer_spans[i] if renderer_spans and i < len(renderer_spans) else None - ) - renderer_end = renderer_span[1] if renderer_span else renderer_start trace.nodes.append( MessageNode.model_construct( parent=parent, tools=turn.tools if parent is None else [], message=msg, token_ids=node_tokens, - renderer_token_ids=renderer_prompt_ids[renderer_start:renderer_end], mask=[False] * len(node_tokens), is_content=is_content[start:end] if has_is_content else [], ) @@ -933,15 +766,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: idx[key] = parent new_node_ids.append(parent) cursor = end - renderer_cursor = renderer_end comp_ids = tokens.completion_ids if tokens else [] gen_start = path_len if cursor is None else cursor gen_prompt = prompt_ids[gen_start:] - renderer_gen_start = ( - renderer_path_len if renderer_cursor is None else renderer_cursor - ) - renderer_gen_prompt = renderer_prompt_ids[renderer_gen_start:] trace.nodes.append( MessageNode.model_construct( parent=parent, @@ -949,7 +777,6 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: message=response.message, sampled=True, token_ids=[*gen_prompt, *comp_ids], - renderer_token_ids=[*renderer_gen_prompt, *comp_ids], mask=[False] * len(gen_prompt) + [True] * len(comp_ids), is_content=([False] * len(gen_prompt) + [True] * len(comp_ids)) if has_is_content @@ -966,7 +793,7 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: # nodes in creation order, then the assistant node), each getting the routing for its tokens. # The prefix goes in too, so the position the previous turn could only pad can be corrected. _attribute_routed_experts( - trace, prefix, new_node_ids, path_len, tokens.routed_experts if tokens else None + trace, prefix, new_node_ids, path_len, routed_experts, routing_start ) # Sampling masks are completion-aligned, so only the sampled node carries them. diff --git a/verifiers/v1/harnesses/hermes_agent/harness.py b/verifiers/v1/harnesses/hermes_agent/harness.py index 734ed9b5bf..9dd528191c 100644 --- a/verifiers/v1/harnesses/hermes_agent/harness.py +++ b/verifiers/v1/harnesses/hermes_agent/harness.py @@ -84,9 +84,8 @@ async def prepare_acp( "api": endpoint, "api_key": secret, "discover_models": False, + "transport": "${HERMES_INTERCEPT_TRANSPORT}", } - if ctx.client.type == "eval": - provider["transport"] = "${HERMES_INTERCEPT_TRANSPORT}" config = { "model": model, # The ACP client already approves tool requests. Avoid routing Hermes' diff --git a/verifiers/v1/interception/server.py b/verifiers/v1/interception/server.py index 9232c5ebde..52df775a1b 100644 --- a/verifiers/v1/interception/server.py +++ b/verifiers/v1/interception/server.py @@ -38,10 +38,10 @@ from pydantic_core import PydanticSerializationError, from_json, to_json from verifiers.v1 import graph -from verifiers.v1.clients import Client, resolve_client +from verifiers.v1.clients import Client from verifiers.v1.clients.base import join_url from verifiers.v1.configs.client import ( - BaseClientConfig, + ClientConfig, resolve_api_key, resolve_headers, ) @@ -263,7 +263,7 @@ def load(self) -> int: """Rollouts currently registered — what the pools balance on.""" return len(self.sessions) - def _client(self, config: BaseClientConfig) -> Client: + def _client(self, config: ClientConfig) -> Client: """The server-owned client for `config` — one per distinct endpoint config, shared by every session registered under it, so the rollouts this server multiplexes reuse one bounded keepalive pool instead of each opening (and tearing down) their own @@ -271,7 +271,7 @@ def _client(self, config: BaseClientConfig) -> Client: key = config.model_dump_json() client = self.clients.get(key) if client is None: - client = self.clients[key] = resolve_client(config) + client = self.clients[key] = Client(config) self.stack.push_async_callback(client.close) return client @@ -647,10 +647,8 @@ async def sample(body: dict, model_request: Request) -> ReplayResponse: call_response, events = await session.client._complete( dialect, body, - session.ctx.sampling, headers=upstream_headers, session_id=session.trace.id, - turn=turn, ) logger.debug( "intercept turn: id=%s tools=%d", diff --git a/verifiers/v1/serve/pool.py b/verifiers/v1/serve/pool.py index 625e6a827f..7f32ead370 100644 --- a/verifiers/v1/serve/pool.py +++ b/verifiers/v1/serve/pool.py @@ -1,7 +1,7 @@ """Env-server worker pool: a ROUTER broker over N worker processes. A lone `EnvServer` runs every rollout as an `asyncio.Task` on one event loop, so -CPU-bound work (renderer tokenization, scoring) competes for that loop; the pool +CPU-bound task setup and scoring compete for that loop; the pool spreads that work across processes. A broker binds the client-facing ROUTER (the *same* wire protocol as a lone diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 24fe9ebaa5..045539ef01 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -1,7 +1,7 @@ import base64 from collections.abc import Iterable from dataclasses import dataclass -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, NotRequired import numpy as np from pydantic import ( @@ -109,7 +109,7 @@ class ToolMessage(BaseModel): tool_call_id: str content: MessageContent name: str | None = None - """Needed by templates such as Harmony when bridge tails omit the issuing call.""" + """Tool name for templates such as Harmony when the issuing call is absent.""" Message = Annotated[ @@ -205,11 +205,12 @@ def total_tokens(self) -> int: class RoutedExperts(TypedDict): - """Base64 uint8 `[tokens, layers, top_k]` routing and its prompt offset.""" + """Base64 integer `[tokens, layers, top_k]` routing and its prompt offset.""" data: Any shape: list[int] start: int + dtype: NotRequired[str] @dataclass @@ -253,7 +254,7 @@ def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: class TurnTokens(BaseModel): - """Training tokens from renderer tokenization or provider-returned token IDs.""" + """Exact inference tokens and optional per-message training attribution.""" model_config = ConfigDict(arbitrary_types_allowed=True) @@ -262,28 +263,20 @@ class TurnTokens(BaseModel): completion_ids: list[int] = Field(default_factory=list) completion_logprobs: list[float] = Field(default_factory=list) - # Transient graph-construction metadata, consumed by the turn's commit and excluded - # from serialized responses. - renderer_prompt_ids: list[int] | None = Field(default=None, exclude=True) - """Logical renderer IDs before multimodal expansion, retained for bridge extension.""" - bridged: bool = Field(default=False, exclude=True) - """Whether the renderer constructed this prompt by extending a stored prefix.""" # Transient carrier (excluded): per-message token spans into `prompt_ids` from the renderer, # consumed by the turn's `commit` to attribute tokens per message, then dropped. message_spans: list[tuple[int, int] | None] | None = Field( default=None, exclude=True ) is_content: list[bool] | None = Field(default=None, exclude=True) - # Authoritative effective-prompt ranges returned by vLLM, flattened across modalities. - mm_placeholders: list[tuple[int, int]] | None = Field(default=None, exclude=True) # Transient carrier (excluded): the renderer's special-token id -> modality marker map, # stamped onto `Trace.mm_token_type_id_map` by the turn's `commit`. None unless the # rendering renderer is multimodal. mm_token_type_id_map: dict[int, int] | None = Field(default=None, exclude=True) - # Transient carrier (excluded): the MoE expert-routing data from `generate` (expert ids + # Transient carrier (excluded): the inference server's MoE expert-routing data (expert ids # per token), attributed per node by the turn's `commit` into `MessageNode.routed_experts`, # then dropped. None unless the engine ran with `enable_return_routed_experts`. - routed_experts: RoutedExperts | None = Field(default=None, exclude=True) + routed_experts: RoutedExperts | str | None = Field(default=None, exclude=True) # Transient carrier (excluded): per-completion-token sampling masks, # attributed to the assistant node by the turn's `commit`, then dropped. sampling_mask: SamplingMask | None = Field(default=None, exclude=True)