From 2b0bfb0b45ebf5e37ca54575c058408813bc654f Mon Sep 17 00:00:00 2001 From: jack-scitix-ai Date: Tue, 28 Jul 2026 15:33:59 +0800 Subject: [PATCH 1/2] feat(models): capability-based Model IR + Transport frontends (RFC #25) --- docs/guide/configuration.md | 14 +- pyproject.toml | 4 + sieval/cli/leaderboard/session.py | 55 +- sieval/core/models/__init__.py | 43 + sieval/core/models/capabilities.py | 46 + sieval/core/models/chat_model.py | 294 +---- sieval/core/models/exceptions.py | 13 + sieval/core/models/gen_model.py | 272 +--- sieval/core/models/ir.py | 350 ++++++ sieval/core/models/model.py | 366 +++++- sieval/core/models/sglang_gen_model.py | 347 +----- sieval/core/models/transport.py | 48 + sieval/core/models/transports/__init__.py | 17 + sieval/core/models/transports/openai_chat.py | 280 +++++ .../models/transports/openai_completions.py | 314 +++++ sieval/core/models/transports/sglang.py | 371 ++++++ sieval/core/tasks/task.py | 16 +- tests/README.md | 6 +- tests/conftest.py | 226 ++-- tests/integration/resume/test_advanced.py | 7 +- .../integration/test_model_backward_compat.py | 129 ++ tests/integration/test_multi_task.py | 2 +- tests/integration/test_runner_edge_cases.py | 18 +- tests/unit/cli/leaderboard/test_session.py | 35 +- tests/unit/conftest.py | 9 + tests/unit/core/models/test_chat_model.py | 1106 +---------------- tests/unit/core/models/test_gen_model.py | 785 +----------- tests/unit/core/models/test_ir.py | 216 ++++ tests/unit/core/models/test_model.py | 300 ++++- tests/unit/core/models/test_model_arun.py | 104 ++ .../unit/core/models/test_model_derivation.py | 111 +- .../unit/core/models/test_sglang_gen_model.py | 598 +-------- tests/unit/core/models/transports/__init__.py | 0 .../models/transports/test_openai_chat.py | 677 ++++++++++ .../transports/test_openai_completions.py | 587 +++++++++ .../core/models/transports/test_sglang.py | 504 ++++++++ tests/unit/core/tasks/test_task.py | 80 +- tests/unit/tasks/conftest.py | 31 + .../tasks/test_arc_challenge_kshot_clp.py | 43 +- .../tasks/test_arc_challenge_kshot_ppl.py | 60 +- tests/unit/tasks/test_arc_easy_kshot_clp.py | 36 +- tests/unit/tasks/test_arc_easy_kshot_ppl.py | 47 +- tests/unit/tasks/test_c_eval_kshot_clp.py | 44 +- tests/unit/tasks/test_cmmlu_kshot_clp.py | 57 +- tests/unit/tasks/test_gsm8k_0shot_gen.py | 38 +- tests/unit/tasks/test_gsm8k_kshot_base_gen.py | 39 +- tests/unit/tasks/test_hellaswag_kshot_ppl.py | 55 +- .../test_hendrycks_math_kshot_base_gen.py | 37 +- .../tasks/test_human_eval_0shot_base_gen.py | 43 +- ...odebench_code_generation_kshot_base_gen.py | 42 +- tests/unit/tasks/test_mbpp_kshot_base_gen.py | 45 +- tests/unit/tasks/test_mmlu_0shot_gen.py | 25 +- tests/unit/tasks/test_mmlu_kshot_clp.py | 44 +- tests/unit/tasks/test_mmmlu_kshot_clp.py | 48 +- tests/unit/tasks/test_openbookqa_kshot_gen.py | 29 +- .../tasks/test_simpleqa_verified_0shot_gen.py | 28 +- .../tasks/test_theoremqa_kshot_base_gen.py | 41 +- 57 files changed, 5148 insertions(+), 4034 deletions(-) create mode 100644 sieval/core/models/capabilities.py create mode 100644 sieval/core/models/exceptions.py create mode 100644 sieval/core/models/ir.py create mode 100644 sieval/core/models/transport.py create mode 100644 sieval/core/models/transports/__init__.py create mode 100644 sieval/core/models/transports/openai_chat.py create mode 100644 sieval/core/models/transports/openai_completions.py create mode 100644 sieval/core/models/transports/sglang.py create mode 100644 tests/integration/test_model_backward_compat.py create mode 100644 tests/unit/core/models/test_ir.py create mode 100644 tests/unit/core/models/test_model_arun.py create mode 100644 tests/unit/core/models/transports/__init__.py create mode 100644 tests/unit/core/models/transports/test_openai_chat.py create mode 100644 tests/unit/core/models/transports/test_openai_completions.py create mode 100644 tests/unit/core/models/transports/test_sglang.py create mode 100644 tests/unit/tasks/conftest.py diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index f6b26c55..aad1af00 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -28,9 +28,10 @@ models: args: concurrency_limit: 64 - # Derived model with type conversion (ChatModel -> GenModel) + # A different kind needs its own base model — a derived model always + # inherits its base's kind (cross-kind `type:` removed by RFC #25) gen_model: - base: base_model + name: "gpt-4o" type: "gen" args: concurrency_limit: 32 @@ -53,8 +54,8 @@ tasks: ### Key Concepts -- **Model derivation**: `base: parent_model` inherits client, limiter, and args -- **Type conversion**: `type: "gen"` switches between ChatModel and GenModel +- **Model derivation**: `base: parent_model` inherits client, limiter, args — and always the base's kind +- **Model kind**: `type: "chat"` / `type: "gen"` selects ChatModel or GenModel on a base model only; cross-kind `type:` conversion on derived models was removed by RFC #25 — define a separate base model instead - **Quota allocation**: `concurrency_limit` in `args` reserves capacity from base - **Class resolution**: built-in classes (exported by `sieval.tasks` / `sieval.datasets`) use short names; custom classes must use full module paths (`my_pkg.my_module.MyTask`) @@ -81,17 +82,18 @@ preprocess -> infer -> postprocess -> feedback -> report Hierarchical concurrency control — derive child models from a base and allocate API quotas: ```python -from sieval.core.models import ChatModel, GenModel +from sieval.core.models import ChatModel base = ChatModel("gpt-4o", concurrency_limit=128) math_model = base.with_args(concurrency_limit=64) # reserves 64 code_model = base.with_args(concurrency_limit=32) # reserves 32 -gen_model = base.as_type(GenModel) # same quota, different type # base uses remaining elastic capacity (128 - 64 - 32 = 32) ``` +A derived model always keeps its base's kind. Cross-kind conversion (`as_type`) was removed by RFC #25 — to reach the same endpoint through a `GenModel`, define a separate base model instead. + ## Anomaly Detection Built-in anomaly detection runs automatically after each task and saves to `anomalies.json`. Rules are filtered by task tags — custom rules can be added via `@sieval_detection_rule` (see `sieval/core/tasks/anomaly.py`). diff --git a/pyproject.toml b/pyproject.toml index ad53d496..6a09ce93 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,7 +177,11 @@ also_copy = [ "sieval/tasks", "sieval/infer", "sieval/cli", + "sieval/meta", + "sieval/__init__.py", "sieval/__main__.py", + "sieval/_version.py", + "scripts", ] [tool.pytest] diff --git a/sieval/cli/leaderboard/session.py b/sieval/cli/leaderboard/session.py index 9ff7c92e..5405cf41 100644 --- a/sieval/cli/leaderboard/session.py +++ b/sieval/cli/leaderboard/session.py @@ -24,7 +24,7 @@ from sieval.cli.leaderboard.card import AlignmentCard, load_card from sieval.core.datasets import Dataset -from sieval.core.models import ChatModel, GenModel, Model, SglangGenModel +from sieval.core.models import Capability, ChatModel, GenModel, Model, SglangGenModel from sieval.core.runners import MultiTaskRunner, TaskRunnerConfig from sieval.core.tasks.context import TaskAction from sieval.core.types import JSONValue @@ -1048,41 +1048,28 @@ def _setup_models(self) -> None: # Extract concurrency_limit separately for with_args concurrency_limit = args.pop("concurrency_limit", None) - # Check if type conversion is needed + # RFC #25 dropped cross-kind model conversion (as_type): a + # derived model always inherits its base's kind. `type:` on a + # derived model is accepted only when it matches the base. target_type = cfg.get("type") - if target_type == "gen": - # An sglang-backed base is already "gen"; as_type(GenModel) - # would swap it to the OpenAI /v1/completions path (which - # rejects echo+logprobs), silently defeating engine: sglang. - # Preserve it. A plain GenModel base is unaffected. - if isinstance(base_model, SglangGenModel): - new_model = base_model - else: - new_model = base_model.as_type(GenModel) - logger.info( - "Created derived model '{}' from '{}' " - "with type conversion to '{}'", - name, - base_name, - target_type, - ) - elif target_type == "chat": - new_model = base_model.as_type(ChatModel) - logger.info( - "Created derived model '{}' from '{}' " - "with type conversion to '{}'", - name, - base_name, - target_type, - ) - elif target_type: - raise ValueError( - f"Model '{name}' has invalid type '{target_type}'. " - "Expected 'chat' or 'gen'" + if target_type: + if target_type not in ("chat", "gen"): + raise ValueError( + f"Model '{name}' has invalid type '{target_type}'. " + "Expected 'chat' or 'gen'" + ) + base_kind = ( + "chat" if Capability.Chat in base_model.capabilities else "gen" ) - else: - # No type conversion, just derive - new_model = base_model + if target_type != base_kind: + raise ValueError( + f"Derived model '{name}' cannot convert base " + f"'{base_name}' from '{base_kind}' to " + f"'{target_type}': cross-kind conversion was " + "removed (RFC #25). Define a separate base model " + "with the desired type instead." + ) + new_model = base_model # Apply additional args (including concurrency_limit) if concurrency_limit is not None or args: diff --git a/sieval/core/models/__init__.py b/sieval/core/models/__init__.py index aaa1a3b7..de249f55 100644 --- a/sieval/core/models/__init__.py +++ b/sieval/core/models/__init__.py @@ -1,15 +1,58 @@ +"""Model backends, the provider-agnostic IR, and Transport frontends. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from .capabilities import Capability from .chat_model import ChatModel +from .exceptions import CapabilityError from .gen_model import GenModel +from .ir import ( + Citation, + GroundingChunk, + GroundingMetadata, + InputScoringResult, + ReasoningOutput, + ReasoningParams, + Request, + Response, + SamplingParams, + ServerToolSpec, + ServerToolType, + ServerToolUse, + TokenLogprob, + TopKEntry, + UsageStats, +) from .model import Model, ModelCallMeta, ModelMeta, ModelOutput, ModelUsage from .sglang_gen_model import SglangGenModel +from .transport import Transport __all__ = [ + "Capability", + "CapabilityError", "ChatModel", + "Citation", "GenModel", + "GroundingChunk", + "GroundingMetadata", + "InputScoringResult", "Model", "ModelCallMeta", "ModelMeta", "ModelOutput", "ModelUsage", + "ReasoningOutput", + "ReasoningParams", + "Request", + "Response", + "SamplingParams", + "ServerToolSpec", + "ServerToolType", + "ServerToolUse", "SglangGenModel", + "TokenLogprob", + "TopKEntry", + "Transport", + "UsageStats", ] diff --git a/sieval/core/models/capabilities.py b/sieval/core/models/capabilities.py new file mode 100644 index 00000000..35138e57 --- /dev/null +++ b/sieval/core/models/capabilities.py @@ -0,0 +1,46 @@ +"""Capability catalog for the Model IR. + +Each ``Capability`` represents a feature a Transport can declare support for. +``Model.assert_capability()`` checks these at setup time — a Request using a +feature the Transport lacks is rejected immediately, never silently ignored. + +AI-Generated Code - Claude Opus 4.8 (Anthropic) +""" + +from enum import Flag, auto + + +class Capability(Flag): + # ── Input modality ──────────────────────────────────────── + Completion = auto() # input: str (base model) + Chat = auto() # input: list[Message] + + # ── Tools ───────────────────────────────────────────────── + FunctionCalling = auto() # client-side tool call (formerly Tools) + ServerTools = auto() # provider-hosted tools master switch + WebSearch = auto() + WebFetch = auto() + HostedCodeExecution = auto() + FileSearch = auto() + ComputerUse = auto() + + # ── Reasoning ───────────────────────────────────────────── + Reasoning = auto() + ReasoningOptional = auto() # can be turned off (OpenAI o-series, older Claude) + ReasoningAlwaysOn = auto() # forced on (Claude Fable 5 / Mythos 5) + ReasoningEffort = auto() # supports effort enum + ReasoningBudget = auto() # supports exact budget_tokens + ReasoningMode = auto() # supports mode (OpenAI standard/pro) + ReasoningContext = auto() # supports cross-turn reasoning state + ReasoningTaskBudget = auto() # supports cross-call total budget (Anthropic) + + # ── Logprobs / scoring ──────────────────────────────────── + TopKLogprobs = auto() + InputScoring = auto() # prompt-side logprobs (BPB/PPL) + SampledLogprobs = auto() + SampledLogprobsWithTokenIds = auto() # token_id populated (sglang native) + + # ── Other ───────────────────────────────────────────────── + StructuredOutput = auto() + Prefill = auto() + FIM = auto() diff --git a/sieval/core/models/chat_model.py b/sieval/core/models/chat_model.py index f72d11d0..bf6fb470 100644 --- a/sieval/core/models/chat_model.py +++ b/sieval/core/models/chat_model.py @@ -1,288 +1,26 @@ -"""ChatModel: chat completions API backend with reasoning content support.""" +"""ChatModel: chat completions API backend. + +RFC #25 moved the wire logic (streaming accumulation, reasoning extraction, +logprob parsing) into +:class:`~sieval.core.models.transports.openai_chat.OpenAIChatTransport`; this +class is the backend selector that pairs the shared client/limiter pool with +that transport. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" from collections.abc import Iterable -from typing import override from openai.types.chat import ChatCompletionMessageParam -from .model import Model, ModelOutput, ModelUsage +from .model import Model +from .transport import Transport class ChatModel(Model[str | Iterable[ChatCompletionMessageParam]]): - """Model subclass for the chat completions API (streaming + non-streaming).""" - - @override - async def _agenerate_impl( - self, prompt: str | Iterable[ChatCompletionMessageParam], **kwargs - ) -> ModelOutput: - if isinstance(prompt, str): - messages = [{"role": "user", "content": prompt}] - elif isinstance(prompt, Iterable): - messages = prompt - else: - raise TypeError("ChatModel requires a string or iterable of messages.") - - final_kwargs = {**self._kwargs, **kwargs} - - num_choices_raw = final_kwargs.get("n", 1) - if isinstance(num_choices_raw, bool) or not isinstance(num_choices_raw, int): - raise TypeError( - "n must be an int, got " - f"{type(num_choices_raw).__name__}: {num_choices_raw!r}" - ) - if num_choices_raw < 1: - raise ValueError(f"n must be >= 1, got {num_choices_raw}") - num_choices = num_choices_raw - texts = [""] * num_choices - finish_reasons = [""] * num_choices - reasoning_texts = [""] * num_choices - usage: ModelUsage | None = None - - # Snapshot params for stable meta serialization - request_params = dict(final_kwargs) - if "stream" not in request_params: - request_params["stream"] = True - stream_mode = request_params["stream"] - if not isinstance(stream_mode, bool): - raise TypeError( - "stream must be a bool, got " - f"{type(stream_mode).__name__}: {stream_mode!r}" - ) - if stream_mode and "stream_options" not in request_params: - request_params["stream_options"] = {"include_usage": True} - - resp = await self._client.chat.completions.create( # type: ignore[no-matching-overload] - model=self._model, - messages=messages, - **request_params, - ) - response_model: str | None = None - system_fingerprint: str | None = None - - if stream_mode: - async for chunk in resp: - if response_model is None: - response_model = getattr(chunk, "model", None) - if system_fingerprint is None: - system_fingerprint = getattr(chunk, "system_fingerprint", None) - if chunk.choices: - for choice in chunk.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - finish_reasons[idx] = choice.finish_reason or "" - if choice.delta: - if choice.delta.content: - texts[idx] += choice.delta.content - rc = getattr(choice.delta, "reasoning", None) - if rc: - reasoning_texts[idx] += rc - else: - rc = getattr(choice.delta, "reasoning_content", None) - if rc: - reasoning_texts[idx] += rc - chunk_usage = getattr(chunk, "usage", None) - if chunk_usage is not None: - usage = { - "input_tokens": chunk_usage.prompt_tokens, - "output_tokens": chunk_usage.completion_tokens, - "total_tokens": chunk_usage.total_tokens, - } - else: - response_model = getattr(resp, "model", None) - system_fingerprint = getattr(resp, "system_fingerprint", None) - for choice in resp.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - - message = choice.message - content = message.content - if content is not None: - texts[idx] += content - finish_reasons[idx] = choice.finish_reason or "" - rc = getattr(message, "reasoning", None) - if rc: - reasoning_texts[idx] += rc - else: - rc = getattr(message, "reasoning_content", None) - if rc: - reasoning_texts[idx] += rc - - raw_usage = resp.usage - if raw_usage is not None: - usage = { - "input_tokens": raw_usage.prompt_tokens, - "output_tokens": raw_usage.completion_tokens, - "total_tokens": raw_usage.total_tokens, - } - - return ModelOutput( - model=self.meta(), # Auto-attach model info - texts=texts, - finish_reasons=finish_reasons, - reasoning_texts=reasoning_texts, - usage=usage, - request_params=request_params, - response_model=response_model, - system_fingerprint=system_fingerprint, - ) - - @override - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - messages = [{"role": "user", "content": prompt}] - - final_kwargs = { - **self._kwargs, - **kwargs, - "max_tokens": max_tokens, - "temperature": temperature, - "logprobs": True, - "top_logprobs": logprobs, - } - - num_choices_raw = final_kwargs.get("n", 1) - if isinstance(num_choices_raw, bool) or not isinstance(num_choices_raw, int): - raise TypeError( - "n must be an int, got " - f"{type(num_choices_raw).__name__}: {num_choices_raw!r}" - ) - if num_choices_raw < 1: - raise ValueError(f"n must be >= 1, got {num_choices_raw}") - num_choices = num_choices_raw - if num_choices > 1: - raise ValueError(f"alogprobs only supports n=1; received n={num_choices}") - texts = [""] * num_choices - reasoning_texts = [""] * num_choices - finish_reasons: list[str] = [""] * num_choices - usage: ModelUsage | None = None - tokens: list[str] = [] - token_logprobs: list[float | None] = [] - saw_logprobs = False - - # Snapshot params for stable meta serialization - request_params = dict(final_kwargs) - if "stream" not in request_params: - request_params["stream"] = True - stream_mode = request_params["stream"] - if not isinstance(stream_mode, bool): - raise TypeError( - "stream must be a bool, got " - f"{type(stream_mode).__name__}: {stream_mode!r}" - ) - if stream_mode and "stream_options" not in request_params: - request_params["stream_options"] = {"include_usage": True} - - resp = await self._client.chat.completions.create( # type: ignore[no-matching-overload] - model=self._model, - messages=messages, - **request_params, - ) - response_model: str | None = None - system_fingerprint: str | None = None - - if stream_mode: - async for chunk in resp: - if response_model is None: - response_model = getattr(chunk, "model", None) - if system_fingerprint is None: - system_fingerprint = getattr(chunk, "system_fingerprint", None) - if chunk.choices: - for choice in chunk.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - finish_reasons[idx] = choice.finish_reason or "" - - if choice.delta: - if choice.delta.content: - texts[idx] += choice.delta.content - rc = getattr(choice.delta, "reasoning", None) - if rc: - reasoning_texts[idx] += rc - else: - rc = getattr(choice.delta, "reasoning_content", None) - if rc: - reasoning_texts[idx] += rc - if idx == 0: - logprobs_obj = getattr(choice, "logprobs", None) - if logprobs_obj is not None: - saw_logprobs = True - content = getattr(logprobs_obj, "content", None) or [] - if content: - tokens.extend(item.token for item in content) - token_logprobs.extend( - item.logprob for item in content - ) - - chunk_usage = getattr(chunk, "usage", None) - if chunk_usage is not None: - usage = { - "input_tokens": chunk_usage.prompt_tokens, - "output_tokens": chunk_usage.completion_tokens, - "total_tokens": chunk_usage.total_tokens, - } - else: - response_model = getattr(resp, "model", None) - system_fingerprint = getattr(resp, "system_fingerprint", None) - for choice in resp.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - - message = choice.message - content = message.content - if content is not None: - texts[idx] += content - finish_reasons[idx] = choice.finish_reason or "" - rc = getattr(message, "reasoning", None) - if rc: - reasoning_texts[idx] += rc - else: - rc = getattr(message, "reasoning_content", None) - if rc: - reasoning_texts[idx] += rc - if idx == 0: - logprobs_obj = choice.logprobs - if logprobs_obj is not None: - saw_logprobs = True - if logprobs_obj.content: - tokens.extend(item.token for item in logprobs_obj.content) - token_logprobs.extend( - item.logprob for item in logprobs_obj.content - ) - - raw_usage = resp.usage - if raw_usage is not None: - usage = { - "input_tokens": raw_usage.prompt_tokens, - "output_tokens": raw_usage.completion_tokens, - "total_tokens": raw_usage.total_tokens, - } + """Model backend for the chat completions API (streaming + non-streaming).""" - if not saw_logprobs: - raise RuntimeError( - "Streaming logprobs not supported by server for chat completions." - ) + def _build_default_transport(self) -> Transport: + from .transports.openai_chat import OpenAIChatTransport - return ModelOutput( - model=self.meta(), - texts=texts, - finish_reasons=finish_reasons, - reasoning_texts=reasoning_texts, - logprobs_tokens=tokens, - logprobs=token_logprobs, - usage=usage, - request_params=request_params, - response_model=response_model, - system_fingerprint=system_fingerprint, - ) + return OpenAIChatTransport(client=self._client, model=self._model) diff --git a/sieval/core/models/exceptions.py b/sieval/core/models/exceptions.py new file mode 100644 index 00000000..2973b7d9 --- /dev/null +++ b/sieval/core/models/exceptions.py @@ -0,0 +1,13 @@ +"""Model-layer exceptions. + +AI-Generated Code - Claude Opus 4.8 (Anthropic) +""" + + +class CapabilityError(ValueError): + """Raised at setup when a Request requires a Capability the Transport lacks. + + This replaces the previous pattern of silently ignoring unsupported + parameters (e.g. ``echo=True`` on ChatModel). Fail loud at setup, never + silently return wrong results. + """ diff --git a/sieval/core/models/gen_model.py b/sieval/core/models/gen_model.py index 8debc87b..74e08cef 100644 --- a/sieval/core/models/gen_model.py +++ b/sieval/core/models/gen_model.py @@ -1,266 +1,22 @@ -"""GenModel: text completions API backend.""" +"""GenModel: text completions API backend. -from collections.abc import Mapping, Sequence -from typing import override +RFC #25 moved the wire logic (streaming accumulation, echo split, top-logprob +sanitizing) into +:class:`~sieval.core.models.transports.openai_completions.OpenAICompletionsTransport`; +this class is the backend selector that pairs the shared client/limiter pool +with that transport. -from .model import Model, ModelOutput, ModelUsage +AI-Generated Code - Claude Fable 5 (Anthropic) +""" - -def _completion_top_logprobs(raw: object) -> list[dict[str, float]]: - if not isinstance(raw, Sequence) or isinstance(raw, str | bytes): - return [] - - top_logprobs = [] - for item in raw: - if item is None: - top_logprobs.append({}) - continue - if not isinstance(item, Mapping): - continue - top_logprobs.append( - { - token: float(logprob) - for token, logprob in item.items() - if isinstance(token, str) and isinstance(logprob, int | float) - } - ) - return top_logprobs +from .model import Model +from .transport import Transport class GenModel(Model[str]): - """Model subclass for the completions API (streaming + non-streaming).""" - - @override - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - if not isinstance(prompt, str): - raise TypeError("GenModel requires a string prompt.") - - final_kwargs = {**self._kwargs, **kwargs} - - num_choices_raw = final_kwargs.get("n", 1) - if isinstance(num_choices_raw, bool) or not isinstance(num_choices_raw, int): - raise TypeError( - "n must be an int, got " - f"{type(num_choices_raw).__name__}: {num_choices_raw!r}" - ) - if num_choices_raw < 1: - raise ValueError(f"n must be >= 1, got {num_choices_raw}") - num_choices = num_choices_raw - texts = [""] * num_choices - finish_reasons = [""] * num_choices - usage: ModelUsage | None = None - - # Snapshot params for stable meta serialization - request_params = dict(final_kwargs) - if "stream" not in request_params: - request_params["stream"] = True - stream_mode = request_params["stream"] - if not isinstance(stream_mode, bool): - raise TypeError( - "stream must be a bool, got " - f"{type(stream_mode).__name__}: {stream_mode!r}" - ) - if stream_mode and "stream_options" not in request_params: - request_params["stream_options"] = {"include_usage": True} - - resp = await self._client.completions.create( - model=self._model, - prompt=prompt, - **request_params, - ) - response_model: str | None = None - system_fingerprint: str | None = None - - if stream_mode: - async for chunk in resp: - if response_model is None: - response_model = getattr(chunk, "model", None) - if system_fingerprint is None: - system_fingerprint = getattr(chunk, "system_fingerprint", None) - if chunk.choices: - for choice in chunk.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - texts[idx] += choice.text or "" - finish_reasons[idx] = choice.finish_reason or "" - chunk_usage = getattr(chunk, "usage", None) - if chunk_usage is not None: - usage = { - "input_tokens": chunk_usage.prompt_tokens, - "output_tokens": chunk_usage.completion_tokens, - "total_tokens": chunk_usage.total_tokens, - } - else: - response_model = getattr(resp, "model", None) - system_fingerprint = getattr(resp, "system_fingerprint", None) - for choice in resp.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - texts[idx] += choice.text or "" - finish_reasons[idx] = choice.finish_reason or "" - raw_usage = resp.usage - if raw_usage is not None: - usage = { - "input_tokens": raw_usage.prompt_tokens, - "output_tokens": raw_usage.completion_tokens, - "total_tokens": raw_usage.total_tokens, - } - - # No reasoning_texts for GenModel now - return ModelOutput( - model=self.meta(), # Auto-attach model info - texts=texts, - finish_reasons=finish_reasons, - usage=usage, - request_params=request_params, - response_model=response_model, - system_fingerprint=system_fingerprint, - ) - - @override - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - final_kwargs = { - **self._kwargs, - **kwargs, - "max_tokens": max_tokens, - "logprobs": logprobs, - "echo": echo, - "temperature": temperature, - } - - num_choices_raw = final_kwargs.get("n", 1) - if isinstance(num_choices_raw, bool) or not isinstance(num_choices_raw, int): - raise TypeError( - "n must be an int, got " - f"{type(num_choices_raw).__name__}: {num_choices_raw!r}" - ) - if num_choices_raw < 1: - raise ValueError(f"n must be >= 1, got {num_choices_raw}") - num_choices = num_choices_raw - if num_choices > 1: - raise ValueError(f"alogprobs only supports n=1; received n={num_choices}") - texts = [""] * num_choices - finish_reasons = [""] * num_choices - usage: ModelUsage | None = None - tokens: list[str] = [] - token_logprobs: list[float | None] = [] - top_token_logprobs: list[dict[str, float]] = [] - saw_logprobs = False - - # Snapshot params for stable meta serialization - request_params = dict(final_kwargs) - if "stream" not in request_params: - request_params["stream"] = True - stream_mode = request_params["stream"] - if not isinstance(stream_mode, bool): - raise TypeError( - "stream must be a bool, got " - f"{type(stream_mode).__name__}: {stream_mode!r}" - ) - if stream_mode and "stream_options" not in request_params: - request_params["stream_options"] = {"include_usage": True} - - resp = await self._client.completions.create( # ty: ignore[no-matching-overload] - model=self._model, - prompt=prompt, - **request_params, - ) - response_model: str | None = None - system_fingerprint: str | None = None - - if stream_mode: - async for chunk in resp: - if response_model is None: - response_model = getattr(chunk, "model", None) - if system_fingerprint is None: - system_fingerprint = getattr(chunk, "system_fingerprint", None) - if chunk.choices: - for choice in chunk.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - - texts[idx] += choice.text or "" - finish_reasons[idx] = choice.finish_reason or "" - if idx == 0: - logprobs_obj = getattr(choice, "logprobs", None) - if logprobs_obj is not None: - saw_logprobs = True - if logprobs_obj.tokens: - tokens.extend(logprobs_obj.tokens) - if logprobs_obj.token_logprobs: - token_logprobs.extend(logprobs_obj.token_logprobs) - top_token_logprobs.extend( - _completion_top_logprobs( - getattr(logprobs_obj, "top_logprobs", None) - ) - ) - - chunk_usage = getattr(chunk, "usage", None) - if chunk_usage is not None: - usage = { - "input_tokens": chunk_usage.prompt_tokens, - "output_tokens": chunk_usage.completion_tokens, - "total_tokens": chunk_usage.total_tokens, - } - else: - response_model = getattr(resp, "model", None) - system_fingerprint = getattr(resp, "system_fingerprint", None) - for choice in resp.choices: - idx = choice.index - if not 0 <= idx < num_choices: - continue - - texts[idx] += choice.text or "" - finish_reasons[idx] = choice.finish_reason or "" - if idx == 0: - logprobs_obj = choice.logprobs - if logprobs_obj is not None: - saw_logprobs = True - if logprobs_obj.tokens: - tokens.extend(logprobs_obj.tokens) - if logprobs_obj.token_logprobs: - token_logprobs.extend(logprobs_obj.token_logprobs) - top_token_logprobs.extend( - _completion_top_logprobs( - getattr(logprobs_obj, "top_logprobs", None) - ) - ) - - raw_usage = resp.usage - if raw_usage is not None: - usage = { - "input_tokens": raw_usage.prompt_tokens, - "output_tokens": raw_usage.completion_tokens, - "total_tokens": raw_usage.total_tokens, - } + """Model backend for the completions API (streaming + non-streaming).""" - if not saw_logprobs: - raise RuntimeError( - "Streaming logprobs not supported by server for completions." - ) + def _build_default_transport(self) -> Transport: + from .transports.openai_completions import OpenAICompletionsTransport - # No reasoning_texts for GenModel now - return ModelOutput( - model=self.meta(), # Auto-attach model info - texts=texts, - finish_reasons=finish_reasons, - logprobs_tokens=tokens, - logprobs=token_logprobs, - top_logprobs=top_token_logprobs or None, - usage=usage, - request_params=request_params, - response_model=response_model, - system_fingerprint=system_fingerprint, - ) + return OpenAICompletionsTransport(client=self._client, model=self._model) diff --git a/sieval/core/models/ir.py b/sieval/core/models/ir.py new file mode 100644 index 00000000..6f1e4a8d --- /dev/null +++ b/sieval/core/models/ir.py @@ -0,0 +1,350 @@ +"""Provider-agnostic Model IR: Request, Response, and all supporting types. + +Design notes +------------ +- All dataclasses are frozen (immutable) to allow safe sharing across async tasks. +- Collection fields use ``tuple``, not ``list``, for the same reason. +- ``session_id`` supports stateful multi-turn (the Transport passes + ``previous_response_id`` / ``previous_interaction_id`` internally; the caller + never touches opaque state). +- ``ReasoningParams.opaque_roundtrip`` supports stateless multi-turn: the caller + echoes the value from ``Response.reasoning.opaque_roundtrip`` back into the + next ``Request.reasoning.opaque_roundtrip``; the Transport re-embeds it into + the correct wire item (OpenAI ``encrypted_content`` / Anthropic ``signature`` / + Google ``signature``). + +Persistence +----------- +``Response`` is the persisted record schema (coupled to RFC #24's resume +version gate). Every response-side nested type carries ``@sieval_record`` so it +round-trips back into a typed object rather than a plain dict: ``obj_to_dict`` +only stamps ``__sieval_mod__``/``__sieval_cls__`` markers on records, and a +``tuple`` of records only rehydrates correctly when the element type is itself a +record. Request-side types (``SamplingParams``, ``ReasoningParams``, +``ServerToolSpec``) are not persisted and are intentionally undecorated. + +AI-Generated Code - Claude Opus 4.8 (Anthropic) +""" + +from dataclasses import dataclass +from typing import Any, Literal + +from sieval.core.utils.serialization import sieval_record + +# ── Sampling params ─────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class SamplingParams: + """Provider-agnostic sampling controls.""" + + temperature: float | None = None + top_p: float | None = None + top_k_sampling: int | None = None # sampling top-k; distinct from logprobs top_k + max_tokens: int | None = None + stop: tuple[str, ...] | None = None + stop_token_ids: tuple[int, ...] | None = None + seed: int | None = None + frequency_penalty: float | None = None + presence_penalty: float | None = None + n: int = 1 + + +# ── Reasoning controls (five axes) ───────────────────────────────────────────── + + +@dataclass(frozen=True) +class ReasoningParams: + """Five-axis reasoning control. + + Axis 1 — intensity (transport picks whichever it supports): + effort: semantic level (OpenAI / Anthropic new / Google Gemini 3) + budget_tokens: exact token cap (Anthropic older models / Google Gemini 2.5) + + Axis 2 — execution tier (OpenAI only): + mode: "standard" | "pro" + + Axis 3 — cross-turn reasoning state: + context: "current_turn" | "all_turns" (OpenAI) + + Axis 4 — summary visibility: + summary: "none" | "auto" | "concise" | "detailed" + + Axis 5 — cross-call total budget (Anthropic agentic loop): + task_budget: int + + Stateless multi-turn: + opaque_roundtrip: echo ``Response.reasoning.opaque_roundtrip`` back here on + the next turn. The Transport re-embeds it into the correct wire format + (OpenAI ``encrypted_content`` / Anthropic ``ThinkingBlock.signature`` / + Google ``ThoughtStep.signature``). Do NOT inspect or modify this value. + """ + + effort: Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None = None + budget_tokens: int | None = None + mode: Literal["standard", "pro"] | None = None + context: Literal["current_turn", "all_turns"] | None = None + summary: Literal["none", "auto", "concise", "detailed"] | None = None + task_budget: int | None = None + opaque_roundtrip: str | None = None # stateless multi-turn round-trip + + +# ── Server tools ──────────────────────────────────────────────────────────────── + +ServerToolType = Literal[ + "web_search", + "web_fetch", + "code_execution", + "file_search", + "computer_use", + "image_generation", + "hosted_mcp", +] + + +@dataclass(frozen=True) +class ServerToolSpec: + """Spec for a single provider-hosted tool.""" + + type: ServerToolType + config: dict[str, Any] | None = None + + +# ── Request IR ────────────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class Request: + """Provider-agnostic inference request. + + input: + str → Completion modality (base model) + list[dict] → Chat modality (list of ``{"role": ..., "content": ...}``) + + session_id: + Stateful multi-turn. When set, the Transport passes it as + ``previous_response_id`` (OpenAI) or ``previous_interaction_id`` (Google) + on each ``arun`` call. The caller never needs to manage opaque state. + Obtain the value from ``Response.session_id`` of the previous turn. + + extra_wire_params: + Unrecognized generation kwargs the request builders could not map onto a + first-class IR field. Each Transport's ``lower()`` merges these into the + native wire body last (so an explicit IR field always wins). This is the + escape hatch that keeps ``with_args(**infer_args)`` working end-to-end. + """ + + input: str | list[dict[str, Any]] + + sampling: SamplingParams | None = None + + # logprobs / scoring + return_logprobs: bool = False + top_k: int = 0 + score_input: bool = False # InputScoring: prompt-side logprobs + + # structured output + response_format: dict[str, Any] | None = None + + # tools + tools: list[dict[str, Any]] | None = None # FunctionCalling + server_tools: tuple[ServerToolSpec, ...] | None = None # ServerTools + + # reasoning + reasoning: ReasoningParams | None = None + + # prefill / FIM + prefix: str | None = None + suffix: str | None = None + + # stateful multi-turn + session_id: str | None = None + + # wire scheduling: None → transport default (single-shot). True asks the + # transport to stream internally and accumulate; the caller still receives + # one terminal Response. Pure scheduling — never affects response content. + stream: bool | None = None + + # passthrough for kwargs without a first-class IR field + extra_wire_params: dict[str, Any] | None = None + + +# ── Response sub-types ──────────────────────────────────────────────────────── + + +@sieval_record +@dataclass(frozen=True) +class TokenLogprob: + """Single token with its log-probability. + + ``logprob`` is nullable: the first prompt token in an input-scoring result + has no preceding context and therefore no log-probability (sglang emits + ``None``; this matches the legacy ``ModelOutput.logprobs`` ``float | None`` + contract). + + ``token_id`` is always populated for sglang (native triple), and best-effort + for other transports. Consumers that need ``token_id`` must declare the + ``SampledLogprobsWithTokenIds`` capability. + + Both nullable fields are declared last with defaults so they survive + persistence: ``obj_to_dict`` drops ``None`` fields, so a required nullable + field could not be reconstructed on the path where it is ``None``. + """ + + token: str + logprob: float | None = None + token_id: int | None = None + + +@sieval_record +@dataclass(frozen=True) +class TopKEntry: + """One candidate in a top-k logprob list. + + ``token_id`` is declared last with a default for the same persistence reason + as :class:`TokenLogprob`. + """ + + token: str + logprob: float + token_id: int | None = None + + +@sieval_record +@dataclass(frozen=True) +class InputScoringResult: + """Prompt-side logprobs for BPB / perplexity computation.""" + + token_logprobs: tuple[TokenLogprob, ...] + byte_count: int | None = None + char_count: int | None = None + + +@sieval_record +@dataclass(frozen=True) +class ReasoningOutput: + """Reasoning channel in the response. + + opaque_roundtrip: + Provider-signed value that must be echoed back in + ``Request.reasoning.opaque_roundtrip`` on the next turn (stateless mode). + Source fields by provider: + OpenAI → ResponseReasoningItem.encrypted_content + Anthropic → ThinkingBlock.signature + Google → ThoughtStep.signature + Do NOT inspect or modify this value. + """ + + text: str | None = None + opaque_roundtrip: str | None = None + thinking_tokens: int = 0 + effort_used: str | None = None + + +@sieval_record +@dataclass(frozen=True) +class ServerToolUse: + """Record of one provider-hosted tool invocation and its result.""" + + tool_type: str + tool_use_id: str + input: dict[str, Any] + result: dict[str, Any] | None = None + error_code: str | None = None + + +@sieval_record +@dataclass(frozen=True) +class Citation: + """Web search source attribution.""" + + url: str + title: str | None = None + page_age: str | None = None + + +@sieval_record +@dataclass(frozen=True) +class GroundingChunk: + """A single source chunk backing a Google Search grounding result.""" + + uri: str + title: str | None = None + + +@sieval_record +@dataclass(frozen=True) +class GroundingMetadata: + """Google Search grounding result. + + ``rendered_content`` MUST be preserved and rendered per Google ToS. + A Transport's ``lift()`` must not drop this field. + """ + + chunks: tuple[GroundingChunk, ...] + rendered_content: str | None = None # Google ToS: must be rendered + + +@sieval_record +@dataclass(frozen=True) +class UsageStats: + """Token usage. ``reasoning_tokens`` is billed separately on all three cloud + providers.""" + + input_tokens: int = 0 + output_tokens: int = 0 + reasoning_tokens: int = 0 + cached_tokens: int = 0 + total_tokens: int = 0 + + +# ── Response IR ──────────────────────────────────────────────────────────────── + + +@sieval_record +@dataclass(frozen=True) +class Response: + """Provider-agnostic inference response. + + session_id: + Populated by the Transport in stateful mode. Pass this value as + ``Request.session_id`` on the next turn; the Transport will use it as + ``previous_response_id`` / ``previous_interaction_id``. + + logprobs / top_logprobs empty-vs-absent contract: + ``None`` means the server sent no logprob channel at all; an empty tuple + means the channel was present but carried no entries. Anomaly detection + distinguishes the two (present-but-empty is flagged), so Transports must + not collapse ``()`` to ``None`` when logprobs were requested. + + Provenance fields (``request_params``, ``response_model``, + ``system_fingerprint``) record what was actually sent on the wire and what + the server reported back. They are persisted for reproducibility records + (``build_model_call_meta``) and never branched on. + """ + + texts: tuple[str, ...] + + reasoning: ReasoningOutput | None = None + tool_calls: tuple[dict[str, Any], ...] | None = None + server_tool_uses: tuple[ServerToolUse, ...] | None = None + + logprobs: tuple[TokenLogprob, ...] | None = None + top_logprobs: tuple[tuple[TopKEntry, ...], ...] | None = None + input_scoring: InputScoringResult | None = None + + citations: tuple[Citation, ...] | None = None + grounding: GroundingMetadata | None = None + + # stateful multi-turn: pass to next Request.session_id + session_id: str | None = None + + # None when the server reported no usage (absence ≠ zeros). + usage: UsageStats | None = None + finish_reasons: tuple[str, ...] | None = None + + # provenance: the lowered wire params (prompt/messages excluded) and the + # server-reported model / fingerprint. + request_params: dict[str, Any] | None = None + response_model: str | None = None + system_fingerprint: str | None = None diff --git a/sieval/core/models/model.py b/sieval/core/models/model.py index aeaf3736..ab13625a 100644 --- a/sieval/core/models/model.py +++ b/sieval/core/models/model.py @@ -1,9 +1,18 @@ -"""Abstract model base class and shared types for model backends.""" +"""Abstract model base class and shared types for model backends. + +RFC #25: ``arun(Request) -> Response`` is the one primitive — it acquires both +limiters and delegates to the composed :class:`Transport`. ``agenerate`` and +``alogprobs`` are capability-gated sugar: thin wrappers that build a +:class:`Request` from legacy OpenAI-style kwargs, run it through ``arun``, and +bridge the :class:`Response` back to the legacy :class:`ModelOutput` shape for +existing consumers. New code should call ``arun`` directly. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" import copy -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import NotRequired, Self, TypedDict +from dataclasses import dataclass, replace +from typing import Any, NotRequired, Self, TypedDict, cast import anyio from openai import AsyncOpenAI @@ -12,6 +21,11 @@ from sieval.core.utils.concurrency import CompositeLimiter from sieval.core.utils.serialization import sieval_record +from .capabilities import Capability +from .exceptions import CapabilityError +from .ir import ReasoningParams, Request, Response, SamplingParams +from .transport import Transport + class ModelUsage(TypedDict): """Token usage statistics from a single model API call.""" @@ -81,13 +95,30 @@ class ModelOutput: system_fingerprint: str | None = None -class Model[TModelInput](ABC): - """Abstract base for all model backends. - - Uses an OpenAI-compatible AsyncClient. Provides two-level concurrency - control: ``_parent_limiter`` (shared API quota from a base model) and - ``_limiter`` (this model's reserved quota). Both limiters are acquired - before every ``agenerate`` / ``alogprobs`` call. +# Legacy OpenAI-style kwarg -> SamplingParams field handled by the request +# builders. Anything not listed here (and not a first-class Request field) +# falls through to Request.extra_wire_params. +_SAMPLING_KWARGS: dict[str, str] = { + "max_tokens": "max_tokens", + "max_completion_tokens": "max_tokens", + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k_sampling", + "seed": "seed", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", +} + + +class Model[TModelInput]: + """Base for all model backends. + + Composes a resource pool (two-level ``anyio`` limiters) with a + :class:`Transport` (the provider frontend). ``arun`` is the one primitive; + ``agenerate`` / ``alogprobs`` are backward-compatible sugar over it. + Concrete backends only select a default Transport via + ``_build_default_transport``; a bare ``Model`` has no transport and no + capabilities. """ def __init__( @@ -99,6 +130,7 @@ def __init__( concurrency_limit: int | None = None, parent_limiter: anyio.CapacityLimiter | None = None, extra: dict[str, JSONValue] | None = None, + transport: Transport | None = None, **kwargs, ): self._model = model @@ -125,6 +157,25 @@ def __init__( else None ) + # Provider frontend. When not injected, the subclass builds its default + # transport over the shared client (backward-compatible path). The + # transport is the single source of truth for `capabilities`; a model + # without one simply declares no IR capabilities and cannot `arun`. + self._transport: Transport | None = ( + transport if transport is not None else self._build_default_transport() + ) + + def _build_default_transport(self) -> Transport | None: + """Build this model's default Transport over ``self._client``. + + Overridden by each concrete backend to return its transport. Returns + ``None`` on the base so that a bare ``Model`` subclass remains + constructible (it just exposes no capabilities). Kept as a hook (rather + than an abstract method) so lazily-imported transport modules avoid an + import cycle with the backend subclasses. + """ + return None + def with_args( self, concurrency_limit: int | None = None, @@ -137,6 +188,12 @@ def with_args( model's limiter; ``None`` shares the existing limiter. Multi-level derivation is forbidden. + RFC #25 eventually decomposes this: per-task ``infer_args`` become + plain :class:`Request` fields and the sub-quota an explicit + resource-pool collaborator. That lands together with the task layer + migrating from ``agenerate(**kwargs)`` to ``arun(Request)``; until + then this remains the supported derivation mechanism. + Example:: child = base.with_args(concurrency_limit=64) @@ -164,23 +221,25 @@ def with_args( return new_model - def as_type(self, model_type: type[Self]) -> Self: - """Re-type this model (e.g. GenModel ↔ ChatModel), sharing client and limiters. + @property + def capabilities(self) -> frozenset[Capability]: + """The IR features this model's Transport honours (empty if none).""" + if self._transport is None: + return frozenset() + return self._transport.capabilities - Example:: + def assert_capability(self, *caps: Capability) -> None: + """Raise :class:`CapabilityError` if any requested capability is missing. - chat_model = gen_model.as_type(ChatModel) + This is the setup-time gate that replaces silently ignoring unsupported + parameters (e.g. ``echo=True`` on a chat backend). """ - if not isinstance(model_type, type) or not issubclass(model_type, Model): - raise TypeError(f"model_type must be a Model subclass, got {model_type}") - - # Copy and change class type - # This is safe because GenModel and ChatModel only differ in methods, - # not in instance attributes - new_model = copy.copy(self) - new_model.__class__ = model_type - - return new_model + missing = frozenset(caps) - self.capabilities + if missing: + raise CapabilityError( + f"{type(self._transport).__name__} does not support: " + + ", ".join(sorted(c.name for c in missing if c.name is not None)) + ) def get_available_quota(self) -> int | float: """Return the minimum available tokens across both limiters.""" @@ -246,10 +305,32 @@ def meta(self) -> ModelMeta: result["extra"] = dict(self._extra) return result - async def agenerate(self, prompt: TModelInput, **kwargs) -> ModelOutput: - """Generate text; acquires both limiters first.""" + # ── the primitive ───────────────────────────────────────────────────────── + + async def arun(self, req: Request) -> Response: + """Run one inference through the Transport; acquires both limiters first. + + The provider-agnostic primitive. Transport contract: + 1. Returns a terminal Response (all internal tool loops resolved). + 2. Stateful mode: pass ``Request.session_id``, get ``Response.session_id`` + back. + 3. Stateless mode: echo ``Response.reasoning.opaque_roundtrip`` back as + ``Request.reasoning.opaque_roundtrip`` on the next turn. + """ + if self._transport is None: + raise CapabilityError( + f"{type(self).__name__} has no Transport; arun() is unavailable." + ) async with CompositeLimiter(self._parent_limiter, self._limiter): - return await self._agenerate_impl(prompt, **kwargs) + return await self._transport.arun(req) + + # ── legacy sugar (thin wrappers over arun) ──────────────────────────────── + + async def agenerate(self, prompt: TModelInput, **kwargs) -> ModelOutput: + """Generate text. Sugar over :meth:`arun`, returning ``ModelOutput``.""" + req = self._build_generate_request(prompt, **kwargs) + resp = await self.arun(req) + return self._response_to_model_output(resp) async def alogprobs( self, @@ -261,28 +342,227 @@ async def alogprobs( temperature: float = 0.0, **kwargs, ) -> ModelOutput: - """Extract logprobs; acquires both limiters first.""" - async with CompositeLimiter(self._parent_limiter, self._limiter): - return await self._alogprobs_impl( - prompt, - max_tokens=max_tokens, - logprobs=logprobs, - echo=echo, - temperature=temperature, - **kwargs, + """Extract logprobs. Sugar over :meth:`arun`, returning ``ModelOutput``. + + ``echo=True`` requests prompt-side scoring, which requires the + ``InputScoring`` capability. On a backend that lacks it (e.g. a chat + completions transport) this raises :class:`CapabilityError` at the + call boundary instead of being silently ignored (the historical bug). + """ + if echo: + self.assert_capability(Capability.InputScoring) + req = self._build_logprobs_request( + prompt, + max_tokens=max_tokens, + logprobs=logprobs, + score_input=echo, + temperature=temperature, + **kwargs, + ) + resp = await self.arun(req) + if ( + resp.logprobs is None + and resp.input_scoring is None + and resp.top_logprobs is None + ): + raise RuntimeError("logprobs requested but the server returned none.") + return self._response_to_model_output(resp) + + # ── request builders (legacy kwargs -> IR) ──────────────────────────────── + + @staticmethod + def _validate_n(final_kwargs: dict[str, Any]) -> int: + """Validate and return ``n`` from merged kwargs.""" + n = final_kwargs.get("n", 1) + if isinstance(n, bool) or not isinstance(n, int): + raise TypeError(f"n must be an int, got {type(n).__name__}: {n!r}") + if n < 1: + raise ValueError(f"n must be >= 1, got {n}") + return n + + @staticmethod + def _coerce_input(prompt: Any) -> str | list[dict[str, Any]]: + """Coerce a legacy prompt into ``Request.input``. + + Modality validation stays with the Transport (a completions transport + rejects non-str input); this only normalizes the container shape. + """ + if isinstance(prompt, str): + return prompt + try: + return [dict(m) for m in prompt] + except TypeError: + raise TypeError( + "prompt must be a string or an iterable of messages, " + f"got {type(prompt).__name__}." + ) from None + + def _kwargs_to_request( + self, input_: str | list[dict[str, Any]], final_kwargs: dict[str, Any] + ) -> Request: + """Map merged OpenAI-style kwargs onto Request fields. + + Recognized keys become first-class IR fields; the remainder rides in + ``extra_wire_params`` so ``with_args(**infer_args)`` keeps working for + provider-specific params the IR does not model. + """ + kw = dict(final_kwargs) + + n = self._validate_n(kw) + kw.pop("n", None) + + stream = kw.pop("stream", True) + if not isinstance(stream, bool): + raise TypeError( + f"stream must be a bool, got {type(stream).__name__}: {stream!r}" ) - @abstractmethod - async def _agenerate_impl(self, prompt: TModelInput, **kwargs) -> ModelOutput: ... + sampling_fields: dict[str, Any] = {"n": n} + for src, dst in _SAMPLING_KWARGS.items(): + if src in kw: + value = kw.pop(src) + if value is not None: + sampling_fields.setdefault(dst, value) + if "stop" in kw: + stop = kw.pop("stop") + if stop is not None: + sampling_fields["stop"] = ( + (stop,) if isinstance(stop, str) else tuple(stop) + ) + if "stop_token_ids" in kw: + stop_ids = kw.pop("stop_token_ids") + if stop_ids is not None: + sampling_fields["stop_token_ids"] = tuple(stop_ids) + + # `logprobs` carries both dialects: chat's bool switch and the + # completions-style int top-k count. + return_logprobs = False + top_k = 0 + lp = kw.pop("logprobs", None) + if isinstance(lp, bool): + return_logprobs = lp + elif lp is not None: + return_logprobs = True + top_k = int(lp) + tlp = kw.pop("top_logprobs", None) + if tlp is not None: + top_k = int(tlp) + + response_format = kw.pop("response_format", None) + tools = kw.pop("tools", None) + + reasoning = None + effort = kw.pop("reasoning_effort", None) + if effort is not None: + reasoning = ReasoningParams(effort=cast("Any", effort)) + + return Request( + input=input_, + sampling=SamplingParams(**sampling_fields), + return_logprobs=return_logprobs, + top_k=top_k, + response_format=response_format, + tools=tools, + reasoning=reasoning, + stream=stream, + extra_wire_params=kw or None, + ) + + def _build_generate_request(self, prompt: TModelInput, **kwargs) -> Request: + """Build the Request for :meth:`agenerate` from merged kwargs.""" + final_kwargs = {**self._kwargs, **kwargs} + return self._kwargs_to_request(self._coerce_input(prompt), final_kwargs) - @abstractmethod - async def _alogprobs_impl( + def _build_logprobs_request( self, prompt: str, *, max_tokens: int, logprobs: int, - echo: bool, + score_input: bool, temperature: float, **kwargs, - ) -> ModelOutput: ... + ) -> Request: + """Build the Request for :meth:`alogprobs` from merged kwargs.""" + final_kwargs = { + **self._kwargs, + **kwargs, + "max_tokens": max_tokens, + "temperature": temperature, + } + n = self._validate_n(final_kwargs) + if n > 1: + raise ValueError(f"alogprobs only supports n=1; received n={n}") + req = self._kwargs_to_request(self._coerce_input(prompt), final_kwargs) + return replace( + req, return_logprobs=True, top_k=logprobs, score_input=score_input + ) + + # ── Response -> legacy ModelOutput bridge ───────────────────────────────── + + def _response_to_model_output(self, resp: Response) -> ModelOutput: + """Bridge a Response back to the legacy ``ModelOutput`` shape. + + Deliberate, documented conversions: + - Input scoring is re-flattened: ``input_scoring.token_logprobs`` are + concatenated ahead of the sampled ``logprobs``, restoring the legacy + echo layout that PPL consumers slice at ``usage.input_tokens``. + - ``top_logprobs`` collapses ``TopKEntry`` tuples to ``{token: logprob}`` + dicts, coalescing duplicate normalized token texts by max (sglang + byte-level collisions); ``token_id`` is dropped. On input scoring the + prompt-side top-k is absent by IR design (no consumer reads it). + - ``reasoning_texts`` carries the single IR reasoning channel (first + choice) when present. + """ + logprobs_present = resp.logprobs is not None or resp.input_scoring is not None + logprobs_tokens: list[str] | None = None + logprobs: list[float | None] | None = None + if logprobs_present: + segments = [] + if resp.input_scoring is not None: + segments.extend(resp.input_scoring.token_logprobs) + if resp.logprobs is not None: + segments.extend(resp.logprobs) + logprobs_tokens = [t.token for t in segments] + logprobs = [t.logprob for t in segments] + + top_logprobs: list[dict[str, float]] | None = None + if resp.top_logprobs is not None: + top_logprobs = [] + for per_pos in resp.top_logprobs: + merged: dict[str, float] = {} + for entry in per_pos: + if entry.token not in merged or entry.logprob > merged[entry.token]: + merged[entry.token] = entry.logprob + top_logprobs.append(merged) + top_logprobs = top_logprobs or None + + usage: ModelUsage | None = None + if resp.usage is not None: + usage = { + "input_tokens": resp.usage.input_tokens, + "output_tokens": resp.usage.output_tokens, + "total_tokens": resp.usage.total_tokens, + } + + reasoning_texts: list[str] | None = None + if resp.reasoning is not None and resp.reasoning.text: + reasoning_texts = [resp.reasoning.text] + + return ModelOutput( + model=self.meta(), + texts=list(resp.texts), + finish_reasons=( + list(resp.finish_reasons) if resp.finish_reasons is not None else None + ), + reasoning_texts=reasoning_texts, + logprobs_tokens=logprobs_tokens, + logprobs=logprobs, + top_logprobs=top_logprobs, + usage=usage, + request_params=( + dict(resp.request_params) if resp.request_params is not None else None + ), + response_model=resp.response_model, + system_fingerprint=resp.system_fingerprint, + ) diff --git a/sieval/core/models/sglang_gen_model.py b/sieval/core/models/sglang_gen_model.py index 7b22aaed..9574dba8 100644 --- a/sieval/core/models/sglang_gen_model.py +++ b/sieval/core/models/sglang_gen_model.py @@ -3,343 +3,36 @@ sglang's OpenAI ``/v1/completions`` endpoint rejects ``echo=True`` together with ``logprobs``, so PPL-style scoring (ARC/HellaSwag read the logprob of an answer token appended to the prompt; CMMLU/MMLU-Base read the first output -token's top-k) cannot go through it. This model speaks sglang's native +token's top-k) cannot go through it. This backend speaks sglang's native ``/generate`` protocol for BOTH generation and logprob extraction, so a single object talks one wire protocol end-to-end. +RFC #25 moved the wire logic (param translation, triple parsing, token-text +normalization, radix-cache guard) into +:class:`~sieval.core.models.transports.sglang.SglangTransport`; this class is +the backend selector that pairs the shared client/limiter pool with that +transport. + It extends ``Model[str]`` rather than ``GenModel`` deliberately: the only thing -``GenModel`` would contribute is its OpenAI-completions ``_agenerate_impl``, -which is a different protocol than the ``/generate`` logprob path — incidental -reuse, not coupling. The genuinely shared infrastructure (OpenAI async client, -limiters, ``with_args``/``meta``, the public ``agenerate``/``alogprobs`` -wrappers) lives in ``Model`` and is inherited directly. +``GenModel`` would contribute is its OpenAI-completions transport, which is a +different protocol than the ``/generate`` logprob path — incidental reuse, not +coupling. The genuinely shared infrastructure (OpenAI async client, limiters, +``with_args``/``meta``, the ``arun``/``agenerate``/``alogprobs`` surface) lives +in ``Model`` and is inherited directly. -AI-Generated Code - Claude Opus 4.8 (Anthropic) +AI-Generated Code - Claude Fable 5 (Anthropic) """ -from typing import cast, override - -from sieval.core.types import JSONValue - -from .model import Model, ModelOutput, ModelUsage - -# OpenAI-style generation kwarg -> sglang sampling_params key. Only these are -# forwarded to /generate; unrecognized kwargs (e.g. seed, stream, echo) are -# dropped rather than risk sglang rejecting an unknown sampling param. -_SAMPLING_PARAM_MAP: dict[str, str] = { - "max_tokens": "max_new_tokens", - "temperature": "temperature", - "top_p": "top_p", - "top_k": "top_k", - "min_p": "min_p", - "stop": "stop", - "frequency_penalty": "frequency_penalty", - "presence_penalty": "presence_penalty", - "repetition_penalty": "repetition_penalty", -} - - -def _request_params(body: dict[str, JSONValue]) -> dict[str, JSONValue]: - """Return the persisted request params: the /generate body minus the prompt. - - ``body["text"]`` is the full prompt, already recorded as the sample input — - copying it verbatim into every per-call record would duplicate it. This - shape is sglang-native (``sampling_params`` etc.) and intentionally differs - from the OpenAI-flavoured ``GenModel``/``ChatModel`` request_params; the - client/protocol decoupling that would unify them is tracked in RFC #25. - """ - return {k: v for k, v in body.items() if k != "text"} - - -def _normalize_token_text(text: str | None) -> str: - """Map GPT-2 byte-level BPE markers back to literal whitespace. - - sglang detokenizes when ``return_text_in_logprobs=True``, but some - tokenizers (e.g. Qwen) surface the raw byte-level markers ``Ġ`` (space) - and ``Ċ`` (newline). ``extract_option_logprob`` matches ``" A"`` / - ``A`` and CMMLU keys its top-k on the token text, so an un-normalized - ``"ĠA"`` would silently never match and the prediction would degrade. - Normalize here so downstream scoring is fed the same token text the - OpenAI path would produce. - - ``text`` is ``None`` when the server did not detokenize the logprobs - (a server launched with ``--skip-tokenizer-init`` ignores - ``return_text_in_logprobs``). Letter/option scoring cannot work without - token text, so fail loud with an actionable message rather than crash on - ``None.replace`` or silently degrade every token to ``""``. - - Limitation: only GPT-2 byte-level markers are handled. SentencePiece - (``▁``, U+2581) and other tokenizer conventions pass through unchanged — - add them here if a tokenizer that uses them needs the same contract. - """ - if text is None: - raise RuntimeError( - "sglang returned a logprob entry with no token text; option/letter " - "scoring needs detokenized text. Do not launch sglang with " - "--skip-tokenizer-init (it ignores return_text_in_logprobs)." - ) - return text.replace("Ġ", " ").replace("Ċ", "\n") +from .model import Model +from .transport import Transport class SglangGenModel(Model[str]): - """Model backend reading text and logprobs from sglang native ``/generate``. - - AI-Generated Code - Claude Opus 4.8 (Anthropic) - """ - - def _generate_url(self) -> str: - """Derive the native ``/generate`` URL from the OpenAI ``/v1`` base.""" - base = (self._api_base or "").rstrip("/").removesuffix("/v1").rstrip("/") - return f"{base}/generate" - - async def _post(self, body: dict[str, JSONValue]) -> dict | list: - """POST ``body`` to ``/generate`` via the OpenAI client. - - Reuses the OpenAI SDK's low-level ``self._client.post`` to speak the - native ``/generate`` protocol: this keeps the configured auth and - ``max_retries``, and an absolute URL is required because the client - would otherwise append the path to the ``/v1`` base. It couples us to - an SDK-internal surface — the client/protocol decoupling is tracked in - RFC #25. Returns the parsed JSON (a dict, or a list when - ``sampling_params.n > 1``). - """ - return cast( - "dict | list", - await self._client.post(self._generate_url(), cast_to=object, body=body), - ) - - @staticmethod - def _validate_n(final_kwargs: dict) -> int: - """Validate and return ``n`` (mirrors GenModel's guard).""" - n = final_kwargs.get("n", 1) - if isinstance(n, bool) or not isinstance(n, int): - raise TypeError(f"n must be an int, got {type(n).__name__}: {n!r}") - if n < 1: - raise ValueError(f"n must be >= 1, got {n}") - return n - - @classmethod - def _sampling_params( - cls, final_kwargs: dict, *, temperature: float | None = None - ) -> dict[str, JSONValue]: - """Translate recognized OpenAI-style kwargs into sglang sampling_params.""" - params: dict[str, JSONValue] = {} - for src, dst in _SAMPLING_PARAM_MAP.items(): - if src in final_kwargs and final_kwargs[src] is not None: - params[dst] = final_kwargs[src] - if temperature is not None: - params["temperature"] = temperature - return params - - @staticmethod - def _finish_reason(meta: dict) -> str: - """Extract a flat finish-reason string from sglang ``meta_info``.""" - fr = meta.get("finish_reason") - if isinstance(fr, dict): - return str(fr.get("type", "")) - return str(fr) if fr else "" - - @staticmethod - def _parse_logprobs(meta: dict, echo: bool) -> tuple[list[str], list[float | None]]: - """Flatten sglang ``*_token_logprobs`` into token-text + logprob lists. - - Each entry is ``[logprob, token_id, token_text]`` (first input - logprob is ``None``). With ``echo`` the input segment precedes the - output segment so echoed candidate tokens land at the sequence end. - """ - entries: list[list] = [] - if echo: - entries.extend(meta.get("input_token_logprobs") or []) - entries.extend(meta.get("output_token_logprobs") or []) - - tokens: list[str] = [] - token_logprobs: list[float | None] = [] - for logprob, _token_id, token_text in entries: - tokens.append(_normalize_token_text(token_text)) - token_logprobs.append(logprob) - return tokens, token_logprobs - - @staticmethod - def _parse_top_logprobs(meta: dict, echo: bool) -> list[dict[str, float]] | None: - """Flatten sglang ``*_top_logprobs`` into ``[{token: logprob}, ...]``. - - Aligns index-for-index with the token list from ``_parse_logprobs`` - (input segment first when ``echo``). A ``None``/empty per-token entry - (e.g. the first input token) becomes ``{}``. Returns ``None`` when the - server sent no top-k at all, matching ``ModelOutput.top_logprobs``'s - optional shape. CMMLU keys A/B/C/D off ``top_logprobs[0]``. - - Distinct token ids can normalize to identical text (e.g. a byte-level - ``"ĠA"`` and a literal ``" A"`` both → ``" A"``). Coalescing them by - keeping the highest logprob prevents a low-probability duplicate from - clobbering the real one, matching CMMLU's max-over-strip semantics. - """ - entries: list = [] - if echo: - entries.extend(meta.get("input_top_logprobs") or []) - entries.extend(meta.get("output_top_logprobs") or []) - if not entries: - return None - - result: list[dict[str, float]] = [] - for per_token in entries: - if not per_token: - result.append({}) - continue - merged: dict[str, float] = {} - for logprob, _token_id, token_text in per_token: - key = _normalize_token_text(token_text) - if key not in merged or logprob > merged[key]: - merged[key] = logprob - result.append(merged) - return result - - @staticmethod - def _parse_usage(meta: dict) -> ModelUsage | None: - """Build ``ModelUsage`` from sglang ``meta_info`` token counts.""" - input_tokens = meta.get("prompt_tokens") - output_tokens = meta.get("completion_tokens") - if input_tokens is None or output_tokens is None: - return None - return { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - - @override - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - if not isinstance(prompt, str): - raise TypeError("SglangGenModel requires a string prompt.") - - final_kwargs = {**self._kwargs, **kwargs} - num_choices = self._validate_n(final_kwargs) - - # Cross-engine parity note: with max_tokens unset no max_new_tokens is - # sent, so sglang applies its own default (128) while the vllm/OpenAI - # completions path applies the OpenAI default. Set max_tokens explicitly - # for identical output length when flipping engine: vllm <-> sglang. - sampling = self._sampling_params(final_kwargs) - if num_choices > 1: - sampling["n"] = num_choices - - body: dict[str, JSONValue] = {"text": prompt, "sampling_params": sampling} - raw = await self._post(body) - - # n>1 yields a list of per-sample dicts; n==1 a single dict. - results = raw if isinstance(raw, list) else [raw] - if not results or not all( - isinstance(r, dict) and "meta_info" in r for r in results - ): - raise RuntimeError( - "sglang /generate returned an unexpected response shape " - "(missing meta_info)." - ) - texts = [r.get("text", "") for r in results] - metas = [r["meta_info"] for r in results] - finish_reasons = [self._finish_reason(m) for m in metas] - - # Prompt tokens are shared across samples; completions sum. - input_tokens = metas[0].get("prompt_tokens") - output_tokens = sum(m.get("completion_tokens") or 0 for m in metas) - usage: ModelUsage | None = ( - { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - if input_tokens is not None - else None - ) - - return ModelOutput( - model=self.meta(), - texts=texts, - finish_reasons=finish_reasons, - usage=usage, - request_params=_request_params(body), - response_model=self._model, - ) - - @override - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - final_kwargs = {**self._kwargs, **kwargs} - num_choices = self._validate_n(final_kwargs) - if num_choices > 1: - raise ValueError(f"alogprobs only supports n=1; received n={num_choices}") - - sampling = self._sampling_params(final_kwargs, temperature=temperature) - # sglang rejects max_new_tokens=0; the generated token is ignored for - # scoring but at least one is required. - sampling["max_new_tokens"] = max(max_tokens, 1) - - body: dict[str, JSONValue] = { - "text": prompt, - "sampling_params": sampling, - "return_logprob": True, - # 0 → all echoed input token logprobs; -1 → output only. - "logprob_start_len": 0 if echo else -1, - "top_logprobs_num": logprobs, - "return_text_in_logprobs": True, - } - - data = await self._post(body) - if not isinstance(data, dict): - raise RuntimeError( - f"sglang /generate returned {type(data).__name__}, expected an object." - ) - meta = data["meta_info"] - - # sglang's radix prefix cache does not recompute logprobs for cached - # positions: on a cache hit it truncates input_token_logprobs to - # (prompt_tokens - cached_tokens). echo-based scoring reads the full - # echoed input sequence, so a truncated set would score silently wrong - # (vLLM errors in this case; sglang stays silent). Deliberate stance: - # ANY cache touch — or a response we can't verify against because it - # omitted prompt_tokens — is untrusted, so fail loud. echo-based scoring - # requires launching sglang with --disable-radix-cache. - if echo: - input_lps = meta.get("input_token_logprobs") or [] - prompt_tokens = meta.get("prompt_tokens") - cached_tokens = meta.get("cached_tokens") or 0 - if prompt_tokens is None: - raise RuntimeError( - "sglang response omitted prompt_tokens, so echoed-input " - "completeness cannot be verified; refusing to score silently. " - "Launch sglang with --disable-radix-cache." - ) - if cached_tokens or len(input_lps) != prompt_tokens: - raise RuntimeError( - "sglang returned partial echoed-input logprobs " - f"({len(input_lps)} of {prompt_tokens} prompt tokens, " - f"cached_tokens={cached_tokens}): its radix prefix cache does " - "not recompute logprobs for cached positions, so echo-based " - "scoring would be silently wrong. Launch sglang with " - "--disable-radix-cache." - ) + """Model backend reading text and logprobs from sglang native ``/generate``.""" - tokens, token_logprobs = self._parse_logprobs(meta, echo) - top_logprobs = self._parse_top_logprobs(meta, echo) - if not token_logprobs and not top_logprobs: - raise RuntimeError("sglang /generate returned no logprobs.") + def _build_default_transport(self) -> Transport: + from .transports.sglang import SglangTransport - return ModelOutput( - model=self.meta(), - texts=[data.get("text", "")], - finish_reasons=[self._finish_reason(meta)], - logprobs_tokens=tokens, - logprobs=token_logprobs, - top_logprobs=top_logprobs, - usage=self._parse_usage(meta), - request_params=_request_params(body), - response_model=self._model, + return SglangTransport( + client=self._client, model=self._model, api_base=self._api_base ) diff --git a/sieval/core/models/transport.py b/sieval/core/models/transport.py new file mode 100644 index 00000000..09928080 --- /dev/null +++ b/sieval/core/models/transport.py @@ -0,0 +1,48 @@ +"""Transport Protocol: the provider-frontend abstraction. + +A Transport lowers a :class:`Request` into its native wire format and lifts the +native response back into a :class:`Response`. Each provider frontend is one +Transport, composed into a :class:`~sieval.core.models.model.Model` (a strategy, +not a base class). + +Three behavioural contracts every implementation must honour: + +1. ``arun()`` returns a terminal Response. + All provider-internal loops are resolved before returning, including any + server-tool round-trips. The caller never drives a server-tool loop. + +2. Opaque state handling depends on mode: + - Stateful (``Request.session_id`` set): the Transport passes + ``previous_response_id`` / ``previous_interaction_id`` internally and + carries the new id on ``Response.session_id``. The caller never sees + opaque state. + - Stateless (``Request.session_id`` absent): the Transport embeds + ``Request.reasoning.opaque_roundtrip`` into the correct wire item on + lower(), and extracts the provider-signed value into + ``Response.reasoning.opaque_roundtrip`` on lift(). The caller echoes it + back next turn without modification. + +3. ``Response.grounding.rendered_content`` must not be dropped in lift(). + Google ToS requires it to be rendered. + +AI-Generated Code - Claude Opus 4.8 (Anthropic) +""" + +from typing import Protocol, runtime_checkable + +from .capabilities import Capability +from .ir import Request, Response + + +@runtime_checkable +class Transport(Protocol): + """Provider frontend: lowers a Request to wire form and lifts the reply.""" + + @property + def capabilities(self) -> frozenset[Capability]: + """The IR features this Transport honours.""" + ... + + async def arun(self, req: Request) -> Response: + """Execute *req* and return a terminal :class:`Response`.""" + ... diff --git a/sieval/core/models/transports/__init__.py b/sieval/core/models/transports/__init__.py new file mode 100644 index 00000000..3f1a7977 --- /dev/null +++ b/sieval/core/models/transports/__init__.py @@ -0,0 +1,17 @@ +"""Provider frontends (Transports) for the Model IR. + +Each module here implements the :class:`~sieval.core.models.transport.Transport` +protocol for one provider wire protocol. + +AI-Generated Code - Claude Opus 4.8 (Anthropic) +""" + +from .openai_chat import OpenAIChatTransport +from .openai_completions import OpenAICompletionsTransport +from .sglang import SglangTransport + +__all__ = [ + "OpenAIChatTransport", + "OpenAICompletionsTransport", + "SglangTransport", +] diff --git a/sieval/core/models/transports/openai_chat.py b/sieval/core/models/transports/openai_chat.py new file mode 100644 index 00000000..b10d3c87 --- /dev/null +++ b/sieval/core/models/transports/openai_chat.py @@ -0,0 +1,280 @@ +"""OpenAIChatTransport: OpenAI ``/v1/chat/completions`` frontend for the IR. + +Capabilities: Chat, FunctionCalling, TopKLogprobs, SampledLogprobs, +StructuredOutput. + +InputScoring is NOT supported — ``echo`` has no equivalent in chat completions. +A Request with ``score_input=True`` (or a stateful ``session_id``, which chat +completions cannot honour) is rejected here, never silently ignored. This is the +fix for the historical "echo silently dropped on chat" bug. + +Streaming: ``Request.stream=True`` makes the transport consume the SSE stream +and accumulate deltas internally (ported from the legacy ``ChatModel`` impl); +the caller always receives one terminal :class:`Response`. ``stream=None`` +defaults to a single-shot request. + +``token_id`` is not available on this wire protocol, so it is left ``None``. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from typing import Any + +from ..capabilities import Capability +from ..exceptions import CapabilityError +from ..ir import ( + ReasoningOutput, + Request, + Response, + TokenLogprob, + TopKEntry, + UsageStats, +) + + +def _tool_call_to_dict(tc: Any) -> dict[str, Any]: + """Best-effort conversion of an SDK tool-call object to a plain dict.""" + if hasattr(tc, "model_dump"): + return tc.model_dump() + if isinstance(tc, dict): + return dict(tc) + return dict(vars(tc)) if hasattr(tc, "__dict__") else {} + + +def _usage_stats(raw: Any) -> UsageStats | None: + """Map an OpenAI usage object to :class:`UsageStats` (None when absent).""" + if raw is None: + return None + return UsageStats( + input_tokens=raw.prompt_tokens, + output_tokens=raw.completion_tokens, + total_tokens=raw.total_tokens, + ) + + +def _delta_reasoning(part: Any) -> str: + """Extract reasoning text from a message or delta (either field spelling).""" + rc = getattr(part, "reasoning", None) + if rc: + return str(rc) + rc = getattr(part, "reasoning_content", None) + return str(rc) if rc else "" + + +def _content_to_ir( + content: list, +) -> tuple[tuple[TokenLogprob, ...], tuple[tuple[TopKEntry, ...], ...]]: + """Map chat logprobs content items to IR token / top-k tuples.""" + logprobs = tuple( + TokenLogprob(token=item.token, logprob=item.logprob) for item in content + ) + top_logprobs = tuple( + tuple( + TopKEntry(token=t.token, logprob=t.logprob) + for t in (getattr(item, "top_logprobs", None) or []) + ) + for item in content + ) + return logprobs, top_logprobs + + +class OpenAIChatTransport: + """Transport for OpenAI ``/v1/chat/completions``.""" + + CAPABILITIES: frozenset[Capability] = frozenset( + { + Capability.Chat, + Capability.FunctionCalling, + Capability.TopKLogprobs, + Capability.SampledLogprobs, + Capability.StructuredOutput, + } + ) + + def __init__(self, client: Any, model: str): + self._client = client + self._model = model + + @property + def capabilities(self) -> frozenset[Capability]: + return self.CAPABILITIES + + # ── lower ───────────────────────────────────────────────────────────────── + + def _lower(self, req: Request) -> tuple[list[dict[str, Any]], dict[str, Any]]: + if req.score_input: + raise CapabilityError( + "OpenAIChatTransport does not support InputScoring " + "(echo/score_input has no chat-completions equivalent)." + ) + if req.session_id is not None: + raise CapabilityError( + "OpenAIChatTransport does not support stateful session_id " + "(chat completions has no previous_response_id)." + ) + + if isinstance(req.input, str): + messages: list[dict[str, Any]] = [{"role": "user", "content": req.input}] + else: + messages = list(req.input) + + params: dict[str, Any] = {} + sp = req.sampling + if sp is not None: + if sp.max_tokens is not None: + params["max_tokens"] = sp.max_tokens + if sp.temperature is not None: + params["temperature"] = sp.temperature + if sp.top_p is not None: + params["top_p"] = sp.top_p + if sp.top_k_sampling is not None: + # vLLM extension; upstream OpenAI rejects it, matching the + # legacy behaviour of forwarding top_k verbatim. + params["top_k"] = sp.top_k_sampling + if sp.stop is not None: + params["stop"] = list(sp.stop) + if sp.stop_token_ids is not None: + params["stop_token_ids"] = list(sp.stop_token_ids) + if sp.seed is not None: + params["seed"] = sp.seed + if sp.frequency_penalty is not None: + params["frequency_penalty"] = sp.frequency_penalty + if sp.presence_penalty is not None: + params["presence_penalty"] = sp.presence_penalty + if sp.n != 1: + params["n"] = sp.n + + if req.return_logprobs: + params["logprobs"] = True + if req.top_k > 0: + params["top_logprobs"] = req.top_k + + if req.response_format is not None: + params["response_format"] = req.response_format + if req.tools is not None: + params["tools"] = req.tools + if req.reasoning is not None and req.reasoning.effort is not None: + params["reasoning_effort"] = req.reasoning.effort + + params["stream"] = bool(req.stream) + + if req.extra_wire_params: + for k, v in req.extra_wire_params.items(): + params.setdefault(k, v) + + # Injected default last so an explicit stream_options (via + # extra_wire_params) wins, matching the legacy "if not present" rule. + if params["stream"] and "stream_options" not in params: + params["stream_options"] = {"include_usage": True} + + return messages, params + + # ── lift (single-shot) ──────────────────────────────────────────────────── + + def _lift(self, resp: Any, *, n: int, params: dict[str, Any]) -> Response: + texts = [""] * n + finish_reasons = [""] * n + reasoning_text = "" + tool_calls: tuple[dict[str, Any], ...] | None = None + logprobs: tuple[TokenLogprob, ...] | None = None + top_logprobs: tuple[tuple[TopKEntry, ...], ...] | None = None + + for choice in resp.choices: + idx = choice.index + if not 0 <= idx < n: + continue + message = choice.message + if message.content is not None: + texts[idx] += message.content + finish_reasons[idx] = choice.finish_reason or "" + if idx == 0: + reasoning_text = _delta_reasoning(message) + raw_calls = getattr(message, "tool_calls", None) + if raw_calls: + tool_calls = tuple(_tool_call_to_dict(tc) for tc in raw_calls) + lp_obj = choice.logprobs + if lp_obj is not None: + logprobs, top_logprobs = _content_to_ir( + getattr(lp_obj, "content", None) or [] + ) + + return Response( + texts=tuple(texts), + reasoning=ReasoningOutput(text=reasoning_text) if reasoning_text else None, + tool_calls=tool_calls, + logprobs=logprobs, + top_logprobs=top_logprobs, + usage=_usage_stats(getattr(resp, "usage", None)), + finish_reasons=tuple(finish_reasons), + request_params=dict(params), + response_model=getattr(resp, "model", None), + system_fingerprint=getattr(resp, "system_fingerprint", None), + ) + + # ── lift (streaming) ────────────────────────────────────────────────────── + + async def _lift_stream( + self, stream: Any, *, n: int, params: dict[str, Any] + ) -> Response: + texts = [""] * n + finish_reasons = [""] * n + reasoning_text = "" + usage: UsageStats | None = None + saw_logprobs = False + lp_tokens: list[TokenLogprob] = [] + lp_topk: list[tuple[TopKEntry, ...]] = [] + response_model: str | None = None + system_fingerprint: str | None = None + + async for chunk in stream: + if response_model is None: + response_model = getattr(chunk, "model", None) + if system_fingerprint is None: + system_fingerprint = getattr(chunk, "system_fingerprint", None) + if chunk.choices: + for choice in chunk.choices: + idx = choice.index + if not 0 <= idx < n: + continue + finish_reasons[idx] = choice.finish_reason or "" + if choice.delta: + if choice.delta.content: + texts[idx] += choice.delta.content + if idx == 0: + reasoning_text += _delta_reasoning(choice.delta) + if idx == 0: + lp_obj = getattr(choice, "logprobs", None) + if lp_obj is not None: + saw_logprobs = True + content = getattr(lp_obj, "content", None) or [] + if content: + chunk_lp, chunk_topk = _content_to_ir(content) + lp_tokens.extend(chunk_lp) + lp_topk.extend(chunk_topk) + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage is not None: + usage = _usage_stats(chunk_usage) + + return Response( + texts=tuple(texts), + reasoning=ReasoningOutput(text=reasoning_text) if reasoning_text else None, + logprobs=tuple(lp_tokens) if saw_logprobs else None, + top_logprobs=tuple(lp_topk) if saw_logprobs else None, + usage=usage, + finish_reasons=tuple(finish_reasons), + request_params=dict(params), + response_model=response_model, + system_fingerprint=system_fingerprint, + ) + + # ── arun ────────────────────────────────────────────────────────────────── + + async def arun(self, req: Request) -> Response: + messages, params = self._lower(req) + n = req.sampling.n if req.sampling is not None else 1 + resp = await self._client.chat.completions.create( + model=self._model, messages=messages, **params + ) + if params["stream"]: + return await self._lift_stream(resp, n=n, params=params) + return self._lift(resp, n=n, params=params) diff --git a/sieval/core/models/transports/openai_completions.py b/sieval/core/models/transports/openai_completions.py new file mode 100644 index 00000000..040b3ce4 --- /dev/null +++ b/sieval/core/models/transports/openai_completions.py @@ -0,0 +1,314 @@ +"""OpenAICompletionsTransport: OpenAI ``/v1/completions`` frontend for the IR. + +Capabilities: Completion, InputScoring (via ``echo=True``), TopKLogprobs, +SampledLogprobs. + +InputScoring is implemented via ``echo=True`` — the only place in the entire +codebase where ``echo`` appears. It is an OpenAI-Completions-specific workaround +for the missing native scoring endpoint and must not leak into any other +transport or the IR layer. On lift, the echoed prompt tokens are split off into +``Response.input_scoring`` at the ``prompt_tokens`` boundary; the remaining +tokens are the sampled completion (``Response.logprobs``). + +Streaming: ``Request.stream=True`` makes the transport consume the SSE stream +and accumulate chunks internally (ported from the legacy ``GenModel`` impl); +the caller always receives one terminal :class:`Response`. ``stream=None`` +defaults to a single-shot request. + +``token_id`` is not available on this wire protocol, so it is left ``None``. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from collections.abc import Mapping, Sequence +from typing import Any + +from ..capabilities import Capability +from ..ir import ( + InputScoringResult, + Request, + Response, + TokenLogprob, + TopKEntry, + UsageStats, +) + + +def _completion_top_logprobs(raw: object) -> list[dict[str, float]]: + """Sanitize the completions API's per-position top-logprob dicts.""" + if not isinstance(raw, Sequence) or isinstance(raw, str | bytes): + return [] + + top_logprobs = [] + for item in raw: + if item is None: + top_logprobs.append({}) + continue + if not isinstance(item, Mapping): + continue + top_logprobs.append( + { + token: float(logprob) + for token, logprob in item.items() + if isinstance(token, str) and isinstance(logprob, int | float) + } + ) + return top_logprobs + + +def _usage_stats(raw: Any) -> UsageStats | None: + """Map an OpenAI usage object to :class:`UsageStats` (None when absent).""" + if raw is None: + return None + return UsageStats( + input_tokens=raw.prompt_tokens, + output_tokens=raw.completion_tokens, + total_tokens=raw.total_tokens, + ) + + +class OpenAICompletionsTransport: + """Transport for OpenAI ``/v1/completions``.""" + + CAPABILITIES: frozenset[Capability] = frozenset( + { + Capability.Completion, + Capability.InputScoring, + Capability.TopKLogprobs, + Capability.SampledLogprobs, + } + ) + + def __init__(self, client: Any, model: str): + self._client = client + self._model = model + + @property + def capabilities(self) -> frozenset[Capability]: + return self.CAPABILITIES + + # ── lower ───────────────────────────────────────────────────────────────── + + def _lower(self, req: Request) -> dict[str, Any]: + if not isinstance(req.input, str): + raise TypeError( + "OpenAICompletionsTransport requires str input (Completion modality)." + ) + + params: dict[str, Any] = {} + sp = req.sampling + if sp is not None: + if sp.max_tokens is not None: + params["max_tokens"] = sp.max_tokens + if sp.temperature is not None: + params["temperature"] = sp.temperature + if sp.top_p is not None: + params["top_p"] = sp.top_p + if sp.top_k_sampling is not None: + # vLLM extension; upstream OpenAI rejects it, matching the + # legacy behaviour of forwarding top_k verbatim. + params["top_k"] = sp.top_k_sampling + if sp.stop is not None: + params["stop"] = list(sp.stop) + if sp.stop_token_ids is not None: + params["stop_token_ids"] = list(sp.stop_token_ids) + if sp.seed is not None: + params["seed"] = sp.seed + if sp.frequency_penalty is not None: + params["frequency_penalty"] = sp.frequency_penalty + if sp.presence_penalty is not None: + params["presence_penalty"] = sp.presence_penalty + if sp.n != 1: + params["n"] = sp.n + + # echo appears ONLY here — the InputScoring workaround. + if req.score_input: + params["echo"] = True + + # completions `logprobs` is the count of top alternatives (0 → own token + # logprob only); required to receive any logprobs at all. + if req.return_logprobs or req.score_input: + params["logprobs"] = req.top_k + + params["stream"] = bool(req.stream) + + if req.extra_wire_params: + for k, v in req.extra_wire_params.items(): + params.setdefault(k, v) + + # Injected default last so an explicit stream_options (via + # extra_wire_params) wins, matching the legacy "if not present" rule. + if params["stream"] and "stream_options" not in params: + params["stream_options"] = {"include_usage": True} + + return params + + # ── lift ──────────────────────────────────────────────────────────────── + + def _build_response( + self, + *, + texts: list[str], + finish_reasons: list[str], + tokens: list[str], + token_logprobs: list[float | None], + top_raw: list[dict[str, float]], + saw_logprobs: bool, + score_input: bool, + usage: UsageStats | None, + params: dict[str, Any], + response_model: str | None, + system_fingerprint: str | None, + ) -> Response: + """Assemble the terminal Response shared by both wire paths. + + ``saw_logprobs`` is the presence signal (a logprobs object appeared on + the wire): only then are ``logprobs``/``top_logprobs`` tuples — possibly + empty — attached; otherwise they stay ``None`` (empty-vs-absent + contract on :class:`Response`). + """ + input_scoring: InputScoringResult | None = None + logprobs: tuple[TokenLogprob, ...] | None = None + top_logprobs: tuple[tuple[TopKEntry, ...], ...] | None = None + + if saw_logprobs: + all_tokens = tuple( + TokenLogprob(token=tok, logprob=lp) + for tok, lp in zip(tokens, token_logprobs, strict=False) + ) + all_topk = tuple( + tuple(TopKEntry(token=tok, logprob=lp) for tok, lp in per_token.items()) + for per_token in top_raw + ) + if score_input: + # Split at the prompt/completion boundary. The echoed prompt + # occupies the first prompt_tokens positions. + boundary = usage.input_tokens if usage is not None else 0 + input_scoring = InputScoringResult(token_logprobs=all_tokens[:boundary]) + logprobs = all_tokens[boundary:] + top_logprobs = all_topk[boundary:] + else: + logprobs = all_tokens + top_logprobs = all_topk + + return Response( + texts=tuple(texts), + logprobs=logprobs, + top_logprobs=top_logprobs, + input_scoring=input_scoring, + usage=usage, + finish_reasons=tuple(finish_reasons), + request_params=dict(params), + response_model=response_model, + system_fingerprint=system_fingerprint, + ) + + def _lift( + self, resp: Any, *, n: int, req: Request, params: dict[str, Any] + ) -> Response: + texts = [""] * n + finish_reasons = [""] * n + tokens: list[str] = [] + token_logprobs: list[float | None] = [] + top_raw: list[dict[str, float]] = [] + saw_logprobs = False + + for choice in resp.choices: + idx = choice.index + if not 0 <= idx < n: + continue + texts[idx] += choice.text or "" + finish_reasons[idx] = choice.finish_reason or "" + if idx == 0: + lp_obj = choice.logprobs + if lp_obj is not None: + saw_logprobs = True + if getattr(lp_obj, "tokens", None): + tokens.extend(lp_obj.tokens) + if getattr(lp_obj, "token_logprobs", None): + token_logprobs.extend(lp_obj.token_logprobs) + top_raw.extend( + _completion_top_logprobs(getattr(lp_obj, "top_logprobs", None)) + ) + + return self._build_response( + texts=texts, + finish_reasons=finish_reasons, + tokens=tokens, + token_logprobs=token_logprobs, + top_raw=top_raw, + saw_logprobs=saw_logprobs, + score_input=req.score_input, + usage=_usage_stats(getattr(resp, "usage", None)), + params=params, + response_model=getattr(resp, "model", None), + system_fingerprint=getattr(resp, "system_fingerprint", None), + ) + + async def _lift_stream( + self, stream: Any, *, n: int, req: Request, params: dict[str, Any] + ) -> Response: + texts = [""] * n + finish_reasons = [""] * n + tokens: list[str] = [] + token_logprobs: list[float | None] = [] + top_raw: list[dict[str, float]] = [] + saw_logprobs = False + usage: UsageStats | None = None + response_model: str | None = None + system_fingerprint: str | None = None + + async for chunk in stream: + if response_model is None: + response_model = getattr(chunk, "model", None) + if system_fingerprint is None: + system_fingerprint = getattr(chunk, "system_fingerprint", None) + if chunk.choices: + for choice in chunk.choices: + idx = choice.index + if not 0 <= idx < n: + continue + texts[idx] += choice.text or "" + finish_reasons[idx] = choice.finish_reason or "" + if idx == 0: + lp_obj = getattr(choice, "logprobs", None) + if lp_obj is not None: + saw_logprobs = True + if getattr(lp_obj, "tokens", None): + tokens.extend(lp_obj.tokens) + if getattr(lp_obj, "token_logprobs", None): + token_logprobs.extend(lp_obj.token_logprobs) + top_raw.extend( + _completion_top_logprobs( + getattr(lp_obj, "top_logprobs", None) + ) + ) + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage is not None: + usage = _usage_stats(chunk_usage) + + return self._build_response( + texts=texts, + finish_reasons=finish_reasons, + tokens=tokens, + token_logprobs=token_logprobs, + top_raw=top_raw, + saw_logprobs=saw_logprobs, + score_input=req.score_input, + usage=usage, + params=params, + response_model=response_model, + system_fingerprint=system_fingerprint, + ) + + # ── arun ────────────────────────────────────────────────────────────────── + + async def arun(self, req: Request) -> Response: + params = self._lower(req) + n = req.sampling.n if req.sampling is not None else 1 + resp = await self._client.completions.create( + model=self._model, prompt=req.input, **params + ) + if params["stream"]: + return await self._lift_stream(resp, n=n, req=req, params=params) + return self._lift(resp, n=n, req=req, params=params) diff --git a/sieval/core/models/transports/sglang.py b/sieval/core/models/transports/sglang.py new file mode 100644 index 00000000..50ba41c0 --- /dev/null +++ b/sieval/core/models/transports/sglang.py @@ -0,0 +1,371 @@ +"""SglangTransport: native sglang ``/generate`` frontend for the Model IR. + +sglang's OpenAI ``/v1/completions`` endpoint rejects ``echo=True`` together +with ``logprobs``, so PPL-style scoring cannot go through it. This transport +speaks sglang's native ``/generate`` protocol for BOTH generation and logprob +extraction, lowering a :class:`Request` and lifting a :class:`Response`. + +The token-text normalization (:func:`_normalize_token_text`), finish-reason and +usage extraction, and the radix prefix-cache guard were moved here verbatim from +the legacy ``SglangGenModel`` implementation — they have been validated against +real sglang responses. One improvement over the legacy path: the native +``[logprob, token_id, token_text]`` triples are parsed directly so +``TokenLogprob.token_id`` is populated (``SampledLogprobsWithTokenIds``). + +The IR split is cleaner than the legacy ``echo`` concatenation: +``input_token_logprobs`` → ``Response.input_scoring`` and +``output_token_logprobs`` → ``Response.logprobs``, never merged. sglang's +``input_top_logprobs`` has no IR field and is not lifted (no consumer reads +prompt-side top-k). + +``Request.stream`` is ignored: the native ``/generate`` path has always been a +single POST (pure scheduling, no content impact). Unrecognized +``extra_wire_params`` are mapped through the OpenAI→sglang sampling-param table +when known and otherwise dropped, matching the legacy stance of never risking +sglang rejecting an unknown param. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from typing import Any, cast + +from sieval.core.types import JSONValue + +from ..capabilities import Capability +from ..ir import ( + InputScoringResult, + Request, + Response, + TokenLogprob, + TopKEntry, + UsageStats, +) + +# OpenAI-style generation kwarg -> sglang sampling_params key. Only these are +# forwarded to /generate; unrecognized kwargs (e.g. seed, stream, echo) are +# dropped rather than risk sglang rejecting an unknown sampling param. +_SAMPLING_PARAM_MAP: dict[str, str] = { + "max_tokens": "max_new_tokens", + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k", + "min_p": "min_p", + "stop": "stop", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "repetition_penalty": "repetition_penalty", +} + + +def _request_params(body: dict[str, JSONValue]) -> dict[str, JSONValue]: + """Return the persisted request params: the /generate body minus the prompt. + + ``body["text"]`` is the full prompt, already recorded as the sample input — + copying it verbatim into every per-call record would duplicate it. This + shape is sglang-native (``sampling_params`` etc.) and intentionally differs + from the OpenAI-flavoured transports' request_params. + """ + return {k: v for k, v in body.items() if k != "text"} + + +def _normalize_token_text(text: str | None) -> str: + """Map GPT-2 byte-level BPE markers back to literal whitespace. + + sglang detokenizes when ``return_text_in_logprobs=True``, but some + tokenizers (e.g. Qwen) surface the raw byte-level markers ``Ġ`` (space) + and ``Ċ`` (newline). ``extract_option_logprob`` matches ``" A"`` / + ``A`` and CMMLU keys its top-k on the token text, so an un-normalized + ``"ĠA"`` would silently never match and the prediction would degrade. + Normalize here so downstream scoring is fed the same token text the + OpenAI path would produce. + + ``text`` is ``None`` when the server did not detokenize the logprobs + (a server launched with ``--skip-tokenizer-init`` ignores + ``return_text_in_logprobs``). Letter/option scoring cannot work without + token text, so fail loud with an actionable message rather than crash on + ``None.replace`` or silently degrade every token to ``""``. + + Limitation: only GPT-2 byte-level markers are handled. SentencePiece + (``▁``, U+2581) and other tokenizer conventions pass through unchanged — + add them here if a tokenizer that uses them needs the same contract. + """ + if text is None: + raise RuntimeError( + "sglang returned a logprob entry with no token text; option/letter " + "scoring needs detokenized text. Do not launch sglang with " + "--skip-tokenizer-init (it ignores return_text_in_logprobs)." + ) + return text.replace("Ġ", " ").replace("Ċ", "\n") + + +def _finish_reason(meta: dict) -> str: + """Extract a flat finish-reason string from sglang ``meta_info``.""" + fr = meta.get("finish_reason") + if isinstance(fr, dict): + return str(fr.get("type", "")) + return str(fr) if fr else "" + + +class SglangTransport: + """Transport for sglang native ``/generate``. + + ``token_id`` is always populated (sglang returns ``[logprob, token_id, + token_text]`` triples). + """ + + CAPABILITIES: frozenset[Capability] = frozenset( + { + Capability.Completion, + Capability.InputScoring, + Capability.SampledLogprobs, + Capability.SampledLogprobsWithTokenIds, + Capability.TopKLogprobs, + Capability.Prefill, + } + ) + + def __init__(self, client: Any, model: str, api_base: str | None = None): + self._client = client + self._model = model + self._api_base = api_base + + @property + def capabilities(self) -> frozenset[Capability]: + return self.CAPABILITIES + + # ── wire helpers ────────────────────────────────────────────────────────── + + def _generate_url(self) -> str: + """Derive the native ``/generate`` URL from the OpenAI ``/v1`` base.""" + base = (self._api_base or "").rstrip("/").removesuffix("/v1").rstrip("/") + return f"{base}/generate" + + async def _post(self, body: dict[str, JSONValue]) -> dict | list: + """POST ``body`` to ``/generate`` via the OpenAI client. + + Reuses the OpenAI SDK's low-level ``post`` to speak the native + ``/generate`` protocol: this keeps the configured auth and + ``max_retries``, and an absolute URL is required because the client + would otherwise append the path to the ``/v1`` base. Returns the parsed + JSON (a dict, or a list when ``sampling_params.n > 1``). + """ + return cast( + "dict | list", + await self._client.post(self._generate_url(), cast_to=object, body=body), + ) + + # ── lower ───────────────────────────────────────────────────────────────── + + def _lower(self, req: Request) -> dict[str, JSONValue]: + if not isinstance(req.input, str): + raise TypeError("SglangTransport requires str input (Completion modality).") + + sampling: dict[str, JSONValue] = {} + sp = req.sampling + n = 1 + if sp is not None: + if sp.max_tokens is not None: + sampling["max_new_tokens"] = sp.max_tokens + if sp.temperature is not None: + sampling["temperature"] = sp.temperature + if sp.top_p is not None: + sampling["top_p"] = sp.top_p + if sp.top_k_sampling is not None: + sampling["top_k"] = sp.top_k_sampling + if sp.stop is not None: + sampling["stop"] = list(sp.stop) + if sp.frequency_penalty is not None: + sampling["frequency_penalty"] = sp.frequency_penalty + if sp.presence_penalty is not None: + sampling["presence_penalty"] = sp.presence_penalty + n = sp.n + if n > 1: + sampling["n"] = n + + # sglang rejects max_new_tokens=0; scoring still needs at least one. + if req.score_input and not sampling.get("max_new_tokens"): + sampling["max_new_tokens"] = 1 + + # Prefill capability: sglang accepts a forced prefill on sampling_params. + if req.prefix is not None: + sampling["prefill"] = req.prefix + + # Legacy-parity passthrough: only kwargs with a known sglang + # sampling-param equivalent are forwarded; the rest are dropped. + if req.extra_wire_params: + for k, v in req.extra_wire_params.items(): + dst = _SAMPLING_PARAM_MAP.get(k) + if dst is not None and v is not None: + sampling.setdefault(dst, v) + + body: dict[str, JSONValue] = {"text": req.input, "sampling_params": sampling} + + if req.return_logprobs or req.score_input: + body["return_logprob"] = True + # 0 → all echoed input token logprobs; -1 → output only. + body["logprob_start_len"] = 0 if req.score_input else -1 + body["top_logprobs_num"] = req.top_k + body["return_text_in_logprobs"] = True + + return body + + # ── lift ──────────────────────────────────────────────────────────────── + + @staticmethod + def _triples_to_tokens(entries: list) -> tuple[TokenLogprob, ...]: + """Map sglang ``[logprob, token_id, token_text]`` triples to TokenLogprobs. + + Unlike the legacy parser, ``token_id`` is preserved. + """ + return tuple( + TokenLogprob( + token=_normalize_token_text(token_text), + logprob=logprob, + token_id=token_id, + ) + for logprob, token_id, token_text in entries + ) + + @staticmethod + def _triples_to_topk( + entries: list, + ) -> tuple[tuple[TopKEntry, ...], ...] | None: + """Map per-token sglang top-k triple lists to tuples of TopKEntry. + + Returns ``None`` when the server sent no top-k at all, matching the + legacy optional shape. A ``None``/empty per-token entry becomes ``()``. + """ + if not entries: + return None + result: list[tuple[TopKEntry, ...]] = [] + for per_token in entries: + if not per_token: + result.append(()) + continue + result.append( + tuple( + TopKEntry( + token=_normalize_token_text(token_text), + logprob=logprob, + token_id=token_id, + ) + for logprob, token_id, token_text in per_token + ) + ) + return tuple(result) + + @staticmethod + def _parse_usage(metas: list[dict]) -> UsageStats | None: + """Build usage from sglang ``meta_info`` token counts. + + Prompt tokens are shared across n samples; completions sum. + """ + input_tokens = metas[0].get("prompt_tokens") + if input_tokens is None: + return None + output_tokens = sum(m.get("completion_tokens") or 0 for m in metas) + return UsageStats( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) + + def _guard_radix_cache(self, meta: dict) -> None: + """Reject partial echoed-input logprobs from the radix prefix cache. + + sglang's radix prefix cache does not recompute logprobs for cached + positions: on a cache hit it truncates ``input_token_logprobs`` to + ``prompt_tokens - cached_tokens``. Echo-based scoring reads the full + echoed input sequence, so a truncated set would score silently wrong + (vLLM errors in this case; sglang stays silent). Deliberate stance: ANY + cache touch — or a response we can't verify against because it omitted + ``prompt_tokens`` — is untrusted, so fail loud. Input scoring requires + launching sglang with ``--disable-radix-cache``. + """ + input_lps = meta.get("input_token_logprobs") or [] + prompt_tokens = meta.get("prompt_tokens") + cached_tokens = meta.get("cached_tokens") or 0 + if prompt_tokens is None: + raise RuntimeError( + "sglang response omitted prompt_tokens, so echoed-input " + "completeness cannot be verified; refusing to score silently. " + "Launch sglang with --disable-radix-cache." + ) + if cached_tokens or len(input_lps) != prompt_tokens: + raise RuntimeError( + "sglang returned partial echoed-input logprobs " + f"({len(input_lps)} of {prompt_tokens} prompt tokens, " + f"cached_tokens={cached_tokens}): its radix prefix cache does " + "not recompute logprobs for cached positions, so echo-based " + "scoring would be silently wrong. Launch sglang with " + "--disable-radix-cache." + ) + + def _lift( + self, + results: list[dict], + body: dict[str, JSONValue], + *, + score_input: bool, + want_logprobs: bool, + ) -> Response: + metas = [r["meta_info"] for r in results] + texts = tuple(r.get("text", "") for r in results) + finish_reasons = tuple(_finish_reason(m) for m in metas) + usage = self._parse_usage(metas) + + input_scoring: InputScoringResult | None = None + logprobs: tuple[TokenLogprob, ...] | None = None + top_logprobs: tuple[tuple[TopKEntry, ...], ...] | None = None + + if want_logprobs or score_input: + # Logprobs are read from the first sample only (alogprobs enforces + # n=1; direct IR callers sampling n>1 with logprobs get sample 0). + meta = metas[0] + if score_input: + self._guard_radix_cache(meta) + input_scoring = InputScoringResult( + token_logprobs=self._triples_to_tokens( + meta.get("input_token_logprobs") or [] + ) + ) + logprobs = self._triples_to_tokens(meta.get("output_token_logprobs") or []) + top_logprobs = self._triples_to_topk(meta.get("output_top_logprobs") or []) + if ( + not logprobs + and not top_logprobs + and not (input_scoring and input_scoring.token_logprobs) + ): + raise RuntimeError("sglang /generate returned no logprobs.") + + return Response( + texts=texts, + logprobs=logprobs, + top_logprobs=top_logprobs, + input_scoring=input_scoring, + usage=usage, + finish_reasons=finish_reasons, + request_params=_request_params(body), + response_model=self._model, + ) + + # ── arun ────────────────────────────────────────────────────────────────── + + async def arun(self, req: Request) -> Response: + body = self._lower(req) + raw = await self._post(body) + # n>1 yields a list of per-sample dicts; n==1 a single dict. + results = raw if isinstance(raw, list) else [raw] + if not results or not all( + isinstance(r, dict) and "meta_info" in r for r in results + ): + raise RuntimeError( + "sglang /generate returned an unexpected response shape " + "(missing meta_info)." + ) + return self._lift( + results, + body, + score_input=req.score_input, + want_logprobs=req.return_logprobs, + ) diff --git a/sieval/core/tasks/task.py b/sieval/core/tasks/task.py index fef60442..96ebed38 100644 --- a/sieval/core/tasks/task.py +++ b/sieval/core/tasks/task.py @@ -1,4 +1,7 @@ -"""Abstract base class for the five-stage evaluation pipeline.""" +"""Abstract base class for the five-stage evaluation pipeline. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" import re from abc import ABC, abstractmethod @@ -6,7 +9,7 @@ from typing import ClassVar, Literal from sieval.core.datasets import Dataset -from sieval.core.models import Model +from sieval.core.models import Capability, Model from .context import TaskContext @@ -44,10 +47,17 @@ class Task[ *cli/session*). tags: Free-form tag set describing the task (e.g. ``{"gen", "zero_shot"}``). Used by the anomaly-detection framework to decide which rules apply. + requires: IR :class:`~sieval.core.models.Capability` set the task needs + from its model's Transport. Checked at construction via + :meth:`~sieval.core.models.Model.assert_capability`, so a model that + cannot honour a required feature (e.g. a chat backend asked for + ``InputScoring``) fails loud at setup rather than silently producing + wrong results. Defaults to the empty set (no IR requirement). """ model_type: ClassVar[Literal["chat", "gen"] | None] = None tags: ClassVar[AbstractSet[str]] = frozenset() # override in subclasses + requires: ClassVar[frozenset[Capability]] = frozenset() # override in subclasses def __init__( self, dataset: Dataset[TRawSample], model: Model, name: str | None = None @@ -58,6 +68,8 @@ def __init__( if self.model_type is not None: self._validate_model_type() + if self.requires: + self._model.assert_capability(*self.requires) def _validate_model_type(self) -> None: """Raise ``TypeError`` if the model's kind does not match :attr:`model_type`.""" diff --git a/tests/README.md b/tests/README.md index 6384be00..16f32a58 100644 --- a/tests/README.md +++ b/tests/README.md @@ -146,6 +146,8 @@ python -m pytest tests/integration/resume/test_basic.py::TestResumePartialComple All shared test infrastructure lives here — available to every test layer without any explicit import. +Mock models stub at the **Transport seam** (RFC #25): each overrides `_build_default_transport()` to return a `HandlerTransport` bound to its `async _stub_arun(Request) -> Response` handler, so `agenerate` / `alogprobs` exercise the real request builders and the `Response -> ModelOutput` bridge while the wire layer stays canned. Subclass mocks override `_stub_arun` and chain via `await super()._stub_arun(req)`. + ### Unit / Integration mocks | Class | Description | @@ -154,10 +156,12 @@ All shared test infrastructure lives here — available to every test layer with | `MockChatModel(answers={...})` | Deterministic chat model | | `MockGenModel(logprob_scores={...})` | Deterministic gen model (alogprobs) | | `MockJudgeModel(verdict="yes")` | LLM-as-judge mock | -| `MockCountingChatModel(answers={...})` | `MockChatModel` that counts `_agenerate_impl` calls | +| `MockCountingChatModel(answers={...})` | `MockChatModel` that counts transport hits (`_stub_arun` calls) | | `MockAlwaysFailModel()` | Always raises an exception | | `MockFailingChatModel(fail_count=1)` | Fails N times then succeeds | | `MockSelectiveFailModel(fail_samples={...})` | Fails on first call for specific prompts | +| `HandlerTransport(handler, capabilities)` | Transport double; records every lowered `Request` in `.requests` | +| `prompt_of(req)` / `n_of(req)` | Extract the flat prompt text / sample count from a `Request` | | `make_config(tmp_path, **overrides)` | `TaskRunnerConfig` for unit/integration tests | ### Performance / Acceptance infrastructure diff --git a/tests/conftest.py b/tests/conftest.py index 12ae6fc5..2fd47c76 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,9 +32,20 @@ from loguru import logger as _loguru_logger from sieval.core.datasets import Dataset -from sieval.core.models import ModelOutput +from sieval.core.models import ( + Capability, + ModelOutput, + Request, + Response, + TokenLogprob, + UsageStats, +) from sieval.core.models.chat_model import ChatModel from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import ( + OpenAIChatTransport, + OpenAICompletionsTransport, +) from sieval.core.runners.runner import TaskRunnerConfig from sieval.core.tasks.context import TaskContext, TaskStage from sieval.core.tasks.saver import TaskSaver @@ -122,7 +133,47 @@ def load(self, name_or_path: str, **kwargs) -> HFDatasetDict: # =================================================================== # Mock Models +# +# RFC #25: mocks stub at the Transport seam. Each mock model overrides +# ``_build_default_transport`` to return a ``HandlerTransport`` bound to +# its ``_stub_arun(Request) -> Response`` handler, so ``agenerate`` / +# ``alogprobs`` exercise the real request builders and Response bridge +# while the wire layer stays canned. Subclass mocks override +# ``_stub_arun`` and chain via ``super()``. # =================================================================== +class HandlerTransport: + """Transport double: forwards ``arun`` to a handler coroutine. + + Records every Request in ``self.requests`` so tests can assert on the + lowered IR instead of legacy kwargs. + """ + + def __init__(self, handler, capabilities: frozenset[Capability]): + self._handler = handler + self._capabilities = frozenset(capabilities) + self.requests: list[Request] = [] + + @property + def capabilities(self) -> frozenset[Capability]: + return self._capabilities + + async def arun(self, req: Request) -> Response: + self.requests.append(req) + return await self._handler(req) + + +def prompt_of(req: Request) -> str: + """Extract the flat question text from a Request input (str or messages).""" + if isinstance(req.input, str): + return req.input + return req.input[-1]["content"] if req.input else "" + + +def n_of(req: Request) -> int: + """Number of samples requested.""" + return req.sampling.n if req.sampling is not None else 1 + + class MockChatModel(ChatModel): """ChatModel that returns deterministic answers without calling any API.""" @@ -132,37 +183,30 @@ def __init__( default_answer: str = "unknown", **kwargs, ): - super().__init__(model="mock-chat", api_key="fake", **kwargs) self._answers = answers or {} self._default_answer = default_answer + super().__init__(model="mock-chat", api_key="fake", **kwargs) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - # Extract question from messages - if isinstance(prompt, str): - q = prompt - else: - msgs = list(prompt) - q = msgs[-1]["content"] if msgs else "" + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) + async def _stub_arun(self, req: Request) -> Response: + q = prompt_of(req) answer = self._answers.get(q, self._default_answer) - n = kwargs.get("n", 1) + n = n_of(req) if isinstance(answer, list): texts = answer[:n] if len(answer) >= n else answer else: texts = [answer] * n - return ModelOutput( - model=self.meta(), - texts=texts, - finish_reasons=["stop"] * len(texts), - usage={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + return Response( + texts=tuple(texts), + finish_reasons=("stop",) * len(texts), + usage=UsageStats(input_tokens=10, output_tokens=2, total_tokens=12), request_params={"model": "mock-chat", "n": n}, ) - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class MockGenModel(GenModel): """GenModel that supports alogprobs without calling any API.""" @@ -173,31 +217,35 @@ def __init__( default_answer: str = "unknown", **kwargs, ): - super().__init__(model="mock-gen", api_key="fake", **kwargs) self._logprob_scores = logprob_scores or {} self._default_answer = default_answer + super().__init__(model="mock-gen", api_key="fake", **kwargs) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - return ModelOutput( - model=self.meta(), - texts=[self._default_answer], - finish_reasons=["stop"], - usage={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES ) - async def _alogprobs_impl(self, prompt: str, **kwargs) -> ModelOutput: + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response( + texts=(self._default_answer,), + finish_reasons=("stop",), + usage=UsageStats(input_tokens=10, output_tokens=2, total_tokens=12), + ) + # Extract the last character as the option label + prompt = prompt_of(req) option_label = prompt.rstrip()[-1] if prompt.strip() else "A" score = self._logprob_scores.get(option_label, -10.0) - - return ModelOutput( - model=self.meta(), - texts=[""], - finish_reasons=["stop"], - logprobs_tokens=[f" {option_label}"], - logprobs=[score], - usage={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, - request_params={"max_tokens": kwargs.get("max_tokens", 1)}, + max_tokens = req.sampling.max_tokens if req.sampling is not None else None + + return Response( + texts=("",), + finish_reasons=("stop",), + logprobs=(TokenLogprob(token=f" {option_label}", logprob=score),), + usage=UsageStats(input_tokens=10, output_tokens=1, total_tokens=11), + request_params={"max_tokens": max_tokens}, ) @@ -205,70 +253,68 @@ class MockJudgeModel(ChatModel): """ChatModel that acts as a judge, returning configurable verdicts.""" def __init__(self, verdict: str = "yes", **kwargs): - super().__init__(model="mock-judge", api_key="fake", **kwargs) self._verdict = verdict + super().__init__(model="mock-judge", api_key="fake", **kwargs) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - return ModelOutput( - model=self.meta(), - texts=[self._verdict], - finish_reasons=["stop"], - usage={"input_tokens": 20, "output_tokens": 1, "total_tokens": 21}, + async def _stub_arun(self, req: Request) -> Response: + return Response( + texts=(self._verdict,), + finish_reasons=("stop",), + usage=UsageStats(input_tokens=20, output_tokens=1, total_tokens=21), request_params={"model": "mock-judge"}, ) - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class MockFailingChatModel(ChatModel): """ChatModel that fails for specified number of calls, then succeeds.""" def __init__(self, fail_count: int = 1, success_answer: str = "42", **kwargs): - super().__init__(model="mock-failing", api_key="fake", **kwargs) self._call_count = 0 self._fail_count = fail_count self._success_answer = success_answer + super().__init__(model="mock-failing", api_key="fake", **kwargs) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: + async def _stub_arun(self, req: Request) -> Response: self._call_count += 1 if self._call_count <= self._fail_count: raise TimeoutError(f"Simulated failure #{self._call_count}") - return ModelOutput( - model=self.meta(), - texts=[self._success_answer], - finish_reasons=["stop"], - usage={"input_tokens": 5, "output_tokens": 1, "total_tokens": 6}, + return Response( + texts=(self._success_answer,), + finish_reasons=("stop",), + usage=UsageStats(input_tokens=5, output_tokens=1, total_tokens=6), ) - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class MockAlwaysFailModel(ChatModel): """ChatModel that always raises an exception.""" def __init__(self, error: type[Exception] = TimeoutError, **kwargs): - super().__init__(model="mock-always-fail", api_key="fake", **kwargs) self._error = error + super().__init__(model="mock-always-fail", api_key="fake", **kwargs) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - raise self._error("Always fails") + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError + async def _stub_arun(self, req: Request) -> Response: + raise self._error("Always fails") class MockCountingChatModel(MockChatModel): - """MockChatModel that counts how many times _agenerate_impl is called.""" + """MockChatModel that counts how many times the transport is hit.""" def __init__(self, **kwargs): - super().__init__(**kwargs) self.call_count = 0 + super().__init__(**kwargs) - async def _agenerate_impl(self, prompt, **kwargs): + async def _stub_arun(self, req: Request) -> Response: self.call_count += 1 - return await super()._agenerate_impl(prompt, **kwargs) + return await super()._stub_arun(req) class MockSelectiveFailModel(ChatModel): @@ -281,14 +327,17 @@ def __init__( default_answer: str = "42", **kwargs, ): - super().__init__(model="mock-selective", api_key="fake", **kwargs) self._fail_samples = fail_samples or set() self._answers = answers or {} self._default_answer = default_answer self._call_counts: dict[str, int] = {} + super().__init__(model="mock-selective", api_key="fake", **kwargs) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - q = prompt if isinstance(prompt, str) else list(prompt)[-1]["content"] + async def _stub_arun(self, req: Request) -> Response: + q = prompt_of(req) self._call_counts[q] = self._call_counts.get(q, 0) + 1 # Fail on first call if prompt matches any fail pattern @@ -296,16 +345,12 @@ async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: raise TimeoutError(f"Simulated first-time failure for: {q}") answer = self._answers.get(q, self._default_answer) - return ModelOutput( - model=self.meta(), - texts=[answer], - finish_reasons=["stop"], - usage={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + return Response( + texts=(answer,), + finish_reasons=("stop",), + usage=UsageStats(input_tokens=10, output_tokens=2, total_tokens=12), ) - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - # =================================================================== # Factory Functions @@ -428,28 +473,27 @@ def __init__( self._output_size = output_size self._output_text = default_answer or ("x" * max(1, output_size)) - async def _agenerate_impl(self, prompt: Any, **kwargs: Any) -> ModelOutput: + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) + + async def _stub_arun(self, req: Request) -> Response: jitter = random.uniform(-self._latency_jitter, self._latency_jitter) await anyio.sleep(max(0, self._latency_s + jitter)) - n = kwargs.get("n", 1) - texts = [self._output_text] * n - input_tokens = max(1, len(prompt) // 4) if isinstance(prompt, str) else 10 + n = n_of(req) + prompt = prompt_of(req) + input_tokens = max(1, len(prompt) // 4) if prompt else 10 output_tokens = max(1, self._output_size) - return ModelOutput( - model=self.meta(), - texts=texts, - finish_reasons=["stop"] * n, - usage={ - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - }, + return Response( + texts=(self._output_text,) * n, + finish_reasons=("stop",) * n, + usage=UsageStats( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ), request_params={"model": "mock-latency", "n": n}, ) - async def _alogprobs_impl(self, prompt: Any, **kwargs: Any) -> ModelOutput: - raise NotImplementedError - @classmethod def from_profile(cls, profile: IOProfile, **kwargs: Any) -> "LatencyMockChatModel": return cls( diff --git a/tests/integration/resume/test_advanced.py b/tests/integration/resume/test_advanced.py index 8ad61529..30565e23 100644 --- a/tests/integration/resume/test_advanced.py +++ b/tests/integration/resume/test_advanced.py @@ -14,6 +14,7 @@ MockChatModel, MockDataset, make_config, + prompt_of, ) from .conftest import ( @@ -173,12 +174,12 @@ async def test_resume_from_iteration_boundary(self, tmp_path): call_counts = {} class FailOnSecondIterModel(MockChatModel): - async def _agenerate_impl(self, prompt, **kwargs): - q = prompt if isinstance(prompt, str) else list(prompt)[-1]["content"] + async def _stub_arun(self, req): + q = prompt_of(req) call_counts[q] = call_counts.get(q, 0) + 1 if call_counts[q] == 2: raise RuntimeError("Fail on iteration 1") - return await super()._agenerate_impl(prompt, **kwargs) + return await super()._stub_arun(req) model1 = FailOnSecondIterModel(answers={"I1": "A1", "I2": "A2"}) task1 = IterativeTask(dataset=dataset, model=model1, name="iter_resume") diff --git a/tests/integration/test_model_backward_compat.py b/tests/integration/test_model_backward_compat.py new file mode 100644 index 00000000..5528c878 --- /dev/null +++ b/tests/integration/test_model_backward_compat.py @@ -0,0 +1,129 @@ +""" +RFC #25 Phase-4 regression suite: legacy public Model surface on the new routing. + +``Model._agenerate_impl`` / ``_alogprobs_impl`` are gone — ``agenerate`` and +``alogprobs`` are now sugar over ``arun(Request) -> Response`` through a +Transport, with the Response bridged back to the legacy ``ModelOutput`` shape. +These tests pin the observable legacy contract on transport-stubbed conftest +mocks (no real HTTP): + + - ``agenerate`` output shape (texts / finish_reasons / usage, n-sampling) + - ``alogprobs`` output shape and the ``echo=True`` InputScoring gate + (loud ``CapabilityError`` instead of the historical silent ignore) + - ``with_args`` derivation: generation still works, resource pool is shared + - ``meta()`` / ``get_quota_info()`` introspection dict shapes + - ``Model.as_type`` no longer exists + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +import pytest + +from sieval.core.models import Capability, CapabilityError, Model, ModelOutput +from tests.conftest import MockChatModel, MockGenModel + + +class TestAgenerateLegacySurface: + @pytest.mark.anyio + async def test_agenerate_returns_legacy_model_output(self): + out = await MockChatModel().agenerate("hello") + + assert isinstance(out, ModelOutput) + assert out.texts == ["unknown"] + assert out.finish_reasons == ["stop"] + assert out.usage == { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + } + + @pytest.mark.anyio + async def test_agenerate_n_sampling_returns_n_texts(self): + out = await MockChatModel().agenerate("hello", n=3) + + assert isinstance(out, ModelOutput) + assert len(out.texts) == 3 + assert out.texts == ["unknown"] * 3 + assert out.finish_reasons == ["stop"] * 3 + + +class TestAlogprobsLegacySurface: + @pytest.mark.anyio + async def test_alogprobs_no_echo_returns_logprobs(self): + out = await MockGenModel().alogprobs("prompt A", echo=False) + + assert isinstance(out, ModelOutput) + assert out.logprobs_tokens == [" A"] + assert out.logprobs == [-10.0] + + @pytest.mark.anyio + async def test_alogprobs_echo_works_on_gen_model(self): + model = MockGenModel() + assert Capability.InputScoring in model.capabilities + + out = await model.alogprobs("prompt A", echo=True) + + assert isinstance(out, ModelOutput) + assert out.logprobs_tokens == [" A"] + assert out.logprobs == [-10.0] + + @pytest.mark.anyio + async def test_alogprobs_echo_on_chat_model_raises_capability_error(self): + # Historically echo=True was silently ignored on chat backends; + # RFC #25 makes the missing InputScoring capability loud. + model = MockChatModel() + assert Capability.InputScoring not in model.capabilities + + with pytest.raises(CapabilityError, match="InputScoring"): + await model.alogprobs("prompt A", echo=True) + + +class TestWithArgsDerivation: + @pytest.mark.anyio + async def test_derived_model_generates_and_shares_pool(self): + base = MockChatModel(concurrency_limit=8) + derived = base.with_args(concurrency_limit=4) + + out = await derived.agenerate("hello") + assert isinstance(out, ModelOutput) + assert out.texts == ["unknown"] + + # Two-level pool: the base's limiter becomes the derived model's parent. + assert base._limiter is not None + assert derived._parent_limiter is base._limiter + + +class TestIntrospectionSurface: + def test_meta_returns_model_meta_dict_shape(self): + meta = MockChatModel().meta() + + assert set(meta) >= {"model", "api_base", "default_params"} + assert meta["model"] == "mock-chat" + assert meta["api_base"] is None + assert meta["default_params"] == {} + + def test_get_quota_info_returns_quota_dict_shape(self): + base = MockChatModel(concurrency_limit=8) + info = base.get_quota_info() + + assert set(info) == {"available", "total", "parent", "child"} + assert info["available"] == 8 + assert info["total"] == 8 + assert info["parent"] is None + assert info["child"] == {"available": 8, "total": 8} + + def test_get_quota_info_on_derived_model_reports_parent(self): + base = MockChatModel(concurrency_limit=8) + derived = base.with_args(concurrency_limit=4) + info = derived.get_quota_info() + + assert info["parent"] == {"available": 8, "total": 8} + assert info["child"] == {"available": 4, "total": 4} + assert info["available"] == 4 + + +class TestRemovedSurface: + def test_as_type_no_longer_exists(self): + base = MockChatModel() + assert not hasattr(base, "as_type") + assert not hasattr(Model, "as_type") diff --git a/tests/integration/test_multi_task.py b/tests/integration/test_multi_task.py index d61b8baf..c61e17cc 100644 --- a/tests/integration/test_multi_task.py +++ b/tests/integration/test_multi_task.py @@ -110,7 +110,7 @@ async def test_one_task_fails_other_succeeds(self, tmp_path): good_model = MockChatModel(answers={"B1?": "X", "B2?": "Y"}) class FailingModel(MockChatModel): - async def _agenerate_impl(self, prompt, **kwargs): + async def _stub_arun(self, req): raise RuntimeError("Task A always fails") bad_model = FailingModel() diff --git a/tests/integration/test_runner_edge_cases.py b/tests/integration/test_runner_edge_cases.py index c4247e5e..c97abc3e 100644 --- a/tests/integration/test_runner_edge_cases.py +++ b/tests/integration/test_runner_edge_cases.py @@ -14,13 +14,13 @@ AI-Generated Code - Claude Sonnet 4.6 (Anthropic) """ +import dataclasses import json from pathlib import Path from typing import ClassVar import pytest -from sieval.core.models import ModelOutput from sieval.core.runners.runner import TaskRunner from sieval.core.tasks.consts import TaskStage from sieval.core.tasks.task import Task @@ -269,15 +269,11 @@ async def test_anomaly_report_includes_pre_resume_samples(self, tmp_path): # --- Mock model that returns truncated outputs --- class MockTruncatedModel(MockChatModel): - async def _agenerate_impl(self, prompt, **kwargs): - result = await super()._agenerate_impl(prompt, **kwargs) + async def _stub_arun(self, req): + resp = await super()._stub_arun(req) # Return truncated finish_reason - return ModelOutput( - model=result.model, - texts=result.texts, - finish_reasons=["length"] * len(result.texts), - usage=result.usage, - request_params=result.request_params, + return dataclasses.replace( + resp, finish_reasons=("length",) * len(resp.texts) ) # --- First run: complete 2 samples with truncated outputs --- @@ -286,12 +282,12 @@ def __init__(self, **kwargs): super().__init__(**kwargs) self._call_count = 0 - async def _agenerate_impl(self, prompt, **kwargs): + async def _stub_arun(self, req): self._call_count += 1 # Fail on third sample (E3) if self._call_count >= 3: raise TimeoutError("Simulated failure on E3") - return await super()._agenerate_impl(prompt, **kwargs) + return await super()._stub_arun(req) model1 = PartialFailTruncatedModel(answers=EDGE_ANSWERS) task1 = EdgeTask(dataset=dataset, model=model1, name="anomaly_resume_test") diff --git a/tests/unit/cli/leaderboard/test_session.py b/tests/unit/cli/leaderboard/test_session.py index 7f8f164a..d7e7dedb 100644 --- a/tests/unit/cli/leaderboard/test_session.py +++ b/tests/unit/cli/leaderboard/test_session.py @@ -1179,8 +1179,8 @@ def test_derived_model_no_type_conversion(self): # child should share parent_limiter = base._limiter assert child._parent_limiter is base._limiter - def test_derived_model_with_type_conversion(self): - """A derived model with 'type' should call as_type on the base.""" + def test_derived_model_matching_type_is_accepted(self): + """`type:` on a derived model is a no-op when it matches the base kind.""" base = MockChatModel(concurrency_limit=64) runner = self._make_runner( { @@ -1196,23 +1196,18 @@ def test_derived_model_with_type_conversion(self): ) runner.models["base"] = base - mock_converted = MagicMock() - mock_converted.with_args.return_value = mock_converted - - with ( - patch.object(base, "as_type", return_value=mock_converted) as mock_as_type, - patch( - "sieval.cli.leaderboard.session.ChatModel", - return_value=base, - ) as mock_chat_cls, + with patch( + "sieval.cli.leaderboard.session.ChatModel", + return_value=base, ): runner._setup_models() - mock_as_type.assert_called_once_with(mock_chat_cls) - assert runner.models["child"] is mock_converted + # No conversion machinery: the child derives directly from the base + # (no args → no with_args fork either). + assert runner.models["child"] is base - def test_derived_model_with_gen_type_conversion(self): - """Derived model with type='gen' should request GenModel conversion.""" + def test_derived_model_cross_kind_type_raises(self): + """RFC #25 dropped as_type: cross-kind derived `type:` is a config error.""" base = MockChatModel(concurrency_limit=64) runner = self._make_runner( { @@ -1228,23 +1223,15 @@ def test_derived_model_with_gen_type_conversion(self): ) runner.models["base"] = base - mock_converted = MagicMock() - mock_converted.with_args.return_value = mock_converted - with ( - patch.object(base, "as_type", return_value=mock_converted) as mock_as_type, patch( "sieval.cli.leaderboard.session.ChatModel", return_value=base, ), + pytest.raises(ValueError, match="cross-kind conversion was removed"), ): runner._setup_models() - from sieval.core.models.gen_model import GenModel - - mock_as_type.assert_called_once_with(GenModel) - assert runner.models["child"] is mock_converted - def test_derived_model_validation_errors(self): """Derived model should reject unknown base and invalid type.""" runner = self._make_runner( diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 69afef51..35a757a8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,11 +4,20 @@ AI-Generated Code - Claude Opus 4.6 (Anthropic) """ +import os + import pytest from sieval.core.models.model import ModelMeta, ModelOutput, ModelUsage from sieval.core.tasks.context import TaskContext, TaskStageMeta +# scripts/ tests are repo-hygiene checks (preflight, layer imports) that run +# tooling against the LIVE repo tree. Inside mutmut's mutants/ sandbox that +# tree is a partial copy, so they fail spuriously — and they exercise no +# sieval/core mutants. Skip collecting them under mutation testing only +# (mutmut sets MUTANT_UNDER_TEST for every pytest run it drives). +collect_ignore = ["scripts"] if os.environ.get("MUTANT_UNDER_TEST") else [] + @pytest.fixture def sample_model_meta() -> ModelMeta: diff --git a/tests/unit/core/models/test_chat_model.py b/tests/unit/core/models/test_chat_model.py index fd4a56af..94d1307e 100644 --- a/tests/unit/core/models/test_chat_model.py +++ b/tests/unit/core/models/test_chat_model.py @@ -1,1088 +1,28 @@ +"""Shell tests: backend selector wiring for ChatModel. + +RFC #25 moved the chat-completions wire logic (streaming accumulation, +reasoning extraction, logprob parsing) into ``OpenAIChatTransport``; the +former ``_agenerate_impl`` / ``_alogprobs_impl`` coverage moved with it to +tests/unit/core/models/transports/test_openai_chat.py, and the request-builder +validation (n/stream types, alogprobs n=1) lives on ``Model`` in +tests/unit/core/models/test_model.py. What remains here is the selector +contract: ChatModel pairs the shared client with that transport, which +supplies the model's capabilities. + +AI-Generated Code - Claude Fable 5 (Anthropic) """ -Unit tests for sieval/core/models/chat_model.py. -Covers: _agenerate_impl (string prompt, message list prompt, streaming -accumulation, n>1 choices, usage, reasoning_content), _alogprobs_impl -(streaming logprobs, no-logprobs raises, usage). +from sieval.core.models import Capability, ChatModel +from sieval.core.models.transports.openai_chat import OpenAIChatTransport -All OpenAI client calls are mocked — no real API traffic. -AI-Generated Code - Claude Opus 4.6 (Anthropic) -""" - -from types import SimpleNamespace -from typing import Any -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from sieval.core.models.chat_model import ChatModel -from sieval.core.models.model import ModelOutput - - -# --------------------------------------------------------------------------- -# Async streaming helpers -# --------------------------------------------------------------------------- -def _make_chunk( - index: int = 0, - content: str = "", - finish_reason: str = "", - usage=None, - reasoning: str = "", -): - """Build a minimal streaming chunk object.""" - chunk = MagicMock() - chunk.usage = usage - - if content or finish_reason is not None or reasoning: - choice = MagicMock() - choice.index = index - choice.finish_reason = finish_reason - delta = MagicMock() - delta.content = content or None - # reasoning_content attribute - reasoning_attr = reasoning or None - delta.reasoning = None - delta.reasoning_content = reasoning_attr - choice.delta = delta - chunk.choices = [choice] - else: - chunk.choices = [] - - return chunk - - -def _make_usage_chunk(prompt_tokens=10, completion_tokens=5): - """Chunk that carries usage but no choices.""" - chunk = MagicMock() - chunk.choices = [] - chunk.usage = MagicMock() - chunk.usage.prompt_tokens = prompt_tokens - chunk.usage.completion_tokens = completion_tokens - chunk.usage.total_tokens = prompt_tokens + completion_tokens - return chunk - - -def _make_non_stream_response( - *, - text: str = "", - finish_reason: str | None = "stop", - reasoning: str | None = None, - reasoning_content: str | None = None, - usage=None, -): - resp = MagicMock() - choice = MagicMock() - choice.index = 0 - choice.finish_reason = finish_reason - message = MagicMock() - message.content = text - message.reasoning = reasoning - message.reasoning_content = reasoning_content - choice.message = message - resp.choices = [choice] - resp.usage = usage - return resp - - -def _make_non_stream_choice( - *, - index: int = 0, - text: str | list[object] = "", - finish_reason: str | None = "stop", - reasoning: str | None = None, - reasoning_content: str | None = None, - logprob_items: list[tuple[str, float]] | None = None, -): - choice = MagicMock() - choice.index = index - choice.finish_reason = finish_reason - message = MagicMock() - message.content = text - message.reasoning = reasoning - message.reasoning_content = reasoning_content - choice.message = message - if logprob_items is not None: - logprobs_obj = MagicMock() - logprobs_obj.content = [ - SimpleNamespace(token=token, logprob=logprob) - for token, logprob in logprob_items - ] - choice.logprobs = logprobs_obj - else: - choice.logprobs = None - return choice - - -def _make_non_stream_response_from_choices(choices: list[object], usage=None): - resp = MagicMock() - resp.choices = choices - resp.usage = usage - return resp - - -def _make_usage(prompt_tokens=10, completion_tokens=5): - usage = MagicMock() - usage.prompt_tokens = prompt_tokens - usage.completion_tokens = completion_tokens - usage.total_tokens = prompt_tokens + completion_tokens - return usage - - -class _AsyncIterator: - """Wraps a list into an async iterator.""" - - def __init__(self, items): - self._items = iter(items) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self._items) - except StopIteration as e: - raise StopAsyncIteration from e - - -# --------------------------------------------------------------------------- -# Concrete ChatModel (delegates to parent _agenerate_impl) -# --------------------------------------------------------------------------- -class _TestChatModel(ChatModel): - """ChatModel that calls the real parent _agenerate_impl / _alogprobs_impl.""" - - # No override — let parent do the work so we test actual implementation. - pass - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- -@pytest.fixture -def model(): - return _TestChatModel(model="test-chat", api_key="fake") - - -def _patch_create(model: _TestChatModel, chunks): - """Patch model._client.chat.completions.create to return async iterator.""" - mock_create = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.chat.completions - target.create = mock_create # type: ignore[invalid-assignment] - return mock_create - - -# =================================================================== -# _agenerate_impl — string prompt -# =================================================================== -class TestAGenerateString: - @pytest.mark.anyio - async def test_basic_string_prompt(self, model): - chunks = [ - _make_chunk(content="Hello"), - _make_chunk(content=" world", finish_reason="stop"), - _make_usage_chunk(10, 3), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("Hi") - assert isinstance(out, ModelOutput) - assert out.texts == ["Hello world"] - assert out.finish_reasons == ["stop"] - - @pytest.mark.anyio - async def test_usage_captured(self, model): - chunks = [ - _make_chunk(content="ok", finish_reason="stop"), - _make_usage_chunk(8, 2), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.usage is not None - assert out.usage["input_tokens"] == 8 - assert out.usage["output_tokens"] == 2 - assert out.usage["total_tokens"] == 10 - - @pytest.mark.anyio - async def test_no_usage_chunk(self, model): - chunks = [_make_chunk(content="hi", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.usage is None - - @pytest.mark.anyio - async def test_model_meta_attached(self, model): - chunks = [_make_chunk(content="x", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.model["model"] == "test-chat" - - @pytest.mark.anyio - async def test_request_params_captured(self, model): - chunks = [_make_chunk(content="x", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt", temperature=0.7) - assert out.request_params is not None - assert out.request_params.get("temperature") == 0.7 - - @pytest.mark.anyio - async def test_non_stream_mode_supported(self, model): - response = _make_non_stream_response( - text="done", - finish_reason="stop", - usage=_make_usage(9, 4), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - - assert out.texts == ["done"] - assert out.finish_reasons == ["stop"] - assert out.usage is not None - assert out.usage["input_tokens"] == 9 - call_kwargs = mock_create.call_args[1] - assert call_kwargs["stream"] is False - - @pytest.mark.anyio - async def test_non_stream_missing_finish_reason_defaults_to_empty(self, model): - response = _make_non_stream_response( - text="done", - finish_reason=None, - usage=_make_usage(2, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.texts == ["done"] - assert out.finish_reasons == [""] - - @pytest.mark.anyio - async def test_stream_options_can_be_overridden(self, model): - chunks = [_make_chunk(content="x", finish_reason="stop")] - mock_create = _patch_create(model, chunks) - out = await model._agenerate_impl( - "prompt", - stream_options={"include_usage": False}, - ) - assert out.texts == ["x"] - call_kwargs = mock_create.call_args[1] - assert call_kwargs["stream_options"] == {"include_usage": False} - - @pytest.mark.anyio - async def test_non_stream_invalid_content_type_raises(self, model): - response = _make_non_stream_response_from_choices( - choices=[ - _make_non_stream_choice( - text=[ - "A", - SimpleNamespace(text="B"), - SimpleNamespace(text=123), - ] - ) - ], - usage=_make_usage(3, 2), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - with pytest.raises(TypeError): - await model._agenerate_impl("prompt", stream=False) - - @pytest.mark.anyio - async def test_non_stream_ignores_out_of_range_choice(self, model): - response = _make_non_stream_response_from_choices( - choices=[_make_non_stream_choice(index=5, text="ignored")], - usage=_make_usage(7, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", n=1, stream=False) - assert out.texts == [""] - assert out.finish_reasons == [""] - assert out.usage is not None - assert out.usage["input_tokens"] == 7 - - @pytest.mark.anyio - async def test_stream_ignores_out_of_range_choice(self, model): - chunk = _make_chunk(index=5, content="ignored", finish_reason="stop") - _patch_create(model, [chunk, _make_usage_chunk(4, 2)]) - out = await model._agenerate_impl("prompt", n=1) - assert out.texts == [""] - assert out.finish_reasons == [""] - assert out.usage is not None - assert out.usage["total_tokens"] == 6 - - -# =================================================================== -# _agenerate_impl — message list prompt -# =================================================================== -class TestAGenerateMessageList: - @pytest.mark.anyio - async def test_message_list_prompt(self, model): - messages = [{"role": "user", "content": "hello"}] - chunks = [_make_chunk(content="reply", finish_reason="stop")] - mock = _patch_create(model, chunks) - out = await model._agenerate_impl(messages) - assert out.texts == ["reply"] - # Verify messages forwarded - call_kwargs = mock.call_args[1] - assert call_kwargs["messages"] == messages - - -# =================================================================== -# _agenerate_impl — n > 1 choices -# =================================================================== -class TestAGenerateMultipleChoices: - @pytest.mark.anyio - async def test_two_choices_in_single_chunk(self, model): - def _multi_choice_chunk(pairs, finish_reason=""): - chunk = MagicMock() - chunk.usage = None - chunk.choices = [] - for index, content in pairs: - choice = MagicMock() - choice.index = index - choice.finish_reason = finish_reason - delta = MagicMock() - delta.content = content - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - chunk.choices.append(choice) - return chunk - - chunks = [ - _multi_choice_chunk([(0, "A1"), (1, "B1")]), - _multi_choice_chunk([(0, "A2"), (1, "B2")], finish_reason="stop"), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt", n=2) - assert out.texts[0] == "A1A2" - assert out.texts[1] == "B1B2" - - -# =================================================================== -# _agenerate_impl — reasoning_content -# =================================================================== -class TestAGenerateReasoning: - @pytest.mark.anyio - @pytest.mark.parametrize( - "reasoning, reasoning_content, expected", - [ - (None, "think", "think"), # reasoning_content fallback - ("primary", "secondary", "primary"), # reasoning takes priority - ], - ids=["reasoning_content-fallback", "reasoning-priority"], - ) - async def test_stream_reasoning_extraction( - self, model, reasoning, reasoning_content, expected - ): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "answer" - delta.reasoning = reasoning - delta.reasoning_content = reasoning_content - choice.delta = delta - chunk.choices = [choice] - - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == expected - - @pytest.mark.anyio - async def test_reasoning_accumulated_multi_chunk(self, model): - """Reasoning tokens accumulate across multiple streaming chunks.""" - chunks = [ - _make_chunk(content="a", reasoning="think1"), - _make_chunk(content="b", reasoning="think2", finish_reason="stop"), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.texts == ["ab"] - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == "think1think2" - - @pytest.mark.anyio - @pytest.mark.parametrize( - "reasoning, reasoning_content, expected", - [ - ("step by step", None, "step by step"), # reasoning field - (None, "fallback", "fallback"), # reasoning_content fallback - ("primary", "secondary", "primary"), # priority - ], - ids=["reasoning-field", "reasoning_content-fallback", "priority"], - ) - async def test_non_stream_reasoning_extraction( - self, model, reasoning, reasoning_content, expected - ): - """Non-streaming: reasoning extraction of chat_model.py.""" - response = _make_non_stream_response( - text="answer", - finish_reason="stop", - reasoning=reasoning, - reasoning_content=reasoning_content, - usage=_make_usage(10, 5), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.texts == ["answer"] - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == expected - - -# =================================================================== -# _alogprobs_impl — reasoning in logprobs mode -# =================================================================== -class TestALogprobsReasoning: - """Cover reasoning extraction paths in _alogprobs_impl.""" - - def _make_stream_logprobs_chunk(self, reasoning, reasoning_content): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "answer" - delta.reasoning = reasoning - delta.reasoning_content = reasoning_content - choice.delta = delta - lp_item = MagicMock() - lp_item.token = "A" - lp_item.logprob = -0.1 - logprobs_obj = MagicMock() - logprobs_obj.content = [lp_item] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - return chunk - - def _patch_logprobs_create(self, model: _TestChatModel, chunks): - mock_create = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.chat.completions - target.create = mock_create # type: ignore[invalid-assignment] - return mock_create - - @pytest.mark.anyio - @pytest.mark.parametrize( - "reasoning, reasoning_content, expected", - [ - ("think", None, "think"), # reasoning field - (None, "fallback", "fallback"), # reasoning_content fallback - ], - ids=["reasoning-field", "reasoning_content-fallback"], - ) - async def test_stream_logprobs_reasoning_extraction( - self, model, reasoning, reasoning_content, expected - ): - chunk = self._make_stream_logprobs_chunk(reasoning, reasoning_content) - self._patch_logprobs_create(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == expected - - @pytest.mark.anyio - @pytest.mark.parametrize( - "reasoning, reasoning_content, expected", - [ - ("thinking", None, "thinking"), # reasoning field - (None, "secondary", "secondary"), # reasoning_content fallback - ("primary", "secondary", "primary"), # priority - ], - ids=["reasoning-field", "reasoning_content-fallback", "priority"], - ) - async def test_non_stream_logprobs_reasoning_extraction( - self, model, reasoning, reasoning_content, expected - ): - response = _make_non_stream_response_from_choices( - choices=[ - _make_non_stream_choice( - text="answer", - reasoning=reasoning, - reasoning_content=reasoning_content, - logprob_items=[("A", -0.1)], - ) - ], - usage=_make_usage(6, 2), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == expected - - -# =================================================================== -# _alogprobs_impl — logprobs -# =================================================================== -class TestALogprobs: - def _make_logprobs_chunk(self, token, logprob, index=0, finish_reason=""): - chunk = MagicMock() - chunk.usage = None - - choice = MagicMock() - choice.index = index - choice.finish_reason = finish_reason - delta = MagicMock() - delta.content = token - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - - lp_item = MagicMock() - lp_item.token = token - lp_item.logprob = logprob - logprobs_obj = MagicMock() - logprobs_obj.content = [lp_item] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - return chunk - - def _patch_logprobs_create(self, model: _TestChatModel, chunks): - mock_create = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.chat.completions - target.create = mock_create # type: ignore[invalid-assignment] - return mock_create - - @pytest.mark.anyio - async def test_n_gt_1_raises(self, model): - mock = self._patch_logprobs_create( - model, [self._make_logprobs_chunk("A", -0.1, finish_reason="stop")] - ) - with pytest.raises(ValueError, match="only supports n=1"): - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5, n=2) - mock.assert_not_called() - - @pytest.mark.anyio - async def test_logprobs_extracted(self, model): - chunks = [ - self._make_logprobs_chunk("A", -0.1, finish_reason="stop"), - self._make_logprobs_chunk("B", -0.5), - ] - self._patch_logprobs_create(model, chunks) - out = await model._alogprobs_impl("prompt", max_tokens=2, logprobs=5) - assert out.logprobs_tokens is not None - assert "A" in out.logprobs_tokens - assert out.logprobs is not None - assert -0.1 in out.logprobs - - @pytest.mark.anyio - async def test_no_logprobs_raises(self, model): - # Chunk with no logprobs attribute - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "x" - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - choice.logprobs = None - chunk.choices = [choice] - - self._patch_logprobs_create(model, [chunk]) - with pytest.raises(RuntimeError, match="Streaming logprobs not supported"): - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - - @pytest.mark.anyio - async def test_logprobs_usage_captured(self, model): - lp_chunk = self._make_logprobs_chunk("A", -0.1, finish_reason="stop") - usage_chunk = _make_usage_chunk(5, 1) - self._patch_logprobs_create(model, [lp_chunk, usage_chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.usage is not None - assert out.usage["input_tokens"] == 5 - - @pytest.mark.anyio - async def test_logprobs_reasoning_content_accumulated(self, model): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "answer" - delta.reasoning = None - delta.reasoning_content = "think" - choice.delta = delta - lp_item = MagicMock() - lp_item.token = "A" - lp_item.logprob = -0.1 - logprobs_obj = MagicMock() - logprobs_obj.content = [lp_item] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - self._patch_logprobs_create(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.reasoning_texts is not None - assert out.reasoning_texts[0] == "think" - - @pytest.mark.anyio - async def test_logprobs_non_stream_supported(self, model): - response = _make_non_stream_response_from_choices( - choices=[ - _make_non_stream_choice( - text="answer", - logprob_items=[("A", -0.1), ("B", -0.2)], - ) - ], - usage=_make_usage(6, 2), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._alogprobs_impl( - "prompt", - max_tokens=1, - logprobs=5, - stream=False, - ) - assert out.texts == ["answer"] - assert out.logprobs_tokens == ["A", "B"] - assert out.logprobs == [-0.1, -0.2] - assert out.usage is not None - assert out.usage["input_tokens"] == 6 - - @pytest.mark.anyio - async def test_logprobs_non_stream_missing_logprobs_raises(self, model): - response = _make_non_stream_response_from_choices( - choices=[_make_non_stream_choice(text="answer", logprob_items=None)], - usage=_make_usage(1, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - with pytest.raises(RuntimeError, match="Streaming logprobs not supported"): - await model._alogprobs_impl( - "prompt", - max_tokens=1, - logprobs=5, - stream=False, - ) - - @pytest.mark.anyio - async def test_logprobs_stream_out_of_range_choice_index_skipped(self, model): - """Out-of-range choice.index should be ignored when n=1.""" - oob_chunk = self._make_logprobs_chunk("BAD", -9.0, index=5) - valid_chunk = self._make_logprobs_chunk( - "A", -0.1, index=0, finish_reason="stop" - ) - self._patch_logprobs_create(model, [oob_chunk, valid_chunk]) - - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - - assert out.texts == ["A"] - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [-0.1] - - @pytest.mark.anyio - async def test_logprobs_non_stream_out_of_range_choice_index_skipped(self, model): - """Non-stream mode should also ignore out-of-range choice.index.""" - response = _make_non_stream_response_from_choices( - choices=[ - _make_non_stream_choice( - index=5, - text="BAD", - logprob_items=[("BAD", -9.0)], - ), - _make_non_stream_choice( - index=0, - text="ok", - logprob_items=[("A", -0.1)], - ), - ], - usage=_make_usage(2, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - - assert out.texts == ["ok"] - assert "BAD" not in out.texts[0] - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [-0.1] - - -# =================================================================== -# Response metadata capture (response_model, system_fingerprint) -# =================================================================== -class TestResponseMetadata: - @pytest.mark.anyio - async def test_agenerate_captures_response_model_streaming(self, model): - """Streaming: response_model captured from first chunk.""" - chunk = _make_chunk(content="ok", finish_reason="stop") - chunk.model = "actual-model-v2" - chunk.system_fingerprint = "fp_abc123" - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.response_model == "actual-model-v2" - assert out.system_fingerprint == "fp_abc123" - - @pytest.mark.anyio - async def test_agenerate_captures_response_model_non_stream(self, model): - """Non-streaming: response_model captured from response object.""" - response = _make_non_stream_response( - text="ok", finish_reason="stop", usage=_make_usage(3, 1) - ) - response.model = "actual-model-v2" - response.system_fingerprint = "fp_xyz789" - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - out = await model._agenerate_impl("prompt", stream=False) - assert out.response_model == "actual-model-v2" - assert out.system_fingerprint == "fp_xyz789" - - @pytest.mark.anyio - async def test_agenerate_streaming_captures_from_first_chunk_only(self, model): - """Streaming: only the first chunk's metadata is used.""" - chunk1 = _make_chunk(content="a", finish_reason="") - chunk1.model = "model-v1" - chunk1.system_fingerprint = "fp_first" - chunk2 = _make_chunk(content="b", finish_reason="stop") - chunk2.model = "model-v2" - chunk2.system_fingerprint = "fp_second" - _patch_create(model, [chunk1, chunk2]) - out = await model._agenerate_impl("prompt") - assert out.response_model == "model-v1" - assert out.system_fingerprint == "fp_first" - - @pytest.mark.anyio - async def test_agenerate_streaming_none_metadata(self, model): - """Streaming: missing model/fingerprint attrs → None.""" - chunk = _make_chunk(content="ok", finish_reason="stop") - del chunk.model - del chunk.system_fingerprint - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.response_model is None - assert out.system_fingerprint is None - - @pytest.mark.anyio - async def test_alogprobs_captures_response_model_streaming(self, model): - """Logprobs streaming: response_model captured from first chunk.""" - chunk = MagicMock() - chunk.usage = None - chunk.model = "logprobs-model-v1" - chunk.system_fingerprint = "fp_lp_stream" - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "A" - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - lp_item = MagicMock() - lp_item.token = "A" - lp_item.logprob = -0.1 - logprobs_obj = MagicMock() - logprobs_obj.content = [lp_item] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - _patch_create(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.response_model == "logprobs-model-v1" - assert out.system_fingerprint == "fp_lp_stream" - - @pytest.mark.anyio - async def test_alogprobs_captures_response_model_non_stream(self, model): - """Logprobs non-streaming: response_model captured from response object.""" - response = _make_non_stream_response_from_choices( - choices=[ - _make_non_stream_choice( - text="A", - logprob_items=[("A", -0.2)], - ) - ], - usage=_make_usage(4, 1), - ) - response.model = "logprobs-model-v2" - response.system_fingerprint = "fp_lp_nonstream" - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.response_model == "logprobs-model-v2" - assert out.system_fingerprint == "fp_lp_nonstream" - - -# =================================================================== -# _agenerate_impl — invalid prompt type -# =================================================================== -class TestAGenerateInvalidPrompt: - @pytest.mark.anyio - async def test_non_string_non_iterable_raises(self, model): - with pytest.raises(TypeError, match="string or iterable"): - await model._agenerate_impl(12345) - - -class TestParamValidation: - @pytest.mark.anyio - async def test_invalid_n_type_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._agenerate_impl("prompt", n="1") - - @pytest.mark.anyio - async def test_bool_n_rejected(self, model): - """bool is a subclass of int in Python; the guard must catch it.""" - with pytest.raises(TypeError, match="n must be an int"): - await model._agenerate_impl("prompt", n=True) - - @pytest.mark.anyio - async def test_invalid_n_value_raises(self, model): - with pytest.raises(ValueError, match="n must be >= 1"): - await model._agenerate_impl("prompt", n=0) - - @pytest.mark.anyio - async def test_stream_non_bool_raises(self, model): - with pytest.raises(TypeError, match="stream must be a bool"): - await model._agenerate_impl("prompt", stream="false") - - @pytest.mark.anyio - async def test_logprobs_bool_n_rejected(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n=True) - - @pytest.mark.anyio - async def test_logprobs_invalid_n_type_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n="1") - - @pytest.mark.anyio - async def test_logprobs_invalid_n_value_raises(self, model): - with pytest.raises(ValueError, match="n must be >= 1"): - await model._alogprobs_impl("prompt", n=0) - - @pytest.mark.anyio - async def test_logprobs_stream_non_bool_raises(self, model): - with pytest.raises(TypeError, match="stream must be a bool"): - await model._alogprobs_impl("prompt", stream="false") - - -# =================================================================== -# Branch coverage: null/empty edge cases -# =================================================================== -class TestAGenerateNullBranches: - """Cover streaming/non-stream branches where delta, content, or usage is None.""" - - @pytest.mark.anyio - async def test_stream_delta_none(self, model): - """choice.delta is None → no content/reasoning accumulated.""" - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - choice.delta = None - chunk.choices = [choice] - - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.texts == [""] - - @pytest.mark.anyio - async def test_stream_delta_content_none(self, model): - """choice.delta.content is None → text stays empty.""" - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = None - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - chunk.choices = [choice] - - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.texts == [""] - - @pytest.mark.anyio - async def test_non_stream_content_none(self, model): - """Non-stream message.content is None → text stays empty.""" - response = _make_non_stream_response(text="") - response.choices[0].message.content = None - response.usage = _make_usage(3, 1) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.texts == [""] - - @pytest.mark.anyio - async def test_non_stream_no_usage(self, model): - """Non-stream resp.usage is None → usage stays None.""" - response = _make_non_stream_response(text="ok", usage=None) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.usage is None - - -class TestALogprobsNullBranches: - """Cover logprobs-mode branches where delta, content, or usage is None.""" - - def _make_lp_chunk(self, *, delta=True, content="A", logprob=-0.1): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - if delta: - d = MagicMock() - d.content = content - d.reasoning = None - d.reasoning_content = None - choice.delta = d - else: - choice.delta = None - lp_item = MagicMock() - lp_item.token = "A" - lp_item.logprob = logprob - logprobs_obj = MagicMock() - logprobs_obj.content = [lp_item] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - return chunk - - def _patch(self, model, chunks): - mock = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.chat.completions - target.create = mock - return mock - - @pytest.mark.anyio - async def test_stream_delta_none(self, model): - """Logprobs stream: delta=None → no text, but logprobs still collected.""" - chunk = self._make_lp_chunk(delta=False) - self._patch(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.texts == [""] - assert out.logprobs_tokens == ["A"] - - @pytest.mark.anyio - async def test_stream_delta_content_none(self, model): - """Logprobs stream: delta.content=None → text empty.""" - chunk = self._make_lp_chunk(content=None) - self._patch(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.texts == [""] - assert out.logprobs_tokens == ["A"] - - @pytest.mark.anyio - async def test_stream_empty_logprobs_content(self, model): - """Logprobs stream: logprobs_obj.content is empty list → no tokens.""" - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - delta = MagicMock() - delta.content = "x" - delta.reasoning = None - delta.reasoning_content = None - choice.delta = delta - logprobs_obj = MagicMock() - logprobs_obj.content = [] - choice.logprobs = logprobs_obj - chunk.choices = [choice] - - # Need a second chunk with actual logprobs so saw_logprobs=True - chunk2 = self._make_lp_chunk() - self._patch(model, [chunk, chunk2]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.logprobs_tokens == ["A"] - - @pytest.mark.anyio - async def test_non_stream_content_none(self, model): - """Logprobs non-stream: message.content=None → text empty.""" - choice = _make_non_stream_choice(text="", logprob_items=[("A", -0.1)]) - choice.message.content = None - response = _make_non_stream_response_from_choices( - choices=[choice], usage=_make_usage(3, 1) - ) - mock = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.texts == [""] - assert out.logprobs_tokens == ["A"] - - @pytest.mark.anyio - async def test_non_stream_empty_logprobs_content(self, model): - """Logprobs non-stream: logprobs_obj.content is empty → no tokens.""" - choice = MagicMock() - choice.index = 0 - choice.finish_reason = "stop" - message = MagicMock() - message.content = "x" - message.reasoning = None - message.reasoning_content = None - choice.message = message - logprobs_obj = MagicMock() - logprobs_obj.content = [] - choice.logprobs = logprobs_obj - - response = _make_non_stream_response_from_choices( - choices=[choice], usage=_make_usage(3, 1) - ) - mock = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock - - # saw_logprobs=True (logprobs_obj is not None) but content empty - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.logprobs_tokens == [] - - @pytest.mark.anyio - async def test_non_stream_no_usage(self, model): - """Logprobs non-stream: resp.usage=None → usage stays None.""" - response = _make_non_stream_response_from_choices( - choices=[_make_non_stream_choice(text="ok", logprob_items=[("A", -0.1)])], - usage=None, - ) - mock = AsyncMock(return_value=response) - target: Any = model._client.chat.completions - target.create = mock +class TestDefaultTransport: + def test_builds_openai_chat_transport(self): + m = ChatModel(model="m", api_key="k") + assert isinstance(m._transport, OpenAIChatTransport) + assert Capability.Chat in m.capabilities - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.usage is None + def test_transport_bound_to_shared_client_and_model(self): + m = ChatModel(model="m", api_key="k") + assert m._transport._client is m._client + assert m._transport._model == "m" diff --git a/tests/unit/core/models/test_gen_model.py b/tests/unit/core/models/test_gen_model.py index 7fde9475..b3604580 100644 --- a/tests/unit/core/models/test_gen_model.py +++ b/tests/unit/core/models/test_gen_model.py @@ -1,765 +1,30 @@ +"""Shell tests: backend selector wiring for GenModel. + +RFC #25 moved the completions wire logic (streaming accumulation, echo split, +``_completion_top_logprobs`` sanitizing) into ``OpenAICompletionsTransport``; +the former ``_agenerate_impl`` / ``_alogprobs_impl`` coverage moved with it to +tests/unit/core/models/transports/test_openai_completions.py, and the +request-builder validation (n/stream types, alogprobs n=1) lives on ``Model`` +in tests/unit/core/models/test_model.py. What remains here is the selector +contract: GenModel pairs the shared client with that transport, which supplies +the model's capabilities. + +AI-Generated Code - Claude Fable 5 (Anthropic) """ -Unit tests for sieval/core/models/gen_model.py. -Covers: _agenerate_impl (string prompt, non-string raises, streaming -accumulation, n>1 choices, usage), _alogprobs_impl (streaming logprobs, -no-logprobs raises, usage). +from sieval.core.models import Capability, GenModel +from sieval.core.models.transports.openai_completions import ( + OpenAICompletionsTransport, +) -All OpenAI client calls are mocked — no real API traffic. -AI-Generated Code - Claude Opus 4.6 (Anthropic) -""" - -from typing import Any -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from sieval.core.models.gen_model import GenModel, _completion_top_logprobs -from sieval.core.models.model import ModelOutput - - -# --------------------------------------------------------------------------- -# Async streaming helpers -# --------------------------------------------------------------------------- -class _AsyncIterator: - def __init__(self, items): - self._items = iter(items) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self._items) - except StopIteration as e: - raise StopAsyncIteration from e - - -def _make_completion_chunk(index=0, text="", finish_reason=""): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = index - choice.text = text - choice.finish_reason = finish_reason - chunk.choices = [choice] - return chunk - - -def _make_usage_chunk(prompt_tokens=10, completion_tokens=5): - chunk = MagicMock() - chunk.choices = [] - chunk.usage = MagicMock() - chunk.usage.prompt_tokens = prompt_tokens - chunk.usage.completion_tokens = completion_tokens - chunk.usage.total_tokens = prompt_tokens + completion_tokens - return chunk - - -def _make_usage(prompt_tokens=10, completion_tokens=5): - usage = MagicMock() - usage.prompt_tokens = prompt_tokens - usage.completion_tokens = completion_tokens - usage.total_tokens = prompt_tokens + completion_tokens - return usage - - -def _make_non_stream_response( - *, - text: str | None = "", - finish_reason: str | None = "stop", - usage=None, - with_logprobs: bool = False, - logprob: float = -0.1, -): - resp = MagicMock() - choice = MagicMock() - choice.index = 0 - choice.text = text - choice.finish_reason = finish_reason - if with_logprobs: - lp_obj = MagicMock() - lp_obj.tokens = [text] - lp_obj.token_logprobs = [logprob] - choice.logprobs = lp_obj - else: - choice.logprobs = None - resp.choices = [choice] - resp.usage = usage - return resp - - -def _make_non_stream_choice( - *, - index: int, - text: str, - token_logprobs: list[float] | None, - tokens: list[str] | None = None, -): - choice = MagicMock() - choice.index = index - choice.text = text - choice.finish_reason = "stop" - if token_logprobs is None: - choice.logprobs = None - else: - lp_obj = MagicMock() - lp_obj.tokens = tokens or [text] - lp_obj.token_logprobs = token_logprobs - choice.logprobs = lp_obj - return choice - - -# --------------------------------------------------------------------------- -# Concrete GenModel (delegates to parent _agenerate_impl) -# --------------------------------------------------------------------------- -@pytest.fixture -def model(): - return GenModel(model="test-gen", api_key="fake") - - -def _patch_create(model: GenModel, chunks): - mock_create = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.completions - target.create = mock_create # type: ignore[invalid-assignment] - return mock_create - - -# =================================================================== -# _agenerate_impl -# =================================================================== -class TestGenAGenerate: - @pytest.mark.anyio - async def test_basic_string_prompt(self, model): - chunks = [ - _make_completion_chunk(text="Hello"), - _make_completion_chunk(text=" world", finish_reason="stop"), - _make_usage_chunk(5, 2), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("What is 1+1?") - assert isinstance(out, ModelOutput) - assert out.texts == ["Hello world"] - assert out.finish_reasons == ["stop"] - - @pytest.mark.anyio - async def test_non_string_prompt_raises(self, model): - with pytest.raises(TypeError, match="GenModel requires a string"): - await model._agenerate_impl(["not", "a", "string"]) - - @pytest.mark.anyio - async def test_usage_captured(self, model): - chunks = [ - _make_completion_chunk(text="ok", finish_reason="stop"), - _make_usage_chunk(7, 3), - ] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.usage is not None - assert out.usage["input_tokens"] == 7 - assert out.usage["output_tokens"] == 3 - assert out.usage["total_tokens"] == 10 - - @pytest.mark.anyio - async def test_no_usage_chunk(self, model): - chunks = [_make_completion_chunk(text="hi", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.usage is None - - @pytest.mark.anyio - async def test_model_meta_attached(self, model): - chunks = [_make_completion_chunk(text="x", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt") - assert out.model["model"] == "test-gen" - - @pytest.mark.anyio - async def test_request_params_captured(self, model): - chunks = [_make_completion_chunk(text="x", finish_reason="stop")] - _patch_create(model, chunks) - out = await model._agenerate_impl("prompt", temperature=0.5) - assert out.request_params is not None - assert out.request_params.get("temperature") == 0.5 - - @pytest.mark.anyio - async def test_two_choices_in_single_chunk(self, model): - chunk1 = MagicMock() - chunk1.usage = None - chunk1.choices = [ - _make_completion_chunk(index=0, text="A1").choices[0], - _make_completion_chunk(index=1, text="B1").choices[0], - ] - - chunk2 = MagicMock() - chunk2.usage = None - chunk2.choices = [ - _make_completion_chunk(index=0, text="A2", finish_reason="stop").choices[0], - _make_completion_chunk(index=1, text="B2", finish_reason="stop").choices[0], - ] - - _patch_create(model, [chunk1, chunk2]) - out = await model._agenerate_impl("prompt", n=2) - assert out.texts[0] == "A1A2" - assert out.texts[1] == "B1B2" - - @pytest.mark.anyio - async def test_non_stream_mode_supported(self, model): - response = _make_non_stream_response( - text="done", - finish_reason="stop", - usage=_make_usage(6, 2), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - - assert out.texts == ["done"] - assert out.finish_reasons == ["stop"] - assert out.usage is not None - assert out.usage["input_tokens"] == 6 - call_kwargs = mock_create.call_args[1] - assert call_kwargs["stream"] is False - - @pytest.mark.anyio - async def test_non_stream_missing_finish_reason_defaults_to_empty(self, model): - response = _make_non_stream_response( - text="done", - finish_reason=None, - usage=_make_usage(3, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.texts == ["done"] - assert out.finish_reasons == [""] - - @pytest.mark.anyio - async def test_non_stream_missing_text_defaults_to_empty(self, model): - response = _make_non_stream_response( - text=None, - finish_reason="stop", - usage=_make_usage(3, 1), - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", stream=False) - assert out.texts == [""] - assert out.finish_reasons == ["stop"] - - @pytest.mark.anyio - async def test_stream_options_can_be_overridden(self, model): - chunks = [_make_completion_chunk(text="x", finish_reason="stop")] - mock_create = _patch_create(model, chunks) - out = await model._agenerate_impl( - "prompt", - stream_options={"include_usage": False}, - ) - assert out.texts == ["x"] - call_kwargs = mock_create.call_args[1] - assert call_kwargs["stream_options"] == {"include_usage": False} - - @pytest.mark.anyio - async def test_stream_out_of_bounds_choice_index_skipped(self, model): - """Streaming chunks whose choice.index >= n are silently skipped.""" - # n=1 but chunk has choice.index=5 (out of bounds) - valid_chunk = _make_completion_chunk(index=0, text="ok", finish_reason="stop") - oob_chunk = MagicMock() - oob_chunk.usage = None - oob_choice = MagicMock() - oob_choice.index = 5 # out of bounds for n=1 - oob_choice.text = "SHOULD_NOT_APPEAR" - oob_choice.finish_reason = "" - oob_chunk.choices = [oob_choice] - - _patch_create(model, [valid_chunk, oob_chunk]) - out = await model._agenerate_impl("prompt", n=1) - assert out.texts == ["ok"] - assert "SHOULD_NOT_APPEAR" not in out.texts[0] - - @pytest.mark.anyio - async def test_non_stream_out_of_bounds_choice_index_skipped(self, model): - """Non-streaming responses whose choice.index >= n are silently skipped.""" - resp = MagicMock() - valid_choice = MagicMock() - valid_choice.index = 0 - valid_choice.text = "ok" - valid_choice.finish_reason = "stop" - valid_choice.logprobs = None - - oob_choice = MagicMock() - oob_choice.index = 3 # out of bounds for n=1 - oob_choice.text = "SHOULD_NOT_APPEAR" - oob_choice.finish_reason = "stop" - oob_choice.logprobs = None - - resp.choices = [valid_choice, oob_choice] - resp.usage = _make_usage(5, 2) - - mock_create = AsyncMock(return_value=resp) - target = model._client.completions - target.create = mock_create - - out = await model._agenerate_impl("prompt", n=1, stream=False) - assert out.texts == ["ok"] - assert "SHOULD_NOT_APPEAR" not in out.texts[0] - - -# =================================================================== -# _alogprobs_impl -# =================================================================== -class TestGenALogprobs: - def _make_logprobs_chunk( - self, token, logprob, index=0, finish_reason="", top_logprobs=None - ): - chunk = MagicMock() - chunk.usage = None - - choice = MagicMock() - choice.index = index - choice.text = token - choice.finish_reason = finish_reason - - lp_obj = MagicMock() - lp_obj.tokens = [token] - lp_obj.token_logprobs = [logprob] - lp_obj.top_logprobs = top_logprobs - choice.logprobs = lp_obj - chunk.choices = [choice] - return chunk - - def _patch_logprobs_create(self, model: GenModel, chunks): - mock_create = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.completions - target.create = mock_create # type: ignore[invalid-assignment] - return mock_create - - @pytest.mark.anyio - async def test_n_gt_1_raises(self, model): - mock = self._patch_logprobs_create( - model, [self._make_logprobs_chunk("A", -0.1, finish_reason="stop")] - ) - with pytest.raises(ValueError, match="only supports n=1"): - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5, n=2) - mock.assert_not_called() - - @pytest.mark.anyio - async def test_logprobs_extracted(self, model): - chunks = [ - self._make_logprobs_chunk( - "A", - -0.1, - finish_reason="stop", - top_logprobs=[{"A": -0.1, "B": -0.5}], - ), - self._make_logprobs_chunk("B", -0.5), - ] - self._patch_logprobs_create(model, chunks) - out = await model._alogprobs_impl("prompt", max_tokens=2, logprobs=5) - assert out.logprobs_tokens is not None - assert "A" in out.logprobs_tokens - assert out.logprobs is not None - assert -0.1 in out.logprobs - assert out.top_logprobs == [{"A": -0.1, "B": -0.5}] - - @pytest.mark.anyio - async def test_no_logprobs_raises(self, model): - # Chunk with logprobs=None - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.text = "x" - choice.finish_reason = "stop" - choice.logprobs = None - chunk.choices = [choice] - - self._patch_logprobs_create(model, [chunk]) - with pytest.raises(RuntimeError, match="Streaming logprobs not supported"): - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - - @pytest.mark.anyio - async def test_logprobs_usage_captured(self, model): - lp_chunk = self._make_logprobs_chunk("A", -0.1, finish_reason="stop") - usage_chunk = _make_usage_chunk(4, 1) - self._patch_logprobs_create(model, [lp_chunk, usage_chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.usage is not None - assert out.usage["input_tokens"] == 4 - assert out.usage["total_tokens"] == 5 - - @pytest.mark.anyio - async def test_default_params_forwarded(self, model): - mock = self._patch_logprobs_create( - model, [self._make_logprobs_chunk("A", -0.1, finish_reason="stop")] - ) - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=3, echo=True) - call_kwargs = mock.call_args[1] - assert call_kwargs.get("logprobs") == 3 - assert call_kwargs.get("echo") is True - - @pytest.mark.anyio - async def test_non_stream_logprobs_supported(self, model): - response = _make_non_stream_response( - text="A", - finish_reason="stop", - usage=_make_usage(4, 1), - with_logprobs=True, - logprob=-0.2, - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - - out = await model._alogprobs_impl( - "prompt", - max_tokens=1, - logprobs=5, - stream=False, - ) - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [-0.2] - call_kwargs = mock_create.call_args[1] - assert call_kwargs["stream"] is False - - @pytest.mark.anyio - async def test_non_stream_missing_logprobs_raises(self, model): - response = _make_non_stream_response( - text="A", - finish_reason="stop", - usage=_make_usage(4, 1), - with_logprobs=False, - ) - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - - with pytest.raises(RuntimeError, match="Streaming logprobs not supported"): - await model._alogprobs_impl( - "prompt", - max_tokens=1, - logprobs=5, - stream=False, - ) - - @pytest.mark.anyio - async def test_stream_out_of_range_choice_index_skipped(self, model): - """Out-of-range choice.index should be ignored when n=1.""" - oob_chunk = self._make_logprobs_chunk("BAD", -9.0, index=5) - valid_chunk = self._make_logprobs_chunk( - "A", -0.1, index=0, finish_reason="stop" - ) - self._patch_logprobs_create(model, [oob_chunk, valid_chunk]) - - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - - assert out.texts == ["A"] - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [-0.1] - - @pytest.mark.anyio - async def test_non_stream_out_of_range_choice_index_skipped(self, model): - """Non-stream mode should also ignore out-of-range choice.index.""" - resp = MagicMock() - resp.choices = [ - _make_non_stream_choice(index=5, text="BAD", token_logprobs=[-9.0]), - _make_non_stream_choice( - index=0, - text="ok", - tokens=["A"], - token_logprobs=[-0.1], - ), - ] - resp.usage = _make_usage(2, 1) - mock_create = AsyncMock(return_value=resp) - target: Any = model._client.completions - target.create = mock_create - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - - assert out.texts == ["ok"] - assert "BAD" not in out.texts[0] - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [-0.1] - - -# =================================================================== -# Response metadata capture (response_model, system_fingerprint) -# =================================================================== -class TestResponseMetadata: - @pytest.mark.anyio - async def test_agenerate_captures_response_model_streaming(self, model): - """Streaming: response_model captured from first chunk.""" - chunk = _make_completion_chunk(text="ok", finish_reason="stop") - chunk.model = "actual-model-v2" - chunk.system_fingerprint = "fp_abc123" - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.response_model == "actual-model-v2" - assert out.system_fingerprint == "fp_abc123" - - @pytest.mark.anyio - async def test_agenerate_captures_response_model_non_stream(self, model): - """Non-streaming: response_model captured from response object.""" - response = _make_non_stream_response( - text="ok", finish_reason="stop", usage=_make_usage(3, 1) - ) - response.model = "actual-model-v2" - response.system_fingerprint = "fp_xyz789" - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - out = await model._agenerate_impl("prompt", stream=False) - assert out.response_model == "actual-model-v2" - assert out.system_fingerprint == "fp_xyz789" - - @pytest.mark.anyio - async def test_agenerate_streaming_captures_from_first_chunk_only(self, model): - """Streaming: only the first chunk's metadata is used.""" - chunk1 = _make_completion_chunk(text="a", finish_reason="") - chunk1.model = "model-v1" - chunk1.system_fingerprint = "fp_first" - chunk2 = _make_completion_chunk(text="b", finish_reason="stop") - chunk2.model = "model-v2" - chunk2.system_fingerprint = "fp_second" - _patch_create(model, [chunk1, chunk2]) - out = await model._agenerate_impl("prompt") - assert out.response_model == "model-v1" - assert out.system_fingerprint == "fp_first" - - @pytest.mark.anyio - async def test_agenerate_streaming_none_metadata(self, model): - """Streaming: missing model/fingerprint attrs → None.""" - chunk = _make_completion_chunk(text="ok", finish_reason="stop") - # Explicitly delete auto-created MagicMock attrs - del chunk.model - del chunk.system_fingerprint - _patch_create(model, [chunk]) - out = await model._agenerate_impl("prompt") - assert out.response_model is None - assert out.system_fingerprint is None - - @pytest.mark.anyio - async def test_alogprobs_captures_response_model_streaming(self, model): - """Logprobs streaming: response_model captured from first chunk.""" - chunk = MagicMock() - chunk.usage = None - chunk.model = "logprobs-model-v1" - chunk.system_fingerprint = "fp_lp_stream" - choice = MagicMock() - choice.index = 0 - choice.text = "A" - choice.finish_reason = "stop" - lp_obj = MagicMock() - lp_obj.tokens = ["A"] - lp_obj.token_logprobs = [-0.1] - choice.logprobs = lp_obj - chunk.choices = [choice] - _patch_create(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.response_model == "logprobs-model-v1" - assert out.system_fingerprint == "fp_lp_stream" - - @pytest.mark.anyio - async def test_alogprobs_captures_response_model_non_stream(self, model): - """Logprobs non-streaming: response_model captured from response object.""" - response = _make_non_stream_response( - text="A", - finish_reason="stop", - usage=_make_usage(4, 1), - with_logprobs=True, - logprob=-0.2, - ) - response.model = "logprobs-model-v2" - response.system_fingerprint = "fp_lp_nonstream" - mock_create = AsyncMock(return_value=response) - target: Any = model._client.completions - target.create = mock_create - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.response_model == "logprobs-model-v2" - assert out.system_fingerprint == "fp_lp_nonstream" - - -class TestParamValidation: - # Parameter validation logic (n type/value, stream type) is identical between - # ChatModel and GenModel. These tests verify GenModel's implementation. - - @pytest.mark.anyio - async def test_invalid_n_type_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._agenerate_impl("prompt", n="2") - - @pytest.mark.anyio - async def test_bool_n_rejected(self, model): - """bool is a subclass of int in Python; the guard must catch it.""" - with pytest.raises(TypeError, match="n must be an int"): - await model._agenerate_impl("prompt", n=True) - - @pytest.mark.anyio - async def test_invalid_n_value_raises(self, model): - with pytest.raises(ValueError, match="n must be >= 1"): - await model._agenerate_impl("prompt", n=0) - - @pytest.mark.anyio - async def test_stream_non_bool_raises(self, model): - with pytest.raises(TypeError, match="stream must be a bool"): - await model._agenerate_impl("prompt", stream="false") - - @pytest.mark.anyio - async def test_logprobs_bool_n_rejected(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n=True) - - @pytest.mark.anyio - async def test_logprobs_invalid_n_type_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n="2") - - @pytest.mark.anyio - async def test_logprobs_invalid_n_value_raises(self, model): - with pytest.raises(ValueError, match="n must be >= 1"): - await model._alogprobs_impl("prompt", n=0) - - @pytest.mark.anyio - async def test_logprobs_stream_non_bool_raises(self, model): - with pytest.raises(TypeError, match="stream must be a bool"): - await model._alogprobs_impl("prompt", stream="false") - - -# =================================================================== -# Branch coverage: null/empty edge cases -# =================================================================== -class TestGenLogprobsNullBranches: - """Cover logprobs branches where tokens/token_logprobs are empty or usage=None.""" - - def _make_lp_chunk(self, *, tokens=None, token_logprobs=None): - chunk = MagicMock() - chunk.usage = None - choice = MagicMock() - choice.index = 0 - choice.text = "x" - choice.finish_reason = "stop" - lp_obj = MagicMock() - lp_obj.tokens = tokens - lp_obj.token_logprobs = token_logprobs - choice.logprobs = lp_obj - chunk.choices = [choice] - return chunk - - def _patch(self, model, chunks): - mock = AsyncMock(return_value=_AsyncIterator(chunks)) - target: Any = model._client.completions - target.create = mock - return mock - - @pytest.mark.anyio - async def test_stream_empty_tokens(self, model): - """logprobs_obj.tokens is empty → no tokens collected.""" - chunk = self._make_lp_chunk(tokens=[], token_logprobs=[-0.1]) - self._patch(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.logprobs_tokens == [] - assert out.logprobs == [-0.1] - - @pytest.mark.anyio - async def test_stream_empty_token_logprobs(self, model): - """logprobs_obj.token_logprobs is empty → no logprobs collected.""" - chunk = self._make_lp_chunk(tokens=["A"], token_logprobs=[]) - self._patch(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [] - - @pytest.mark.anyio - async def test_stream_no_usage(self, model): - """No usage chunk → usage stays None.""" - chunk = self._make_lp_chunk(tokens=["A"], token_logprobs=[-0.1]) - self._patch(model, [chunk]) - out = await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5) - assert out.usage is None - - @pytest.mark.anyio - async def test_non_stream_empty_tokens(self, model): - """Non-stream: logprobs_obj.tokens is empty.""" - resp = _make_non_stream_response( - text="x", - finish_reason="stop", - usage=_make_usage(3, 1), - with_logprobs=True, - ) - resp.choices[0].logprobs.tokens = [] - resp.choices[0].logprobs.token_logprobs = [-0.1] - mock = AsyncMock(return_value=resp) - target: Any = model._client.completions - target.create = mock - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.logprobs_tokens == [] - assert out.logprobs == [-0.1] - - @pytest.mark.anyio - async def test_non_stream_empty_token_logprobs(self, model): - """Non-stream: logprobs_obj.token_logprobs is empty.""" - resp = _make_non_stream_response( - text="x", - finish_reason="stop", - usage=_make_usage(3, 1), - with_logprobs=True, - ) - resp.choices[0].logprobs.tokens = ["A"] - resp.choices[0].logprobs.token_logprobs = [] - mock = AsyncMock(return_value=resp) - target: Any = model._client.completions - target.create = mock - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.logprobs_tokens == ["A"] - assert out.logprobs == [] - - @pytest.mark.anyio - async def test_non_stream_no_usage(self, model): - """Non-stream: resp.usage=None → usage stays None.""" - resp = _make_non_stream_response( - text="x", - finish_reason="stop", - usage=None, - with_logprobs=True, - ) - mock = AsyncMock(return_value=resp) - target: Any = model._client.completions - target.create = mock - - out = await model._alogprobs_impl( - "prompt", max_tokens=1, logprobs=5, stream=False - ) - assert out.usage is None - - -class TestCompletionTopLogprobs: - """Cover _completion_top_logprobs parsing and its defensive branches.""" - - def test_non_sequence_returns_empty(self): - assert _completion_top_logprobs(None) == [] - assert _completion_top_logprobs("AB") == [] +class TestDefaultTransport: + def test_builds_openai_completions_transport(self): + m = GenModel(model="m", api_key="k") + assert isinstance(m._transport, OpenAICompletionsTransport) + assert Capability.InputScoring in m.capabilities - def test_parses_and_skips_malformed_entries(self): - # Mix of: a valid dict; a None entry (→ {}); a non-Mapping entry - # (skipped); and a dict whose non-numeric value ("no") and non-str key - # (2) are both filtered out, leaving only the valid pair. - raw = [{"A": -0.1}, None, "x", {"y": -0.5, "1": "no", 2: -0.3}] - assert _completion_top_logprobs(raw) == [{"A": -0.1}, {}, {"y": -0.5}] + def test_transport_bound_to_shared_client_and_model(self): + m = GenModel(model="m", api_key="k") + assert m._transport._client is m._client + assert m._transport._model == "m" diff --git a/tests/unit/core/models/test_ir.py b/tests/unit/core/models/test_ir.py new file mode 100644 index 00000000..6711ad5d --- /dev/null +++ b/tests/unit/core/models/test_ir.py @@ -0,0 +1,216 @@ +"""Unit tests for the provider-agnostic Model IR. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +import dataclasses + +import pytest + +from sieval.core.models.ir import ( + Citation, + GroundingChunk, + GroundingMetadata, + InputScoringResult, + ReasoningOutput, + ReasoningParams, + Request, + Response, + SamplingParams, + ServerToolSpec, + ServerToolUse, + TokenLogprob, + TopKEntry, + UsageStats, +) +from sieval.core.utils.serialization import ( + dict_to_obj, + global_type_registry, + obj_to_dict, +) + + +class TestRequestConstruction: + def test_minimal_completion_request(self): + req = Request(input="hello") + assert req.input == "hello" + assert req.sampling is None + assert req.return_logprobs is False + assert req.top_k == 0 + assert req.score_input is False + assert req.session_id is None + assert req.extra_wire_params is None + + def test_chat_request_input_is_message_list(self): + req = Request(input=[{"role": "user", "content": "hi"}]) + assert isinstance(req.input, list) + assert req.input[0]["role"] == "user" + + def test_extra_wire_params_passthrough_field(self): + req = Request(input="x", extra_wire_params={"logit_bias": {"1": -100}}) + assert req.extra_wire_params == {"logit_bias": {"1": -100}} + + def test_stream_defaults_to_transport_choice(self): + # None → the transport picks (single-shot); pure scheduling knob. + assert Request(input="x").stream is None + assert Request(input="x", stream=True).stream is True + + def test_sampling_params_defaults(self): + sp = SamplingParams() + assert sp.temperature is None + assert sp.n == 1 + assert sp.stop is None + + def test_reasoning_params_opaque_roundtrip_is_str_or_none(self): + rp = ReasoningParams(opaque_roundtrip="sig-abc") + assert rp.opaque_roundtrip == "sig-abc" + assert ReasoningParams().opaque_roundtrip is None + + def test_server_tool_spec(self): + spec = ServerToolSpec(type="web_search", config={"max_uses": 3}) + req = Request(input="x", server_tools=(spec,)) + assert req.server_tools is not None + assert req.server_tools[0].type == "web_search" + + +class TestResponseConstruction: + def test_minimal_response(self): + resp = Response(texts=("hi",)) + assert resp.texts == ("hi",) + assert resp.reasoning is None + assert resp.logprobs is None + assert resp.session_id is None + assert resp.finish_reasons is None + + def test_provenance_fields_default_none(self): + resp = Response(texts=("hi",)) + assert resp.request_params is None + assert resp.response_model is None + assert resp.system_fingerprint is None + + def test_usage_defaults_to_none(self): + # Absence != zeros: a server that reported no usage must not be + # recorded as zero tokens (the bridge relies on this distinction). + resp = Response(texts=("hi",)) + assert resp.usage is None + + def test_usage_stats_fields_default_zero(self): + usage = UsageStats() + assert usage.input_tokens == 0 + assert usage.output_tokens == 0 + assert usage.reasoning_tokens == 0 + assert usage.cached_tokens == 0 + assert usage.total_tokens == 0 + + def test_grounding_metadata_rendered_content_optional(self): + gm = GroundingMetadata(chunks=(GroundingChunk(uri="http://x"),)) + assert gm.rendered_content is None + assert dataclasses.fields(GroundingMetadata) # field exists + gm2 = GroundingMetadata(chunks=(), rendered_content="
") + assert gm2.rendered_content == "
" + + +class TestImmutability: + @pytest.mark.parametrize( + ("obj", "field_name", "value"), + [ + (Request(input="x"), "input", "y"), + (Response(texts=("a",)), "texts", ("b",)), + (SamplingParams(), "temperature", 0.5), + (ReasoningParams(), "effort", "high"), + (UsageStats(), "input_tokens", 5), + (TokenLogprob(token_id=1, token="a", logprob=-0.1), "logprob", 0.0), + ], + ) + def test_frozen_assignment_raises(self, obj, field_name, value): + with pytest.raises(dataclasses.FrozenInstanceError): + setattr(obj, field_name, value) + + +class TestSerializationRoundTrip: + """Response is the persisted record schema; nested records must rehydrate + back into typed objects, not plain dicts.""" + + def test_response_round_trips_to_typed_object(self): + data = obj_to_dict(Response(texts=("hi", "there")), add_type=True) + back = dict_to_obj(data, global_type_registry) + assert isinstance(back, Response) + assert back.texts == ("hi", "there") + + def test_nested_reasoning_and_usage_rehydrate_as_records(self): + resp = Response( + texts=("answer",), + reasoning=ReasoningOutput(text="think", thinking_tokens=12), + usage=UsageStats(input_tokens=3, output_tokens=4, total_tokens=7), + ) + back = dict_to_obj(obj_to_dict(resp, add_type=True), global_type_registry) + assert isinstance(back.reasoning, ReasoningOutput) + assert back.reasoning.text == "think" + assert back.reasoning.thinking_tokens == 12 + assert isinstance(back.usage, UsageStats) + assert back.usage.total_tokens == 7 + + def test_tuple_of_token_logprobs_rehydrates_element_type(self): + # The gap-4 case: a tuple of records only round-trips when the element + # type is itself @sieval_record. + resp = Response( + texts=("x",), + logprobs=( + TokenLogprob(token_id=10, token=" A", logprob=-0.5), + TokenLogprob(token_id=None, token="B", logprob=-1.5), + ), + ) + back = dict_to_obj(obj_to_dict(resp, add_type=True), global_type_registry) + assert isinstance(back.logprobs, tuple) + assert all(isinstance(t, TokenLogprob) for t in back.logprobs) + assert back.logprobs[0].token_id == 10 + assert back.logprobs[0].token == " A" + assert back.logprobs[1].token_id is None + + def test_nested_tuple_of_topk_entries_rehydrates(self): + resp = Response( + texts=("x",), + top_logprobs=( + ( + TopKEntry(token_id=1, token="A", logprob=-0.1), + TopKEntry(token_id=2, token="B", logprob=-2.0), + ), + ), + ) + back = dict_to_obj(obj_to_dict(resp, add_type=True), global_type_registry) + assert isinstance(back.top_logprobs[0][0], TopKEntry) + assert back.top_logprobs[0][1].token == "B" + + def test_input_scoring_and_grounding_and_tool_use_rehydrate(self): + resp = Response( + texts=("x",), + input_scoring=InputScoringResult( + token_logprobs=(TokenLogprob(token_id=5, token="q", logprob=-0.2),), + byte_count=3, + char_count=1, + ), + citations=(Citation(url="http://a", title="A"),), + grounding=GroundingMetadata( + chunks=(GroundingChunk(uri="http://c", title="C"),), + rendered_content="", + ), + server_tool_uses=( + ServerToolUse( + tool_type="web_search", + tool_use_id="t1", + input={"q": "x"}, + result={"ok": True}, + ), + ), + ) + back = dict_to_obj(obj_to_dict(resp, add_type=True), global_type_registry) + assert isinstance(back.input_scoring, InputScoringResult) + assert back.input_scoring.byte_count == 3 + assert isinstance(back.input_scoring.token_logprobs[0], TokenLogprob) + assert isinstance(back.citations[0], Citation) + assert isinstance(back.grounding, GroundingMetadata) + # Google ToS: rendered_content must survive the round-trip. + assert back.grounding.rendered_content == "" + assert isinstance(back.grounding.chunks[0], GroundingChunk) + assert isinstance(back.server_tool_uses[0], ServerToolUse) + assert back.server_tool_uses[0].result == {"ok": True} diff --git a/tests/unit/core/models/test_model.py b/tests/unit/core/models/test_model.py index f2e8cb59..1a43e063 100644 --- a/tests/unit/core/models/test_model.py +++ b/tests/unit/core/models/test_model.py @@ -1,17 +1,27 @@ """ Focused tests for non-overlapping Model behaviors. -Most with_args/as_type/meta branches are covered in test_model_derivation.py. -This file keeps only unique checks plus runtime concurrency path tests. +with_args/meta branches are covered in test_model_derivation.py; the IR +primitive (arun/capabilities) in test_model_arun.py; wire lowering/lifting in +tests/unit/core/models/transports/. This file keeps unique checks: quota, +runtime concurrency paths, the legacy-kwargs request builders, and the +Response -> ModelOutput bridge. -AI-Generated Code - GPT-5.3-Codex (OpenAI) +AI-Generated Code - Claude Fable 5 (Anthropic) """ -from unittest.mock import AsyncMock - import pytest -from sieval.core.models import ChatModel, GenModel, ModelOutput +from sieval.core.models import ( + GenModel, + InputScoringResult, + ModelOutput, + ReasoningOutput, + Response, + TokenLogprob, + TopKEntry, + UsageStats, +) # --------------------------------------------------------------------------- @@ -42,10 +52,9 @@ def test_derived_overrides_kwargs(self, gen_model): # Parent unchanged assert gen_model._kwargs == {} - def test_as_type_preserves_kwargs(self): - model = GenModel(model="m", api_key="k", temperature=0.5) - chat = model.as_type(ChatModel) - assert chat._kwargs == {"temperature": 0.5} + def test_as_type_is_gone(self, gen_model): + """RFC #25 dropped cross-kind model conversion.""" + assert not hasattr(gen_model, "as_type") # =================================================================== @@ -112,7 +121,7 @@ def test_quota_info_no_limiter(self, unlimited_model): # =================================================================== # agenerate / alogprobs concurrency paths -# (covers model.py lines 206, 208-209, 223-255: parent_limiter branches) +# (covers the parent_limiter branches around arun) # =================================================================== def _build_chat_model_for_path(path): from tests.conftest import MockChatModel @@ -212,18 +221,267 @@ async def test_alogprobs_paths(self, path, prompt): assert result.logprobs_tokens is not None and len(result.logprobs_tokens) == 1 @pytest.mark.anyio - async def test_alogprobs_forwards_echo_flag(self, monkeypatch): + async def test_alogprobs_lowers_echo_to_score_input(self): + """alogprobs args land on the Request the transport receives.""" from tests.conftest import MockGenModel model = MockGenModel() - alogprobs_impl = AsyncMock(side_effect=model._alogprobs_impl) - monkeypatch.setattr(model, "_alogprobs_impl", alogprobs_impl) - await model.alogprobs("A", echo=False, max_tokens=2, logprobs=3) - assert alogprobs_impl.await_count == 1 - call = alogprobs_impl.await_args - assert call is not None - assert call.kwargs["echo"] is False - assert call.kwargs["max_tokens"] == 2 - assert call.kwargs["logprobs"] == 3 + req = model._transport.requests[0] + assert req.score_input is False + assert req.return_logprobs is True + assert req.top_k == 3 + assert req.sampling.max_tokens == 2 + assert req.sampling.temperature == 0.0 + + @pytest.mark.anyio + async def test_alogprobs_echo_true_sets_score_input(self): + from tests.conftest import MockGenModel + + model = MockGenModel() + await model.alogprobs("A", echo=True) + + req = model._transport.requests[0] + assert req.score_input is True + + +# =================================================================== +# Request builders (legacy OpenAI-style kwargs -> IR) +# =================================================================== +class TestBuildGenerateRequest: + def _model(self, **kwargs): + return GenModel(model="m", api_key="k", **kwargs) + + def test_sampling_kwargs_map_to_sampling_params(self): + req = self._model()._build_generate_request( + "p", + max_tokens=64, + temperature=0.7, + top_p=0.9, + seed=42, + frequency_penalty=0.1, + presence_penalty=0.2, + n=3, + ) + sp = req.sampling + assert sp.max_tokens == 64 + assert sp.temperature == 0.7 + assert sp.top_p == 0.9 + assert sp.seed == 42 + assert sp.frequency_penalty == 0.1 + assert sp.presence_penalty == 0.2 + assert sp.n == 3 + assert req.extra_wire_params is None + + def test_max_completion_tokens_aliases_max_tokens(self): + req = self._model()._build_generate_request("p", max_completion_tokens=32) + assert req.sampling.max_tokens == 32 + + def test_max_tokens_wins_over_alias(self): + req = self._model()._build_generate_request( + "p", max_tokens=8, max_completion_tokens=32 + ) + assert req.sampling.max_tokens == 8 + + def test_stop_string_becomes_tuple(self): + req = self._model()._build_generate_request("p", stop="\n\n") + assert req.sampling.stop == ("\n\n",) + + def test_stop_list_becomes_tuple(self): + req = self._model()._build_generate_request("p", stop=["\n\n", "Q:"]) + assert req.sampling.stop == ("\n\n", "Q:") + + def test_top_k_kwarg_is_sampling_top_k(self): + """`top_k` is the vLLM/sglang sampling knob, not the logprobs count.""" + req = self._model()._build_generate_request("p", top_k=40) + assert req.sampling.top_k_sampling == 40 + assert req.top_k == 0 + + def test_logprobs_bool_is_chat_switch(self): + req = self._model()._build_generate_request("p", logprobs=True) + assert req.return_logprobs is True + assert req.top_k == 0 + + def test_logprobs_int_is_completions_count(self): + req = self._model()._build_generate_request("p", logprobs=5) + assert req.return_logprobs is True + assert req.top_k == 5 + + def test_top_logprobs_sets_top_k(self): + req = self._model()._build_generate_request("p", logprobs=True, top_logprobs=7) + assert req.return_logprobs is True + assert req.top_k == 7 + + def test_model_kwargs_merge_with_call_kwargs(self): + model = self._model(temperature=0.5, max_tokens=10) + req = model._build_generate_request("p", max_tokens=99) + assert req.sampling.temperature == 0.5 + assert req.sampling.max_tokens == 99 # call kwargs win + + def test_unknown_kwargs_ride_in_extra_wire_params(self): + req = self._model()._build_generate_request("p", min_p=0.05, echo=True) + assert req.extra_wire_params == {"min_p": 0.05, "echo": True} + + def test_stream_defaults_true_and_is_poppable(self): + assert self._model()._build_generate_request("p").stream is True + assert self._model()._build_generate_request("p", stream=False).stream is False + + def test_reasoning_effort_maps_to_reasoning_params(self): + req = self._model()._build_generate_request("p", reasoning_effort="high") + assert req.reasoning is not None + assert req.reasoning.effort == "high" + + def test_response_format_and_tools(self): + fmt = {"type": "json_object"} + tools = [{"type": "function", "function": {"name": "f"}}] + req = self._model()._build_generate_request( + "p", response_format=fmt, tools=tools + ) + assert req.response_format == fmt + assert req.tools == tools + + def test_message_input_is_materialized(self): + from sieval.core.models import ChatModel + + model = ChatModel(model="m", api_key="k") + + def gen(): + yield {"role": "user", "content": "hi"} + + req = model._build_generate_request(gen()) + assert req.input == [{"role": "user", "content": "hi"}] + + +class TestBuilderValidation: + def _model(self): + return GenModel(model="m", api_key="k") + + def test_n_must_be_int(self): + with pytest.raises(TypeError, match="n must be an int"): + self._model()._build_generate_request("p", n="3") + + def test_n_bool_rejected(self): + with pytest.raises(TypeError, match="n must be an int"): + self._model()._build_generate_request("p", n=True) + + def test_n_below_one_rejected(self): + with pytest.raises(ValueError, match="n must be >= 1"): + self._model()._build_generate_request("p", n=0) + + def test_stream_must_be_bool(self): + with pytest.raises(TypeError, match="stream must be a bool"): + self._model()._build_generate_request("p", stream="yes") + + def test_non_iterable_prompt_rejected(self): + with pytest.raises(TypeError, match="prompt must be a string"): + self._model()._build_generate_request(123) + + def test_alogprobs_rejects_n_gt_1(self): + with pytest.raises(ValueError, match="alogprobs only supports n=1"): + self._model()._build_logprobs_request( + "p", max_tokens=1, logprobs=5, score_input=True, temperature=0.0, n=2 + ) + + def test_logprobs_request_forces_logprob_fields(self): + model = GenModel(model="m", api_key="k", logprobs=99) + req = model._build_logprobs_request( + "p", max_tokens=1, logprobs=5, score_input=True, temperature=0.0 + ) + # explicit alogprobs args override any _kwargs-borne logprobs config + assert req.return_logprobs is True + assert req.top_k == 5 + assert req.score_input is True + assert req.sampling.max_tokens == 1 + + +# =================================================================== +# Response -> ModelOutput bridge +# =================================================================== +class TestResponseBridge: + def _bridge(self, resp: Response) -> ModelOutput: + return GenModel(model="m", api_key="k")._response_to_model_output(resp) + + def test_basic_fields(self): + out = self._bridge( + Response( + texts=("a", "b"), + finish_reasons=("stop", "length"), + usage=UsageStats(input_tokens=3, output_tokens=4, total_tokens=7), + request_params={"max_tokens": 4}, + response_model="served-model", + system_fingerprint="fp", + ) + ) + assert isinstance(out, ModelOutput) + assert out.texts == ["a", "b"] + assert out.finish_reasons == ["stop", "length"] + assert out.usage == {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7} + assert out.request_params == {"max_tokens": 4} + assert out.response_model == "served-model" + assert out.system_fingerprint == "fp" + assert out.model["model"] == "m" + + def test_usage_absent_stays_none(self): + """Absence != zeros: a zero-filled usage dict would silently corrupt + the ARC echoed-logprob slice.""" + out = self._bridge(Response(texts=("t",))) + assert out.usage is None + + def test_input_scoring_is_flattened_ahead_of_sampled_logprobs(self): + """The legacy echo layout: prompt tokens first, then the completion.""" + resp = Response( + texts=("t",), + input_scoring=InputScoringResult( + token_logprobs=( + TokenLogprob(token="Hello", logprob=None), + TokenLogprob(token=" world", logprob=-0.5), + ) + ), + logprobs=(TokenLogprob(token=" !", logprob=-1.5),), + ) + out = self._bridge(resp) + assert out.logprobs_tokens == ["Hello", " world", " !"] + assert out.logprobs == [None, -0.5, -1.5] + + def test_logprobs_absent_stays_none(self): + out = self._bridge(Response(texts=("t",))) + assert out.logprobs_tokens is None + assert out.logprobs is None + assert out.top_logprobs is None + + def test_logprobs_present_but_empty_is_empty_list(self): + """Anomaly detection distinguishes present-but-empty from absent.""" + out = self._bridge(Response(texts=("t",), logprobs=())) + assert out.logprobs_tokens == [] + assert out.logprobs == [] + + def test_top_logprobs_coalesce_duplicates_by_max(self): + """Distinct token ids can normalize to identical text (sglang Ġ); + keep the highest logprob, matching legacy CMMLU semantics.""" + resp = Response( + texts=("t",), + top_logprobs=( + ( + TopKEntry(token=" A", logprob=-2.0, token_id=1), + TopKEntry(token=" A", logprob=-0.5, token_id=2), + TopKEntry(token=" B", logprob=-1.0, token_id=3), + ), + ), + ) + out = self._bridge(resp) + assert out.top_logprobs == [{" A": -0.5, " B": -1.0}] + + def test_top_logprobs_empty_collapses_to_none(self): + out = self._bridge(Response(texts=("t",), top_logprobs=())) + assert out.top_logprobs is None + + def test_reasoning_text_becomes_reasoning_texts(self): + out = self._bridge( + Response(texts=("t",), reasoning=ReasoningOutput(text="thinking...")) + ) + assert out.reasoning_texts == ["thinking..."] + + def test_empty_reasoning_stays_none(self): + out = self._bridge(Response(texts=("t",), reasoning=ReasoningOutput(text=""))) + assert out.reasoning_texts is None diff --git a/tests/unit/core/models/test_model_arun.py b/tests/unit/core/models/test_model_arun.py new file mode 100644 index 00000000..14bce5a0 --- /dev/null +++ b/tests/unit/core/models/test_model_arun.py @@ -0,0 +1,104 @@ +"""Tests for the IR primitive on Model: arun, capabilities, assert_capability, +and the alogprobs echo gate. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +import pytest + +from sieval.core.models import ( + Capability, + CapabilityError, + ChatModel, + GenModel, + Request, + Response, + UsageStats, +) + + +class _StubTransport: + def __init__(self, caps: frozenset[Capability], response: Response): + self._caps = caps + self._response = response + self.calls: list[Request] = [] + + @property + def capabilities(self) -> frozenset[Capability]: + return self._caps + + async def arun(self, req: Request) -> Response: + self.calls.append(req) + return self._response + + +class TestCapabilities: + def test_chat_model_capabilities_from_transport(self): + m = ChatModel(model="c", api_key="k") + assert Capability.Chat in m.capabilities + assert Capability.InputScoring not in m.capabilities + + def test_gen_model_has_input_scoring(self): + m = GenModel(model="g", api_key="k") + assert Capability.InputScoring in m.capabilities + assert Capability.Completion in m.capabilities + + def test_assert_capability_passes_when_present(self): + m = GenModel(model="g", api_key="k") + m.assert_capability(Capability.Completion, Capability.InputScoring) + + def test_assert_capability_raises_when_missing(self): + m = ChatModel(model="c", api_key="k") + with pytest.raises(CapabilityError, match="InputScoring"): + m.assert_capability(Capability.InputScoring) + + +class TestArun: + @pytest.mark.anyio + async def test_arun_delegates_to_transport(self): + resp = Response( + texts=("hi",), usage=UsageStats(output_tokens=1, total_tokens=1) + ) + stub = _StubTransport(frozenset({Capability.Chat}), resp) + m = ChatModel(model="c", api_key="k", transport=stub) + out = await m.arun(Request(input="prompt")) + assert out is resp + assert len(stub.calls) == 1 + assert stub.calls[0].input == "prompt" + + @pytest.mark.anyio + async def test_injected_transport_supplies_capabilities(self): + stub = _StubTransport(frozenset({Capability.Chat}), Response(texts=())) + m = ChatModel(model="c", api_key="k", transport=stub) + assert m.capabilities == frozenset({Capability.Chat}) + + +class TestAlogprobsEchoGate: + @pytest.mark.anyio + async def test_echo_true_on_chat_raises_capability_error(self): + """The historical bug: echo=True on chat was silently ignored.""" + from tests.conftest import MockChatModel + + m = MockChatModel() + with pytest.raises(CapabilityError, match="InputScoring"): + await m.alogprobs("prompt", echo=True) + + @pytest.mark.anyio + async def test_echo_true_on_gen_works(self): + from tests.conftest import MockGenModel + + m = MockGenModel() + out = await m.alogprobs("prompt", echo=True) + assert out.logprobs is not None + + @pytest.mark.anyio + async def test_echo_false_on_chat_skips_the_gate(self): + from tests.conftest import MockChatModel + + m = MockChatModel() + # echo=False must skip the InputScoring gate and reach the transport. + # The chat stub serves no logprob channel, so the post-arun contract + # check fires — proving the gate was bypassed and the request ran. + with pytest.raises(RuntimeError, match="server returned none"): + await m.alogprobs("prompt", echo=False) + assert m._transport.requests[0].score_input is False diff --git a/tests/unit/core/models/test_model_derivation.py b/tests/unit/core/models/test_model_derivation.py index d91e278a..d9762772 100644 --- a/tests/unit/core/models/test_model_derivation.py +++ b/tests/unit/core/models/test_model_derivation.py @@ -1,38 +1,52 @@ """ -Unit tests for Model.with_args, Model.as_type, and meta() derivation logic. +Unit tests for Model.with_args and meta() derivation logic. -Covers parent limiter wiring, type conversion, nested derivation, -and meta() field presence — paths not exercised by test_model.py. +Covers parent limiter wiring, nested derivation, and meta() field presence — +paths not exercised by test_model.py. (The as_type coverage that used to live +here was removed with the method itself; RFC #25 dropped cross-kind model +conversion.) -AI-Generated Code - Claude Sonnet 4.6 (Anthropic) +AI-Generated Code - Claude Fable 5 (Anthropic) """ import pytest -from sieval.core.models import ChatModel, GenModel +from sieval.core.models import ChatModel, GenModel, Request, Response +from sieval.core.models.transports import ( + OpenAIChatTransport, + OpenAICompletionsTransport, +) # --------------------------------------------------------------------------- # Stub implementations — no real API calls # --------------------------------------------------------------------------- -class StubGenModel(GenModel): - """GenModel stub that never hits a real API.""" +class _NeverCallTransport: + """Transport double that fails loudly if any wire call is attempted.""" - async def _agenerate_impl(self, prompt, **kwargs): - raise RuntimeError("Stub must not be called in unit tests") + def __init__(self, capabilities): + self._capabilities = frozenset(capabilities) + + @property + def capabilities(self): + return self._capabilities - async def _alogprobs_impl(self, prompt, **kwargs): + async def arun(self, req: Request) -> Response: raise RuntimeError("Stub must not be called in unit tests") +class StubGenModel(GenModel): + """GenModel stub that never hits a real API.""" + + def _build_default_transport(self): + return _NeverCallTransport(OpenAICompletionsTransport.CAPABILITIES) + + class StubChatModel(ChatModel): """ChatModel stub that never hits a real API.""" - async def _agenerate_impl(self, prompt, **kwargs): - raise RuntimeError("Stub must not be called in unit tests") - - async def _alogprobs_impl(self, prompt, **kwargs): - raise RuntimeError("Stub must not be called in unit tests") + def _build_default_transport(self): + return _NeverCallTransport(OpenAIChatTransport.CAPABILITIES) # --------------------------------------------------------------------------- @@ -50,12 +64,6 @@ def base_gen_no_limit(): return StubGenModel(model="base-gen-unlimited", api_key="fake") -@pytest.fixture -def base_chat(): - """ChatModel with a concurrency limiter (total_tokens=32).""" - return StubChatModel(model="base-chat", api_key="fake", concurrency_limit=32) - - # =================================================================== # TestModelDerivation # =================================================================== @@ -97,56 +105,12 @@ def test_with_args_new_limiter_total_tokens(self, base_gen): assert child._limiter.total_tokens == 16 - # ------------------------------------------------------------------ - # as_type — type conversion - # ------------------------------------------------------------------ - - def test_as_type_conversion_chat_to_gen(self, base_chat): - """as_type(GenModel) returns a GenModel instance.""" - gen = base_chat.as_type(GenModel) - - assert isinstance(gen, GenModel) - assert not isinstance(gen, ChatModel) - - def test_as_type_conversion_gen_to_chat(self, base_gen): - """as_type(ChatModel) returns a ChatModel instance.""" - chat = base_gen.as_type(ChatModel) - - assert isinstance(chat, ChatModel) - assert not isinstance(chat, GenModel) - - def test_as_type_preserves_parent_limiter(self, base_gen): - """as_type does not alter _parent_limiter.""" - child = base_gen.with_args(concurrency_limit=8) - converted = child.as_type(ChatModel) - - # The converted model must carry the same parent limiter reference - assert converted._parent_limiter is child._parent_limiter - assert converted._parent_limiter is base_gen._limiter - - def test_as_type_preserves_own_limiter(self, base_gen): - """as_type does not alter _limiter.""" - child = base_gen.with_args(concurrency_limit=8) - converted = child.as_type(ChatModel) - - assert converted._limiter is child._limiter - assert converted._limiter.total_tokens == 8 - - def test_as_type_preserves_model_name(self, base_gen): - """as_type does not change the model name.""" - chat = base_gen.as_type(ChatModel) - - assert chat._model == base_gen._model - - def test_as_type_invalid_type_raises(self, base_gen): - """as_type with a non-Model type raises TypeError.""" - with pytest.raises(TypeError, match="Model subclass"): - base_gen.as_type(str) + def test_with_args_shares_transport(self, base_gen): + """Derived models reuse the same Transport (same client, same wire).""" + child = base_gen.with_args(temperature=0.7) - def test_as_type_invalid_non_type_raises(self, base_gen): - """as_type with a non-type value raises TypeError.""" - with pytest.raises(TypeError, match="Model subclass"): - base_gen.as_type(42) + assert child._transport is base_gen._transport + assert child.capabilities == base_gen.capabilities # ------------------------------------------------------------------ # Nested derivation @@ -290,13 +254,6 @@ def test_extra_not_in_meta_default_params(self): model = StubGenModel(model="test", api_key="fake", extra=extra) assert "extra" not in model.meta()["default_params"] - def test_as_type_preserves_extra(self): - """as_type() must preserve extra.""" - extra = {"sequence_wrappers": {"dna": "{seq}"}} - model = StubGenModel(model="test", api_key="fake", extra=extra) - chat = model.as_type(StubChatModel) - assert chat.extra == extra - def test_with_args_extra_not_in_child_kwargs(self): """with_args(extra=...) must not leak into child _kwargs.""" model = StubGenModel(model="test", api_key="fake") diff --git a/tests/unit/core/models/test_sglang_gen_model.py b/tests/unit/core/models/test_sglang_gen_model.py index 7013ff0c..0709b244 100644 --- a/tests/unit/core/models/test_sglang_gen_model.py +++ b/tests/unit/core/models/test_sglang_gen_model.py @@ -1,579 +1,29 @@ +"""Shell tests: backend selector wiring for SglangGenModel. + +RFC #25 moved the native /generate wire logic (URL derivation, sampling-param +translation, triple parsing, ``_normalize_token_text``, the radix-cache guard) +into ``SglangTransport``; the former ``_agenerate_impl`` / ``_alogprobs_impl`` +coverage moved with it to tests/unit/core/models/transports/test_sglang.py, +and the request-builder validation (n/stream types, alogprobs n=1) lives on +``Model`` in tests/unit/core/models/test_model.py. What remains here is the +selector contract: SglangGenModel pairs the shared client (and its api_base) +with that transport, which supplies the model's capabilities. + +AI-Generated Code - Claude Fable 5 (Anthropic) """ -Unit tests for sieval/core/models/sglang_gen_model.py. -Covers native /generate generation (_agenerate_impl) and logprob extraction -(_alogprobs_impl): URL derivation, request body, sampling-param translation, -echo→logprob_start_len, input/output token-logprob + top-logprob parsing, -token-text normalization, end-to-end extract_option_logprob / total_logprob / -CMMLU-style top-k consumption, and the n / empty-response guards. The OpenAI -client's public ``post`` is mocked — no real traffic. +from sieval.core.models import Capability, SglangGenModel +from sieval.core.models.transports.sglang import SglangTransport -AI-Generated Code - Claude Opus 4.8 (Anthropic) -""" - -from typing import Any -from unittest.mock import AsyncMock - -import pytest - -from sieval.core.models.model import ModelOutput -from sieval.core.models.sglang_gen_model import ( - SglangGenModel, - _normalize_token_text, -) -from sieval.core.utils.ppl import extract_option_logprob, total_logprob - - -@pytest.fixture -def model(): - return SglangGenModel( - model="test-sglang", api_base="http://host:8000/v1", api_key="local" - ) - - -def _patch_post(model: SglangGenModel, payload): - """Mock the OpenAI client's public ``post`` to return ``payload``.""" - mock_post = AsyncMock(return_value=payload) - target: Any = model._client - target.post = mock_post # type: ignore[invalid-assignment] - return mock_post - - -def _meta( - input_entries=None, output_entries=None, input_top=None, output_top=None, **extra -): - meta: dict[str, Any] = {} - if input_entries is not None: - meta["input_token_logprobs"] = input_entries - if output_entries is not None: - meta["output_token_logprobs"] = output_entries - if input_top is not None: - meta["input_top_logprobs"] = input_top - if output_top is not None: - meta["output_top_logprobs"] = output_top - meta.update(extra) - return meta - - -# =================================================================== -# URL derivation -# =================================================================== -class TestGenerateUrl: - def test_strips_v1_suffix(self, model): - assert model._generate_url() == "http://host:8000/generate" - - def test_trailing_slash_base(self): - m = SglangGenModel(model="x", api_base="http://host:8000/v1/", api_key="local") - assert m._generate_url() == "http://host:8000/generate" - - def test_no_v1_suffix(self): - m = SglangGenModel(model="x", api_base="http://host:8000", api_key="local") - assert m._generate_url() == "http://host:8000/generate" - - def test_none_base(self): - m = SglangGenModel(model="x", api_key="local") - assert m._generate_url() == "/generate" - - -# =================================================================== -# Token text normalization -# =================================================================== -class TestNormalizeTokenText: - def test_space_marker(self): - assert _normalize_token_text("ĠA") == " A" - - def test_newline_marker(self): - assert _normalize_token_text("Ċ") == "\n" - - def test_plain_unchanged(self): - assert _normalize_token_text(" A") == " A" - - def test_none_text_raises(self): - # Server without detokenization (--skip-tokenizer-init) returns no text; - # fail loud rather than crash on None.replace or degrade to "". - with pytest.raises(RuntimeError, match="no token text"): - _normalize_token_text(None) - - -# =================================================================== -# _agenerate_impl (native /generate generation) -# =================================================================== -class TestAgenerate: - @pytest.mark.anyio - async def test_basic_generation(self, model): - post = _patch_post( - model, - { - "text": "hello world", - "meta_info": _meta(prompt_tokens=5, completion_tokens=2), - }, - ) - out = await model._agenerate_impl("hi") - assert isinstance(out, ModelOutput) - assert out.texts == ["hello world"] - assert out.usage == {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} - # posts to /generate with the text; no return_logprob for plain generation. - assert post.call_args[0][0] == "http://host:8000/generate" - body = post.call_args[1]["body"] - assert body["text"] == "hi" - assert "return_logprob" not in body - - @pytest.mark.anyio - async def test_non_string_prompt_raises(self, model): - with pytest.raises(TypeError, match="requires a string"): - await model._agenerate_impl(["not", "a", "string"]) - - @pytest.mark.anyio - async def test_sampling_param_translation(self, model): - post = _patch_post(model, {"text": "x", "meta_info": _meta()}) - await model._agenerate_impl( - "hi", max_tokens=64, temperature=0.7, top_p=0.9, seed=123 - ) - sp = post.call_args[1]["body"]["sampling_params"] - assert sp["max_new_tokens"] == 64 - assert sp["temperature"] == 0.7 - assert sp["top_p"] == 0.9 - # unmapped kwargs (seed) are dropped, not forwarded to sglang. - assert "seed" not in sp - - @pytest.mark.anyio - async def test_n_gt_1_list_response(self, model): - post = _patch_post( - model, - [ - {"text": "a", "meta_info": _meta(prompt_tokens=4, completion_tokens=1)}, - {"text": "b", "meta_info": _meta(prompt_tokens=4, completion_tokens=2)}, - ], - ) - out = await model._agenerate_impl("hi", n=2) - assert out.texts == ["a", "b"] - # prompt tokens counted once, completions summed - assert out.usage == {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7} - assert post.call_args[1]["body"]["sampling_params"]["n"] == 2 - - @pytest.mark.anyio - async def test_finish_reason_extracted(self, model): - _patch_post( - model, - {"text": "x", "meta_info": _meta(finish_reason={"type": "length"})}, - ) - out = await model._agenerate_impl("hi") - assert out.finish_reasons == ["length"] - - @pytest.mark.anyio - async def test_n_non_int_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._agenerate_impl("hi", n="2") - - @pytest.mark.anyio - async def test_n_lt_1_raises(self, model): - with pytest.raises(ValueError, match="n must be >= 1"): - await model._agenerate_impl("hi", n=0) - - @pytest.mark.anyio - async def test_invalid_response_missing_meta_info_raises(self, model): - """A response without meta_info fails loud instead of a bare KeyError.""" - _patch_post(model, {"text": "hi"}) # no meta_info - with pytest.raises(RuntimeError, match="missing meta_info"): - await model._agenerate_impl("hi") - - @pytest.mark.anyio - async def test_request_params_excludes_prompt(self, model): - post = _patch_post( - model, - {"text": "x", "meta_info": _meta(prompt_tokens=1, completion_tokens=1)}, - ) - out = await model._agenerate_impl("secret prompt", temperature=0.0) - # the raw prompt is sent on the wire... - assert post.call_args[1]["body"]["text"] == "secret prompt" - # ...but is NOT persisted into per-call request_params. - assert out.request_params is not None - assert "text" not in out.request_params - assert "sampling_params" in out.request_params - - -# =================================================================== -# _alogprobs_impl request body -# =================================================================== -class TestRequestBody: - @pytest.mark.anyio - async def test_echo_true_request_body(self, model): - post = _patch_post( - model, - { - "text": "", - "meta_info": _meta(input_entries=[[-0.1, 1, " A"]], prompt_tokens=1), - }, - ) - await model._alogprobs_impl("prompt", max_tokens=1, logprobs=5, echo=True) - body = post.call_args[1]["body"] - assert body["text"] == "prompt" - assert body["return_logprob"] is True - assert body["logprob_start_len"] == 0 - assert body["top_logprobs_num"] == 5 - assert body["return_text_in_logprobs"] is True - assert body["sampling_params"]["max_new_tokens"] == 1 - assert body["sampling_params"]["temperature"] == 0.0 - # routed through the public client with an absolute URL. - assert post.call_args[0][0] == "http://host:8000/generate" - assert post.call_args[1]["cast_to"] is object - - @pytest.mark.anyio - async def test_echo_false_sets_start_len_minus_one(self, model): - post = _patch_post( - model, {"text": "", "meta_info": _meta(output_entries=[[-0.1, 1, "x"]])} - ) - await model._alogprobs_impl("prompt", echo=False) - assert post.call_args[1]["body"]["logprob_start_len"] == -1 - - @pytest.mark.anyio - async def test_max_tokens_floored_to_one(self, model): - post = _patch_post( - model, - { - "text": "", - "meta_info": _meta(input_entries=[[-0.1, 1, " A"]], prompt_tokens=1), - }, - ) - await model._alogprobs_impl("prompt", max_tokens=0) - assert post.call_args[1]["body"]["sampling_params"]["max_new_tokens"] == 1 - - @pytest.mark.anyio - async def test_request_params_excludes_prompt(self, model): - _patch_post( - model, - { - "text": "", - "meta_info": _meta(input_entries=[[-0.1, 1, " A"]], prompt_tokens=1), - }, - ) - out = await model._alogprobs_impl("secret prompt", echo=True) - assert out.request_params is not None - assert "text" not in out.request_params - assert out.request_params["return_logprob"] is True - - -# =================================================================== -# Chosen-token logprob parsing -# =================================================================== -class TestParsing: - @pytest.mark.anyio - async def test_input_logprobs_to_tokens_and_logprobs(self, model): - meta = _meta( - input_entries=[[None, 1, "The"], [-0.5, 2, " cat"], [-0.1, 3, " A"]], - prompt_tokens=3, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt") - assert out.logprobs_tokens == ["The", " cat", " A"] - assert out.logprobs == [None, -0.5, -0.1] - - @pytest.mark.anyio - async def test_token_text_normalized(self, model): - meta = _meta( - input_entries=[[None, 1, "ĠThe"], [-0.1, 2, "ĠA"]], prompt_tokens=2 - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt") - assert out.logprobs_tokens == [" The", " A"] - - @pytest.mark.anyio - async def test_array_ordering_input_then_output(self, model): - meta = _meta( - input_entries=[[None, 1, "Q"], [-0.2, 2, " B"]], - output_entries=[[-0.3, 3, " gen"]], - prompt_tokens=2, - ) - _patch_post(model, {"text": " gen", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=True) - assert out.logprobs_tokens == ["Q", " B", " gen"] - assert out.logprobs == [None, -0.2, -0.3] - - @pytest.mark.anyio - async def test_finish_reason_on_logprobs(self, model): - meta = _meta( - input_entries=[[-0.1, 1, " A"]], - prompt_tokens=1, - finish_reason={"type": "length"}, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt") - assert out.finish_reasons == ["length"] - - @pytest.mark.anyio - async def test_usage_parsed(self, model): - # input count must equal prompt_tokens under echo (cold cache). - meta = _meta( - input_entries=[[None, 1, "Q"], [-0.1, 2, " A"]], - prompt_tokens=2, - completion_tokens=1, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt") - assert out.usage == {"input_tokens": 2, "output_tokens": 1, "total_tokens": 3} - - @pytest.mark.anyio - async def test_usage_none_when_counts_absent(self, model): - # echo=False so the radix guard (which requires prompt_tokens) is skipped; - # this isolates _parse_usage returning None when counts are absent. - meta = _meta(output_entries=[[-0.1, 1, " A"]]) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False) - assert out.usage is None - - -# =================================================================== -# Top-k logprob parsing (CMMLU / MMLU-Base consumption) -# =================================================================== -class TestTopLogprobs: - @pytest.mark.anyio - async def test_output_top_logprobs_echo_false(self, model): - """CMMLU shape: echo=False, first output token's top-k as {token: logprob}.""" - meta = _meta( - output_entries=[[-0.7, 100, " A"]], - output_top=[[[-0.7, 100, " A"], [-1.2, 101, " B"], [-3.0, 102, " C"]]], - ) - _patch_post(model, {"text": " A", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False) - assert out.top_logprobs == [{" A": -0.7, " B": -1.2, " C": -3.0}] - - @pytest.mark.anyio - async def test_top_logprobs_normalized_keys(self, model): - meta = _meta( - output_entries=[[-0.7, 100, "ĠA"]], - output_top=[[[-0.7, 100, "ĠA"], [-1.2, 101, "ĠB"]]], - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False) - assert out.top_logprobs == [{" A": -0.7, " B": -1.2}] - - @pytest.mark.anyio - async def test_top_logprobs_echo_aligns_input_first(self, model): - """echo=True: input top-k precede output; None/empty first entry → {}.""" - meta = _meta( - input_entries=[[None, 1, "Q"], [-0.2, 2, " B"]], - output_entries=[[-0.3, 3, " g"]], - input_top=[None, [[-0.2, 2, " B"], [-0.9, 9, " C"]]], - output_top=[[[-0.3, 3, " g"]]], - prompt_tokens=2, - ) - _patch_post(model, {"text": " g", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=True) - assert out.top_logprobs == [{}, {" B": -0.2, " C": -0.9}, {" g": -0.3}] - - @pytest.mark.anyio - async def test_top_logprobs_none_when_absent(self, model): - meta = _meta(input_entries=[[-0.1, 1, " A"]], prompt_tokens=1) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt") - assert out.top_logprobs is None - - @pytest.mark.anyio - async def test_cmmlu_style_scoring(self, model): - """Map first output token's top-k onto A/B/C/D (echo=False, CMMLU shape).""" - meta = _meta( - output_entries=[[-0.7, 100, " B"]], - output_top=[ - [[-2.0, 1, " A"], [-0.7, 2, " B"], [-3.0, 3, " C"], [-2.5, 4, " D"]] - ], - ) - _patch_post(model, {"text": " B", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False, logprobs=100) - scores = {tok.strip(): lp for tok, lp in (out.top_logprobs or [{}])[0].items()} - assert max(scores, key=lambda k: scores[k]) == "B" - - @pytest.mark.anyio - async def test_distinct_tokens_stripping_to_same_letter_both_kept(self, model): - """Two tokens stripping to the same letter (' B' vs '\\tB') stay distinct. - - Observed live on Qwen2.5-72B: the greedy ' B' (high logprob) and a - rare '\\tB' (very low) both strip to 'B'. They have different - normalized text, so they must remain SEPARATE dict entries — that is - what lets CMMLU's ``max``-over-strip recover the high logprob rather - than clobbering it with the low one. - """ - meta = _meta( - output_entries=[[-0.007, 425, " B"]], - output_top=[[[-0.007, 425, " B"], [-11.94, 12791, "\tB"]]], - ) - _patch_post(model, {"text": " B", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False, logprobs=100) - assert out.top_logprobs == [{" B": -0.007, "\tB": -11.94}] - # A max-over-strip consumer (CMMLU) recovers the high logprob. - best = max( - (lp for tok, lp in out.top_logprobs[0].items() if tok.strip() == "B") - ) - assert best == -0.007 - - @pytest.mark.anyio - async def test_duplicate_normalized_tokens_keep_max(self, model): - """Two token ids normalizing to the SAME text keep the highest logprob. - - A byte-level "ĠA" (high) and a literal " A" (low) both normalize to - " A"; the dict must not let the later, lower entry clobber the real - one — otherwise CMMLU would score the option at the wrong logprob. - """ - meta = _meta( - output_entries=[[-0.05, 100, "ĠA"]], - output_top=[[[-0.05, 100, "ĠA"], [-9.9, 55, " A"]]], - ) - _patch_post(model, {"text": " A", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False, logprobs=100) - assert out.top_logprobs == [{" A": -0.05}] - - @pytest.mark.anyio - async def test_none_token_text_in_top_raises(self, model): - """A top-k entry with no token text (no detokenization) fails loud.""" - meta = _meta( - output_entries=[[-0.1, 1, " A"]], - output_top=[[[-0.1, 1, None]]], - ) - _patch_post(model, {"text": "", "meta_info": meta}) - with pytest.raises(RuntimeError, match="no token text"): - await model._alogprobs_impl("prompt", echo=False) - - -# =================================================================== -# End-to-end consumption by echo-based ppl utilities -# =================================================================== -class TestPplConsumption: - @pytest.mark.anyio - async def test_extract_option_logprob_finds_letter(self, model): - meta = _meta( - input_entries=[ - [None, 1, "Question:"], - [-2.0, 2, " text"], - [-0.7, 3, " A"], - ], - prompt_tokens=3, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("Question: text A", echo=True) - assert extract_option_logprob(out.logprobs_tokens, out.logprobs, "A") == -0.7 - - @pytest.mark.anyio - async def test_total_logprob_sums_continuation(self, model): - meta = _meta( - input_entries=[[None, 1, "Ctx"], [-1.0, 2, " the"], [-2.0, 3, " end"]], - prompt_tokens=3, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("Ctx the end", echo=True) - total, count = total_logprob(out.logprobs_tokens, out.logprobs) - assert total == pytest.approx(-3.0) - assert count == 2 - - -# =================================================================== -# Guards -# =================================================================== -class TestGuards: - @pytest.mark.anyio - async def test_n_gt_1_raises(self, model): - post = _patch_post(model, {"text": "", "meta_info": _meta()}) - with pytest.raises(ValueError, match="only supports n=1"): - await model._alogprobs_impl("prompt", n=2) - post.assert_not_called() - - @pytest.mark.anyio - async def test_n_non_int_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n="2") - - @pytest.mark.anyio - async def test_n_bool_raises(self, model): - with pytest.raises(TypeError, match="n must be an int"): - await model._alogprobs_impl("prompt", n=True) - - @pytest.mark.anyio - async def test_empty_response_raises(self, model): - """No token logprobs AND no top logprobs → raise (retryable failure). - - prompt_tokens=0 so the empty input matches (passes the radix guard) and - we reach the no-logprobs check. - """ - _patch_post( - model, {"text": "", "meta_info": _meta(input_entries=[], prompt_tokens=0)} - ) - with pytest.raises(RuntimeError, match="no logprobs"): - await model._alogprobs_impl("prompt") - - @pytest.mark.anyio - async def test_meta_attached(self, model): - _patch_post( - model, - { - "text": "", - "meta_info": _meta(input_entries=[[-0.1, 1, " A"]], prompt_tokens=1), - }, - ) - out = await model._alogprobs_impl("prompt") - assert out.model["model"] == "test-sglang" - assert out.response_model == "test-sglang" - - @pytest.mark.anyio - async def test_non_dict_response_raises(self, model): - """A list response (only valid for n>1 generation) is rejected here.""" - _patch_post(model, [{"text": "", "meta_info": _meta()}]) - with pytest.raises(RuntimeError, match="expected an object"): - await model._alogprobs_impl("prompt") - - -# =================================================================== -# Radix-cache truncation guard (echo=True completeness) -# =================================================================== -class TestEchoCompletenessGuard: - """echo=True must fail loud when sglang's prefix cache truncates input logprobs.""" - - @pytest.mark.anyio - async def test_cached_tokens_nonzero_raises(self, model): - # Cache hit: input_token_logprobs truncated to the uncached tail. - meta = _meta( - input_entries=[[-0.1, 1, " star"]], - prompt_tokens=5, - cached_tokens=4, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - with pytest.raises(RuntimeError, match="disable-radix-cache"): - await model._alogprobs_impl("prompt", echo=True) - - @pytest.mark.anyio - async def test_count_mismatch_raises(self, model): - # cached_tokens field absent, but returned count < prompt_tokens. - meta = _meta(input_entries=[[-0.1, 1, " star"]], prompt_tokens=5) - _patch_post(model, {"text": "", "meta_info": meta}) - with pytest.raises(RuntimeError, match="partial echoed-input"): - await model._alogprobs_impl("prompt", echo=True) - - @pytest.mark.anyio - async def test_missing_prompt_tokens_raises(self, model): - # Without prompt_tokens the count can't be verified → fail loud rather - # than let possibly-truncated logprobs through. - meta = _meta(input_entries=[[None, 1, "a"], [-0.1, 2, " b"]]) - _patch_post(model, {"text": "", "meta_info": meta}) - with pytest.raises(RuntimeError, match="omitted prompt_tokens"): - await model._alogprobs_impl("prompt", echo=True) - @pytest.mark.anyio - async def test_full_input_passes(self, model): - meta = _meta( - input_entries=[[None, 1, "a"], [-0.1, 2, " b"]], - prompt_tokens=2, - cached_tokens=0, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=True) - assert out.logprobs_tokens == ["a", " b"] +class TestDefaultTransport: + def test_builds_sglang_transport(self): + m = SglangGenModel(model="m", api_base="http://host:8000/v1", api_key="local") + assert isinstance(m._transport, SglangTransport) + assert Capability.SampledLogprobsWithTokenIds in m.capabilities - @pytest.mark.anyio - async def test_echo_false_ignores_cache(self, model): - """echo=False (CMMLU) reads output only — cache truncation is irrelevant.""" - meta = _meta( - output_entries=[[-0.1, 1, " A"]], - output_top=[[[-0.1, 1, " A"]]], - prompt_tokens=5, - cached_tokens=4, - ) - _patch_post(model, {"text": "", "meta_info": meta}) - out = await model._alogprobs_impl("prompt", echo=False) - assert out.top_logprobs == [{" A": -0.1}] + def test_transport_bound_to_shared_client_model_and_api_base(self): + m = SglangGenModel(model="m", api_base="http://host:8000/v1", api_key="local") + assert m._transport._client is m._client + assert m._transport._model == "m" + assert m._transport._api_base == "http://host:8000/v1" diff --git a/tests/unit/core/models/transports/__init__.py b/tests/unit/core/models/transports/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/core/models/transports/test_openai_chat.py b/tests/unit/core/models/transports/test_openai_chat.py new file mode 100644 index 00000000..9ea56c47 --- /dev/null +++ b/tests/unit/core/models/transports/test_openai_chat.py @@ -0,0 +1,677 @@ +"""Unit tests for OpenAIChatTransport (lower/lift + InputScoring rejection). + +Streaming accumulation coverage moved here from the legacy ChatModel tests +(tests/unit/core/models/test_chat_model.py) when RFC #25 relocated the wire +logic into the transport: everything is driven through +``Transport.arun(Request(...))`` with the OpenAI client mocked — no real +API traffic. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from sieval.core.models.exceptions import CapabilityError +from sieval.core.models.ir import ( + ReasoningParams, + Request, + SamplingParams, + TokenLogprob, + UsageStats, +) +from sieval.core.models.transports.openai_chat import OpenAIChatTransport + + +def _make_transport(response: object) -> tuple[OpenAIChatTransport, AsyncMock]: + client = MagicMock() + create = AsyncMock(return_value=response) + client.chat.completions.create = create + return OpenAIChatTransport(client=client, model="m"), create + + +def _make_response( + *, + content="", + finish_reason="stop", + reasoning=None, + reasoning_content=None, + logprob_content=None, + tool_calls=None, + prompt_tokens=1, + completion_tokens=1, +): + resp = MagicMock() + resp.model = None + resp.system_fingerprint = None + choice = MagicMock() + choice.index = 0 + choice.finish_reason = finish_reason + msg = MagicMock() + msg.content = content + msg.reasoning = reasoning + msg.reasoning_content = reasoning_content + msg.tool_calls = tool_calls + choice.message = msg + if logprob_content is None: + choice.logprobs = None + else: + lp = MagicMock() + lp.content = logprob_content + choice.logprobs = lp + resp.choices = [choice] + usage = MagicMock() + usage.prompt_tokens = prompt_tokens + usage.completion_tokens = completion_tokens + usage.total_tokens = prompt_tokens + completion_tokens + resp.usage = usage + return resp + + +def _lp_item(token, logprob, top=None): + item = MagicMock() + item.token = token + item.logprob = logprob + item.top_logprobs = top or [] + return item + + +def _top(token, logprob): + t = MagicMock() + t.token = token + t.logprob = logprob + return t + + +# --------------------------------------------------------------------------- +# Async streaming helpers (ported from the legacy ChatModel tests) +# --------------------------------------------------------------------------- +class _AsyncIterator: + """Wraps a list into an async iterator.""" + + def __init__(self, items): + self._items = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._items) + except StopIteration as e: + raise StopAsyncIteration from e + + +def _make_chunk( + index=0, + content=None, + finish_reason="", + usage=None, + reasoning=None, + reasoning_content=None, + logprob_content=None, + model=None, + system_fingerprint=None, +): + """Build a minimal streaming chunk carrying one choice.""" + chunk = MagicMock() + chunk.usage = usage + chunk.model = model + chunk.system_fingerprint = system_fingerprint + choice = MagicMock() + choice.index = index + choice.finish_reason = finish_reason + delta = MagicMock() + delta.content = content + delta.reasoning = reasoning + delta.reasoning_content = reasoning_content + choice.delta = delta + if logprob_content is None: + choice.logprobs = None + else: + lp = MagicMock() + lp.content = logprob_content + choice.logprobs = lp + chunk.choices = [choice] + return chunk + + +def _make_multi_chunk(pairs, finish_reason=""): + """One chunk carrying several choices: ``pairs`` of (index, content).""" + chunk = MagicMock() + chunk.usage = None + chunk.model = None + chunk.system_fingerprint = None + chunk.choices = [ + _make_chunk(index=index, content=content, finish_reason=finish_reason).choices[ + 0 + ] + for index, content in pairs + ] + return chunk + + +def _make_usage_chunk(prompt_tokens=10, completion_tokens=5): + """Final chunk carrying usage but no choices.""" + chunk = MagicMock() + chunk.choices = [] + chunk.model = None + chunk.system_fingerprint = None + usage = MagicMock() + usage.prompt_tokens = prompt_tokens + usage.completion_tokens = completion_tokens + usage.total_tokens = prompt_tokens + completion_tokens + chunk.usage = usage + return chunk + + +class TestLowerRejections: + def test_score_input_raises_capability_error(self): + t, _ = _make_transport(_make_response()) + with pytest.raises(CapabilityError, match="InputScoring"): + t._lower(Request(input="hi", score_input=True)) + + def test_session_id_raises_capability_error(self): + t, _ = _make_transport(_make_response()) + with pytest.raises(CapabilityError, match="session_id"): + t._lower(Request(input="hi", session_id="resp_123")) + + +class TestLower: + def test_str_input_wrapped_as_user_message(self): + t, _ = _make_transport(_make_response()) + messages, _ = t._lower(Request(input="hi")) + assert messages == [{"role": "user", "content": "hi"}] + + def test_message_list_passthrough(self): + t, _ = _make_transport(_make_response()) + msgs = [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}] + messages, _ = t._lower(Request(input=msgs)) + assert messages == msgs + + def test_return_logprobs_and_top_k(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", return_logprobs=True, top_k=5)) + assert params["logprobs"] is True + assert params["top_logprobs"] == 5 + + def test_top_k_zero_omits_top_logprobs(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", return_logprobs=True, top_k=0)) + assert params["logprobs"] is True + assert "top_logprobs" not in params + + def test_reasoning_effort_passthrough(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request(input="hi", reasoning=ReasoningParams(effort="high")) + ) + assert params["reasoning_effort"] == "high" + + def test_response_format_and_tools(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request( + input="hi", + response_format={"type": "json_object"}, + tools=[{"type": "function"}], + ) + ) + assert params["response_format"] == {"type": "json_object"} + assert params["tools"] == [{"type": "function"}] + + def test_sampling_maps(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request(input="hi", sampling=SamplingParams(max_tokens=4, temperature=0.3)) + ) + assert params["max_tokens"] == 4 + assert params["temperature"] == 0.3 + + +class TestLowerStream: + """Wire scheduling: Request.stream lowering + stream_options injection.""" + + def test_stream_true_sets_stream_and_injects_stream_options(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", stream=True)) + assert params["stream"] is True + assert params["stream_options"] == {"include_usage": True} + + def test_explicit_stream_options_wins_over_injected_default(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request( + input="hi", + stream=True, + extra_wire_params={"stream_options": {"include_usage": False}}, + ) + ) + assert params["stream"] is True + assert params["stream_options"] == {"include_usage": False} + + def test_stream_none_defaults_to_single_shot(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", stream=None)) + assert params["stream"] is False + assert "stream_options" not in params + + def test_stream_false_no_stream_options(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", stream=False)) + assert params["stream"] is False + assert "stream_options" not in params + + +class TestLift: + @pytest.mark.anyio + async def test_text_and_usage(self): + t, _ = _make_transport( + _make_response(content="hello", prompt_tokens=3, completion_tokens=2) + ) + out = await t.arun(Request(input="hi")) + assert out.texts == ("hello",) + assert out.usage == UsageStats(input_tokens=3, output_tokens=2, total_tokens=5) + + @pytest.mark.anyio + async def test_reasoning_extracted(self): + t, _ = _make_transport(_make_response(content="a", reasoning="because")) + out = await t.arun(Request(input="hi")) + assert out.reasoning is not None + assert out.reasoning.text == "because" + + @pytest.mark.anyio + async def test_reasoning_content_fallback(self): + t, _ = _make_transport( + _make_response(content="a", reasoning=None, reasoning_content="fallback") + ) + out = await t.arun(Request(input="hi")) + assert out.reasoning is not None + assert out.reasoning.text == "fallback" + + @pytest.mark.anyio + async def test_reasoning_takes_priority_over_reasoning_content(self): + t, _ = _make_transport( + _make_response( + content="a", reasoning="primary", reasoning_content="secondary" + ) + ) + out = await t.arun(Request(input="hi")) + assert out.reasoning is not None + assert out.reasoning.text == "primary" + + @pytest.mark.anyio + async def test_no_reasoning_channel_is_none(self): + t, _ = _make_transport(_make_response(content="a")) + out = await t.arun(Request(input="hi")) + assert out.reasoning is None + + @pytest.mark.anyio + async def test_logprobs_and_top_logprobs_mapped(self): + content = [ + _lp_item("A", -0.1, [_top("A", -0.1), _top("B", -2.0)]), + ] + t, _ = _make_transport(_make_response(content="A", logprob_content=content)) + out = await t.arun(Request(input="hi", return_logprobs=True, top_k=2)) + assert out.logprobs is not None + assert out.logprobs[0].token == "A" + assert out.logprobs[0].token_id is None + assert out.top_logprobs is not None + assert out.top_logprobs[0][1].token == "B" + assert out.top_logprobs[0][1].logprob == -2.0 + + @pytest.mark.anyio + async def test_logprobs_object_with_empty_content_is_empty_tuple(self): + t, _ = _make_transport(_make_response(content="x", logprob_content=[])) + out = await t.arun(Request(input="hi", return_logprobs=True)) + assert out.logprobs == () + assert out.top_logprobs == () + + @pytest.mark.anyio + async def test_absent_logprobs_object_is_none(self): + t, _ = _make_transport(_make_response(content="x")) + out = await t.arun(Request(input="hi", return_logprobs=True)) + assert out.logprobs is None + assert out.top_logprobs is None + + @pytest.mark.anyio + async def test_tool_calls_captured(self): + tc = MagicMock() + tc.model_dump.return_value = {"id": "call_1", "type": "function"} + t, _ = _make_transport(_make_response(content="", tool_calls=[tc])) + out = await t.arun(Request(input="hi", tools=[{"type": "function"}])) + assert out.tool_calls is not None + assert out.tool_calls[0] == {"id": "call_1", "type": "function"} + + @pytest.mark.anyio + async def test_usage_absent_is_none(self): + resp = _make_response(content="hi") + resp.usage = None + t, _ = _make_transport(resp) + out = await t.arun(Request(input="hi")) + assert out.usage is None + + @pytest.mark.anyio + async def test_response_model_and_fingerprint_captured(self): + resp = _make_response(content="ok") + resp.model = "served-model-v2" + resp.system_fingerprint = "fp_xyz789" + t, _ = _make_transport(resp) + out = await t.arun(Request(input="hi")) + assert out.response_model == "served-model-v2" + assert out.system_fingerprint == "fp_xyz789" + + @pytest.mark.anyio + async def test_request_params_is_lowered_params_dict(self): + t, create = _make_transport(_make_response(content="ok")) + out = await t.arun( + Request(input="hi", sampling=SamplingParams(max_tokens=4, temperature=0.3)) + ) + assert out.request_params == { + "max_tokens": 4, + "temperature": 0.3, + "stream": False, + } + # model/messages ride separately on the wire, never in request_params. + call_kwargs = dict(create.call_args.kwargs) + assert call_kwargs.pop("model") == "m" + assert call_kwargs.pop("messages") == [{"role": "user", "content": "hi"}] + assert out.request_params == call_kwargs + + @pytest.mark.anyio + async def test_missing_finish_reason_defaults_to_empty(self): + t, _ = _make_transport(_make_response(content="done", finish_reason=None)) + out = await t.arun(Request(input="hi")) + assert out.texts == ("done",) + assert out.finish_reasons == ("",) + + @pytest.mark.anyio + async def test_content_none_stays_empty(self): + t, _ = _make_transport(_make_response(content=None)) + out = await t.arun(Request(input="hi")) + assert out.texts == ("",) + + @pytest.mark.anyio + async def test_out_of_range_choice_index_ignored(self): + resp = _make_response(content="ignored") + resp.choices[0].index = 5 + t, _ = _make_transport(resp) + out = await t.arun(Request(input="hi")) + assert out.texts == ("",) + assert out.finish_reasons == ("",) + + +class TestLiftStream: + """Streaming accumulation, driven through arun(Request(stream=True)).""" + + @pytest.mark.anyio + async def test_delta_content_stitched(self): + chunks = [ + _make_chunk(content="Hello"), + _make_chunk(content=" world", finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.texts == ("Hello world",) + assert resp.finish_reasons == ("stop",) + + @pytest.mark.anyio + async def test_n_gt_1_routes_by_choice_index(self): + chunks = [ + _make_multi_chunk([(0, "A1"), (1, "B1")]), + _make_multi_chunk([(0, "A2"), (1, "B2")], finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun( + Request(input="hi", sampling=SamplingParams(n=2), stream=True) + ) + assert resp.texts == ("A1A2", "B1B2") + assert resp.finish_reasons == ("stop", "stop") + + @pytest.mark.anyio + async def test_out_of_range_choice_index_ignored(self): + chunks = [ + _make_chunk(index=5, content="ignored", finish_reason="stop"), + _make_chunk(index=0, content="kept", finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.texts == ("kept",) + assert resp.finish_reasons == ("stop",) + + @pytest.mark.anyio + async def test_reasoning_accumulated_from_delta_reasoning(self): + chunks = [ + _make_chunk(content="a", reasoning="think1"), + _make_chunk(content="b", reasoning="think2", finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.texts == ("ab",) + assert resp.reasoning is not None + assert resp.reasoning.text == "think1think2" + + @pytest.mark.anyio + async def test_reasoning_content_fallback(self): + chunks = [ + _make_chunk(content="a", reasoning_content="fallback", finish_reason="stop") + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.reasoning is not None + assert resp.reasoning.text == "fallback" + + @pytest.mark.anyio + async def test_reasoning_takes_priority_over_reasoning_content(self): + chunks = [ + _make_chunk( + content="a", + reasoning="primary", + reasoning_content="secondary", + finish_reason="stop", + ) + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.reasoning is not None + assert resp.reasoning.text == "primary" + + @pytest.mark.anyio + async def test_no_reasoning_deltas_channel_is_none(self): + chunks = [_make_chunk(content="a", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.reasoning is None + + @pytest.mark.anyio + async def test_usage_from_final_chunk(self): + chunks = [ + _make_chunk(content="ok", finish_reason="stop"), + _make_usage_chunk(8, 2), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.usage == UsageStats( + input_tokens=8, output_tokens=2, total_tokens=10 + ) + + @pytest.mark.anyio + async def test_no_usage_chunk_usage_is_none(self): + chunks = [_make_chunk(content="hi", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.usage is None + + @pytest.mark.anyio + async def test_response_metadata_captured_from_first_chunk(self): + chunks = [ + _make_chunk(content="a", model="model-v1", system_fingerprint="fp_first"), + _make_chunk( + content="b", + finish_reason="stop", + model="model-v2", + system_fingerprint="fp_second", + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.response_model == "model-v1" + assert resp.system_fingerprint == "fp_first" + + @pytest.mark.anyio + async def test_missing_response_metadata_is_none(self): + chunks = [_make_chunk(content="ok", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.response_model is None + assert resp.system_fingerprint is None + + @pytest.mark.anyio + async def test_streaming_logprobs_collected_for_choice_zero(self): + chunks = [ + _make_chunk( + content="A", + logprob_content=[ + _lp_item("A", -0.1, [_top("A", -0.1), _top("B", -2.0)]) + ], + ), + _make_chunk( + content="B", + finish_reason="stop", + logprob_content=[_lp_item("B", -0.5)], + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun( + Request(input="hi", return_logprobs=True, top_k=2, stream=True) + ) + assert resp.logprobs == ( + TokenLogprob(token="A", logprob=-0.1), + TokenLogprob(token="B", logprob=-0.5), + ) + assert resp.top_logprobs is not None + assert resp.top_logprobs[0][1].token == "B" + assert resp.top_logprobs[0][1].logprob == -2.0 + assert resp.top_logprobs[1] == () + + @pytest.mark.anyio + async def test_logprobs_only_collected_from_choice_zero(self): + chunks = [ + _make_chunk(index=1, content="b", logprob_content=[_lp_item("X", -9.0)]), + _make_chunk( + index=0, + content="a", + finish_reason="stop", + logprob_content=[_lp_item("A", -0.1)], + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun( + Request( + input="hi", + sampling=SamplingParams(n=2), + return_logprobs=True, + stream=True, + ) + ) + assert resp.logprobs == (TokenLogprob(token="A", logprob=-0.1),) + + @pytest.mark.anyio + async def test_no_logprobs_object_on_any_chunk_is_none(self): + chunks = [_make_chunk(content="x", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", return_logprobs=True, stream=True)) + assert resp.logprobs is None + assert resp.top_logprobs is None + + @pytest.mark.anyio + async def test_logprobs_object_with_empty_content_is_empty_tuple(self): + chunks = [_make_chunk(content="x", finish_reason="stop", logprob_content=[])] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", return_logprobs=True, stream=True)) + assert resp.logprobs == () + assert resp.top_logprobs == () + + @pytest.mark.anyio + async def test_delta_none_accumulates_nothing(self): + chunk = _make_chunk(finish_reason="stop") + chunk.choices[0].delta = None + t, _ = _make_transport(_AsyncIterator([chunk])) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.texts == ("",) + + @pytest.mark.anyio + async def test_delta_content_none_stays_empty(self): + chunks = [_make_chunk(content=None, finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.texts == ("",) + + @pytest.mark.anyio + async def test_request_params_include_stream_and_stream_options(self): + chunks = [_make_chunk(content="x", finish_reason="stop")] + t, create = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="hi", stream=True)) + assert resp.request_params == { + "stream": True, + "stream_options": {"include_usage": True}, + } + assert create.call_args.kwargs["stream"] is True + assert create.call_args.kwargs["stream_options"] == {"include_usage": True} + + +class TestLowerBranches: + def test_all_sampling_params_map(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request( + input="hi", + sampling=SamplingParams( + top_p=0.8, + stop=("X",), + seed=7, + frequency_penalty=0.1, + presence_penalty=0.2, + n=2, + ), + ) + ) + assert params["top_p"] == 0.8 + assert params["stop"] == ["X"] + assert params["seed"] == 7 + assert params["frequency_penalty"] == 0.1 + assert params["presence_penalty"] == 0.2 + assert params["n"] == 2 + + def test_extra_wire_params_passthrough(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower(Request(input="hi", extra_wire_params={"user": "u1"})) + assert params["user"] == "u1" + + def test_explicit_ir_field_wins_over_extra_wire_params(self): + t, _ = _make_transport(_make_response()) + _, params = t._lower( + Request( + input="hi", + sampling=SamplingParams(max_tokens=4), + extra_wire_params={"max_tokens": 99}, + ) + ) + assert params["max_tokens"] == 4 + + +class TestToolCallConversion: + def test_dict_tool_call_passthrough(self): + from sieval.core.models.transports.openai_chat import _tool_call_to_dict + + assert _tool_call_to_dict({"id": "x"}) == {"id": "x"} + + def test_plain_object_tool_call_via_vars(self): + from sieval.core.models.transports.openai_chat import _tool_call_to_dict + + class _Obj: + def __init__(self): + self.id = "y" + + assert _tool_call_to_dict(_Obj()) == {"id": "y"} diff --git a/tests/unit/core/models/transports/test_openai_completions.py b/tests/unit/core/models/transports/test_openai_completions.py new file mode 100644 index 00000000..7d8cacaf --- /dev/null +++ b/tests/unit/core/models/transports/test_openai_completions.py @@ -0,0 +1,587 @@ +"""Unit tests for OpenAICompletionsTransport (echo/InputScoring split). + +Streaming accumulation and ``_completion_top_logprobs`` coverage moved here +from the legacy GenModel tests (tests/unit/core/models/test_gen_model.py) when +RFC #25 relocated the wire logic into the transport: everything is driven +through ``Transport.arun(Request(...))`` with the OpenAI client mocked — no +real API traffic. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from sieval.core.models.ir import Request, SamplingParams, TokenLogprob, UsageStats +from sieval.core.models.transports.openai_completions import ( + OpenAICompletionsTransport, + _completion_top_logprobs, +) + + +def _make_transport(response: object) -> tuple[OpenAICompletionsTransport, AsyncMock]: + client = MagicMock() + create = AsyncMock(return_value=response) + client.completions.create = create + return OpenAICompletionsTransport(client=client, model="m"), create + + +def _make_response( + *, + text="", + tokens=None, + token_logprobs=None, + top_logprobs=None, + prompt_tokens=0, + completion_tokens=0, + finish_reason="stop", +): + resp = MagicMock() + resp.model = None + resp.system_fingerprint = None + choice = MagicMock() + choice.index = 0 + choice.text = text + choice.finish_reason = finish_reason + if tokens is None and token_logprobs is None: + choice.logprobs = None + else: + lp = MagicMock() + lp.tokens = tokens or [] + lp.token_logprobs = token_logprobs or [] + lp.top_logprobs = top_logprobs + choice.logprobs = lp + resp.choices = [choice] + usage = MagicMock() + usage.prompt_tokens = prompt_tokens + usage.completion_tokens = completion_tokens + usage.total_tokens = prompt_tokens + completion_tokens + resp.usage = usage + return resp + + +# --------------------------------------------------------------------------- +# Async streaming helpers (ported from the legacy GenModel tests) +# --------------------------------------------------------------------------- +class _AsyncIterator: + """Wraps a list into an async iterator.""" + + def __init__(self, items): + self._items = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._items) + except StopIteration as e: + raise StopAsyncIteration from e + + +def _make_chunk( + index=0, + text="", + finish_reason="", + tokens=None, + token_logprobs=None, + top_logprobs=None, + model=None, + system_fingerprint=None, +): + """Build a minimal completions streaming chunk carrying one choice.""" + chunk = MagicMock() + chunk.usage = None + chunk.model = model + chunk.system_fingerprint = system_fingerprint + choice = MagicMock() + choice.index = index + choice.text = text + choice.finish_reason = finish_reason + if tokens is None and token_logprobs is None and top_logprobs is None: + choice.logprobs = None + else: + lp = MagicMock() + lp.tokens = tokens or [] + lp.token_logprobs = token_logprobs or [] + lp.top_logprobs = top_logprobs + choice.logprobs = lp + chunk.choices = [choice] + return chunk + + +def _make_usage_chunk(prompt_tokens=10, completion_tokens=5): + """Final chunk carrying usage but no choices.""" + chunk = MagicMock() + chunk.choices = [] + chunk.model = None + chunk.system_fingerprint = None + usage = MagicMock() + usage.prompt_tokens = prompt_tokens + usage.completion_tokens = completion_tokens + usage.total_tokens = prompt_tokens + completion_tokens + chunk.usage = usage + return chunk + + +class TestLower: + def test_score_input_true_sets_echo(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", score_input=True)) + assert params["echo"] is True + + def test_score_input_false_omits_echo(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", return_logprobs=True, top_k=3)) + assert "echo" not in params + + def test_logprobs_count_from_top_k(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", return_logprobs=True, top_k=5)) + assert params["logprobs"] == 5 + + def test_sampling_maps_through(self): + t, _ = _make_transport(_make_response()) + params = t._lower( + Request(input="hi", sampling=SamplingParams(max_tokens=8, temperature=0.7)) + ) + assert params["max_tokens"] == 8 + assert params["temperature"] == 0.7 + + def test_all_sampling_params_map(self): + t, _ = _make_transport(_make_response()) + params = t._lower( + Request( + input="hi", + sampling=SamplingParams( + top_p=0.8, + stop=("X",), + seed=42, + frequency_penalty=0.3, + presence_penalty=0.4, + n=2, + ), + ) + ) + assert params["top_p"] == 0.8 + assert params["stop"] == ["X"] + assert params["seed"] == 42 + assert params["frequency_penalty"] == 0.3 + assert params["presence_penalty"] == 0.4 + assert params["n"] == 2 + + def test_extra_wire_params_passthrough(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", extra_wire_params={"logit_bias": {}})) + assert params["logit_bias"] == {} + + def test_explicit_ir_field_wins_over_extra_wire_params(self): + t, _ = _make_transport(_make_response()) + params = t._lower( + Request( + input="hi", + sampling=SamplingParams(max_tokens=8), + extra_wire_params={"max_tokens": 99}, + ) + ) + assert params["max_tokens"] == 8 + + def test_non_str_input_raises(self): + t, _ = _make_transport(_make_response()) + with pytest.raises(TypeError, match="str input"): + t._lower(Request(input=[{"role": "user", "content": "x"}])) + + +class TestLowerStream: + """Wire scheduling: Request.stream lowering + stream_options injection.""" + + def test_stream_true_sets_stream_and_injects_stream_options(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", stream=True)) + assert params["stream"] is True + assert params["stream_options"] == {"include_usage": True} + + def test_explicit_stream_options_wins_over_injected_default(self): + t, _ = _make_transport(_make_response()) + params = t._lower( + Request( + input="hi", + stream=True, + extra_wire_params={"stream_options": {"include_usage": False}}, + ) + ) + assert params["stream"] is True + assert params["stream_options"] == {"include_usage": False} + + def test_stream_none_defaults_to_single_shot(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", stream=None)) + assert params["stream"] is False + assert "stream_options" not in params + + def test_stream_false_no_stream_options(self): + t, _ = _make_transport(_make_response()) + params = t._lower(Request(input="hi", stream=False)) + assert params["stream"] is False + assert "stream_options" not in params + + +class TestLiftEchoSplit: + @pytest.mark.anyio + async def test_prompt_completion_boundary_split(self): + # 2 prompt tokens echoed + 1 completion token. + resp = _make_response( + text=" g", + tokens=["a", "b", " g"], + token_logprobs=[None, -1.0, -0.5], + prompt_tokens=2, + completion_tokens=1, + ) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="ab", score_input=True, top_k=0)) + assert out.input_scoring is not None + assert len(out.input_scoring.token_logprobs) == 2 # == prompt_tokens + assert out.input_scoring.token_logprobs[0].logprob is None + assert out.logprobs is not None + assert len(out.logprobs) == 1 + assert out.logprobs[0].token == " g" + assert out.logprobs[0].logprob == -0.5 + + @pytest.mark.anyio + async def test_missing_usage_boundary_zero(self): + # No usage → the echo boundary defaults to 0: input_scoring stays + # empty and every token lands in the sampled-completion channel. + resp = _make_response( + text=" g", + tokens=["a", " g"], + token_logprobs=[None, -0.5], + ) + resp.usage = None + t, _ = _make_transport(resp) + out = await t.arun(Request(input="a", score_input=True, top_k=0)) + assert out.input_scoring is not None + assert out.input_scoring.token_logprobs == () + assert out.logprobs is not None + assert len(out.logprobs) == 2 + + @pytest.mark.anyio + async def test_no_score_input_all_tokens_in_logprobs(self): + resp = _make_response( + text="xy", + tokens=["x", "y"], + token_logprobs=[-0.1, -0.2], + prompt_tokens=1, + completion_tokens=2, + ) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p", return_logprobs=True, top_k=0)) + assert out.input_scoring is None + assert out.logprobs is not None + assert len(out.logprobs) == 2 + + @pytest.mark.anyio + async def test_token_id_is_none_for_completions(self): + resp = _make_response( + text="x", + tokens=["x"], + token_logprobs=[-0.1], + prompt_tokens=0, + completion_tokens=1, + ) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p", return_logprobs=True)) + assert out.logprobs is not None + assert out.logprobs[0].token_id is None + + @pytest.mark.anyio + async def test_no_logprobs_response(self): + resp = _make_response(text="hi", prompt_tokens=1, completion_tokens=1) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p")) + assert out.logprobs is None + assert out.top_logprobs is None + assert out.texts == ("hi",) + + @pytest.mark.anyio + async def test_usage_absent_is_none(self): + resp = _make_response(text="hi") + resp.usage = None + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p")) + assert out.usage is None + + +class TestLift: + """Non-streaming lift of the remaining Response fields.""" + + @pytest.mark.anyio + async def test_usage_present_maps_to_usage_stats(self): + resp = _make_response(text="ok", prompt_tokens=6, completion_tokens=2) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p")) + assert out.usage == UsageStats(input_tokens=6, output_tokens=2, total_tokens=8) + + @pytest.mark.anyio + async def test_response_model_and_fingerprint_captured(self): + resp = _make_response(text="ok") + resp.model = "served-model-v2" + resp.system_fingerprint = "fp_xyz789" + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p")) + assert out.response_model == "served-model-v2" + assert out.system_fingerprint == "fp_xyz789" + + @pytest.mark.anyio + async def test_request_params_is_lowered_params_dict(self): + t, create = _make_transport(_make_response(text="ok")) + out = await t.arun( + Request(input="p", sampling=SamplingParams(max_tokens=8, temperature=0.7)) + ) + assert out.request_params == { + "max_tokens": 8, + "temperature": 0.7, + "stream": False, + } + # model/prompt ride separately on the wire, never in request_params. + call_kwargs = dict(create.call_args.kwargs) + assert call_kwargs.pop("model") == "m" + assert call_kwargs.pop("prompt") == "p" + assert out.request_params == call_kwargs + + @pytest.mark.anyio + async def test_logprobs_object_with_empty_channels_is_empty_tuple(self): + resp = _make_response(text="x", tokens=[], token_logprobs=[]) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p", return_logprobs=True)) + assert out.logprobs == () + assert out.top_logprobs == () + + @pytest.mark.anyio + async def test_top_logprobs_sanitized_per_position(self): + resp = _make_response( + text="A", + tokens=["A"], + token_logprobs=[-0.1], + top_logprobs=[{"A": -0.1, "B": -2.0}], + prompt_tokens=0, + completion_tokens=1, + ) + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p", return_logprobs=True, top_k=2)) + assert out.top_logprobs is not None + assert len(out.top_logprobs) == 1 + assert {(e.token, e.logprob) for e in out.top_logprobs[0]} == { + ("A", -0.1), + ("B", -2.0), + } + + @pytest.mark.anyio + async def test_out_of_range_choice_index_ignored(self): + resp = _make_response(text="ignored") + resp.choices[0].index = 5 + t, _ = _make_transport(resp) + out = await t.arun(Request(input="p")) + assert out.texts == ("",) + assert out.finish_reasons == ("",) + + +class TestLiftStream: + """Streaming accumulation, driven through arun(Request(stream=True)).""" + + @pytest.mark.anyio + async def test_text_concatenated_across_chunks(self): + chunks = [ + _make_chunk(text="Hello"), + _make_chunk(text=" world", finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.texts == ("Hello world",) + assert resp.finish_reasons == ("stop",) + + @pytest.mark.anyio + async def test_n_gt_1_routes_by_choice_index(self): + chunk1 = MagicMock() + chunk1.usage = None + chunk1.model = None + chunk1.system_fingerprint = None + chunk1.choices = [ + _make_chunk(index=0, text="A1").choices[0], + _make_chunk(index=1, text="B1").choices[0], + ] + chunk2 = MagicMock() + chunk2.usage = None + chunk2.model = None + chunk2.system_fingerprint = None + chunk2.choices = [ + _make_chunk(index=0, text="A2", finish_reason="stop").choices[0], + _make_chunk(index=1, text="B2", finish_reason="stop").choices[0], + ] + t, _ = _make_transport(_AsyncIterator([chunk1, chunk2])) + resp = await t.arun( + Request(input="p", sampling=SamplingParams(n=2), stream=True) + ) + assert resp.texts == ("A1A2", "B1B2") + assert resp.finish_reasons == ("stop", "stop") + + @pytest.mark.anyio + async def test_out_of_range_choice_index_ignored(self): + chunks = [ + _make_chunk(index=5, text="ignored", finish_reason="stop"), + _make_chunk(index=0, text="kept", finish_reason="stop"), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.texts == ("kept",) + + @pytest.mark.anyio + async def test_logprobs_accumulated_for_choice_zero(self): + chunks = [ + _make_chunk( + text="A", + tokens=["A"], + token_logprobs=[-0.1], + top_logprobs=[{"A": -0.1, "B": -2.0}], + ), + _make_chunk( + text="B", + finish_reason="stop", + tokens=["B"], + token_logprobs=[-0.5], + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun( + Request(input="p", return_logprobs=True, top_k=2, stream=True) + ) + assert resp.logprobs == ( + TokenLogprob(token="A", logprob=-0.1), + TokenLogprob(token="B", logprob=-0.5), + ) + assert resp.top_logprobs is not None + assert {(e.token, e.logprob) for e in resp.top_logprobs[0]} == { + ("A", -0.1), + ("B", -2.0), + } + + @pytest.mark.anyio + async def test_usage_from_final_chunk(self): + chunks = [ + _make_chunk(text="ok", finish_reason="stop"), + _make_usage_chunk(7, 3), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.usage == UsageStats( + input_tokens=7, output_tokens=3, total_tokens=10 + ) + + @pytest.mark.anyio + async def test_no_usage_chunk_usage_is_none(self): + chunks = [_make_chunk(text="hi", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.usage is None + + @pytest.mark.anyio + async def test_response_metadata_captured_from_first_chunk(self): + chunks = [ + _make_chunk(text="a", model="model-v1", system_fingerprint="fp_first"), + _make_chunk( + text="b", + finish_reason="stop", + model="model-v2", + system_fingerprint="fp_second", + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.response_model == "model-v1" + assert resp.system_fingerprint == "fp_first" + + @pytest.mark.anyio + async def test_echo_split_at_usage_boundary(self): + # Echoed prompt tokens arrive first, then the sampled completion, + # then the usage chunk that defines the split boundary. + chunks = [ + _make_chunk(text="", tokens=["a", "b"], token_logprobs=[None, -1.0]), + _make_chunk( + text=" g", + finish_reason="stop", + tokens=[" g"], + token_logprobs=[-0.5], + ), + _make_usage_chunk(prompt_tokens=2, completion_tokens=1), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="ab", score_input=True, top_k=0, stream=True)) + assert resp.input_scoring is not None + assert [tl.token for tl in resp.input_scoring.token_logprobs] == ["a", "b"] + assert resp.input_scoring.token_logprobs[0].logprob is None + assert resp.logprobs == (TokenLogprob(token=" g", logprob=-0.5),) + + @pytest.mark.anyio + async def test_echo_split_missing_usage_boundary_zero(self): + chunks = [ + _make_chunk( + text=" g", + finish_reason="stop", + tokens=["a", " g"], + token_logprobs=[None, -0.5], + ), + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="a", score_input=True, top_k=0, stream=True)) + assert resp.input_scoring is not None + assert resp.input_scoring.token_logprobs == () + assert resp.logprobs is not None + assert len(resp.logprobs) == 2 + + @pytest.mark.anyio + async def test_no_logprobs_object_on_any_chunk_is_none(self): + chunks = [_make_chunk(text="x", finish_reason="stop")] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", return_logprobs=True, stream=True)) + assert resp.logprobs is None + assert resp.top_logprobs is None + + @pytest.mark.anyio + async def test_logprobs_object_with_empty_channels_is_empty_tuple(self): + chunks = [ + _make_chunk(text="x", finish_reason="stop", tokens=[], token_logprobs=[]) + ] + t, _ = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", return_logprobs=True, stream=True)) + assert resp.logprobs == () + assert resp.top_logprobs == () + + @pytest.mark.anyio + async def test_request_params_include_stream_and_stream_options(self): + chunks = [_make_chunk(text="x", finish_reason="stop")] + t, create = _make_transport(_AsyncIterator(chunks)) + resp = await t.arun(Request(input="p", stream=True)) + assert resp.request_params == { + "stream": True, + "stream_options": {"include_usage": True}, + } + assert create.call_args.kwargs["stream"] is True + assert create.call_args.kwargs["stream_options"] == {"include_usage": True} + + +class TestCompletionTopLogprobs: + """Cover _completion_top_logprobs parsing and its defensive branches. + + Moved from gen_model.py to transports/openai_completions.py by RFC #25. + """ + + def test_non_sequence_returns_empty(self): + assert _completion_top_logprobs(None) == [] + assert _completion_top_logprobs("AB") == [] + + def test_parses_and_skips_malformed_entries(self): + # Mix of: a valid dict; a None entry (→ {}); a non-Mapping entry + # (skipped); and a dict whose non-numeric value ("no") and non-str key + # (2) are both filtered out, leaving only the valid pair. + raw = [{"A": -0.1}, None, "x", {"y": -0.5, "1": "no", 2: -0.3}] + assert _completion_top_logprobs(raw) == [{"A": -0.1}, {}, {"y": -0.5}] diff --git a/tests/unit/core/models/transports/test_sglang.py b/tests/unit/core/models/transports/test_sglang.py new file mode 100644 index 00000000..f722e157 --- /dev/null +++ b/tests/unit/core/models/transports/test_sglang.py @@ -0,0 +1,504 @@ +"""Unit tests for SglangTransport (lower/lift of the native /generate protocol). + +URL derivation, token-text normalization, triple parsing, the radix-cache +guard, and the n>1 list-response handling moved here from the legacy +SglangGenModel tests (tests/unit/core/models/test_sglang_gen_model.py) when +RFC #25 relocated the wire logic into the transport: everything is driven +through ``Transport.arun(Request(...))`` with the OpenAI client's ``post`` +mocked — no real traffic. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from sieval.core.models.ir import ( + Request, + SamplingParams, + TokenLogprob, + TopKEntry, + UsageStats, +) +from sieval.core.models.transports.sglang import SglangTransport, _normalize_token_text + + +def _make_transport(payload: dict | list) -> tuple[SglangTransport, AsyncMock]: + client = MagicMock() + post = AsyncMock(return_value=payload) + client.post = post + t = SglangTransport(client=client, model="m", api_base="http://host/v1") + return t, post + + +def _meta(**overrides: Any) -> dict: + meta = { + "prompt_tokens": 0, + "completion_tokens": 1, + "cached_tokens": 0, + "finish_reason": {"type": "stop"}, + } + meta.update(overrides) + return meta + + +# =================================================================== +# URL derivation (moved from the legacy SglangGenModel tests) +# =================================================================== +class TestGenerateUrl: + def test_strips_v1_suffix(self): + t = SglangTransport( + client=MagicMock(), model="x", api_base="http://host:8000/v1" + ) + assert t._generate_url() == "http://host:8000/generate" + + def test_trailing_slash_base(self): + t = SglangTransport( + client=MagicMock(), model="x", api_base="http://host:8000/v1/" + ) + assert t._generate_url() == "http://host:8000/generate" + + def test_no_v1_suffix(self): + t = SglangTransport(client=MagicMock(), model="x", api_base="http://host:8000") + assert t._generate_url() == "http://host:8000/generate" + + def test_none_base(self): + t = SglangTransport(client=MagicMock(), model="x", api_base=None) + assert t._generate_url() == "/generate" + + +# =================================================================== +# Token text normalization (moved from the legacy SglangGenModel tests) +# =================================================================== +class TestNormalizeTokenText: + def test_space_marker(self): + assert _normalize_token_text("ĠA") == " A" + + def test_newline_marker(self): + assert _normalize_token_text("Ċ") == "\n" + + def test_plain_unchanged(self): + assert _normalize_token_text(" A") == " A" + + def test_none_text_raises(self): + # Server without detokenization (--skip-tokenizer-init) returns no text; + # fail loud rather than crash on None.replace or degrade to "". + with pytest.raises(RuntimeError, match="no token text"): + _normalize_token_text(None) + + +class TestLower: + def test_max_tokens_maps_to_max_new_tokens(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", sampling=SamplingParams(max_tokens=7))) + assert body["sampling_params"]["max_new_tokens"] == 7 + + def test_score_input_sets_logprob_start_len_0(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", score_input=True, top_k=5)) + assert body["logprob_start_len"] == 0 + assert body["return_logprob"] is True + assert body["top_logprobs_num"] == 5 + assert body["return_text_in_logprobs"] is True + + def test_no_score_input_sets_logprob_start_len_minus_1(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", return_logprobs=True, top_k=3)) + assert body["logprob_start_len"] == -1 + + def test_top_k_maps_to_top_logprobs_num(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", return_logprobs=True, top_k=4)) + assert body["top_logprobs_num"] == 4 + + def test_no_logprobs_request_omits_logprob_fields(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi")) + assert "return_logprob" not in body + assert "logprob_start_len" not in body + + def test_echo_never_appears_in_wire_body(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", score_input=True)) + assert "echo" not in body + assert "echo" not in body["sampling_params"] + + def test_prefix_maps_to_prefill(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", prefix="PRE")) + assert body["sampling_params"]["prefill"] == "PRE" + + def test_score_input_forces_min_one_new_token(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower(Request(input="hi", score_input=True)) + assert body["sampling_params"]["max_new_tokens"] == 1 + + def test_all_sampling_params_map(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower( + Request( + input="hi", + sampling=SamplingParams( + temperature=0.7, + top_p=0.9, + top_k_sampling=40, + stop=("",), + frequency_penalty=0.1, + presence_penalty=0.2, + n=3, + ), + ) + ) + sp = body["sampling_params"] + assert sp["temperature"] == 0.7 + assert sp["top_p"] == 0.9 + assert sp["top_k"] == 40 + assert sp["stop"] == [""] + assert sp["frequency_penalty"] == 0.1 + assert sp["presence_penalty"] == 0.2 + assert sp["n"] == 3 + + def test_non_str_input_raises(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + with pytest.raises(TypeError, match="str input"): + t._lower(Request(input=[{"role": "user", "content": "x"}])) + + +class TestExtraWireParams: + """extra_wire_params go through the OpenAI→sglang table or are dropped.""" + + def test_known_key_mapped_into_sampling_params(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower( + Request(input="hi", extra_wire_params={"repetition_penalty": 1.1}) + ) + assert body["sampling_params"]["repetition_penalty"] == 1.1 + assert "repetition_penalty" not in body # nested, never top-level + + def test_unknown_keys_dropped(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower( + Request(input="hi", extra_wire_params={"seed": 1, "custom_flag": True}) + ) + assert "seed" not in body + assert "seed" not in body["sampling_params"] + assert "custom_flag" not in body + assert "custom_flag" not in body["sampling_params"] + + def test_explicit_ir_field_wins_over_extra_wire_params(self): + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + body = t._lower( + Request( + input="hi", + sampling=SamplingParams(max_tokens=7), + extra_wire_params={"max_tokens": 99}, + ) + ) + assert body["sampling_params"]["max_new_tokens"] == 7 + + +# =================================================================== +# Wire plumbing: URL, cast_to, body, provenance +# =================================================================== +class TestWirePlumbing: + @pytest.mark.anyio + async def test_posts_to_generate_url_with_body(self): + t, post = _make_transport({"text": "out", "meta_info": _meta(prompt_tokens=1)}) + await t.arun(Request(input="hi")) + assert post.call_args[0][0] == "http://host/generate" + assert post.call_args[1]["cast_to"] is object + body = post.call_args[1]["body"] + assert body["text"] == "hi" + assert "return_logprob" not in body + + @pytest.mark.anyio + async def test_stream_flag_is_ignored_single_post(self): + # Native /generate is always one POST; Request.stream never reaches + # the wire body (pure scheduling, no content impact). + t, post = _make_transport({"text": "x", "meta_info": _meta(prompt_tokens=1)}) + resp = await t.arun(Request(input="hi", stream=True)) + assert post.call_count == 1 + assert "stream" not in post.call_args[1]["body"] + assert resp.texts == ("x",) + + @pytest.mark.anyio + async def test_request_params_excludes_prompt_text(self): + t, post = _make_transport( + {"text": "x", "meta_info": _meta(prompt_tokens=1, completion_tokens=1)} + ) + resp = await t.arun( + Request(input="secret prompt", sampling=SamplingParams(temperature=0.0)) + ) + # the raw prompt is sent on the wire... + assert post.call_args[1]["body"]["text"] == "secret prompt" + # ...but is NOT persisted into per-call request_params. + assert resp.request_params is not None + assert "text" not in resp.request_params + assert resp.request_params["sampling_params"] == {"temperature": 0.0} + + @pytest.mark.anyio + async def test_response_model_is_configured_model(self): + t, _ = _make_transport({"text": "x", "meta_info": _meta(prompt_tokens=1)}) + resp = await t.arun(Request(input="hi")) + assert resp.response_model == "m" + + +# =================================================================== +# n>1: /generate returns a list of per-sample dicts +# =================================================================== +class TestNSamples: + @pytest.mark.anyio + async def test_list_response_yields_per_sample_texts(self): + t, post = _make_transport( + [ + { + "text": "a", + "meta_info": _meta(prompt_tokens=4, completion_tokens=1), + }, + { + "text": "b", + "meta_info": _meta( + prompt_tokens=4, + completion_tokens=2, + finish_reason={"type": "length"}, + ), + }, + { + "text": "c", + "meta_info": _meta(prompt_tokens=4, completion_tokens=3), + }, + ] + ) + resp = await t.arun(Request(input="hi", sampling=SamplingParams(n=3))) + assert resp.texts == ("a", "b", "c") + assert resp.finish_reasons == ("stop", "length", "stop") + # prompt tokens counted once, completions summed. + assert resp.usage == UsageStats( + input_tokens=4, output_tokens=6, total_tokens=10 + ) + assert post.call_args[1]["body"]["sampling_params"]["n"] == 3 + + +class TestLift: + @pytest.mark.anyio + async def test_triple_populates_token_id(self): + meta = _meta( + prompt_tokens=1, + output_token_logprobs=[[-0.5, 42, "hello"]], + ) + t, _ = _make_transport({"text": "hello", "meta_info": meta}) + resp = await t.arun(Request(input="x", return_logprobs=True, top_k=0)) + assert resp.logprobs is not None + assert resp.logprobs[0].token_id == 42 + assert resp.logprobs[0].token == "hello" + assert resp.logprobs[0].logprob == -0.5 + + @pytest.mark.anyio + async def test_normalize_token_text_called(self): + meta = _meta(output_token_logprobs=[[-0.1, 7, "ĠA"]]) # "ĠA" + t, _ = _make_transport({"text": "", "meta_info": meta}) + resp = await t.arun(Request(input="x", return_logprobs=True)) + assert resp.logprobs is not None + assert resp.logprobs[0].token == " A" # Ġ normalized to space + + @pytest.mark.anyio + async def test_newline_marker_normalized_in_tokens(self): + meta = _meta(output_token_logprobs=[[-0.2, 9, "Ċ"]]) + t, _ = _make_transport({"text": "\n", "meta_info": meta}) + resp = await t.arun(Request(input="x", return_logprobs=True)) + assert resp.logprobs == (TokenLogprob(token="\n", logprob=-0.2, token_id=9),) + + @pytest.mark.anyio + async def test_input_token_logprobs_go_to_input_scoring(self): + meta = _meta( + prompt_tokens=2, + input_token_logprobs=[[None, 1, "a"], [-1.2, 2, "b"]], + output_token_logprobs=[[-0.3, 3, "c"]], + cached_tokens=0, + ) + t, _ = _make_transport({"text": "c", "meta_info": meta}) + resp = await t.arun(Request(input="ab", score_input=True, top_k=0)) + assert resp.input_scoring is not None + assert len(resp.input_scoring.token_logprobs) == 2 + assert resp.input_scoring.token_logprobs[0].logprob is None # first token + assert resp.input_scoring.token_logprobs[1].token_id == 2 + # output segment lands in logprobs, never merged with the input segment + assert resp.logprobs is not None + assert len(resp.logprobs) == 1 + assert resp.logprobs[0].token == "c" + + @pytest.mark.anyio + async def test_usage_and_finish_reason(self): + meta = _meta( + prompt_tokens=4, completion_tokens=6, finish_reason={"type": "length"} + ) + t, _ = _make_transport({"text": "out", "meta_info": meta}) + resp = await t.arun(Request(input="x")) + assert resp.usage == UsageStats( + input_tokens=4, output_tokens=6, total_tokens=10 + ) + assert resp.finish_reasons == ("length",) + assert resp.texts == ("out",) + + @pytest.mark.anyio + async def test_usage_none_when_prompt_tokens_absent(self): + # echo=False path (no radix guard); _parse_usage returns None when the + # server omitted the token counts — absence is not zeros. + meta = { + "completion_tokens": 1, + "finish_reason": "stop", + "output_token_logprobs": [[-0.1, 1, " A"]], + } + t, _ = _make_transport({"text": " A", "meta_info": meta}) + resp = await t.arun(Request(input="p", return_logprobs=True)) + assert resp.usage is None + + +# =================================================================== +# Top-k triple parsing +# =================================================================== +class TestTopKParsing: + @pytest.mark.anyio + async def test_top_k_preserves_token_id_and_normalizes_text(self): + meta = _meta( + output_token_logprobs=[[-0.7, 100, "ĠB"]], + output_top_logprobs=[[[-0.7, 100, "ĠB"], [-1.2, 101, "Ċ"]]], + ) + t, _ = _make_transport({"text": " B", "meta_info": meta}) + resp = await t.arun(Request(input="p", return_logprobs=True, top_k=2)) + assert resp.top_logprobs == ( + ( + TopKEntry(token=" B", logprob=-0.7, token_id=100), + TopKEntry(token="\n", logprob=-1.2, token_id=101), + ), + ) + + @pytest.mark.anyio + async def test_empty_per_token_top_entry_becomes_empty_tuple(self): + meta = _meta( + prompt_tokens=2, + input_token_logprobs=[[None, 1, "Q"], [-0.2, 2, " B"]], + output_token_logprobs=[[-0.3, 3, " g"]], + output_top_logprobs=[None, [[-0.3, 3, " g"]]], + ) + t, _ = _make_transport({"text": " g", "meta_info": meta}) + resp = await t.arun(Request(input="Q B", score_input=True, top_k=1)) + assert resp.top_logprobs == ( + (), + (TopKEntry(token=" g", logprob=-0.3, token_id=3),), + ) + + @pytest.mark.anyio + async def test_top_logprobs_none_when_absent(self): + meta = _meta(output_token_logprobs=[[-0.1, 1, " A"]]) + t, _ = _make_transport({"text": " A", "meta_info": meta}) + resp = await t.arun(Request(input="p", return_logprobs=True)) + assert resp.top_logprobs is None + + @pytest.mark.anyio + async def test_none_token_text_in_top_raises(self): + """A top-k entry with no token text (no detokenization) fails loud.""" + meta = _meta( + output_token_logprobs=[[-0.1, 1, " A"]], + output_top_logprobs=[[[-0.1, 1, None]]], + ) + t, _ = _make_transport({"text": "", "meta_info": meta}) + with pytest.raises(RuntimeError, match="no token text"): + await t.arun(Request(input="p", return_logprobs=True, top_k=1)) + + +# =================================================================== +# Guards: response shape, radix cache, empty logprobs +# =================================================================== +class TestGuards: + @pytest.mark.anyio + async def test_radix_cache_hit_with_score_input_raises(self): + meta = _meta( + prompt_tokens=3, + input_token_logprobs=[[None, 1, "a"]], # truncated (cache hit) + cached_tokens=2, + ) + t, _ = _make_transport({"text": "", "meta_info": meta}) + with pytest.raises(RuntimeError, match="radix prefix cache"): + await t.arun(Request(input="abc", score_input=True)) + + @pytest.mark.anyio + async def test_truncated_input_logprobs_raises(self): + # cached_tokens reads 0, but the returned count < prompt_tokens. + meta = _meta( + prompt_tokens=5, + input_token_logprobs=[[-0.1, 1, " star"]], + cached_tokens=0, + ) + t, _ = _make_transport({"text": "", "meta_info": meta}) + with pytest.raises(RuntimeError, match="partial echoed-input"): + await t.arun(Request(input="abcde", score_input=True)) + + @pytest.mark.anyio + async def test_missing_prompt_tokens_with_score_input_raises(self): + meta = {"completion_tokens": 1, "finish_reason": "stop"} + meta["input_token_logprobs"] = [[None, 1, "a"]] + t, _ = _make_transport({"text": "", "meta_info": meta}) + with pytest.raises(RuntimeError, match="omitted prompt_tokens"): + await t.arun(Request(input="a", score_input=True)) + + @pytest.mark.anyio + async def test_full_echoed_input_passes_guard(self): + meta = _meta( + prompt_tokens=2, + input_token_logprobs=[[None, 1, "a"], [-0.1, 2, " b"]], + cached_tokens=0, + ) + t, _ = _make_transport({"text": "", "meta_info": meta}) + resp = await t.arun(Request(input="ab", score_input=True)) + assert resp.input_scoring is not None + assert [tl.token for tl in resp.input_scoring.token_logprobs] == ["a", " b"] + + @pytest.mark.anyio + async def test_no_score_input_skips_radix_guard(self): + """Output-only logprobs (CMMLU shape) — cache truncation is irrelevant.""" + meta = _meta( + prompt_tokens=5, + cached_tokens=4, + output_token_logprobs=[[-0.1, 1, " A"]], + ) + t, _ = _make_transport({"text": " A", "meta_info": meta}) + resp = await t.arun(Request(input="p", return_logprobs=True)) + assert resp.logprobs is not None + assert resp.logprobs[0].token == " A" + + @pytest.mark.anyio + async def test_no_logprobs_channels_raises(self): + """Logprobs requested but every channel came back empty → fail loud.""" + t, _ = _make_transport({"text": "", "meta_info": _meta()}) + with pytest.raises(RuntimeError, match="no logprobs"): + await t.arun(Request(input="x", return_logprobs=True)) + + @pytest.mark.anyio + async def test_empty_echoed_input_raises(self): + # score_input with an empty echoed input passes the radix guard + # (0 == prompt_tokens) but still yields no logprobs at all. + meta = _meta(prompt_tokens=0, input_token_logprobs=[]) + t, _ = _make_transport({"text": "", "meta_info": meta}) + with pytest.raises(RuntimeError, match="no logprobs"): + await t.arun(Request(input="p", score_input=True)) + + @pytest.mark.anyio + async def test_missing_meta_info_raises(self): + t, _ = _make_transport({"text": "hi"}) + with pytest.raises(RuntimeError, match="missing meta_info"): + await t.arun(Request(input="x")) + + @pytest.mark.anyio + async def test_list_with_non_dict_element_raises(self): + t, _ = _make_transport(["nope"]) + with pytest.raises(RuntimeError, match="missing meta_info"): + await t.arun(Request(input="x")) + + @pytest.mark.anyio + async def test_empty_list_response_raises(self): + t, _ = _make_transport([]) + with pytest.raises(RuntimeError, match="missing meta_info"): + await t.arun(Request(input="x")) diff --git a/tests/unit/core/tasks/test_task.py b/tests/unit/core/tasks/test_task.py index 0b7f92a4..db89d32d 100644 --- a/tests/unit/core/tasks/test_task.py +++ b/tests/unit/core/tasks/test_task.py @@ -12,7 +12,7 @@ from datasets import DatasetDict as HFDatasetDict from sieval.core.datasets import Dataset -from sieval.core.models import ModelOutput +from sieval.core.models import Capability, CapabilityError from sieval.core.models.chat_model import ChatModel from sieval.core.models.gen_model import GenModel from sieval.core.models.sglang_gen_model import SglangGenModel @@ -38,37 +38,25 @@ def load(self, name_or_path, **kwargs) -> HFDatasetDict: class _MockChatModel(ChatModel): + """Construction-only mock: Task validation never invokes the wire.""" + def __init__(self): super().__init__(model="mock", api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - return ModelOutput(model=self.meta(), texts=["ok"]) - - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class _MockGenModel(GenModel): + """Construction-only mock: Task validation never invokes the wire.""" + def __init__(self): super().__init__(model="mock-gen", api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - return ModelOutput(model=self.meta(), texts=["ok"]) - - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class _MockSglangGenModel(SglangGenModel): + """Construction-only mock: Task validation never invokes the wire.""" + def __init__(self): super().__init__(model="mock-sglang", api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - return ModelOutput(model=self.meta(), texts=["ok"]) - - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError - class _ConcreteTask(Task): """Fully concrete Task with no model_type restriction.""" @@ -99,6 +87,18 @@ class _GenOnlyTask(_ConcreteTask): model_type = "gen" +class _ScoringTask(_ConcreteTask): + """Declares an IR capability requirement (prompt-side scoring).""" + + requires = frozenset({Capability.InputScoring}) + + +class _TokenIdTask(_ConcreteTask): + """Requires native token ids (only the sglang transport supplies them).""" + + requires = frozenset({Capability.SampledLogprobsWithTokenIds}) + + # =================================================================== # name property # =================================================================== @@ -152,20 +152,50 @@ def test_no_model_type_restriction_accepts_both(self): def test_unrecognized_model_type_raises(self): """A model that is neither ChatModel nor GenModel should raise TypeError.""" - from sieval.core.models.model import Model, ModelOutput + from sieval.core.models.model import Model class _CustomModel(Model): - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - return ModelOutput(model=self.meta(), texts=["ok"]) - - async def _alogprobs_impl(self, prompt, **kwargs) -> ModelOutput: - raise NotImplementedError + """Bare Model subclass: no transport, no kind.""" custom = _CustomModel(model="custom", api_key="fake") with pytest.raises(TypeError, match="requires a ChatModel or GenModel"): _ChatOnlyTask(_SimpleDataset(), custom) +# =================================================================== +# requires (IR capability gate) +# =================================================================== +class TestRequiresCapabilityGate: + def test_no_requires_accepts_any_model(self): + """The default empty `requires` gates nothing.""" + _ConcreteTask(_SimpleDataset(), _MockChatModel()) + _ConcreteTask(_SimpleDataset(), _MockGenModel()) + + def test_input_scoring_task_with_gen_model_ok(self): + """GenModel's completions transport supplies InputScoring.""" + _ScoringTask(_SimpleDataset(), _MockGenModel()) + + def test_input_scoring_task_with_sglang_model_ok(self): + _ScoringTask(_SimpleDataset(), _MockSglangGenModel()) + + def test_input_scoring_task_with_chat_model_raises(self): + """Chat completions cannot score the prompt — fail loud at construction.""" + with pytest.raises(CapabilityError, match="InputScoring"): + _ScoringTask(_SimpleDataset(), _MockChatModel()) + + def test_token_id_task_with_sglang_model_ok(self): + """Only the sglang transport populates token ids.""" + _TokenIdTask(_SimpleDataset(), _MockSglangGenModel()) + + def test_token_id_task_with_gen_model_raises(self): + with pytest.raises(CapabilityError, match="SampledLogprobsWithTokenIds"): + _TokenIdTask(_SimpleDataset(), _MockGenModel()) + + def test_token_id_task_with_chat_model_raises(self): + with pytest.raises(CapabilityError, match="SampledLogprobsWithTokenIds"): + _TokenIdTask(_SimpleDataset(), _MockChatModel()) + + # =================================================================== # make_context # =================================================================== diff --git a/tests/unit/tasks/conftest.py b/tests/unit/tasks/conftest.py new file mode 100644 index 00000000..030b9d9c --- /dev/null +++ b/tests/unit/tasks/conftest.py @@ -0,0 +1,31 @@ +"""Task-test collection tweaks for mutation testing. + +The ``test_import_does_not_pull_`` convention spawns a FRESH python +interpreter (``subprocess.run([sys.executable, "-c", "import sieval..."])``) +to prove an optional dependency is lazy-imported. Under mutation testing the +subprocess inherits the environment pointing at mutmut's instrumented +``mutants/`` tree, whose trampolines require mutmut's in-process runtime +config — unavailable in a bare interpreter — so the import crashes for +reasons unrelated to what the test verifies. These tests also exercise no +``sieval/core`` mutants, so skipping them under mutation testing loses no +kill power. ``MUTANT_UNDER_TEST`` is set by mutmut for every pytest run it +drives and is absent otherwise. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + +import os + +import pytest + + +def pytest_collection_modifyitems(items): + if not os.environ.get("MUTANT_UNDER_TEST"): + return + skip = pytest.mark.skip( + reason="fresh-interpreter lazy-import check cannot run against " + "mutmut's instrumented tree" + ) + for item in items: + if item.name.startswith("test_import_does_not_pull_"): + item.add_marker(skip) diff --git a/tests/unit/tasks/test_arc_challenge_kshot_clp.py b/tests/unit/tasks/test_arc_challenge_kshot_clp.py index 472a17c0..dabb6c09 100644 --- a/tests/unit/tasks/test_arc_challenge_kshot_clp.py +++ b/tests/unit/tasks/test_arc_challenge_kshot_clp.py @@ -8,8 +8,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TopKEntry from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import EvalMode, TaskContext from sieval.core.tasks.meta import get_task_meta from sieval.datasets.arc_challenge import ( @@ -17,38 +18,34 @@ ARCChallengeDatasetSample, ) from sieval.tasks.arc_challenge_kshot_clp import ARCChallengeFewShotClpTask +from tests.conftest import HandlerTransport class _TopLogprobsGenModel(GenModel): """Returns a fixed next-token top_logprobs map; records the prompt + echo.""" def __init__(self, top: dict[str, float]): - super().__init__(model="mock-gen", api_key="fake") self._top = top self.prompts: list[str] = [] self.echo_flags: list[bool] = [] + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (max_tokens, logprobs, temperature, kwargs) - self.prompts.append(prompt) - self.echo_flags.append(echo) - return ModelOutput( - model=self.meta(), - texts=["A"], - top_logprobs=[dict(self._top)], + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + self.prompts.append(req.input) + self.echo_flags.append(req.score_input) + return Response( + texts=("A",), + top_logprobs=( + tuple(TopKEntry(token=t, logprob=lp) for t, lp in self._top.items()), + ), ) diff --git a/tests/unit/tasks/test_arc_challenge_kshot_ppl.py b/tests/unit/tasks/test_arc_challenge_kshot_ppl.py index afba654b..a3858c25 100644 --- a/tests/unit/tasks/test_arc_challenge_kshot_ppl.py +++ b/tests/unit/tasks/test_arc_challenge_kshot_ppl.py @@ -8,8 +8,15 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ( + ModelOutput, + Request, + Response, + TokenLogprob, + UsageStats, +) from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import EvalMode, TaskContext from sieval.core.tasks.meta import get_task_meta from sieval.datasets.arc_challenge import ( @@ -18,6 +25,7 @@ ) from sieval.tasks._arc import ARC_UNCOND_CONTEXT, echoed_logprob from sieval.tasks.arc_challenge_kshot_ppl import ARCChallengeFewShotPplTask +from tests.conftest import HandlerTransport class _ScriptedGenModel(GenModel): @@ -30,37 +38,33 @@ class _ScriptedGenModel(GenModel): """ def __init__(self, scores: dict[str, tuple[float, float]]): - super().__init__(model="mock-gen", api_key="fake") self._scores = scores self.prompts: list[str] = [] + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (max_tokens, temperature, kwargs) - assert echo is True - assert logprobs == 0 # ppl requests no top-k (matches hellaswag sibling) - self.prompts.append(prompt) - option = prompt.split("Answer:")[-1].strip() + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + assert req.score_input is True + assert req.top_k == 0 # ppl requests no top-k (matches hellaswag sibling) + self.prompts.append(req.input) + option = req.input.split("Answer:")[-1].strip() cond_lp, uncond_lp = self._scores[option] - value = uncond_lp if prompt.startswith(ARC_UNCOND_CONTEXT) else cond_lp - return ModelOutput( - model=self.meta(), - texts=[""], - logprobs_tokens=["_ctx", "_opt", "_generated"], - logprobs=[None, value, -99.0], - usage={"input_tokens": 2, "output_tokens": 1, "total_tokens": 3}, + value = uncond_lp if req.input.startswith(ARC_UNCOND_CONTEXT) else cond_lp + return Response( + texts=("",), + logprobs=( + TokenLogprob(token="_ctx", logprob=None), + TokenLogprob(token="_opt", logprob=value), + TokenLogprob(token="_generated", logprob=-99.0), + ), + usage=UsageStats(input_tokens=2, output_tokens=1, total_tokens=3), ) diff --git a/tests/unit/tasks/test_arc_easy_kshot_clp.py b/tests/unit/tasks/test_arc_easy_kshot_clp.py index 0480ec79..f1c935bd 100644 --- a/tests/unit/tasks/test_arc_easy_kshot_clp.py +++ b/tests/unit/tasks/test_arc_easy_kshot_clp.py @@ -8,36 +8,34 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TopKEntry from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import EvalMode, TaskContext from sieval.core.tasks.meta import get_task_meta from sieval.datasets.arc_easy import ARCEasyDataset, ARCEasyDatasetSample from sieval.tasks.arc_easy_kshot_clp import ARCEasyFewShotClpTask +from tests.conftest import HandlerTransport class _TopLogprobsGenModel(GenModel): def __init__(self, top: dict[str, float]): - super().__init__(model="mock-gen", api_key="fake") self._top = top + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput( - model=self.meta(), texts=["B"], top_logprobs=[dict(self._top)] + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + return Response( + texts=("B",), + top_logprobs=( + tuple(TopKEntry(token=t, logprob=lp) for t, lp in self._top.items()), + ), ) diff --git a/tests/unit/tasks/test_arc_easy_kshot_ppl.py b/tests/unit/tasks/test_arc_easy_kshot_ppl.py index f387bdf5..943027bf 100644 --- a/tests/unit/tasks/test_arc_easy_kshot_ppl.py +++ b/tests/unit/tasks/test_arc_easy_kshot_ppl.py @@ -8,44 +8,41 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TokenLogprob, UsageStats from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import EvalMode, TaskContext from sieval.core.tasks.meta import get_task_meta from sieval.datasets.arc_easy import ARCEasyDataset, ARCEasyDatasetSample from sieval.tasks._arc import ARC_UNCOND_CONTEXT from sieval.tasks.arc_easy_kshot_ppl import ARCEasyFewShotPplTask +from tests.conftest import HandlerTransport class _ScriptedGenModel(GenModel): def __init__(self, scores: dict[str, tuple[float, float]]): - super().__init__(model="mock-gen", api_key="fake") self._scores = scores + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (max_tokens, logprobs, temperature, echo, kwargs) - option = prompt.split("Answer:")[-1].strip() + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + option = req.input.split("Answer:")[-1].strip() cond_lp, uncond_lp = self._scores[option] - value = uncond_lp if prompt.startswith(ARC_UNCOND_CONTEXT) else cond_lp - return ModelOutput( - model=self.meta(), - texts=[""], - logprobs_tokens=["_", "_"], - logprobs=[None, value], - usage={"input_tokens": 2, "output_tokens": 0, "total_tokens": 2}, + value = uncond_lp if req.input.startswith(ARC_UNCOND_CONTEXT) else cond_lp + return Response( + texts=("",), + logprobs=( + TokenLogprob(token="_", logprob=None), + TokenLogprob(token="_", logprob=value), + ), + usage=UsageStats(input_tokens=2, output_tokens=0, total_tokens=2), ) diff --git a/tests/unit/tasks/test_c_eval_kshot_clp.py b/tests/unit/tasks/test_c_eval_kshot_clp.py index f0272b6d..80c554f4 100644 --- a/tests/unit/tasks/test_c_eval_kshot_clp.py +++ b/tests/unit/tasks/test_c_eval_kshot_clp.py @@ -7,39 +7,43 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, TopKEntry from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.c_eval import CEvalDataset, CEvalDatasetSample from sieval.tasks.c_eval_kshot_clp import CEvalFewShotCLPTask +from tests.conftest import HandlerTransport class _ScriptedGenModel(GenModel): """Returns a fixed next-token top_logprobs map for every alogprobs call.""" def __init__(self, top_logprobs: dict[str, float]): - super().__init__(model="mock-gen", api_key="fake") self._top_logprobs = top_logprobs self.calls = 0 self.prompts: list[str] = [] + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt, **kwargs): # pragma: no cover - raise AssertionError("clp task must not call agenerate") - - async def _alogprobs_impl( - self, - prompt, - *, - max_tokens=1, - logprobs=100, - echo=False, - temperature=0.0, - **kwargs, - ) -> ModelOutput: + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): # pragma: no cover + raise AssertionError("clp task must not call agenerate") + assert isinstance(req.input, str) self.calls += 1 - self.prompts.append(prompt) - return ModelOutput( - model=self.meta(), texts=[""], top_logprobs=[dict(self._top_logprobs)] + self.prompts.append(req.input) + return Response( + texts=("",), + top_logprobs=( + tuple( + TopKEntry(token=t, logprob=lp) + for t, lp in self._top_logprobs.items() + ), + ), ) @@ -116,8 +120,8 @@ async def test_postprocess_raises_when_option_token_missing(): @pytest.mark.anyio async def test_infer_does_not_generate(): - # _agenerate_impl raises if touched; reaching the assert proves only - # alogprobs ran and top_logprobs are returned. + # The stub's generation branch raises if touched; reaching the assert + # proves only alogprobs ran and top_logprobs are returned. model = _ScriptedGenModel({"A": -0.1, "B": -1.0, "C": -1.0, "D": -1.0}) task = _task(model) raw = _sample("law", "A") diff --git a/tests/unit/tasks/test_cmmlu_kshot_clp.py b/tests/unit/tasks/test_cmmlu_kshot_clp.py index 5f46ea3e..ee2c1bbd 100644 --- a/tests/unit/tasks/test_cmmlu_kshot_clp.py +++ b/tests/unit/tasks/test_cmmlu_kshot_clp.py @@ -8,8 +8,15 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ( + ModelOutput, + Request, + Response, + TokenLogprob, + TopKEntry, +) from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.cmmlu import CMMLUDataset, CMMLUDatasetSample from sieval.tasks.cmmlu_kshot_clp import ( @@ -18,36 +25,36 @@ CMMLU_SUBJECT_DISPLAY_NAMES, CMMLUFewShotClpTask, ) +from tests.conftest import HandlerTransport class _DummyGenModel(GenModel): def __init__(self): - super().__init__(model="mock-gen", api_key="fake") self.logprob_prompts: list[str] = [] + super().__init__(model="mock-gen", api_key="fake") - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (max_tokens, logprobs, temperature, kwargs) - self.logprob_prompts.append(prompt) - assert echo is False - return ModelOutput( - model=self.meta(), - texts=["B"], - logprobs_tokens=["B"], - logprobs=[-0.1], - top_logprobs=[{"A": -1.0, "B": -0.1, "C": -2.0, "D": -3.0}], + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + self.logprob_prompts.append(req.input) + assert req.score_input is False + return Response( + texts=("B",), + logprobs=(TokenLogprob(token="B", logprob=-0.1),), + top_logprobs=( + ( + TopKEntry(token="A", logprob=-1.0), + TopKEntry(token="B", logprob=-0.1), + TopKEntry(token="C", logprob=-2.0), + TopKEntry(token="D", logprob=-3.0), + ), + ), ) diff --git a/tests/unit/tasks/test_gsm8k_0shot_gen.py b/tests/unit/tasks/test_gsm8k_0shot_gen.py index 9cf78d27..2c09813f 100644 --- a/tests/unit/tasks/test_gsm8k_0shot_gen.py +++ b/tests/unit/tasks/test_gsm8k_0shot_gen.py @@ -7,8 +7,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.chat_model import ChatModel +from sieval.core.models.transports import OpenAIChatTransport from sieval.core.tasks import TaskContext from sieval.datasets.gsm8k import GSM8KDataset, GSM8KDatasetSample from sieval.tasks.gsm8k_0shot_gen import ( @@ -16,31 +17,21 @@ GSM8KZeroShotGenTask, _gold_answer, ) +from tests.conftest import HandlerTransport class _CapturingChatModel(ChatModel): def __init__(self, text: str): - super().__init__(model="mock-chat", api_key="fake") - self.last_kwargs: dict[str, object] = {} + self.last_req: Request | None = None self._text = text + super().__init__(model="mock-chat", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=[self._text]) - - async def _alogprobs_impl( - self, - prompt, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=(self._text,)) def _sample(answer: str = "Solution.\n#### 42") -> GSM8KDatasetSample: @@ -176,5 +167,8 @@ async def test_infer_injects_no_decode_params(): _sample(), TaskContext(sample_id=0, raw_sample=_sample()) ) await task.infer(pre, TaskContext(sample_id=0, raw_sample=_sample())) - for forbidden in ("temperature", "top_p", "max_tokens", "n", "stop"): - assert forbidden not in model.last_kwargs + req = model.last_req + assert req is not None + # No decode params injected: default sampling (n=1, everything else unset). + assert req.sampling == SamplingParams() + assert req.extra_wire_params is None diff --git a/tests/unit/tasks/test_gsm8k_kshot_base_gen.py b/tests/unit/tasks/test_gsm8k_kshot_base_gen.py index 2b3e1ae5..552b3e9d 100644 --- a/tests/unit/tasks/test_gsm8k_kshot_base_gen.py +++ b/tests/unit/tasks/test_gsm8k_kshot_base_gen.py @@ -7,8 +7,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.gsm8k import GSM8KDataset, GSM8KDatasetSample from sieval.tasks.gsm8k_kshot_base_gen import ( @@ -17,30 +18,22 @@ _extract_answer, _extract_flexible_match, ) +from tests.conftest import HandlerTransport class _CapturingGenModel(GenModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-gen", api_key="fake") - self.last_kwargs: dict[str, object] = {} - - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=[" Work shown.\n#### 42"]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=(" Work shown.\n#### 42",)) def _sample(answer: str = "Solution.\n#### 42") -> GSM8KDatasetSample: @@ -75,7 +68,11 @@ async def test_infer_only_forwards_prompt_coupled_stop(): await task.infer("prompt", TaskContext(sample_id=0, raw_sample=_sample())) - assert model.last_kwargs == {"stop": list(STOP_SEQUENCES)} + req = model.last_req + assert req is not None + # Only the prompt-coupled stop is forwarded — no other sampling params. + assert req.sampling == SamplingParams(stop=STOP_SEQUENCES) + assert req.extra_wire_params is None @pytest.mark.anyio diff --git a/tests/unit/tasks/test_hellaswag_kshot_ppl.py b/tests/unit/tasks/test_hellaswag_kshot_ppl.py index 87bc0d95..71e1889d 100644 --- a/tests/unit/tasks/test_hellaswag_kshot_ppl.py +++ b/tests/unit/tasks/test_hellaswag_kshot_ppl.py @@ -7,8 +7,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TokenLogprob from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.hellaswag import HellaSwagDataset from sieval.tasks.hellaswag_kshot_ppl import ( @@ -16,6 +17,7 @@ _argmax, _continuation_logprob, ) +from tests.conftest import HandlerTransport QUERY = "Activity label: a person" # Chosen so raw-LL argmax (acc) and length-normalized argmax (acc_norm) DIVERGE. @@ -50,9 +52,7 @@ def _doc(activity: str, ctx_a: str, ctx_b: str, endings: list[str], label: str) RENDERED_B = "Driving: He starts the car. The man drives away." -def _crafted_output( - model: GenModel, prompt: str, choice: str, ll: float -) -> ModelOutput: +def _crafted_response(prompt: str, choice: str, ll: float) -> Response: """Echoed logprobs for *prompt* (= context + ' ' + choice) summing to *ll*. Layout stresses the helper: a non-empty leading BOS token, the (arbitrary, @@ -65,43 +65,38 @@ def _crafted_output( part1, part2 = cont[:mid], cont[mid:] tokens = ["", context, part1, part2, ""] logprobs = [None, -42.0, ll / 2, ll / 2, -99.0] - return ModelOutput( - model=model.meta(), - texts=[""], - logprobs_tokens=tokens, - logprobs=logprobs, + return Response( + texts=("",), + logprobs=tuple( + TokenLogprob(token=t, logprob=lp) + for t, lp in zip(tokens, logprobs, strict=True) + ), ) class _PPLMockModel(GenModel): def __init__(self, choices: list[str], lls: list[float]): - super().__init__(model="mock-gen", api_key="fake") self._ll_by_choice: dict[str, float] = dict(zip(choices, lls, strict=True)) self.calls: list[str] = [] self.echos: list[bool] = [] self.logprobs_args: list[int] = [] + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (max_tokens, temperature, kwargs) - self.calls.append(prompt) - self.echos.append(echo) - self.logprobs_args.append(logprobs) + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + self.calls.append(req.input) + self.echos.append(req.score_input) + self.logprobs_args.append(req.top_k) # CHOICES are mutually non-suffix, so exactly one endswith-matches. - choice = next(c for c in self._ll_by_choice if prompt.endswith(c)) - return _crafted_output(self, prompt, choice, self._ll_by_choice[choice]) + choice = next(c for c in self._ll_by_choice if req.input.endswith(c)) + return _crafted_response(req.input, choice, self._ll_by_choice[choice]) def _make_task( diff --git a/tests/unit/tasks/test_hendrycks_math_kshot_base_gen.py b/tests/unit/tasks/test_hendrycks_math_kshot_base_gen.py index 3e8a1faf..a1c7eb49 100644 --- a/tests/unit/tasks/test_hendrycks_math_kshot_base_gen.py +++ b/tests/unit/tasks/test_hendrycks_math_kshot_base_gen.py @@ -7,8 +7,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.hendrycks_math import ( HendrycksMathDataset, @@ -18,34 +19,24 @@ N_SHOT, HendrycksMathFewShotBaseGenTask, ) +from tests.conftest import HandlerTransport _FA = "\nFinal Answer: The final answer is ${}$. I hope it is correct." class _CapturingGenModel(GenModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-gen", api_key="fake") - self.last_kwargs: dict[str, object] = {} - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput( - model=self.meta(), texts=[f"$\\boxed{{16}}${_FA.format('16')}"] + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES ) - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=(f"$\\boxed{{16}}${_FA.format('16')}",)) def _sample( @@ -93,9 +84,11 @@ async def test_preprocess_is_deepseek_minerva_prompt(): async def test_infer_forwards_deepseek_stop_only(): task, model = _task() await task.infer("prompt", TaskContext(sample_id=0, raw_sample=_sample())) - assert model.last_kwargs == {"stop": ["\nProblem:"]} - assert "temperature" not in model.last_kwargs - assert "max_tokens" not in model.last_kwargs + req = model.last_req + assert req is not None + # Only the DeepSeek stop is forwarded — no temperature / max_tokens / etc. + assert req.sampling == SamplingParams(stop=("\nProblem:",)) + assert req.extra_wire_params is None @pytest.mark.anyio diff --git a/tests/unit/tasks/test_human_eval_0shot_base_gen.py b/tests/unit/tasks/test_human_eval_0shot_base_gen.py index 00733f76..16ca57a0 100644 --- a/tests/unit/tasks/test_human_eval_0shot_base_gen.py +++ b/tests/unit/tasks/test_human_eval_0shot_base_gen.py @@ -7,39 +7,31 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.human_eval import HumanEvalDataset, HumanEvalDatasetSample from sieval.tasks.human_eval_0shot_base_gen import ( STOP_SEQUENCES, HumanEvalZeroShotBaseGenTask, ) +from tests.conftest import HandlerTransport class _CapturingGenModel(GenModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-gen", api_key="fake") - self.last_prompt = "" - self.last_kwargs: dict[str, object] = {} - - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - self.last_prompt = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=[" return x + 1"]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=(" return x + 1",)) def _sample() -> HumanEvalDatasetSample: @@ -75,12 +67,13 @@ async def test_preprocess_and_infer_use_base_completion_prompt(): await task.infer(pre, TaskContext(sample_id=0, raw_sample=raw)) assert pre == raw["prompt"] - assert model.last_prompt == raw["prompt"] - assert model.last_kwargs["n"] == 2 - assert model.last_kwargs["stop"] == ["\nclass"] + req = model.last_req + assert req is not None + assert req.input == raw["prompt"] # Decoding params (max_tokens, temperature, top_p) are owned by the # model config / infer_args, never injected by the task layer. - assert "max_tokens" not in model.last_kwargs + assert req.sampling == SamplingParams(stop=("\nclass",), n=2) + assert req.extra_wire_params is None finally: await task.shutdown() diff --git a/tests/unit/tasks/test_livecodebench_code_generation_kshot_base_gen.py b/tests/unit/tasks/test_livecodebench_code_generation_kshot_base_gen.py index 00942f7c..d2495ac8 100644 --- a/tests/unit/tasks/test_livecodebench_code_generation_kshot_base_gen.py +++ b/tests/unit/tasks/test_livecodebench_code_generation_kshot_base_gen.py @@ -12,8 +12,9 @@ get_base_model_question_template_answer, get_base_model_target_block, ) -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.livecodebench_code_generation import LiveCodeBenchDataset from sieval.tasks.livecodebench_code_generation_kshot_base_gen import ( @@ -21,33 +22,25 @@ STOP_SEQUENCES, LiveCodeBenchCodeGenerationFewShotBaseGenTask, ) +from tests.conftest import HandlerTransport _STARTER = "class Solution:\n def solve(self) -> int:\n " class _CapturingGenModel(GenModel): def __init__(self, texts: list[str] | None = None): - super().__init__(model="mock-gen", api_key="fake") - self.last_kwargs: dict[str, object] = {} + self.last_req: Request | None = None self._texts = texts if texts is not None else ["print(1)"] + super().__init__(model="mock-gen", api_key="fake") + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=list(self._texts)) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=tuple(self._texts)) def _raw(starter_code: str = "") -> dict: @@ -168,10 +161,11 @@ async def test_infer_forwards_only_stop_and_n_not_decoding_params(): finally: await task.shutdown() - assert model.last_kwargs["stop"] == ["###"] - assert model.last_kwargs["n"] == 4 - assert "max_tokens" not in model.last_kwargs - assert "temperature" not in model.last_kwargs + req = model.last_req + assert req is not None + # Exactly stop + n — no max_tokens / temperature injected by the task. + assert req.sampling == SamplingParams(stop=("###",), n=4) + assert req.extra_wire_params is None @pytest.mark.anyio diff --git a/tests/unit/tasks/test_mbpp_kshot_base_gen.py b/tests/unit/tasks/test_mbpp_kshot_base_gen.py index 9ec0b5a0..d2cdae90 100644 --- a/tests/unit/tasks/test_mbpp_kshot_base_gen.py +++ b/tests/unit/tasks/test_mbpp_kshot_base_gen.py @@ -1,36 +1,34 @@ +"""Unit tests for the MBPP k-shot base-gen task. + +AI-Generated Code - Claude Fable 5 (Anthropic) +""" + import pytest from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.mbpp import MBPPDataset, MBPPDatasetSample from sieval.tasks.mbpp_kshot_base_gen import MBPPFewShotBaseGenTask +from tests.conftest import HandlerTransport class _CapturingGenModel(GenModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-gen", api_key="fake") - self.last_kwargs: dict[str, object] = {} - - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=["def f():\n pass\n[DONE]"]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=("def f():\n pass\n[DONE]",)) def _sample() -> MBPPDatasetSample: @@ -99,10 +97,11 @@ async def test_infer_forwards_n_and_stop_but_not_decoding_params(): await task.shutdown() assert result.texts == ["def f():\n pass\n[DONE]"] - assert model.last_kwargs["n"] == 3 - assert model.last_kwargs["stop"] == ["[DONE]"] + req = model.last_req + assert req is not None # Decoding params stay in the model layer; the task must not inject them. - assert "max_tokens" not in model.last_kwargs + assert req.sampling == SamplingParams(stop=("[DONE]",), n=3) + assert req.extra_wire_params is None def test_k_above_n_raises(): diff --git a/tests/unit/tasks/test_mmlu_0shot_gen.py b/tests/unit/tasks/test_mmlu_0shot_gen.py index c4cd711a..f0725ce3 100644 --- a/tests/unit/tasks/test_mmlu_0shot_gen.py +++ b/tests/unit/tasks/test_mmlu_0shot_gen.py @@ -7,33 +7,24 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response from sieval.core.models.chat_model import ChatModel +from sieval.core.models.transports import OpenAIChatTransport from sieval.core.tasks import TaskContext from sieval.datasets.mmlu import MMLUDataset, MMLUDatasetSample from sieval.tasks.mmlu_0shot_gen import MMLUZeroShotGenTask +from tests.conftest import HandlerTransport class _StubChatModel(ChatModel): def __init__(self): super().__init__(model="mock-chat", api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=["Answer: A"]) - - async def _alogprobs_impl( - self, - prompt, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) + + async def _stub_arun(self, req: Request) -> Response: + return Response(texts=("Answer: A",)) def _sample(subject: str = "anatomy", answer: int = 0) -> MMLUDatasetSample: diff --git a/tests/unit/tasks/test_mmlu_kshot_clp.py b/tests/unit/tasks/test_mmlu_kshot_clp.py index 8dff367a..3f079516 100644 --- a/tests/unit/tasks/test_mmlu_kshot_clp.py +++ b/tests/unit/tasks/test_mmlu_kshot_clp.py @@ -7,8 +7,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TopKEntry from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext from sieval.datasets.mmlu import MMLUDataset, MMLUDatasetSample from sieval.tasks.mmlu_kshot_clp import ( @@ -17,6 +18,7 @@ _format_example, _format_subject, ) +from tests.conftest import HandlerTransport class _ScriptedGenModel(GenModel): @@ -26,34 +28,34 @@ class _ScriptedGenModel(GenModel): """ def __init__(self, winner: str = "A", drop: str | None = None): - super().__init__(model="mock-gen", api_key="fake") self._winner = winner self._drop = drop self.prompts: list[str] = [] self.call_count = 0 + super().__init__(model="mock-gen", api_key="fake") - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - self.prompts.append(prompt) + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + assert isinstance(req.input, str) + self.prompts.append(req.input) self.call_count += 1 - assert echo is False # clp reads the next-token distribution, no echo - dist = { - f" {label}": (-0.1 if label == self._winner else -5.0) + # clp reads the next-token distribution, no echo + assert req.score_input is False + dist = tuple( + TopKEntry( + token=f" {label}", + logprob=-0.1 if label == self._winner else -5.0, + ) for label in CHOICES if label != self._drop - } - return ModelOutput(model=self.meta(), texts=[""], top_logprobs=[dist]) + ) + return Response(texts=("",), top_logprobs=(dist,)) def _sample( diff --git a/tests/unit/tasks/test_mmmlu_kshot_clp.py b/tests/unit/tasks/test_mmmlu_kshot_clp.py index 0ac98fd3..f553e95a 100644 --- a/tests/unit/tasks/test_mmmlu_kshot_clp.py +++ b/tests/unit/tasks/test_mmmlu_kshot_clp.py @@ -9,8 +9,9 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, TokenLogprob, TopKEntry from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks import TaskContext, TaskStageOutput from sieval.datasets.mmmlu import MMMLUDataset, MMMLUDatasetSample from sieval.tasks.mmmlu_kshot_clp import ( @@ -18,6 +19,7 @@ MMMLUKShotClpTask, OfficialScores, ) +from tests.conftest import HandlerTransport _FinalCtx = TaskContext[ MMMLUDatasetSample, str, TaskStageOutput[OfficialScores], str, Feedback @@ -25,30 +27,26 @@ class _TopLogprobGenModel(GenModel): - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = (prompt, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput( - model=self.meta(), - texts=[" C"], - finish_reasons=["length"], - logprobs_tokens=[" C"], - logprobs=[-0.1], - top_logprobs=[ - {" A": -3.0, " B": -2.0, " C": -0.1, " D": -4.0}, - ], + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + if not (req.return_logprobs or req.score_input): + return Response(texts=("",)) + return Response( + texts=(" C",), + finish_reasons=("length",), + logprobs=(TokenLogprob(token=" C", logprob=-0.1),), + top_logprobs=( + ( + TopKEntry(token=" A", logprob=-3.0), + TopKEntry(token=" B", logprob=-2.0), + TopKEntry(token=" C", logprob=-0.1), + TopKEntry(token=" D", logprob=-4.0), + ), + ), ) diff --git a/tests/unit/tasks/test_openbookqa_kshot_gen.py b/tests/unit/tasks/test_openbookqa_kshot_gen.py index c940895c..7cb6e9f4 100644 --- a/tests/unit/tasks/test_openbookqa_kshot_gen.py +++ b/tests/unit/tasks/test_openbookqa_kshot_gen.py @@ -8,25 +8,29 @@ from datasets import DatasetDict as HFDatasetDict from sieval.community.openbookqa import OBQA_PROMPT_TEMPLATE -from sieval.core.models import ModelOutput +from sieval.core.models import ModelOutput, Request, Response, SamplingParams from sieval.core.models.chat_model import ChatModel +from sieval.core.models.transports import OpenAIChatTransport from sieval.core.tasks import TaskContext from sieval.datasets.openbookqa import OpenBookQADataset, OpenBookQADatasetSample from sieval.tasks.openbookqa_kshot_gen import ( STOP_SEQUENCES, OpenBookQAFewShotGenTask, ) +from tests.conftest import HandlerTransport class _CapturingChatModel(ChatModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-chat", api_key="fake") - self.last_kwargs: dict[str, object] = {} - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=["The answer is A."]) + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) + + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=("The answer is A.",)) def _sample(stem: str, answer_key: str = "A") -> OpenBookQADatasetSample: @@ -127,8 +131,11 @@ async def test_infer_does_not_forward_decoding_params(): TaskContext(sample_id=0, raw_sample=raw), ) - for forbidden in ("temperature", "top_p", "max_tokens", "n", "stop"): - assert forbidden not in model.last_kwargs + req = model.last_req + assert req is not None + # No decoding params forwarded: default sampling (n=1, everything else unset). + assert req.sampling == SamplingParams() + assert req.extra_wire_params is None def test_stop_sequences_pinned(): @@ -147,7 +154,8 @@ async def test_infer_bounds_generation_at_kshot_but_not_zero_shot(): [{"role": "user", "content": "x"}], TaskContext(sample_id=0, raw_sample=_sample("q-test")), ) - assert model_k.last_kwargs.get("stop") == list(STOP_SEQUENCES) + assert model_k.last_req is not None + assert model_k.last_req.sampling == SamplingParams(stop=STOP_SEQUENCES) # k=0: no stop — preserves upstream 0-shot parity. model_0 = _CapturingChatModel() @@ -156,7 +164,8 @@ async def test_infer_bounds_generation_at_kshot_but_not_zero_shot(): [{"role": "user", "content": "x"}], TaskContext(sample_id=0, raw_sample=_sample("q-test")), ) - assert "stop" not in model_0.last_kwargs + assert model_0.last_req is not None + assert model_0.last_req.sampling == SamplingParams() @pytest.mark.anyio diff --git a/tests/unit/tasks/test_simpleqa_verified_0shot_gen.py b/tests/unit/tasks/test_simpleqa_verified_0shot_gen.py index 8dda1a32..946766e6 100644 --- a/tests/unit/tasks/test_simpleqa_verified_0shot_gen.py +++ b/tests/unit/tasks/test_simpleqa_verified_0shot_gen.py @@ -8,8 +8,9 @@ from datasets import DatasetDict as HFDatasetDict from sieval.community.simpleqa_verified import aggregate_metrics, parse_grade -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response from sieval.core.models.chat_model import ChatModel +from sieval.core.models.transports import OpenAIChatTransport from sieval.core.tasks import TaskContext from sieval.datasets.simpleqa_verified import ( SimpleQAVerifiedDataset, @@ -19,26 +20,23 @@ GradeFeedback, SimpleQAVerifiedZeroShotGenTask, ) +from tests.conftest import HandlerTransport class _ScriptedChatModel(ChatModel): - """ChatModel returning a fixed reply, recording the last agenerate kwargs.""" + """ChatModel returning a fixed reply, recording the last Request.""" def __init__(self, reply: str, model: str = "mock"): - super().__init__(model=model, api_key="fake") self._reply = reply - self.last_kwargs: dict[str, object] = {} + self.last_req: Request | None = None + super().__init__(model=model, api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=[self._reply]) + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _alogprobs_impl( - self, prompt, *, max_tokens=1, logprobs=5, echo=True, temperature=0.0, **kwargs - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + return Response(texts=(self._reply,)) def _sample() -> SimpleQAVerifiedDatasetSample: @@ -105,7 +103,9 @@ async def test_infer_forwards_n(): grader = _ScriptedChatModel(reply="A", model="grader") task = SimpleQAVerifiedZeroShotGenTask(dataset, model, grader=grader, n=3) await task.infer([{"role": "user", "content": "q"}], TaskContext(sample_id=0)) - assert model.last_kwargs.get("n") == 3 + assert model.last_req is not None + assert model.last_req.sampling is not None + assert model.last_req.sampling.n == 3 # --- feedback: grades each answer via the grader, records provenance --- diff --git a/tests/unit/tasks/test_theoremqa_kshot_base_gen.py b/tests/unit/tasks/test_theoremqa_kshot_base_gen.py index ea947623..cd052647 100644 --- a/tests/unit/tasks/test_theoremqa_kshot_base_gen.py +++ b/tests/unit/tasks/test_theoremqa_kshot_base_gen.py @@ -12,9 +12,11 @@ from datasets import Dataset as HFDataset from datasets import DatasetDict as HFDatasetDict -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response, SamplingParams from sieval.core.models.gen_model import GenModel +from sieval.core.models.transports import OpenAICompletionsTransport from sieval.core.tasks.context import TaskContext +from tests.conftest import HandlerTransport _TASK_MODULE = "sieval.tasks.theoremqa_kshot_base_gen" _DATASET_MODULE = "sieval.datasets.theoremqa" @@ -67,26 +69,19 @@ def _preserve_registries(): class _MockGenModel(GenModel): def __init__(self): + self.last_req: Request | None = None super().__init__(model="mock-gen", api_key="fake") - self.last_kwargs: dict[str, object] = {} - - async def _agenerate_impl(self, prompt: str, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=["The answer is 4"]) - - async def _alogprobs_impl( - self, - prompt: str, - *, - max_tokens: int = 1, - logprobs: int = 5, - echo: bool = True, - temperature: float = 0.0, - **kwargs, - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - raise NotImplementedError + + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport( + self._stub_arun, OpenAICompletionsTransport.CAPABILITIES + ) + + async def _stub_arun(self, req: Request) -> Response: + if req.return_logprobs or req.score_input: + raise NotImplementedError # the gen task never requests logprobs + self.last_req = req + return Response(texts=("The answer is 4",)) def _task_module(): @@ -158,7 +153,11 @@ async def test_infer_only_forwards_prompt_coupled_stop(): TaskContext(sample_id=0, raw_sample={"Question": "What is 2+2?"}), ) - assert model.last_kwargs == {"stop": task_module._STOP_TOKENS} + req = model.last_req + assert req is not None + # Only the prompt-coupled stop is forwarded — no other sampling params. + assert req.sampling == SamplingParams(stop=tuple(task_module._STOP_TOKENS)) + assert req.extra_wire_params is None @pytest.mark.anyio From 5006b0a7fbccbf31372eade2ce1a35868891db1a Mon Sep 17 00:00:00 2001 From: jack-scitix-ai Date: Tue, 28 Jul 2026 16:15:52 +0800 Subject: [PATCH 2/2] test(models): fix ty narrowing + migrate browsecomp mock to transport seam --- tests/unit/core/models/test_chat_model.py | 1 + tests/unit/core/models/test_gen_model.py | 1 + tests/unit/core/models/test_model.py | 8 +++-- tests/unit/core/models/test_model_arun.py | 3 +- .../unit/core/models/test_sglang_gen_model.py | 1 + .../core/models/transports/test_sglang.py | 26 +++++++++++------ tests/unit/tasks/test_browsecomp_0shot_gen.py | 29 ++++++++++--------- 7 files changed, 43 insertions(+), 26 deletions(-) diff --git a/tests/unit/core/models/test_chat_model.py b/tests/unit/core/models/test_chat_model.py index 94d1307e..ef917e9c 100644 --- a/tests/unit/core/models/test_chat_model.py +++ b/tests/unit/core/models/test_chat_model.py @@ -24,5 +24,6 @@ def test_builds_openai_chat_transport(self): def test_transport_bound_to_shared_client_and_model(self): m = ChatModel(model="m", api_key="k") + assert isinstance(m._transport, OpenAIChatTransport) assert m._transport._client is m._client assert m._transport._model == "m" diff --git a/tests/unit/core/models/test_gen_model.py b/tests/unit/core/models/test_gen_model.py index b3604580..77b38d63 100644 --- a/tests/unit/core/models/test_gen_model.py +++ b/tests/unit/core/models/test_gen_model.py @@ -26,5 +26,6 @@ def test_builds_openai_completions_transport(self): def test_transport_bound_to_shared_client_and_model(self): m = GenModel(model="m", api_key="k") + assert isinstance(m._transport, OpenAICompletionsTransport) assert m._transport._client is m._client assert m._transport._model == "m" diff --git a/tests/unit/core/models/test_model.py b/tests/unit/core/models/test_model.py index 1a43e063..bce282d7 100644 --- a/tests/unit/core/models/test_model.py +++ b/tests/unit/core/models/test_model.py @@ -223,25 +223,28 @@ async def test_alogprobs_paths(self, path, prompt): @pytest.mark.anyio async def test_alogprobs_lowers_echo_to_score_input(self): """alogprobs args land on the Request the transport receives.""" - from tests.conftest import MockGenModel + from tests.conftest import HandlerTransport, MockGenModel model = MockGenModel() await model.alogprobs("A", echo=False, max_tokens=2, logprobs=3) + assert isinstance(model._transport, HandlerTransport) req = model._transport.requests[0] assert req.score_input is False assert req.return_logprobs is True assert req.top_k == 3 + assert req.sampling is not None assert req.sampling.max_tokens == 2 assert req.sampling.temperature == 0.0 @pytest.mark.anyio async def test_alogprobs_echo_true_sets_score_input(self): - from tests.conftest import MockGenModel + from tests.conftest import HandlerTransport, MockGenModel model = MockGenModel() await model.alogprobs("A", echo=True) + assert isinstance(model._transport, HandlerTransport) req = model._transport.requests[0] assert req.score_input is True @@ -392,6 +395,7 @@ def test_logprobs_request_forces_logprob_fields(self): assert req.return_logprobs is True assert req.top_k == 5 assert req.score_input is True + assert req.sampling is not None assert req.sampling.max_tokens == 1 diff --git a/tests/unit/core/models/test_model_arun.py b/tests/unit/core/models/test_model_arun.py index 14bce5a0..83c6971f 100644 --- a/tests/unit/core/models/test_model_arun.py +++ b/tests/unit/core/models/test_model_arun.py @@ -93,7 +93,7 @@ async def test_echo_true_on_gen_works(self): @pytest.mark.anyio async def test_echo_false_on_chat_skips_the_gate(self): - from tests.conftest import MockChatModel + from tests.conftest import HandlerTransport, MockChatModel m = MockChatModel() # echo=False must skip the InputScoring gate and reach the transport. @@ -101,4 +101,5 @@ async def test_echo_false_on_chat_skips_the_gate(self): # check fires — proving the gate was bypassed and the request ran. with pytest.raises(RuntimeError, match="server returned none"): await m.alogprobs("prompt", echo=False) + assert isinstance(m._transport, HandlerTransport) assert m._transport.requests[0].score_input is False diff --git a/tests/unit/core/models/test_sglang_gen_model.py b/tests/unit/core/models/test_sglang_gen_model.py index 0709b244..2c269d29 100644 --- a/tests/unit/core/models/test_sglang_gen_model.py +++ b/tests/unit/core/models/test_sglang_gen_model.py @@ -24,6 +24,7 @@ def test_builds_sglang_transport(self): def test_transport_bound_to_shared_client_model_and_api_base(self): m = SglangGenModel(model="m", api_base="http://host:8000/v1", api_key="local") + assert isinstance(m._transport, SglangTransport) assert m._transport._client is m._client assert m._transport._model == "m" assert m._transport._api_base == "http://host:8000/v1" diff --git a/tests/unit/core/models/transports/test_sglang.py b/tests/unit/core/models/transports/test_sglang.py index f722e157..251d3642 100644 --- a/tests/unit/core/models/transports/test_sglang.py +++ b/tests/unit/core/models/transports/test_sglang.py @@ -44,6 +44,14 @@ def _meta(**overrides: Any) -> dict: return meta +def _sampling(body: dict) -> dict: + """Narrow ``body["sampling_params"]`` (typed JSONValue) to a dict for the + type checker; the sglang lower() always nests a dict there.""" + sp = body["sampling_params"] + assert isinstance(sp, dict) + return sp + + # =================================================================== # URL derivation (moved from the legacy SglangGenModel tests) # =================================================================== @@ -93,7 +101,7 @@ class TestLower: def test_max_tokens_maps_to_max_new_tokens(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) body = t._lower(Request(input="hi", sampling=SamplingParams(max_tokens=7))) - assert body["sampling_params"]["max_new_tokens"] == 7 + assert _sampling(body)["max_new_tokens"] == 7 def test_score_input_sets_logprob_start_len_0(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) @@ -123,17 +131,17 @@ def test_echo_never_appears_in_wire_body(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) body = t._lower(Request(input="hi", score_input=True)) assert "echo" not in body - assert "echo" not in body["sampling_params"] + assert "echo" not in _sampling(body) def test_prefix_maps_to_prefill(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) body = t._lower(Request(input="hi", prefix="PRE")) - assert body["sampling_params"]["prefill"] == "PRE" + assert _sampling(body)["prefill"] == "PRE" def test_score_input_forces_min_one_new_token(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) body = t._lower(Request(input="hi", score_input=True)) - assert body["sampling_params"]["max_new_tokens"] == 1 + assert _sampling(body)["max_new_tokens"] == 1 def test_all_sampling_params_map(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) @@ -151,7 +159,7 @@ def test_all_sampling_params_map(self): ), ) ) - sp = body["sampling_params"] + sp = _sampling(body) assert sp["temperature"] == 0.7 assert sp["top_p"] == 0.9 assert sp["top_k"] == 40 @@ -174,7 +182,7 @@ def test_known_key_mapped_into_sampling_params(self): body = t._lower( Request(input="hi", extra_wire_params={"repetition_penalty": 1.1}) ) - assert body["sampling_params"]["repetition_penalty"] == 1.1 + assert _sampling(body)["repetition_penalty"] == 1.1 assert "repetition_penalty" not in body # nested, never top-level def test_unknown_keys_dropped(self): @@ -183,9 +191,9 @@ def test_unknown_keys_dropped(self): Request(input="hi", extra_wire_params={"seed": 1, "custom_flag": True}) ) assert "seed" not in body - assert "seed" not in body["sampling_params"] + assert "seed" not in _sampling(body) assert "custom_flag" not in body - assert "custom_flag" not in body["sampling_params"] + assert "custom_flag" not in _sampling(body) def test_explicit_ir_field_wins_over_extra_wire_params(self): t, _ = _make_transport({"text": "", "meta_info": _meta()}) @@ -196,7 +204,7 @@ def test_explicit_ir_field_wins_over_extra_wire_params(self): extra_wire_params={"max_tokens": 99}, ) ) - assert body["sampling_params"]["max_new_tokens"] == 7 + assert _sampling(body)["max_new_tokens"] == 7 # =================================================================== diff --git a/tests/unit/tasks/test_browsecomp_0shot_gen.py b/tests/unit/tasks/test_browsecomp_0shot_gen.py index 2ea1d0f1..633651cd 100644 --- a/tests/unit/tasks/test_browsecomp_0shot_gen.py +++ b/tests/unit/tasks/test_browsecomp_0shot_gen.py @@ -8,8 +8,9 @@ from datasets import DatasetDict as HFDatasetDict from sieval.community.browsecomp import aggregate_metrics, parse_grade -from sieval.core.models import ModelOutput +from sieval.core.models import Request, Response from sieval.core.models.chat_model import ChatModel +from sieval.core.models.transports import OpenAIChatTransport from sieval.core.tasks import TaskContext from sieval.datasets.browsecomp import ( BrowseCompDataset, @@ -19,26 +20,24 @@ BrowseCompZeroShotGenTask, GradeFeedback, ) +from tests.conftest import HandlerTransport, n_of class _ScriptedChatModel(ChatModel): - """ChatModel returning a fixed reply, recording the last agenerate kwargs.""" + """ChatModel returning a fixed reply, recording the last Request seen.""" def __init__(self, reply: str, model: str = "mock"): - super().__init__(model=model, api_key="fake") self._reply = reply - self.last_kwargs: dict[str, object] = {} + self.last_req: Request | None = None + super().__init__(model=model, api_key="fake") - async def _agenerate_impl(self, prompt, **kwargs) -> ModelOutput: - _ = prompt - self.last_kwargs = dict(kwargs) - return ModelOutput(model=self.meta(), texts=[self._reply]) + def _build_default_transport(self) -> HandlerTransport: + return HandlerTransport(self._stub_arun, OpenAIChatTransport.CAPABILITIES) - async def _alogprobs_impl( - self, prompt, *, max_tokens=1, logprobs=5, echo=True, temperature=0.0, **kwargs - ) -> ModelOutput: - _ = (prompt, max_tokens, logprobs, echo, temperature, kwargs) - return ModelOutput(model=self.meta(), texts=[""]) + async def _stub_arun(self, req: Request) -> Response: + self.last_req = req + n = n_of(req) + return Response(texts=(self._reply,) * n, finish_reasons=("stop",) * n) def _sample() -> BrowseCompDatasetSample: @@ -107,7 +106,9 @@ async def test_infer_forwards_n(): grader = _ScriptedChatModel(reply="correct: yes", model="grader") task = BrowseCompZeroShotGenTask(dataset, model, grader=grader, n=3) await task.infer([{"role": "user", "content": "q"}], TaskContext(sample_id=0)) - assert model.last_kwargs.get("n") == 3 + assert model.last_req is not None + assert model.last_req.sampling is not None + assert model.last_req.sampling.n == 3 # --- feedback: yes/no grading + confidence + provenance ---