diff --git a/src/marsys/agents/exceptions.py b/src/marsys/agents/exceptions.py index fd134598..47a1343b 100644 --- a/src/marsys/agents/exceptions.py +++ b/src/marsys/agents/exceptions.py @@ -839,7 +839,13 @@ def from_provider_response( # Provider-specific error parsing based on status code # This needs to work even when response is None (just using status_code) elif status_code: - if provider == "openai": + # Azure OpenAI serves the Responses API and returns the same error + # envelope, so it classifies identically to first-party OpenAI — + # sharing the branch keeps one behaviour for one wire contract, the + # way ``bedrock`` shares Anthropic's below. Without the share an + # Azure 429 falls past every branch unclassified and is treated as + # not retryable. + if provider in ("openai", "azure"): error_data = raw_response.get("error", {}) if raw_response else {} if error_data: message = error_data.get("message", message) diff --git a/src/marsys/models/adapters/__init__.py b/src/marsys/models/adapters/__init__.py index 07a154f5..7ace4d5a 100644 --- a/src/marsys/models/adapters/__init__.py +++ b/src/marsys/models/adapters/__init__.py @@ -4,6 +4,7 @@ from marsys.models.adapters.openai import OpenAIAdapter, AsyncOpenAIAdapter from marsys.models.adapters.openrouter import OpenRouterAdapter, AsyncOpenRouterAdapter from marsys.models.adapters.anthropic import AnthropicAdapter, AsyncAnthropicAdapter +from marsys.models.adapters.azure import AsyncAzureOpenAIAdapter, AzureOpenAIAdapter from marsys.models.adapters.bedrock import AsyncBedrockAdapter, BedrockAdapter from marsys.models.adapters.google import GoogleAdapter, AsyncGoogleAdapter from marsys.models.adapters.openai_oauth import OpenAIOAuthAdapter, AsyncOpenAIOAuthAdapter @@ -33,6 +34,9 @@ # Bedrock (Claude on Amazon Bedrock) "BedrockAdapter", "AsyncBedrockAdapter", + # Azure OpenAI (OpenAI models on an Azure resource) + "AzureOpenAIAdapter", + "AsyncAzureOpenAIAdapter", # Google "GoogleAdapter", "AsyncGoogleAdapter", diff --git a/src/marsys/models/adapters/anthropic.py b/src/marsys/models/adapters/anthropic.py index e3af976a..eed52c06 100644 --- a/src/marsys/models/adapters/anthropic.py +++ b/src/marsys/models/adapters/anthropic.py @@ -86,11 +86,35 @@ def _anthropic_model_requires_adaptive_thinking(model_name: str) -> bool: CACHE_CONTROL_EPHEMERAL = {"type": "ephemeral"} - -def mark_conversation_tail_for_cache(messages: List[Dict[str, Any]]) -> None: +# A caller marks a message row with this key to say "my content here changes every +# request; do not put the cache breakpoint on me". Neutral and per-item, riding the +# caller's own message dict — the `defer_loading` shape, which is this codebase's +# established way for a caller to signal request structure without a new request +# parameter. Stripped during conversion; it never reaches the wire. +# +# Why a caller needs this: the breakpoint's value is that the NEXT request can read +# the entry this one writes, which requires the entry's hashed prefix to consist of +# bytes the next request still contains. A row whose text is regenerated per request +# (a clock, a budget figure, anything derived from "now") is by construction absent +# from the next request, so an entry written at or after it is unreadable forever — +# each turn writes a fresh entry and reads none. Measured on Bedrock/Opus 5, single- +# step turns, tools present: marker on the volatile row → turn 2 `read=0`; marker on +# the last durable row → turn 2 `read=8425`. +CACHE_EXEMPT_KEY = "cache_exempt" + + +def mark_conversation_tail_for_cache( + messages: List[Dict[str, Any]], *, volatile_tail: int = 0 +) -> None: """Place ONE prompt-cache breakpoint on the last content block of the last message, in place on ``messages`` — the platform's multi-turn caching pattern. + ``volatile_tail`` excludes that many trailing messages from carrying the marker, + for a caller that appends per-request content after the durable conversation (see + ``CACHE_EXEMPT_KEY``). The marker then lands on the last DURABLE row, which the + next request still contains verbatim, so the entry stays readable. Default 0 keeps + the payload byte-identical to the unparameterized form for every other caller. + Adapter-owned and unconditional, matching the only other `cache_control` site in this codebase (the OAuth adapter's static Claude-Code prefix block). Only the payload builder knows the rendered block layout, and caching is prefix-match @@ -118,6 +142,11 @@ def mark_conversation_tail_for_cache(messages: List[Dict[str, Any]]) -> None: the model's cacheable minimum silently writes nothing and costs nothing, so no size check is needed here. """ + if volatile_tail: + # Step back past the per-request rows. Every row is volatile (a caller that + # marked the whole list) → nothing durable to anchor an entry to, so no marker: + # writing one would cost a fresh entry per request and read none. + messages = messages[:-volatile_tail] if not messages: return last = messages[-1] @@ -150,6 +179,125 @@ def mark_conversation_tail_for_cache(messages: List[Dict[str, Any]]) -> None: return +# --- token counting (POST /v1/messages/count_tokens) ------------------------- +# +# The endpoint takes the SAME body as message creation minus the generation +# controls, and answers ``{"input_tokens": N}`` with the model's own tokenizer. +# It is free of charge, rate-limited independently of message creation, and does +# not participate in prompt caching, so a caller may count the exact payload it is +# about to send without perturbing anything. +# +# Both Anthropic legs (api-key and OAuth) build their payload through their own +# ``format_request_payload`` — that is the point: the OAuth leg injects a required +# system block and renames tools, so a count assembled any other way would not be a +# count of what the request actually sends. Only the generation-side keys come off. +_COUNT_TOKENS_REJECTED_KEYS = ( + "max_tokens", + "stream", + "temperature", + "top_p", + "top_k", + "output_config", +) +# Endpoints (by URL) that answered "there is no such route". Structural and +# permanent for the process; a credential-shaped refusal (401/403) is deliberately +# NOT recorded here, because the OAuth token file has several writers and a refresh +# in flight looks exactly like a rejection for one request. +_COUNT_TOKENS_UNSUPPORTED: "set[str]" = set() + + +def count_tokens_url_for(messages_url: str) -> str: + """The count endpoint beside a Messages endpoint, query string preserved (the + OAuth leg's URL carries ``?beta=true`` and the count has to keep it).""" + base, sep, query = messages_url.partition("?") + return f"{base.rstrip('/')}/count_tokens" + (f"{sep}{query}" if sep else "") + + +def strip_for_count_tokens(payload: Dict[str, Any]) -> Dict[str, Any]: + """A Messages payload reduced to a LEGAL count request. + + Two things happen here, and the second is not cosmetic. The generation controls + come off because the endpoint rejects them. And the final assistant message is + right-trimmed because the API rejects one that ends in whitespace — a rule the + messages call rarely meets (something is almost always appended after the last + assistant turn) and a count meets constantly, since counting a conversation means + presenting its last row as final. A settled conversation ends with an assistant + reply by definition, and models end replies with a newline, so without this a + fold's post-fold count fails on exactly the conversations most likely to fold. + """ + reduced = {k: v for k, v in payload.items() if k not in _COUNT_TOKENS_REJECTED_KEYS} + messages = reduced.get("messages") + if isinstance(messages, list) and messages: + trimmed = _trim_final_assistant(messages[-1]) + reduced["messages"] = ( + messages[:-1] if trimmed is None else [*messages[:-1], trimmed] + ) + return reduced + + +def _trim_final_assistant(message: Any) -> Optional[Dict[str, Any]]: + """The message with trailing whitespace off its last text block, or ``None`` when + nothing is left of it. Copies rather than mutating: the payload's blocks may be + the caller's own rows.""" + if not isinstance(message, dict) or message.get("role") != "assistant": + return message + content = message.get("content") + if isinstance(content, str): + trimmed = content.rstrip() + return {**message, "content": trimmed} if trimmed else None + if not isinstance(content, list) or not content: + return message + last = content[-1] + if not isinstance(last, dict) or last.get("type") != "text": + return message + text = str(last.get("text", "")).rstrip() + blocks = list(content[:-1]) if not text else [*content[:-1], {**last, "text": text}] + return {**message, "content": blocks} if blocks else None + + +def count_tokens_supported(url: str) -> bool: + return url not in _COUNT_TOKENS_UNSUPPORTED + + +def read_count_tokens_response( + url: str, status: int, body: Any, *, provider: str +) -> Optional[int]: + """One count response → a token count, or ``None``. + + ``None`` is the whole error vocabulary: a count is an optimisation on the + caller's side, never the turn's business, so nothing here raises. A 404/405 + says the route does not exist on this endpoint and is remembered for the + process; every other failure is treated as this-request-only. + """ + if status == 200 and isinstance(body, dict): + count = body.get("input_tokens") + if isinstance(count, int) and count >= 0: + return count + logger.warning("%s count_tokens returned no input_tokens: %r", provider, body) + return None + if status in (404, 405): + _COUNT_TOKENS_UNSUPPORTED.add(url) + logger.warning( + "%s does not serve %s (HTTP %s); token counting is off for this process", + provider, url, status, + ) + return None + # The provider's own words, not just the status: a bare "HTTP 400" on a request + # nobody sees is undiagnosable, and this one answers in a sentence. + logger.warning( + "%s count_tokens failed with HTTP %s: %s", provider, status, _error_text(body) + ) + return None + + +def _error_text(body: Any) -> str: + if isinstance(body, dict): + error = body.get("error") + if isinstance(error, dict) and error.get("message"): + return str(error["message"]) + return "no error body" + + class AnthropicAdapter(APIProviderAdapter): """Adapter for Anthropic Claude API""" @@ -317,6 +465,19 @@ def format_request_payload(self, messages: List[Dict], **kwargs) -> Dict[str, An system_message = None user_messages = [] + # How many TRAILING rows the caller marked as per-request (``CACHE_EXEMPT_KEY``), + # so the breakpoint below lands on the last durable row instead. Counted from the + # end and stopping at the first unmarked row: the exemption is about position (what + # the next request will still contain), so a marked row with durable rows after it + # is not a tail and does not shift the marker. Each of these converts 1:1 into + # ``user_messages``, so the count carries over. The key itself never reaches the + # wire — the regular-message branch rebuilds rows with only role/content. + volatile_tail = 0 + for msg in reversed(messages): + if not msg.get(CACHE_EXEMPT_KEY): + break + volatile_tail += 1 + for msg in messages: if msg.get("role") == "system": system_message = msg.get("content") @@ -533,13 +694,21 @@ def format_request_payload(self, messages: List[Dict], **kwargs) -> Dict[str, An # final message, and the json-mode fallback above may still append a hint # block there. Placing it after every content mutation is what makes "the # tail" mean the actual tail. - mark_conversation_tail_for_cache(user_messages) + mark_conversation_tail_for_cache(user_messages, volatile_tail=volatile_tail) return payload def get_endpoint_url(self) -> str: return f"{self.base_url.rstrip('/')}/messages" + def get_count_tokens_url(self) -> str: + return count_tokens_url_for(self.get_endpoint_url()) + + def format_count_tokens_payload( + self, messages: List[Dict], **kwargs + ) -> Dict[str, Any]: + return strip_for_count_tokens(self.format_request_payload(messages, **kwargs)) + def handle_api_error(self, error: Exception, response=None) -> ErrorResponse: """Enhanced error handling using ModelAPIError classification.""" from marsys.agents.exceptions import ModelAPIError @@ -845,3 +1014,50 @@ async def arun_streaming( provider=self._provider_name() or "anthropic", response=stream_error_payload({"type": "max_retries"}, 0), ) + + async def acount_tokens( + self, + messages: List[Dict], + *, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[str] = None, + ) -> Optional[int]: + """How many input tokens a Messages call with this body would carry, by the + model's own tokenizer, or ``None`` when the endpoint cannot say. + + ``system`` rides in as the leading ``role:"system"`` message the payload + builder already knows how to hoist, so the counted body is assembled by the + same code the real request goes through — including this leg's tool + conversion and cache markers. + + No retry loop and no raise: the caller is sizing something, not producing + the turn's answer, and a count that fails must cost the turn nothing. + """ + import asyncio + + import aiohttp + + url = self.get_count_tokens_url() + if not count_tokens_supported(url): + return None + if system is not None: + messages = [{"role": "system", "content": system}, *messages] + payload = self.format_count_tokens_payload(messages, tools=tools) + headers = {**self.get_headers(), "accept": "application/json"} + session = await self._ensure_session() + try: + async with session.post( + url, + headers=headers, + json=payload, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + status = response.status + try: + body = await response.json(content_type=None) + except ValueError: + body = None + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + logger.warning("anthropic count_tokens transport failure: %s", exc) + return None + return read_count_tokens_response(url, status, body, provider="anthropic") diff --git a/src/marsys/models/adapters/anthropic_oauth.py b/src/marsys/models/adapters/anthropic_oauth.py index d78f4c12..3f844a44 100644 --- a/src/marsys/models/adapters/anthropic_oauth.py +++ b/src/marsys/models/adapters/anthropic_oauth.py @@ -7,9 +7,14 @@ from typing import Any, Dict, List, Optional from marsys.models.adapters.anthropic import ( + CACHE_EXEMPT_KEY, _anthropic_model_rejects_temperature, _anthropic_model_requires_adaptive_thinking, + count_tokens_supported, + count_tokens_url_for, mark_conversation_tail_for_cache, + read_count_tokens_response, + strip_for_count_tokens, ) from marsys.models.adapters.base import APIProviderAdapter, AsyncBaseAPIAdapter from marsys.models.response_models import ( @@ -295,6 +300,14 @@ def get_endpoint_url(self) -> str: """Return Claude API endpoint URL.""" return self.API_URL + def get_count_tokens_url(self) -> str: + return count_tokens_url_for(self.get_endpoint_url()) + + def format_count_tokens_payload( + self, messages: List[Dict], **kwargs + ) -> Dict[str, Any]: + return strip_for_count_tokens(self.format_request_payload(messages, **kwargs)) + def _build_system_array(self, system_message: Optional[str] = None) -> List[Dict]: """ Build system prompt array with required prefix. @@ -443,6 +456,15 @@ def format_request_payload(self, messages: List[Dict], **kwargs) -> Dict[str, An system_message = None converted_messages = [] + # Trailing per-request rows the caller exempted from the breakpoint — twinned with + # the api-key adapter (this class is not a subclass of it, so the two payload + # builders are kept deliberately parallel). See ``CACHE_EXEMPT_KEY``. + volatile_tail = 0 + for msg in reversed(messages): + if not msg.get(CACHE_EXEMPT_KEY): + break + volatile_tail += 1 + for msg in messages: role = msg.get("role") content = msg.get("content") @@ -611,7 +633,7 @@ def format_request_payload(self, messages: List[Dict], **kwargs) -> Dict[str, An # in an OAuth payload — the static Claude-Code prefix block in # ``_build_system_array`` is the first — which keeps the payload two under # the API's four-breakpoint ceiling. - mark_conversation_tail_for_cache(converted_messages) + mark_conversation_tail_for_cache(converted_messages, volatile_tail=volatile_tail) return payload @@ -1138,4 +1160,47 @@ async def arun_streaming( exception=e ) - + async def acount_tokens( + self, + messages: List[Dict], + *, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[str] = None, + ) -> Optional[int]: + """How many input tokens a Messages call with this body would carry, by the + model's own tokenizer, or ``None`` when the endpoint cannot say. + + Built through this leg's own payload builder, which is why the method exists + here rather than once for both Anthropic legs: this one prepends the required + Claude-Code system block and renames tools, so a body assembled any other way + would count something the request never sends. The ``accept`` header is the + one deliberate difference from the messages call — that call is + streaming-always, and a count is a single JSON object. + + No retry loop and no raise: the caller is sizing something, not producing the + turn's answer, and a count that fails must cost the turn nothing. + """ + import httpx + + url = self.get_count_tokens_url() + if not count_tokens_supported(url): + return None + if system is not None: + messages = [{"role": "system", "content": system}, *messages] + await asyncio.to_thread(self._ensure_fresh_token) + payload = self.format_count_tokens_payload(messages, tools=tools) + headers = {**self.get_headers(), "accept": "application/json"} + try: + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.post(url, headers=headers, json=payload) + status = response.status_code + try: + body = response.json() + except ValueError: + body = None + except httpx.HTTPError as exc: + logger.warning("anthropic-oauth count_tokens transport failure: %s", exc) + return None + return read_count_tokens_response( + url, status, body, provider="anthropic-oauth" + ) diff --git a/src/marsys/models/adapters/azure.py b/src/marsys/models/adapters/azure.py new file mode 100644 index 00000000..2911c76b --- /dev/null +++ b/src/marsys/models/adapters/azure.py @@ -0,0 +1,169 @@ +"""Adapter for OpenAI models hosted on Azure OpenAI (AI Foundry). + +Azure re-hosts OpenAI's own models behind its own resource, and its **v1 surface** +serves the Responses API at the identical relative path with the identical request and +response bodies. That is what makes this a thin subclass of :class:`OpenAIAdapter` +instead of a parallel implementation: message conversion, tool conversion, structured +output, SSE streaming, reasoning effort, usage harmonization and error classification +are all inherited unchanged. ``get_endpoint_url()`` is inherited too — against a +``.../openai/v1`` base it already produces the documented +``POST .../openai/v1/responses``. + +The older Azure surface — ``/openai/deployments//responses`` with a pinned +``api-version`` query parameter — is deliberately not used. On the v1 GA API +``api-version`` is no longer required, and the deployment name travels in the request +body's ``model`` field, which is exactly the shape the inherited payload builder already +emits. Targeting the legacy surface would mean overriding endpoint construction to gain +nothing. + +Four deltas from the first-party adapter, each measured against the live resource +rather than assumed: + +* **base_url** — per-resource, so it cannot be a compiled-in constant. Resolved from + the environment at construction, the way :mod:`marsys.models.adapters.bedrock` + resolves its region-dependent host. +* **Auth** — the resource key travels in an ``api-key`` header. (A resource key is also + accepted in ``Authorization: Bearer`` on this surface, measured; ``api-key`` is the + documented spelling for a key, with Bearer reserved for Entra ID tokens, so the + documented one is what this sends.) +* **Reasoning-effort vocabulary** — this surface serves ``none``, ``low``, ``medium``, + ``high``, ``xhigh`` and ``max``, and answers ``minimal`` with a 400 that enumerates + those six. The smallest thinking budget therefore arrives as ``low`` here while the + first-party leg keeps sending ``minimal``. +* **Model ids** — the ``model`` field carries an operator-chosen *deployment* name, and + the response echoes that same deployment name back rather than an underlying model + snapshot. Measured on ``gpt-5.6-sol`` and ``gpt-5.6-terra``: both echo themselves. + A cost table keyed on the echoed id therefore keys on deployment names. + +Two behaviours of this endpoint that are inherited rather than worked around, recorded +because both are silent until they are not: + +* **Prompt caching is implicit and takes no request parameter.** ``prompt_cache_options`` + accepts ``implicit`` (the default) and ``explicit``; a ``prompt_cache_breakpoint`` + field is rejected as an unknown parameter in every position, and ``explicit`` mode + without one caches nothing. So the correct request is one that says nothing about + caching, and a measured pair of identical large calls reports + ``cache_write_tokens: 3395`` then ``cached_tokens: 3395`` with no parameter sent. + The inherited harmonizer reads both figures. +* **Reasoning-capable deployments reject ``temperature``.** The inherited capability + check is a regex over the model name, which the real deployment names + (``gpt-5.6-*``) satisfy. A deployment renamed to something not starting ``gpt-5``+ + would send ``temperature`` and take a 400 from this endpoint. Naming a deployment + after the model it serves keeps that check honest; this adapter cannot know a + deployment's underlying model from its name alone. +""" + +import logging +import os +from typing import Dict, Optional + +from marsys.models.adapters.openai import AsyncOpenAIAdapter, OpenAIAdapter + +logger = logging.getLogger(__name__) + +AZURE_OPENAI_V1_PATH = "/openai/v1" + + +def azure_openai_base_url(endpoint: Optional[str] = None) -> str: + """The v1 base URL for an Azure OpenAI resource, or ``""`` if unconfigured. + + Resolution order: explicit argument, ``AZURE_OPENAI_ENDPOINT``, then + ``FOUNDRY_ENDPOINT``. The second name is read because this stack's voice legs + already resolve the same resource under it — one resource, one credential, read by + whichever name the caller's environment happens to carry, rather than a second + stored copy of the same secret. + + Accepts the several spellings one resource legitimately arrives as and returns the + one the Responses path hangs off: + + * ``https://.services.ai.azure.com`` — the resource host. + * ``https://.openai.azure.com`` — the same resource's other documented host. + * ``https://.services.ai.azure.com/api/projects/`` — the *project* + endpoint, which is what the Azure portal offers for copying and what an escrowed + copy of this value turned out to hold. It addresses a project inside the + resource and does not serve ``/openai/v1``; the resource host does, so the + project suffix is dropped. + * an endpoint already ending in ``/openai/v1`` — returned as-is, so a caller who + configured the full base is not given ``/openai/v1/openai/v1``. + """ + resolved = ( + endpoint + or os.getenv("AZURE_OPENAI_ENDPOINT") + or os.getenv("FOUNDRY_ENDPOINT") + or "" + ).strip() + if not resolved: + # Empty rather than a guess: there is no default Azure resource, and a + # fabricated host would turn a missing setting into a confusing connection + # error instead of an obvious unconfigured one. + return "" + resolved = resolved.rstrip("/") + if resolved.endswith(AZURE_OPENAI_V1_PATH): + return resolved + marker = "/api/projects/" + if marker in resolved: + resolved = resolved[: resolved.index(marker)] + return f"{resolved}{AZURE_OPENAI_V1_PATH}" + + +class AzureOpenAIAdapter(OpenAIAdapter): + """OpenAI models on an Azure OpenAI resource, via the v1 Responses surface.""" + + def __init__( + self, + model_name: str, + api_key: str = "", + base_url: str = "", + max_tokens: int = 1024, + temperature: float = 0.7, + endpoint: Optional[str] = None, + **kwargs, + ): + super().__init__( + model_name=model_name, + api_key=api_key or os.getenv("AZURE_OPENAI_API_KEY", "") or os.getenv("FOUNDRY_API_KEY", ""), + # Both spellings of the same fact go through the one normalizer: it is + # idempotent on an already-normalized base, and a caller handing over the + # portal's project URL or the bare resource host gets a client that can reach + # ``/responses`` instead of one that 404s. An empty pair falls through to the + # environment, which is what lets a caller whose model-construction path has + # no per-resource endpoint to hand still reach the right host. + base_url=azure_openai_base_url(base_url or endpoint), + max_tokens=max_tokens, + temperature=temperature, + **kwargs, + ) + + def get_headers(self) -> Dict[str, str]: + return { + "api-key": self.api_key, + "Content-Type": "application/json", + } + + def _provider_name(self) -> Optional[str]: + # Pinned rather than derived from the class name, so retry/error settings and + # the harmonized response's provider field both resolve under the provider id + # the rest of the stack uses. + return "azure" + + def _served_effort(self, effort: str, model_lower: str) -> str: + # This surface serves none / low / medium / high / xhigh / max and answers + # `minimal` with a 400 (measured on `gpt-5.6-sol`, whose reply enumerates the + # six it takes). Unconditional rather than keyed on the model name: the name + # here is an operator-chosen deployment label and cannot be read as evidence + # about the generation underneath, and the smallest reasoning this endpoint has + # is what a request for the smallest should get either way. + if effort == "minimal": + return "low" + return super()._served_effort(effort, model_lower) + + +class AsyncAzureOpenAIAdapter(AsyncOpenAIAdapter, AzureOpenAIAdapter): + """Async Azure OpenAI adapter. + + Inherits the Responses SSE stream from :class:`AsyncOpenAIAdapter` and the + endpoint/auth/provider deltas from :class:`AzureOpenAIAdapter`. + """ + + def _provider_name(self) -> Optional[str]: + return "azure" diff --git a/src/marsys/models/adapters/factory.py b/src/marsys/models/adapters/factory.py index d2eb7f19..282990e6 100644 --- a/src/marsys/models/adapters/factory.py +++ b/src/marsys/models/adapters/factory.py @@ -4,6 +4,7 @@ from marsys.models.adapters.openai import OpenAIAdapter from marsys.models.adapters.openrouter import OpenRouterAdapter from marsys.models.adapters.anthropic import AnthropicAdapter +from marsys.models.adapters.azure import AzureOpenAIAdapter from marsys.models.adapters.bedrock import BedrockAdapter from marsys.models.adapters.google import GoogleAdapter from marsys.models.adapters.openai_oauth import OpenAIOAuthAdapter @@ -27,6 +28,7 @@ def create_adapter( "openai": OpenAIAdapter, "anthropic": AnthropicAdapter, "bedrock": BedrockAdapter, # Claude on Amazon Bedrock (Messages-API-shaped) + "azure": AzureOpenAIAdapter, # OpenAI models on Azure OpenAI (Responses-API-shaped) "google": GoogleAdapter, "openrouter": OpenRouterAdapter, # OpenRouter with additional headers support "xai": OpenRouterAdapter, # xAI Grok uses OpenAI-compatible /chat/completions diff --git a/src/marsys/models/adapters/openai.py b/src/marsys/models/adapters/openai.py index 45c17fbc..6c57a297 100644 --- a/src/marsys/models/adapters/openai.py +++ b/src/marsys/models/adapters/openai.py @@ -24,6 +24,38 @@ logger = logging.getLogger(__name__) +# The keys a Responses API `input` message item accepts. Anything else in a caller's +# message dict is theirs, not the wire's, and the endpoint rejects it outright. +_RESPONSES_INPUT_ITEM_KEYS = frozenset({"role", "content", "type"}) + +# `thinking_budget` (a token allowance, the Anthropic-family knob every caller in this +# stack already sets) mapped onto `reasoning.effort` (the bucket the Responses API +# takes). Without this the budget is silently inert on every OpenAI-family leg and the +# model runs at the provider default — `medium` — whatever the caller configured. +# Boundaries follow the Anthropic minimum of 1024: below it the caller is asking for +# as little deliberation as the provider offers. +_THINKING_BUDGET_EFFORT_BUCKETS: tuple[tuple[int, str], ...] = ( + (1024, "minimal"), + (4096, "low"), + (16384, "medium"), +) +_MAX_THINKING_EFFORT = "high" + + +def thinking_budget_to_effort(budget: Optional[int]) -> Optional[str]: + """`reasoning.effort` for a token budget, or None to leave the provider default. + + A non-positive budget means the caller turned thinking off, which is not the same + request as "think as little as possible" — it maps to None so the parameter is + omitted rather than pinned to `minimal`. + """ + if budget is None or budget <= 0: + return None + for ceiling, effort in _THINKING_BUDGET_EFFORT_BUCKETS: + if budget < ceiling: + return effort + return _MAX_THINKING_EFFORT + class OpenAIAdapter(APIProviderAdapter): """Adapter for OpenAI and OpenAI-compatible APIs (OpenRouter, Groq)""" @@ -147,17 +179,38 @@ def convert_content_types(content): "call_id": msg.get("tool_call_id"), "output": msg.get("content", "") }) - # Regular messages - convert content types and ensure content is not None + # Regular messages - rebuild from the keys this endpoint accepts. + # + # Rebuilt from an allow-list rather than copied-and-pruned. A caller's + # message dict is its own working object and routinely carries keys that + # mean something upstream and nothing to a provider — provenance tags, + # routing hints, cache-control markers. Copying the dict forwards all of + # them: the Responses API answers an unrecognized per-item key with + # `400 unknown_parameter: input[0].`, so a caller that annotates its + # messages cannot talk to this endpoint at all. Pruning known offenders + # one at a time only defers that to the next key someone adds, which is + # why `name` was already being popped here by hand. + # + # The Anthropic adapter reached the same shape from the same 400 and + # rebuilds `{role, content}` only; `type` is kept here because the + # Responses input array uses it to discriminate item kinds. else: - cleaned_msg = msg.copy() + cleaned_msg = { + key: value + for key, value in msg.items() + if key in _RESPONSES_INPUT_ITEM_KEYS + } + dropped = set(msg) - set(cleaned_msg) + if dropped: + logger.debug( + "Dropped non-wire message keys before send: %s", + sorted(dropped), + ) if cleaned_msg.get("content") is None: cleaned_msg["content"] = "" else: # Convert content types (text -> input_text, image_url -> input_image) cleaned_msg["content"] = convert_content_types(cleaned_msg["content"]) - # Remove 'name' field - not supported in Responses API - # (was supported in Chat Completions for multi-user/multi-persona dialogues) - cleaned_msg.pop("name", None) converted_messages.append(cleaned_msg) payload = { @@ -265,14 +318,16 @@ def convert_content_types(content): converted_tools.append({"type": "tool_search"}) payload["tools"] = converted_tools - # Handle OpenAI reasoning (effort-based for all models via Responses API) + # Handle OpenAI reasoning (effort-based for all models via Responses API). + # An explicit `reasoning_effort` wins; failing that, a caller's thinking budget + # selects the bucket, so the one knob this stack exposes reaches this leg too. reasoning_effort = kwargs.get("reasoning_effort") + if not reasoning_effort: + reasoning_effort = thinking_budget_to_effort(kwargs.get("thinking_budget")) if reasoning_effort and reasoning_effort.lower() in ["minimal", "low", "medium", "high"]: - effort_value = reasoning_effort.lower() - # Codex models don't support 'minimal' - map to 'low' - if "codex" in model_lower and effort_value == "minimal": - effort_value = "low" - payload["reasoning"] = {"effort": effort_value} + payload["reasoning"] = { + "effort": self._served_effort(reasoning_effort.lower(), model_lower) + } # Only accept known OpenAI Responses API parameters - warn about unknown ones # Based on: https://platform.openai.com/docs/api-reference/responses/create @@ -334,6 +389,24 @@ def convert_content_types(content): return payload + def _served_effort(self, effort: str, model_lower: str) -> str: + """The nearest effort THIS surface will serve for the one the caller asked for. + + `minimal` is not universal. Codex models reject it; so does GPT-5.6 on Azure's + re-hosted surface, which answers a request for it with + `Unsupported value: 'minimal' is not supported ... Supported values are: 'none', + 'low', 'medium', 'high', 'xhigh', and 'max'` — a 400 on every call, for a caller + who only configured a small thinking budget. + + `low` is the substitution and `none` is not: a positive budget means "think a + little", and turning that into no reasoning at all would answer a different + question than the one asked. A budget of zero never reaches here (it maps to no + parameter at all). + """ + if effort == "minimal" and "codex" in model_lower: + return "low" + return effort + def get_endpoint_url(self) -> str: # Migrate to OpenAI Responses API (unified endpoint for all models) # Supports reasoning parameter for GPT-5, o-series, and all future models @@ -344,7 +417,12 @@ def handle_api_error(self, error: Exception, response=None) -> ErrorResponse: from marsys.agents.exceptions import ModelAPIError # Create classified API error - api_error = ModelAPIError.from_provider_response(provider="openai", response=response, exception=error) + # The subclass's own provider id, matching what the retry path below already + # reports — a hardcoded "openai" mislabels an Azure-hosted failure. The + # classifier shares one branch for both, since they share one error envelope. + api_error = ModelAPIError.from_provider_response( + provider=self._provider_name() or "openai", response=response, exception=error + ) # For critical errors, raise the exception to stop execution if api_error.is_critical(): @@ -460,25 +538,56 @@ def harmonize_response( usage = None if usage_data: output_details = usage_data.get("output_tokens_details") or {} + input_details = usage_data.get("input_tokens_details") or {} + # Cache accounting, converted from this API's convention to UsageInfo's. + # + # The two conventions are inverses and both are internally consistent, so a + # naive field-to-field mapping produces numbers that look plausible and are + # wrong. Here, `cached_tokens` and `cache_write_tokens` are SLICES OF + # `input_tokens` — a measured call reads `input_tokens: 3398` with + # `cached_tokens: 3395` inside it, and `3398 + output 5 == total 3403`. + # `UsageInfo.prompt_tokens` means the opposite: the uncached REMAINDER, with + # the cache counts sitting beside it and `full_prompt_tokens` summing all + # three. Mapping `input_tokens` straight onto `prompt_tokens` therefore + # counts the cached slice twice — once inside the prompt figure and once + # again as a cache field — inflating the billable prompt by up to the whole + # cached prefix and charging that slice at the fresh-input rate on top. + # + # Subtracting here, once, keeps every downstream reading correct without a + # provider conditional: `full_prompt_tokens` recovers `input_tokens` + # exactly, and a price split over (fresh, read, write) sums to the same + # whole. Clamped because a vendor's slices must not exceed the whole they + # come from, and a negative prompt count would walk a spend ledger + # backwards. + reported_input = ( + usage_data.get("input_tokens") + or usage_data.get("prompt_tokens") + or 0 + ) + cached_tokens = input_details.get("cached_tokens") or 0 + cache_write_tokens = input_details.get("cache_write_tokens") or 0 + uncached_input = max(0, reported_input - cached_tokens - cache_write_tokens) usage = UsageInfo( - prompt_tokens=( - usage_data.get("input_tokens") - or usage_data.get("prompt_tokens") - ), + prompt_tokens=uncached_input, completion_tokens=( usage_data.get("output_tokens") or usage_data.get("completion_tokens") ), total_tokens=usage_data.get("total_tokens"), + # A subset of `completion_tokens`, billed as output. Recorded for + # visibility; a consumer that adds it to the completion count is + # double-counting. reasoning_tokens=( output_details.get("reasoning_tokens") or usage_data.get("reasoning_tokens") ), + cache_read_input_tokens=cached_tokens or None, + cache_creation_input_tokens=cache_write_tokens or None, ) # Build metadata metadata = ResponseMetadata( - provider="openai", + provider=self._provider_name() or "openai", model=raw_response.get("model", self.model_name), request_id=raw_response.get("id"), created=raw_response.get("created") or raw_response.get("created_at"), diff --git a/src/marsys/models/models.py b/src/marsys/models/models.py index 413598a6..fcf827e2 100644 --- a/src/marsys/models/models.py +++ b/src/marsys/models/models.py @@ -50,6 +50,9 @@ # Bedrock — resolved by name for the async twin, so it must be in scope here BedrockAdapter, AsyncBedrockAdapter, + # Azure OpenAI — same reason: the async twin is looked up by name in this module + AzureOpenAIAdapter, + AsyncAzureOpenAIAdapter, # Google GoogleAdapter, AsyncGoogleAdapter, @@ -83,6 +86,43 @@ def _bedrock_default_base_url() -> str: return bedrock_base_url() +def _azure_default_base_url() -> str: + """Azure OpenAI's base URL names a customer's own resource, so there is no + literal to write. Resolved from the environment the way Bedrock's region is, + and empty when nothing is configured — a caller that supplies the endpoint + per instance (the normal path) never reads this entry.""" + from marsys.models.adapters.azure import azure_openai_base_url + + return azure_openai_base_url() + + +# The providers whose endpoint is not vendor-wide: Bedrock's carries an AWS region, +# Azure OpenAI's names a customer's own resource. Both are read from the environment, +# which is why they need to be re-readable rather than table literals. +_PER_RESOURCE_BASE_URL_RESOLVERS = { + "bedrock": _bedrock_default_base_url, + "azure": _azure_default_base_url, +} + + +def default_base_url(provider: str) -> Optional[str]: + """The provider's default endpoint, re-resolving the per-resource ones. + + ``PROVIDER_BASE_URLS`` is built once at import, so for the two providers above + its entry is only a snapshot of the environment as it stood then. A process that + learns its region or its resource *after* this module was imported — the usual + order, since configuration is injected at startup and this module is imported by + the package — would read that stale snapshot forever, and for Azure the snapshot + is typically the empty string. The cost of re-reading is one environment lookup; + the cost of not re-reading is an import-order dependency that is invisible until + it silently sends a request nowhere. + """ + resolver = _PER_RESOURCE_BASE_URL_RESOLVERS.get(provider) + if resolver is not None: + return resolver() or PROVIDER_BASE_URLS.get(provider) + return PROVIDER_BASE_URLS.get(provider) + + # Define the provider base URLs dictionary PROVIDER_BASE_URLS = { "openai": "https://api.openai.com/v1/", @@ -94,6 +134,11 @@ def _bedrock_default_base_url() -> str: # so this entry is the AWS_REGION-resolved default; the adapter re-resolves # it per instance (see adapters/bedrock.bedrock_base_url). "bedrock": _bedrock_default_base_url(), + # OpenAI models on an Azure OpenAI resource. Per-resource rather than + # region-dependent, so this entry is only the environment-resolved default and + # is empty on a host that has none; the adapter re-resolves per instance (see + # adapters/azure.azure_openai_base_url). + "azure": _azure_default_base_url(), "openai-oauth": "https://chatgpt.com/backend-api/codex/responses", # ChatGPT OAuth endpoint "anthropic-oauth": "https://api.anthropic.com/v1/messages?beta=true", # Claude OAuth endpoint } @@ -115,7 +160,7 @@ class ModelConfig(BaseModel): description="Model identifier (e.g., 'gpt-4o', 'mistralai/Mistral-7B-Instruct-v0.1')", ) provider: Optional[ - Literal["openai", "openrouter", "google", "anthropic", "xai", "bedrock", "openai-oauth", "anthropic-oauth"] + Literal["openai", "openrouter", "google", "anthropic", "xai", "bedrock", "azure", "openai-oauth", "anthropic-oauth"] ] = Field( None, description="API provider name (used to determine base_url if not set)" ) @@ -208,14 +253,23 @@ def _set_base_url_from_provider(cls, data: Any) -> Any: provider = data.get("provider") if provider: # Look up base_url from the dictionary - base_url = PROVIDER_BASE_URLS.get(provider) + base_url = default_base_url(provider) if base_url: data["base_url"] = base_url - else: - # Provider specified but not in our known dictionary + elif provider not in PROVIDER_BASE_URLS: warnings.warn( f"Unknown API provider '{provider}'. 'base_url' must be set explicitly if needed." ) + else: + # Known provider whose endpoint names a customer's own resource + # (Azure OpenAI), unresolved on this host. Distinct from an + # unknown provider, and said so: the fix is to supply the + # endpoint, not to correct the provider name. + warnings.warn( + f"API provider '{provider}' has no default endpoint — its base URL is " + "per-resource. Set 'base_url' explicitly or configure the provider's " + "endpoint environment variable." + ) else: # Raise error only if type is API and neither provider nor base_url is set raise ValueError( @@ -241,6 +295,8 @@ def _validate_api_key(self) -> "ModelConfig": # Bedrock authenticates with a bearer token, not SigV4, on the # Messages-API-shaped endpoint this stack targets. "bedrock": "AWS_BEARER_TOKEN_BEDROCK", + # An Azure OpenAI resource key, sent in the `api-key` header. + "azure": "AZURE_OPENAI_API_KEY", } # Providers that use OAuth or other credential mechanisms (not API keys) oauth_providers = {"openai-oauth", "anthropic-oauth"} @@ -871,6 +927,32 @@ def sync_run(): response = await loop.run_in_executor(None, sync_run) return response + async def acount_tokens( + self, + messages: List[Dict[str, Any]], + *, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[str] = None, + ) -> Optional[int]: + """How many input tokens this model would charge for that body, counted by + the provider rather than estimated, or ``None`` when the provider offers no + such service. + + ``None`` is the default and the only honest answer for a provider without a + counting endpoint: a caller that needs a number can fall back to whatever + estimate it already has, but it must be able to tell an estimate from a + count. Providers that CAN answer implement ``acount_tokens`` on their async + adapter, where the payload rendering and credentials live. + + This is not the per-message ``TokenCounter`` in ``marsys.utils.tokens``: + that protocol returns a count per message from a character heuristic, which + one endpoint call cannot produce. Two different capabilities. + """ + counter = getattr(self.async_adapter, "acount_tokens", None) + if counter is None: + return None + return await counter(messages, tools=tools, system=system) + async def cleanup(self): """Clean up async resources.""" if self.async_adapter and hasattr(self.async_adapter, 'cleanup'): diff --git a/src/marsys/models/serialize.py b/src/marsys/models/serialize.py index 1f434683..d8d782ab 100644 --- a/src/marsys/models/serialize.py +++ b/src/marsys/models/serialize.py @@ -39,6 +39,7 @@ "anthropic", "xai", "bedrock", + "azure", "openai-oauth", "anthropic-oauth", ] diff --git a/tests/models/test_azure_openai_leg.py b/tests/models/test_azure_openai_leg.py new file mode 100644 index 00000000..7612b974 --- /dev/null +++ b/tests/models/test_azure_openai_leg.py @@ -0,0 +1,478 @@ +"""The Azure OpenAI leg: what the wire accepts, and what the meter is told it cost. + +Every assertion here is the offline half of something measured against a real Azure +OpenAI resource. The three that would be cheap to get wrong and expensive to discover +live: + +* an unrecognized per-item key in the ``input`` array is a hard ``400 + unknown_parameter``, so the message sanitize is a connectivity requirement rather + than hygiene; +* ``cached_tokens`` is a SLICE of ``input_tokens`` on this API and the ADDITIVE + complement of ``prompt_tokens`` in ``UsageInfo``, so a field-to-field mapping + double-counts the cached prefix and bills it at the fresh-input rate; +* ``reasoning_tokens`` is a SLICE of ``output_tokens``, so anything that adds the two + overstates output by up to the whole reasoning budget. +""" + +import warnings + +import pytest + +from marsys.models.adapters.azure import ( + AsyncAzureOpenAIAdapter, + AzureOpenAIAdapter, + azure_openai_base_url, +) +from marsys.models.adapters.factory import ProviderAdapterFactory +from marsys.models.adapters.openai import ( + OpenAIAdapter, + thinking_budget_to_effort, +) +from marsys.models.models import PROVIDER_BASE_URLS +from marsys.models.serialize import ApiProvider + +MESSAGES = [{"role": "user", "content": "hi"}] +RESOURCE = "https://marsys-dev-fn-01.services.ai.azure.com" + + +def _azure(model_name: str = "gpt-5.6-sol", **kwargs) -> AzureOpenAIAdapter: + return AzureOpenAIAdapter( + model_name=model_name, + api_key="not-a-real-key", + base_url=f"{RESOURCE}/openai/v1", + max_tokens=4096, + **kwargs, + ) + + +# --- the endpoint ------------------------------------------------------------ + + +@pytest.mark.parametrize( + "given, expected", + [ + (RESOURCE, f"{RESOURCE}/openai/v1"), + (f"{RESOURCE}/", f"{RESOURCE}/openai/v1"), + # The portal's copyable "project endpoint", which is what the operator's + # escrowed copy of this value actually holds. It does not serve /openai/v1. + (f"{RESOURCE}/api/projects/marsys-llm-api-01", f"{RESOURCE}/openai/v1"), + (f"{RESOURCE}/api/projects/marsys-llm-api-01/", f"{RESOURCE}/openai/v1"), + # Idempotent, so a caller who configured the full base is not doubled up. + (f"{RESOURCE}/openai/v1", f"{RESOURCE}/openai/v1"), + (f"{RESOURCE}/openai/v1/", f"{RESOURCE}/openai/v1"), + ("https://marsys-dev-fn-01.openai.azure.com", "https://marsys-dev-fn-01.openai.azure.com/openai/v1"), + (" " + RESOURCE + " ", f"{RESOURCE}/openai/v1"), + ], +) +def test_base_url_normalizes_every_form_one_resource_arrives_as(given, expected): + assert azure_openai_base_url(given) == expected + + +def test_an_explicit_base_url_is_normalized_like_every_other_spelling(monkeypatch): + """The ``base_url`` argument and the ``endpoint`` argument are the same fact arriving + under two names, so one of them cannot skip the normalizer: a caller passing the bare + resource host would otherwise build a client that POSTs to ``/responses`` off the host + root and 404s on every call.""" + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + monkeypatch.delenv("FOUNDRY_ENDPOINT", raising=False) + bare = AzureOpenAIAdapter(model_name="gpt-5.6-sol", api_key="k", base_url=RESOURCE) + assert bare.get_endpoint_url() == f"{RESOURCE}/openai/v1/responses" + project = AzureOpenAIAdapter( + model_name="gpt-5.6-sol", + api_key="k", + base_url=f"{RESOURCE}/api/projects/marsys-llm-api-01", + ) + assert project.get_endpoint_url() == f"{RESOURCE}/openai/v1/responses" + # Still idempotent on the normalized form Spren actually threads in. + already = AzureOpenAIAdapter( + model_name="gpt-5.6-sol", api_key="k", base_url=f"{RESOURCE}/openai/v1" + ) + assert already.get_endpoint_url() == f"{RESOURCE}/openai/v1/responses" + + +def test_base_url_is_empty_when_nothing_is_configured(monkeypatch): + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + monkeypatch.delenv("FOUNDRY_ENDPOINT", raising=False) + assert azure_openai_base_url() == "" + + +def test_base_url_reads_the_vendor_name_first_then_the_house_name(monkeypatch): + monkeypatch.setenv("FOUNDRY_ENDPOINT", "https://house.services.ai.azure.com") + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + assert azure_openai_base_url() == "https://house.services.ai.azure.com/openai/v1" + + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", RESOURCE) + assert azure_openai_base_url() == f"{RESOURCE}/openai/v1" + + +def test_endpoint_is_the_v1_responses_path_with_no_api_version(monkeypatch): + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", RESOURCE) + url = AzureOpenAIAdapter(model_name="gpt-5.6-sol", api_key="k").get_endpoint_url() + assert url == f"{RESOURCE}/openai/v1/responses" + # The legacy surface's two markers. Both absent: no per-deployment path segment, + # and no pinned api-version query — the v1 GA API requires neither. + assert "/deployments/" not in url + assert "api-version" not in url + + +# --- auth and identity ------------------------------------------------------- + + +def test_auth_is_the_api_key_header_not_a_bearer(): + headers = _azure().get_headers() + assert headers["api-key"] == "not-a-real-key" + assert "Authorization" not in headers + assert "x-api-key" not in headers + + +def test_both_adapters_report_the_azure_provider_id(): + assert _azure()._provider_name() == "azure" + assert ( + AsyncAzureOpenAIAdapter( + model_name="gpt-5.6-sol", api_key="k", base_url=f"{RESOURCE}/openai/v1" + )._provider_name() + == "azure" + ) + + +def test_the_deployment_name_travels_in_the_model_body_field(): + """Azure addresses a deployment, not a model snapshot; on the v1 surface the + deployment name is the `model` field's value and the response echoes it back.""" + payload = _azure("gpt-5.6-terra").format_request_payload(MESSAGES) + assert payload["model"] == "gpt-5.6-terra" + + +def test_the_registry_rows_agree_that_azure_exists(): + from typing import get_args + + assert "azure" in get_args(ApiProvider) + assert "azure" in PROVIDER_BASE_URLS + adapter = ProviderAdapterFactory.create_adapter( + provider="azure", model_name="gpt-5.6-sol", api_key="k", base_url=f"{RESOURCE}/openai/v1" + ) + assert isinstance(adapter, AzureOpenAIAdapter) + + +def test_a_config_for_this_provider_validates(monkeypatch): + """Also the regression for the import-order trap: this endpoint is not in the + module's literal table, it is read from the environment, and `PROVIDER_BASE_URLS` + is built once at import — which happens before a daemon injects its configuration. + Snapshotted, the entry stays empty for the life of the process and every request + goes nowhere. The env vars here are deliberately set AFTER the import above.""" + from marsys.models.models import ModelConfig + + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "resource-key") + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", RESOURCE) + assert PROVIDER_BASE_URLS["azure"] == "", "the snapshot is empty — that is the point" + cfg = ModelConfig(type="api", provider="azure", name="gpt-5.6-sol") + assert cfg.api_key == "resource-key" + assert cfg.base_url == f"{RESOURCE}/openai/v1" + + +def test_an_unresolved_resource_says_the_endpoint_is_missing(monkeypatch): + """A known provider with no resolvable endpoint is a different failure from a + misspelled provider name, and the message has to say which: the fix here is to + supply an endpoint, not to correct a spelling. (The misspelling case never + reaches this branch — `provider` is a Literal, so Pydantic rejects it first.)""" + from marsys.models.models import ModelConfig + + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "resource-key") + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + monkeypatch.delenv("FOUNDRY_ENDPOINT", raising=False) + with pytest.warns(UserWarning, match="per-resource"): + ModelConfig(type="api", provider="azure", name="gpt-5.6-sol") + + +# --- error classification ---------------------------------------------------- + + +class _FakeResponse: + def __init__(self, status_code, body, headers=None): + self.status_code = status_code + self._body = body + self.headers = headers or {} + + def json(self): + return self._body + + +def test_an_azure_rate_limit_is_classified_the_way_a_first_party_one_is(): + """The classifier branches on the error *envelope*, and Azure's is OpenAI's. Left + out of that branch, an Azure 429 matches nothing, classifies as UNKNOWN and comes + back not retryable — the request is dropped instead of backed off, which on the + interactive path looks like the model refusing to answer.""" + response = _FakeResponse( + 429, + {"error": {"code": "429", "message": "Rate limit exceeded", "type": "rate_limit_error"}}, + {"retry-after": "13"}, + ) + error = _azure().handle_api_error(RuntimeError("429 Too Many Requests"), response=response) + assert error.classification["category"] == "rate_limit" + assert error.classification["is_retryable"] is True + assert error.classification["retry_after"] == 13 + # And the vendor is named honestly, rather than reported as first-party OpenAI. + assert error.provider == "azure" + + +def test_an_azure_server_error_is_retryable(): + error = _azure().handle_api_error( + RuntimeError("503"), response=_FakeResponse(503, {"error": {"message": "busy"}}) + ) + assert error.classification["category"] == "service_unavailable" + assert error.classification["is_retryable"] is True + + +# --- the wire sanitize ------------------------------------------------------- + + +def test_internal_message_keys_never_reach_the_wire(): + """A caller's message dict is its own working object. This stack's reasoner + annotates messages with routing and provenance fields, and the Responses API + answers an unrecognized per-item key with `400 unknown_parameter: input[0].` + — so forwarding them is not untidy, it is a leg that cannot connect.""" + annotated = [ + { + "role": "user", + "content": "hi", + "kind": "user_message", + "actor": "principal", + "cache_exempt": True, + "name": "someone", + } + ] + payload = _azure().format_request_payload(annotated) + item = payload["input"][0] + assert set(item) == {"role", "content"} + assert item["role"] == "user" + assert item["content"] == "hi" + + +def test_the_item_discriminator_survives_the_sanitize(): + """`type` is not an internal key here — the Responses input array uses it to tell + a message from a function_call, so an allow-list that dropped it would break + tool round-trips.""" + payload = _azure().format_request_payload( + [{"type": "function_call_output", "call_id": "c1", "role": "user", "content": "x"}] + ) + item = payload["input"][0] + assert item["type"] == "function_call_output" + assert "call_id" not in item + + +def test_the_sanitize_is_the_shared_openai_behaviour_not_an_azure_special_case(): + """One wire contract, one implementation. A fix that lived only on the subclass + would leave the first-party leg rejecting the same annotated messages.""" + first_party = OpenAIAdapter( + model_name="gpt-5.5", api_key="k", base_url="https://api.openai.com/v1/" + ) + payload = first_party.format_request_payload( + [{"role": "user", "content": "hi", "kind": "user_message"}] + ) + assert set(payload["input"][0]) == {"role", "content"} + + +def test_no_cache_option_field_is_sent(): + """Prompt caching on this API is implicit and takes no request parameter. A + `prompt_cache_breakpoint` is rejected outright in every position, and `explicit` + mode without one caches nothing — so the correct request says nothing at all, + and the measured pair of identical calls still reports a write then a read.""" + payload = _azure().format_request_payload(MESSAGES) + assert "prompt_cache_options" not in payload + assert "prompt_cache_breakpoint" not in payload + assert not any("cache" in key for key in payload) + + +# --- the thinking knob ------------------------------------------------------- + + +@pytest.mark.parametrize( + "budget, expected", + [ + (None, None), + (0, None), # thinking off is not "think as little as possible" + (-1, None), + (512, "minimal"), + (1024, "low"), + (4095, "low"), + (4096, "medium"), + (8192, "medium"), + (16384, "high"), + (32768, "high"), + ], +) +def test_a_thinking_budget_selects_a_reasoning_effort(budget, expected): + assert thinking_budget_to_effort(budget) == expected + + +def test_the_configured_thinking_budget_reaches_this_leg(): + """The only deliberation knob this stack exposes is a token budget. Unmapped, a + caller's setting is inert and every call runs at the provider default (`medium`), + which is indistinguishable from the knob working.""" + payload = _azure().format_request_payload(MESSAGES, thinking_budget=32768) + assert payload["reasoning"] == {"effort": "high"} + + +def test_an_explicit_effort_beats_the_budget(): + payload = _azure().format_request_payload( + MESSAGES, thinking_budget=32768, reasoning_effort="low" + ) + assert payload["reasoning"] == {"effort": "low"} + + +def test_thinking_off_sends_no_reasoning_block(): + payload = _azure().format_request_payload(MESSAGES, thinking_budget=0) + assert "reasoning" not in payload + + +def test_the_budget_kwarg_does_not_warn_as_unknown(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + _azure().format_request_payload(MESSAGES, thinking_budget=8192) + + +def test_the_smallest_budget_asks_for_an_effort_this_surface_actually_serves(): + """`minimal` is a 400 here — the endpoint's own reply lists none / low / medium / + high / xhigh / max — so the smallest bucket has to arrive as the smallest this + surface has. It must still ask for reasoning: a positive budget is "think a little", + and `none` would answer a question nobody asked.""" + payload = _azure().format_request_payload(MESSAGES, thinking_budget=512) + assert payload["reasoning"] == {"effort": "low"} + + +def test_an_explicit_minimal_is_substituted_too(): + """The caller who names the effort outright is on the same endpoint as the one who + named a budget, and it rejects the value for both of them.""" + payload = _azure().format_request_payload(MESSAGES, reasoning_effort="minimal") + assert payload["reasoning"] == {"effort": "low"} + + +def test_the_first_party_leg_still_sends_minimal(): + """The control, and the scope line: `minimal` is served by OpenAI's own endpoint and + the substitution above belongs to this re-hosting surface, not to the shared payload + builder. A run of this file that changed the first-party leg would be a silent change + to every OpenAI caller in the stack.""" + payload = OpenAIAdapter( + model_name="gpt-5.6", api_key="k", base_url="https://api.openai.com/v1" + ).format_request_payload(MESSAGES, thinking_budget=512) + assert payload["reasoning"] == {"effort": "minimal"} + codex = OpenAIAdapter( + model_name="gpt-5.6-codex", api_key="k", base_url="https://api.openai.com/v1" + ).format_request_payload(MESSAGES, thinking_budget=512) + assert codex["reasoning"] == {"effort": "low"} # the pre-existing codex exception + + +# --- the meter --------------------------------------------------------------- + + +def _harmonized(usage: dict): + raw = { + "id": "resp_1", + "model": "gpt-5.6-sol", + "created_at": 1, + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok"}], + } + ], + "usage": usage, + } + return _azure().harmonize_response(raw, request_start_time=0.0) + + +def test_a_cached_prefix_is_not_counted_twice(): + """Measured on the resource: `input_tokens: 3398` CONTAINING `cached_tokens: 3395`, + with `3398 + output 5 == total 3403`. `UsageInfo.prompt_tokens` is the inverse + convention — the uncached remainder, with the cache figures beside it — so the + slice has to be subtracted once here. Mapped straight across, the cached prefix + would be billed as fresh input on top of being billed as a cache read.""" + usage = _harmonized( + { + "input_tokens": 3398, + "output_tokens": 5, + "total_tokens": 3403, + "input_tokens_details": {"cached_tokens": 3395}, + } + ).metadata.usage + assert usage.prompt_tokens == 3 + assert usage.cache_read_input_tokens == 3395 + assert usage.cache_creation_input_tokens is None + # The whole prompt is recoverable, exactly: nothing lost, nothing doubled. + assert usage.full_prompt_tokens == 3398 + + +def test_a_cache_write_is_split_out_the_same_way(): + """The first of the measured pair: the same 3395 tokens, reported as a write.""" + usage = _harmonized( + { + "input_tokens": 3398, + "output_tokens": 5, + "total_tokens": 3403, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 3395}, + } + ).metadata.usage + assert usage.prompt_tokens == 3 + assert usage.cache_creation_input_tokens == 3395 + assert usage.cache_read_input_tokens is None + assert usage.full_prompt_tokens == 3398 + + +def test_slices_that_exceed_their_whole_do_not_walk_the_ledger_backwards(): + usage = _harmonized( + { + "input_tokens": 100, + "output_tokens": 5, + "input_tokens_details": {"cached_tokens": 400}, + } + ).metadata.usage + assert usage.prompt_tokens == 0 + + +def test_a_provider_that_reports_no_cache_leaves_both_fields_unset(): + usage = _harmonized( + {"input_tokens": 75, "output_tokens": 1186, "total_tokens": 1261} + ).metadata.usage + assert usage.prompt_tokens == 75 + assert usage.cache_read_input_tokens is None + assert usage.cache_creation_input_tokens is None + assert usage.full_prompt_tokens == 75 + + +def test_reasoning_tokens_are_recorded_as_a_subset_of_output(): + """The vendor's own arithmetic proves the containment: `input 75 + output 1186 == + total 1261`, with `reasoning_tokens: 1024` reported inside the 1186. A meter that + adds the two charges the reasoning slice twice — here, 86% high.""" + usage = _harmonized( + { + "input_tokens": 75, + "output_tokens": 1186, + "total_tokens": 1261, + "output_tokens_details": {"reasoning_tokens": 1024}, + } + ).metadata.usage + assert usage.completion_tokens == 1186 + assert usage.reasoning_tokens == 1024 + assert usage.prompt_tokens + usage.completion_tokens == usage.total_tokens == 1261 + + +def test_the_harmonized_response_reports_azure_and_the_echoed_deployment(): + """`metadata.model` is what a cost table is keyed on, and a table miss prices at + zero silently. Azure echoes the DEPLOYMENT name (measured: `gpt-5.6-sol` for a + request made with `gpt-5.6-sol`), so the echo needs no renormalization — but the + provider label must be this leg's, not the first-party one's.""" + response = _harmonized({"input_tokens": 3, "output_tokens": 5, "total_tokens": 8}) + assert response.metadata.provider == "azure" + assert response.metadata.model == "gpt-5.6-sol" + + +def test_chat_completions_usage_names_still_read(): + """The same harmonizer serves any OpenAI-compatible endpoint; the older field + names must not silently become zeros.""" + usage = _harmonized( + {"prompt_tokens": 40, "completion_tokens": 7, "total_tokens": 47} + ).metadata.usage + assert usage.prompt_tokens == 40 + assert usage.completion_tokens == 7 diff --git a/tests/models/test_count_tokens.py b/tests/models/test_count_tokens.py new file mode 100644 index 00000000..24588403 --- /dev/null +++ b/tests/models/test_count_tokens.py @@ -0,0 +1,416 @@ +"""Provider-counted input tokens on the two Anthropic legs. + +No network. What is pinned here: + +* **The counted body is the body that would be SENT.** Each leg builds the count + payload through its own ``format_request_payload``, so the OAuth leg's required + Claude-Code system block and its tool renaming are counted too. A count assembled + any other way measures a request nobody makes. +* **Only the generation controls come off.** ``max_tokens``/``stream``/sampling are + rejected by the count endpoint; ``system``/``messages``/``tools``/``thinking`` are + exactly what decides the number. +* **A failure is ``None``, never an exception.** The caller is sizing something, not + producing an answer. A missing route (404/405) is remembered for the process; a + credential-shaped refusal is not — the OAuth token file has several writers and a + refresh in flight looks exactly like a rejection for one request. +""" + +import json + +import httpx +import pytest + +from marsys.models.adapters import anthropic as anthropic_mod +from marsys.models.adapters.anthropic import ( + AsyncAnthropicAdapter, + count_tokens_url_for, + strip_for_count_tokens, +) +from marsys.models.adapters.anthropic_oauth import AsyncAnthropicOAuthAdapter +from marsys.models.models import BaseAPIModel + +MESSAGES = [{"role": "user", "content": "hi"}] +TOOLS = [ + { + "type": "function", + "function": { + "name": "read_file", + "description": "read a file", + "parameters": {"type": "object", "properties": {}}, + }, + } +] + + +@pytest.fixture(autouse=True) +def _forget_unsupported_endpoints(): + """The unavailable-endpoint memo is process-wide by design; tests must not + inherit each other's.""" + anthropic_mod._COUNT_TOKENS_UNSUPPORTED.clear() + yield + anthropic_mod._COUNT_TOKENS_UNSUPPORTED.clear() + + +class _FakeResponse: + def __init__(self, status: int, body): + self.status = status + self._body = body + + async def json(self, content_type=None): + if self._body is None: + raise ValueError("not json") + return self._body + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +_DEFAULT_BODY = object() # so ``body=None`` can mean "the response is not JSON" + + +class _FakeSession: + """The aiohttp seam ``AsyncBaseAPIAdapter._ensure_session`` hands out.""" + + # ``_ensure_session`` reuses a live session and rebuilds a closed one. + closed = False + + def __init__(self, status: int = 200, body=_DEFAULT_BODY, raises=None): + self.status = status + self.body = {"input_tokens": 1234} if body is _DEFAULT_BODY else body + self.raises = raises + self.calls: list[dict] = [] + + def post(self, url, headers=None, json=None, timeout=None): + self.calls.append({"url": url, "headers": headers, "payload": json}) + if self.raises is not None: + raise self.raises + return _FakeResponse(self.status, self.body) + + +def _api(session: _FakeSession) -> AsyncAnthropicAdapter: + adapter = AsyncAnthropicAdapter( + model_name="claude-opus-5", + api_key="not-a-real-key", + base_url="https://api.anthropic.com/v1", + max_tokens=8192, + ) + adapter._session = session + return adapter + + +def _oauth() -> AsyncAnthropicOAuthAdapter: + """An OAuth adapter that never touches the credentials file on disk.""" + adapter = object.__new__(AsyncAnthropicOAuthAdapter) + adapter.model_name = "claude-sonnet-5" + adapter.max_tokens = 8192 + adapter.temperature = 0.7 + adapter.enable_thinking = False + adapter.thinking_budget = 0 + adapter.auto_refresh = False + adapter.access_token = "not-a-real-token" + adapter.credentials = {"access_token": "not-a-real-token"} + adapter._credentials_path = "unused" + adapter._session = None + return adapter + + +def _oauth_transport(monkeypatch, handler): + real_async_client = httpx.AsyncClient + transport = httpx.MockTransport(handler) + monkeypatch.setattr( + httpx, "AsyncClient", lambda **kw: real_async_client(transport=transport) + ) + + +# ── the URL and the payload ─────────────────────────────────────────────────── + + +def test_the_count_url_sits_beside_the_messages_url_and_keeps_its_query(): + assert ( + count_tokens_url_for("https://api.anthropic.com/v1/messages") + == "https://api.anthropic.com/v1/messages/count_tokens" + ) + # The OAuth leg's endpoint carries a query string; dropping it would send the + # count somewhere the messages call never goes. + assert ( + count_tokens_url_for("https://api.anthropic.com/v1/messages?beta=true") + == "https://api.anthropic.com/v1/messages/count_tokens?beta=true" + ) + + +def test_the_final_assistant_message_is_trimmed_of_trailing_whitespace(): + """The API rejects a final assistant message that ends in whitespace, and a count + is the one request that routinely presents one: a settled conversation ends with + an assistant reply, and models end replies with a newline. Measured live — the + provider answers "final assistant content cannot end with trailing whitespace".""" + payload = { + "model": "m", + "messages": [ + {"role": "user", "content": "ask"}, + {"role": "assistant", "content": [{"type": "text", "text": "answered.\n\n"}]}, + ], + } + stripped = strip_for_count_tokens(payload) + assert stripped["messages"][-1]["content"][-1]["text"] == "answered." + # …and the caller's own rows are untouched: the payload's blocks may be the + # durable conversation's dicts. + assert payload["messages"][-1]["content"][-1]["text"] == "answered.\n\n" + + +def test_a_plain_string_final_assistant_message_is_trimmed_too(): + stripped = strip_for_count_tokens( + {"model": "m", "messages": [{"role": "assistant", "content": "done "}]} + ) + assert stripped["messages"][-1]["content"] == "done" + + +def test_a_whitespace_only_final_assistant_message_is_dropped(): + """Nothing is left of it to count, and an empty text block is itself rejected.""" + stripped = strip_for_count_tokens( + { + "model": "m", + "messages": [ + {"role": "user", "content": "ask"}, + {"role": "assistant", "content": [{"type": "text", "text": " "}]}, + ], + } + ) + assert stripped["messages"] == [{"role": "user", "content": "ask"}] + + +def test_a_final_user_message_is_left_exactly_as_it_is(): + """The rule is about assistant content. Trimming a user row would change what is + being counted for no reason.""" + stripped = strip_for_count_tokens( + {"model": "m", "messages": [{"role": "user", "content": "ask \n"}]} + ) + assert stripped["messages"] == [{"role": "user", "content": "ask \n"}] + + +def test_only_the_generation_controls_are_stripped(): + payload = { + "model": "claude-opus-5", + "max_tokens": 8192, + "stream": True, + "temperature": 0.7, + "top_p": 0.9, + "output_config": {"effort": "high"}, + "system": [{"type": "text", "text": "SYS"}], + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "t", "input_schema": {}}], + "thinking": {"type": "adaptive"}, + } + stripped = strip_for_count_tokens(payload) + assert set(stripped) == {"model", "system", "messages", "tools", "thinking"} + # the survivors are untouched, not rebuilt + assert stripped["system"] == payload["system"] + assert stripped["tools"] == payload["tools"] + + +# ── the api-key leg ─────────────────────────────────────────────────────────── + + +async def test_the_api_key_leg_counts_the_body_it_would_send(): + session = _FakeSession(body={"input_tokens": 4_242}) + adapter = _api(session) + + count = await adapter.acount_tokens(MESSAGES, tools=TOOLS, system="SYS") + + assert count == 4_242 + call = session.calls[0] + assert call["url"] == "https://api.anthropic.com/v1/messages/count_tokens" + # The system string rides as the payload builder's own rendered system field — + # the block shape the messages call sends, not the raw string. + assert call["payload"]["system"][0]["text"] == "SYS" + # …the tools are the converted Anthropic shape the request would carry… + assert call["payload"]["tools"][0]["name"] == "read_file" + assert "input_schema" in call["payload"]["tools"][0] + # …and the generation controls are gone. + assert "max_tokens" not in call["payload"] and "stream" not in call["payload"] + assert call["headers"]["accept"] == "application/json" + assert call["headers"]["x-api-key"] == "not-a-real-key" + + +async def test_a_body_with_no_system_is_counted_as_is(): + session = _FakeSession(body={"input_tokens": 11}) + adapter = _api(session) + assert await adapter.acount_tokens(MESSAGES) == 11 + assert session.calls[0]["payload"]["messages"][0]["role"] == "user" + + +async def test_a_transport_failure_answers_none_rather_than_raising(): + import aiohttp + + session = _FakeSession(raises=aiohttp.ClientError("connection reset")) + adapter = _api(session) + assert await adapter.acount_tokens(MESSAGES) is None + + +async def test_a_body_that_is_not_json_answers_none(): + session = _FakeSession(status=200, body=None) + adapter = _api(session) + assert await adapter.acount_tokens(MESSAGES) is None + + +async def test_a_200_without_input_tokens_answers_none(): + session = _FakeSession(status=200, body={"unexpected": True}) + adapter = _api(session) + assert await adapter.acount_tokens(MESSAGES) is None + + +# ── availability: structural vs credential-shaped ───────────────────────────── + + +@pytest.mark.parametrize("status", [404, 405]) +async def test_a_missing_route_is_remembered_for_the_process(status): + """The endpoint is not there; asking again every turn buys nothing.""" + session = _FakeSession(status=status, body={"error": "not found"}) + adapter = _api(session) + + assert await adapter.acount_tokens(MESSAGES) is None + assert await adapter.acount_tokens(MESSAGES) is None + assert len(session.calls) == 1 # the second call never left the process + + +@pytest.mark.parametrize("status", [401, 403, 429, 500]) +async def test_a_credential_or_transient_refusal_is_never_remembered(status): + """The OAuth token file has several writers, so a refresh in flight looks like a + rejection for one request. Memoising that would disable counting for the life of + a daemon that never restarts.""" + session = _FakeSession(status=status, body={"error": "nope"}) + adapter = _api(session) + + assert await adapter.acount_tokens(MESSAGES) is None + assert await adapter.acount_tokens(MESSAGES) is None + assert len(session.calls) == 2 # asked again, as it must be + + +async def test_the_memo_is_per_endpoint_not_global(): + unavailable = _api(_FakeSession(status=404, body={})) + assert await unavailable.acount_tokens(MESSAGES) is None + + other_session = _FakeSession(body={"input_tokens": 7}) + other = AsyncAnthropicAdapter( + model_name="claude-opus-5", + api_key="k", + base_url="https://other.example.invalid/v1", + max_tokens=1024, + ) + other._session = other_session + assert await other.acount_tokens(MESSAGES) == 7 + + +# ── the OAuth leg ───────────────────────────────────────────────────────────── + + +async def test_the_oauth_leg_counts_what_its_own_payload_builder_renders(monkeypatch): + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["headers"] = dict(request.headers) + seen["payload"] = json.loads(request.content) + return httpx.Response(200, json={"input_tokens": 9_001}) + + _oauth_transport(monkeypatch, handler) + adapter = _oauth() + monkeypatch.setattr(adapter, "_ensure_fresh_token", lambda: None, raising=False) + + count = await adapter.acount_tokens(MESSAGES, tools=TOOLS, system="SYS") + + assert count == 9_001 + assert seen["url"] == "https://api.anthropic.com/v1/messages/count_tokens?beta=true" + # The required Claude-Code prefix block is COUNTED, because the messages call + # always sends it — a count without it undercounts every request. + assert seen["payload"]["system"][0]["text"] == adapter.CLAUDE_CODE_PREFIX + assert seen["payload"]["system"][1]["text"] == "SYS" + # …as is the reserved-name transform this leg applies to tools. + assert seen["payload"]["tools"][0]["name"] == "Read" + assert "max_tokens" not in seen["payload"] and "stream" not in seen["payload"] + # JSON, not the messages call's SSE. + assert seen["headers"]["accept"] == "application/json" + assert seen["headers"]["authorization"].startswith("Bearer ") + assert "oauth-2025-04-20" in seen["headers"]["anthropic-beta"] + + +async def test_the_oauth_leg_refreshes_its_token_before_counting(monkeypatch): + """Same discipline as the messages call: the cached token is a per-request + snapshot of a file other processes rewrite.""" + refreshed: list[bool] = [] + _oauth_transport( + monkeypatch, lambda request: httpx.Response(200, json={"input_tokens": 5}) + ) + adapter = _oauth() + monkeypatch.setattr( + adapter, "_ensure_fresh_token", lambda: refreshed.append(True), raising=False + ) + + assert await adapter.acount_tokens(MESSAGES) == 5 + assert refreshed == [True] + + +async def test_an_unauthorised_oauth_leg_answers_none_and_is_asked_again(monkeypatch): + """If the subscription credential is not accepted here, compaction reports "not + measured" and the turn is untouched — and the next turn asks again. This leg's + token file has several writers, so a refusal is as likely to be a refresh in + flight as a verdict, and a daemon that never restarts would otherwise stop + measuring for days on one blip.""" + calls: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(1) + return httpx.Response(403, json={"error": {"type": "forbidden"}}) + + _oauth_transport(monkeypatch, handler) + adapter = _oauth() + monkeypatch.setattr(adapter, "_ensure_fresh_token", lambda: None, raising=False) + + assert await adapter.acount_tokens(MESSAGES) is None + assert await adapter.acount_tokens(MESSAGES) is None + assert len(calls) == 2 + + +async def test_an_oauth_transport_failure_answers_none(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("no route to host") + + _oauth_transport(monkeypatch, handler) + adapter = _oauth() + monkeypatch.setattr(adapter, "_ensure_fresh_token", lambda: None, raising=False) + + assert await adapter.acount_tokens(MESSAGES) is None + + +# ── the model port ──────────────────────────────────────────────────────────── + + +class _AdapterWithCount: + def __init__(self): + self.seen: dict = {} + + async def acount_tokens(self, messages, *, tools=None, system=None): + self.seen = {"messages": messages, "tools": tools, "system": system} + return 321 + + +async def test_the_model_delegates_the_count_to_its_async_adapter(): + model = object.__new__(BaseAPIModel) + adapter = _AdapterWithCount() + model.async_adapter = adapter + + assert await model.acount_tokens(MESSAGES, tools=TOOLS, system="SYS") == 321 + assert adapter.seen == {"messages": MESSAGES, "tools": TOOLS, "system": "SYS"} + + +async def test_a_provider_that_cannot_count_answers_none_rather_than_guessing(): + """Every non-Anthropic leg takes this path. The caller must be able to tell an + absent count from a small one.""" + model = object.__new__(BaseAPIModel) + model.async_adapter = object() + assert await model.acount_tokens(MESSAGES) is None + + model.async_adapter = None + assert await model.acount_tokens(MESSAGES) is None diff --git a/tests/models/test_prompt_cache_breakpoint.py b/tests/models/test_prompt_cache_breakpoint.py index e7366d91..85008d27 100644 --- a/tests/models/test_prompt_cache_breakpoint.py +++ b/tests/models/test_prompt_cache_breakpoint.py @@ -22,6 +22,7 @@ import pytest from marsys.models.adapters.anthropic import ( + CACHE_EXEMPT_KEY, AnthropicAdapter, mark_conversation_tail_for_cache, ) @@ -502,3 +503,129 @@ def test_the_marker_lands_after_the_json_mode_hint_not_before_it(): assert "JSON" in json.dumps(content) assert content[-1]["cache_control"] == EPHEMERAL assert sum(1 for b in content if "cache_control" in b) == 1 + + +# ── the per-request-content exemption (CACHE_EXEMPT_KEY) ───────────────────────────── +# +# A caller that appends per-request content after the durable conversation (a clock, a +# budget figure — anything derived from "now") needs the breakpoint to stay on the last +# DURABLE row. The entry's value is that the NEXT request can read it, which requires the +# hashed prefix to be bytes the next request still contains; a regenerated row is absent +# from it by construction, so an entry written at or after that row is unreadable forever. +# Measured on Bedrock/Opus 5, single-step turns with tools present: marker on the volatile +# row -> turn 2 read=0; marker on the last durable row -> turn 2 read=8425. + + +def _marked_indices(payload) -> list[int]: + out = [] + for i, msg in enumerate(payload["messages"]): + content = msg.get("content") + blocks = content if isinstance(content, list) else [] + if any(isinstance(b, dict) and b.get("cache_control") for b in blocks): + out.append(i) + return out + + +def test_an_exempt_tail_row_moves_the_marker_to_the_last_durable_row(): + payload = _api().format_request_payload( + [ + {"role": "user", "content": "durable question"}, + {"role": "assistant", "content": "durable answer"}, + {"role": "user", "content": "now: 09:00", CACHE_EXEMPT_KEY: True}, + ] + ) + assert _marked_indices(payload) == [1], "the marker must sit on the last durable row" + assert payload["messages"][2]["content"] == "now: 09:00", "the row still reaches the model" + + +def test_several_exempt_trailing_rows_are_all_stepped_past(): + payload = _api().format_request_payload( + [ + {"role": "user", "content": "durable"}, + {"role": "user", "content": "volatile a", CACHE_EXEMPT_KEY: True}, + {"role": "user", "content": "volatile b", CACHE_EXEMPT_KEY: True}, + ] + ) + assert _marked_indices(payload) == [0] + + +def test_an_exempt_row_with_durable_rows_after_it_does_not_move_the_marker(): + """The exemption is about POSITION — what the next request will still contain. A marked + row that is not part of the trailing run is not a tail, so the marker stays at the end.""" + payload = _api().format_request_payload( + [ + {"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}, + {"role": "user", "content": "durable"}, + ] + ) + assert _marked_indices(payload) == [1] + + +def test_an_all_exempt_list_writes_no_marker_at_all(): + """Nothing durable to anchor an entry to: a marker would cost a fresh entry per request + and read none, so none is written.""" + payload = _api().format_request_payload( + [{"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}] + ) + assert _marked_indices(payload) == [] + + +def test_the_exemption_key_never_reaches_the_wire(): + payload = _api().format_request_payload( + [ + {"role": "user", "content": "durable"}, + {"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}, + ] + ) + assert CACHE_EXEMPT_KEY not in json.dumps(payload) + + +def test_the_no_key_path_is_byte_identical_to_the_unparameterized_form(): + """Every other framework caller must be unaffected: with no exempt row the payload is + exactly what it was before this parameter existed.""" + messages = [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "two"}, + {"role": "user", "content": "three"}, + ] + payload = _api().format_request_payload([dict(m) for m in messages]) + baseline = [dict(m) for m in messages] + mark_conversation_tail_for_cache(baseline) # the default, volatile_tail=0 + assert payload["messages"] == baseline + + +def test_oauth_leg_honors_the_exemption_too(): + """The OAuth adapter is not a subclass, so its payload builder is kept deliberately + twinned. Its static Claude-Code prefix block is the other marker.""" + payload = _oauth().format_request_payload( + [ + {"role": "user", "content": "durable"}, + {"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}, + ] + ) + assert _marked_indices(payload) == [0] + assert CACHE_EXEMPT_KEY not in json.dumps(payload) + + +def test_bedrock_inherits_the_exemption(): + """Bedrock does not override ``format_request_payload``, and it is the production leg + the measured figures come from.""" + payload = BedrockAdapter( + model_name="claude-opus-5", api_key="tok" + ).format_request_payload( + [ + {"role": "user", "content": "durable"}, + {"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}, + ] + ) + assert _marked_indices(payload) == [0] + + +def test_the_exemption_is_deterministic_and_idempotent(): + messages = [ + {"role": "user", "content": "durable"}, + {"role": "user", "content": "volatile", CACHE_EXEMPT_KEY: True}, + ] + first = _api().format_request_payload([dict(m) for m in messages]) + second = _api().format_request_payload([dict(m) for m in messages]) + assert json.dumps(first, sort_keys=True) == json.dumps(second, sort_keys=True)