diff --git a/pyproject.toml b/pyproject.toml index cb962d6..b763b1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,37 @@ name = "nullrun" # the full ``flush_interval`` (5s default). Plus CI hygiene: # pip cache, ``fail-fast`` matrix, ``pytest-xdist -n auto``. No # on-wire change; backends on 1.0.0 keep working unchanged. -version = "0.13.9" +# 0.13.6 (2026-07-11): multi-agent span attachment (parent_trace_id) +# on the langgraph callback; new cost_events.parent_trace_id column +# (backend migration 217). Wire-additive — legacy backends ignore +# the field. Pairs with PR #61. +# 0.13.7 (2026-07-12): wire ``parent_trace_id`` end-to-end on +# ``/track`` (v3 single-event + legacy /track/batch). Pre-fix the +# SDK stamped the field in the langgraph callback but dropped it +# at the runtime._enrich_event / _build_v3_track_payload layers. +# No on-wire change for legacy backends; new column required on +# the v3 path. Pairs with PR #64. +# 0.13.8 (2026-07-12): hotfix #2 for parent_trace_id — the +# runtime._enrich_event parent_trace_id fallback used an +# "if not in enriched" guard that missed whenever langgraph.py +# callback's _active_runs lookup missed (run_id drift between +# auto-injected and user-supplied callbacks, or non-langgraph +# stacks). Switched to override semantics: the chain contextvar +# is the single source of truth; both caller-set and +# contextvar-fallback resolve to the same value (idempotent for +# the happy path, closes the drift in the unhappy path). +# Pairs with PR #66. +# 0.13.9 (2026-07-13): crewai 1.15 compatibility — replace +# step_callback kwargs injection (removed upstream) with a +# crewai_event_bus bridge; gate_cache re-capture for fresh +# server-minted execution_id on cache-hit. Pairs with PR #67. +# 0.13.10 (2026-07-13): close 5 vendor extractor edge cases +# missed in the 0.13.9 audit — Cohere v2 nested tool_calls + +# cached_tokens + UPPERCASE finish_reason; Mistral flat +# num_cached_tokens fallback; Gemini 2.5+ thoughtsTokenCount; +# Anthropic 4.5+ extended-thinking tokens; Bedrock Mistral/Llama +# finish_reason paths. No on-wire change; no SDK_MIN_VERSION bump. +version = "0.13.10" # Kept under the 200-char preview threshold so the full line is visible # without an "expand" click. Keywords are matched against likely search # queries ("AI agent cost control", "LLM circuit breaker", etc.). diff --git a/src/nullrun/__version__.py b/src/nullrun/__version__.py index 012d739..7581faa 100644 --- a/src/nullrun/__version__.py +++ b/src/nullrun/__version__.py @@ -1,5 +1,120 @@ """NullRun Platform SDK. +v3.24 / 0.13.10 (2026-07-13) — close 5 vendor extractor edge cases +missed in the 0.13.9 audit. + + 1. Cohere v2 tool_calls path: the pre-0.13.10 extractor read + top-level payload["tool_calls"], but Cohere v2 nests the field + under message.tool_calls (OpenAI shape). Every v2 Cohere call + shipped with tool_names=[] and the backend's loop detection + could not see Cohere tool use. Fix walks both v1 (top-level) + and v2 (message.tool_calls) paths. Same patch adds + usage.tokens.cached_tokens (cache-hit read was always 0) and + the UPPERCASE finish_reason vocabulary + (COMPLETE | MAX_TOKENS | TOOL_CALL) — the _FINISH_REASON_MAP + already lower-cased both vocabularies; the missing piece was + the test snapshot. + + 2. Mistral num_cached_tokens (flat field on usage, not nested + under prompt_tokens_details.cached_tokens like OpenAI's). The + OpenAI extractor only read the nested shape, so Mistral + customers always saw cache_read_tokens=0 even when the + inference cache hit. Fix reads the Mistral flat field as a + fallback inside the same chain. The _openai_extractor host + map (line 567) already covers Mistral so no host-routing + change was needed. + + 3. Gemini 2.5+ thoughtsTokenCount (reasoning tokens in + usageMetadata) — was hard-coded to 0, so thinking-mode Gemini + calls had no visible reasoning column on the dashboard. + Surfaced as reasoning_tokens while the total stays at + totalTokenCount (reasoning tokens are part of + candidatesTokenCount upstream). + + 4. Anthropic 4.5+ output_tokens_details.thinking_tokens + (extended-thinking mode) — was hard-coded to 0 for the same + reason. The pre-0.13.10 comment ("reasoning tokens are part + of output_tokens") was correct for the non-thinking baseline, + but the thinking-mode field was still readable and was being + dropped. Now we read the breakdown while keeping the total at + input+output (Anthropic bills thinking tokens at the output + rate upstream). + + 5. AWS Bedrock finish_reason for the Mistral-on-Bedrock / + OpenAI-compat and Llama-on-Bedrock adapter shapes. The + pre-0.13.10 extractor only read top-level stopReason / + stop_reason (Anthropic + Llama top-level). Mistral's + OpenAI-compat shape puts the field under + choices[0].finish_reason and was always None. The + matched_shape discriminator (already tracked in the + tool-detection block) tells us which body to read from and + the new branch picks choices[0].finish_reason when + matched_shape == 'openai_choices'. + +The same audit identified the following as should-fix but +deferred to a follow-up PR (none is a billing gap; all are +visibility / observability gaps): + + - Anthropic cache_creation.ephemeral_{1h,5m}_input_tokens + TTL breakdown (different billing rates; Bedrock does not + yet expose the breakdown as of 2026-Q3). + - Anthropic server_tool_use.{web_search_requests, + web_fetch_requests} — server-side tool invocations not + visible to loop detection. + - Gemini multimodal *TokensDetails[] (TEXT vs IMAGE vs + AUDIO) — image-heavy calls mask the real cost driver. + - Cohere billed_units.{search_units, classifications} for + RAG / classify workloads. + - Cohere reasoning models (command-a-reasoning-*). + - Bedrock Converse API (separate envelope from InvokeModel). + +Wire format: unchanged. Backends on 1.0.0 keep working +unchanged. Pinning unchanged: SDK_MIN_VERSION_FOR_V3 = +"0.12.0". Recommended upgrade path: 0.13.9 -> 0.13.10. + +Tests (8 new in tests/test_extractors.py): + + - test_cohere_v2_message_tool_calls_path — v2 nested + message.tool_calls returns the right tool_names. + - test_cohere_v2_cached_tokens — tokens.cached_tokens + surfaces as cache_read_tokens. + - test_cohere_v1_top_level_tool_calls_fallback — v1 + callers (legacy top-level tool_calls) keep working. + - test_openai_mistral_num_cached_tokens — Mistral + usage.num_cached_tokens fallback in the OpenAI extractor. + - test_gemini_2_5_thinking_tokens — thoughtsTokenCount + surfaces as reasoning_tokens while the total stays at + totalTokenCount. + - test_anthropic_extended_thinking_tokens — + output_tokens_details.thinking_tokens surfaces + alongside cache_read_input_tokens / + cache_creation_input_tokens already extracted. + - test_bedrock_mistral_finish_reason_via_choices — + Mistral-on-Bedrock OpenAI-compat finish_reason is now + captured. + - test_bedrock_llama_finish_reason_via_top_level — + Llama-on-Bedrock stop_reason snake_case is captured + (already worked, but had no test snapshot before). + +Verification locally (origin/master + this commit on top): + + * pytest tests/test_extractors.py tests/test_crewai_patch.py + tests/test_runtime.py tests/test_runtime_branches.py + tests/test_track_batch_retry.py + tests/test_track_span_context.py + tests/test_v3_wire_contract.py tests/test_release_polish.py + — 185 passed, 1 skipped (8 new tests net-new from this + commit; no regression on the 177 tests that were green + on master). + * ruff check src/ — "All checks passed!". + * mypy src/ — 11 pre-existing errors (langgraph overload + mismatches at lines 1818, 1821, 1827; same count as + origin/master). No new mypy findings from this release. + +No public API change. No SDK_MIN_VERSION bump. + +--- + v3.23 / 0.13.9 (2026-07-13) — crewai 1.15 compatibility + gate_cache re-capture. @@ -579,5 +694,5 @@ """ -__version__ = "0.13.9" +__version__ = "0.13.10" __platform_version__ = "1.0.0" diff --git a/src/nullrun/instrumentation/auto.py b/src/nullrun/instrumentation/auto.py index 748f8b1..adb53c8 100644 --- a/src/nullrun/instrumentation/auto.py +++ b/src/nullrun/instrumentation/auto.py @@ -147,6 +147,9 @@ def _openai_extractor(body: bytes, status: int) -> ExtractedUsage | None: # Optional nested usage detail blocks. OpenAI added these in 2024 # to expose cache hits (prompt caching) and reasoning tokens (o1). + # Mistral exposes a flat ``num_cached_tokens`` field at the same + # level (no ``prompt_tokens_details`` wrapper); the chained + # ``or`` reads either shape. prompt_details = usage.get("prompt_tokens_details") or {} completion_details = usage.get("completion_tokens_details") or {} @@ -182,7 +185,16 @@ def _openai_extractor(body: bytes, status: int) -> ExtractedUsage | None: # Previously these were reachable only via raw_usage (now # stripped at the wire boundary). Backend gate/budget/loop # detection now sees them as first-class columns. - "cache_read_tokens": int(prompt_details.get("cached_tokens", 0) or 0), + # OpenAI nests cache hits under ``prompt_tokens_details.cached_tokens``; + # Mistral exposes ``usage.num_cached_tokens`` at the same level + # (no nested object). Read either shape so cache-savings metrics + # work for both providers. Mistral also does not have a write + # side, so cache_write_tokens stays 0. + "cache_read_tokens": int( + prompt_details.get("cached_tokens", 0) + or usage.get("num_cached_tokens", 0) + or 0 + ), "cache_write_tokens": 0, # OpenAI does not expose cache creation "reasoning_tokens": int(completion_details.get("reasoning_tokens", 0) or 0), "finish_reason": _normalize_finish_reason(raw_finish), @@ -304,10 +316,19 @@ def _anthropic_extractor(body: bytes, status: int) -> ExtractedUsage | None: "id": payload.get("id"), "cache_read_tokens": int(usage.get("cache_read_input_tokens", 0) or 0), "cache_write_tokens": int(usage.get("cache_creation_input_tokens", 0) or 0), - # Anthropic reasoning tokens are part of output_tokens (they're - # billed at the output rate). Surface 0 here so the backend - # can detect providers that report them separately. - "reasoning_tokens": 0, + # Anthropic 4.5+ extended-thinking surfaces + # ``output_tokens_details.thinking_tokens`` so callers can + # split reasoning from visible output for billing and + # latency dashboards. Native Messages API keeps + # thinking tokens rolled into ``output_tokens`` upstream + # (they're billed at the output rate), so the total + # stays correct without any adjustment; we only surface + # the breakdown. + "reasoning_tokens": int( + (usage.get("output_tokens_details") or {}).get( + "thinking_tokens", 0 + ) or 0 + ), "finish_reason": _normalize_finish_reason(payload.get("stop_reason")), "tool_names": tool_names, } @@ -362,7 +383,13 @@ def _gemini_extractor(body: bytes, status: int) -> ExtractedUsage | None: "id": payload.get("responseId") or payload.get("id"), "cache_read_tokens": int(usage.get("cachedContentTokenCount", 0) or 0), "cache_write_tokens": 0, - "reasoning_tokens": 0, + # Gemini 2.5+ "thinking" models surface + # ``usageMetadata.thoughtsTokenCount`` (reasoning tokens + # are part of ``candidatesTokenCount`` upstream but Gemini + # splits them out for billing/dashboards). Without this + # read the dashboard can't distinguish reasoning vs + # visible output for thinking-mode Gemini calls. + "reasoning_tokens": int(usage.get("thoughtsTokenCount", 0) or 0), "finish_reason": _normalize_finish_reason(raw_finish), "tool_names": tool_names, } @@ -374,6 +401,25 @@ def _cohere_extractor(body: bytes, status: int) -> ExtractedUsage | None: response.usage.{tokens, input_tokens, output_tokens}. Note: Cohere streaming has no usage in stream — only non-streaming responses carry it. Documented in the plan. + + 2026-07-13 (drift fix #N): v2 has THREE schema changes the SDK + silently missed: + + 1. ``tool_calls`` live under ``message.tool_calls`` (not at + the top level). v1 still used top-level ``tool_calls``; + v2 moved them into the assistant message envelope. The + top-level path is preserved as a fallback for v1 + the + rare v2 adapter that lifts the field back up, so neither + version is broken by the new primary path. + + 2. ``usage.tokens.cached_tokens`` is the v2 cache hit counter + (Cohere's inference cache). Previously always read as 0. + + 3. ``finish_reason`` values are UPPERCASE + (``COMPLETE | MAX_TOKENS | STOP_SEQUENCE | TOOL_CALL | + ERROR | TIMEOUT``); the v1 vocabulary was lowercase. The + ``_normalize_finish_reason`` helper lower-cases before + mapping so both vocabularies work. """ if status >= 400 or not body: return None @@ -385,28 +431,52 @@ def _cohere_extractor(body: bytes, status: int) -> ExtractedUsage | None: if not isinstance(usage, dict): return None # v2 uses input_tokens/output_tokens; v1 used prompt_tokens/completion_tokens. + # ``usage.tokens`` in v2 is a nested object (not an int) so we cannot + # use it as the ``total`` value here — that path was wrong in v1 and + # silently crashes on v2. Compute total from the actual int fields + # below; the ``tokens`` nested object is only consulted for the + # ``cached_tokens`` field (its only int it carries in v2). inp = int( usage.get("input_tokens", 0) or usage.get("prompt_tokens", 0) or 0 ) out = int( usage.get("output_tokens", 0) or usage.get("completion_tokens", 0) or 0 ) - total = int(usage.get("tokens", 0) or 0) or (inp + out) - if total == 0 and inp == 0 and out == 0: + if inp == 0 and out == 0: return None + total = inp + out - # Cohere tool_calls are top-level (not under choices[]). The - # schema varies: v2 uses {id, type: "function", function: {name - # arguments}}; some adapters use {name, parameters}. We accept - # both shapes. + # Cache reads — v2 exposes ``usage.tokens.cached_tokens``; v1 + # had no cache concept. The nested path is the new primary; + # the top-level fallback covers a future v1.1 re-introduction. + cache_read = int( + (usage.get("tokens") or {}).get("cached_tokens", 0) + or usage.get("cached_tokens", 0) + or 0 + ) + + # Cohere tool_calls are top-level (v1) OR under + # ``message.tool_calls`` (v2). v2 path is primary because + # current Cohere SDK always puts the field there; the + # top-level path remains so v1 callers keep working. tool_names: list[str] = [] - for tc in payload.get("tool_calls") or []: - if not isinstance(tc, dict): - continue - if isinstance(tc.get("function"), dict) and tc["function"].get("name"): - tool_names.append(tc["function"]["name"]) - elif tc.get("name"): - tool_names.append(tc["name"]) + message = payload.get("message") if isinstance(payload, dict) else None + if isinstance(message, dict): + for tc in message.get("tool_calls") or []: + if not isinstance(tc, dict): + continue + if isinstance(tc.get("function"), dict) and tc["function"].get("name"): + tool_names.append(tc["function"]["name"]) + elif tc.get("name"): + tool_names.append(tc["name"]) + if not tool_names: + for tc in payload.get("tool_calls") or []: + if not isinstance(tc, dict): + continue + if isinstance(tc.get("function"), dict) and tc["function"].get("name"): + tool_names.append(tc["function"]["name"]) + elif tc.get("name"): + tool_names.append(tc["name"]) return { "prompt_tokens": inp, @@ -417,9 +487,13 @@ def _cohere_extractor(body: bytes, status: int) -> ExtractedUsage | None: # surface a stable response id at the top level; rely on # model+provider for disambiguation. See _openai_extractor. "id": payload.get("id") or payload.get("generation_id"), - "cache_read_tokens": 0, + "cache_read_tokens": cache_read, "cache_write_tokens": 0, "reasoning_tokens": 0, + # _normalize_finish_reason lower-cases before mapping so + # both v1 ("stop"/"length"/"tool_calls") and v2 + # ("COMPLETE"/"MAX_TOKENS"/"TOOL_CALL") are normalized + # to the backend's canonical vocabulary. "finish_reason": _normalize_finish_reason(payload.get("finish_reason")), "tool_names": tool_names, } @@ -543,6 +617,23 @@ def _bedrock_extractor(body: bytes, status: int) -> ExtractedUsage | None: or 0 ) + # Finish reason — shape depends on the underlying model. + # Anthropic-on-Bedrock / native Anthropic: top-level + # ``stopReason`` (camelCase, AWS-style). Llama-on-Bedrock: + # top-level ``stop_reason`` (snake_case). Mistral-on-Bedrock + # / OpenAI-compat: ``choices[0].finish_reason`` (per the + # OpenAI shape) — the ``matched_shape`` discriminator we + # already track above tells us which body the response used. + # We capture from all three sources in priority order so the + # backend always sees a finish_reason on Mistral / Llama + # Bedrock calls, not just on Anthropic. + raw_finish: str | None = payload.get("stopReason") or payload.get("stop_reason") + if raw_finish is None and matched_shape == "openai_choices": + for choice in payload.get("choices") or []: + if isinstance(choice, dict) and choice.get("finish_reason"): + raw_finish = choice["finish_reason"] + break + return { "prompt_tokens": inp, "completion_tokens": out, @@ -558,7 +649,7 @@ def _bedrock_extractor(body: bytes, status: int) -> ExtractedUsage | None: "cache_read_tokens": cache_read, "cache_write_tokens": cache_write, "reasoning_tokens": 0, - "finish_reason": _normalize_finish_reason(payload.get("stopReason") or payload.get("stop_reason")), + "finish_reason": _normalize_finish_reason(raw_finish), "tool_names": tool_names, } diff --git a/tests/test_extractors.py b/tests/test_extractors.py index 0a67c39..bfe4973 100644 --- a/tests/test_extractors.py +++ b/tests/test_extractors.py @@ -86,6 +86,29 @@ def test_openai_malformed_json_returns_none(): assert _openai_extractor(b"not-json", 200) is None +def test_openai_mistral_num_cached_tokens(): + """Mistral exposes a flat ``usage.num_cached_tokens`` field at the + same level (no ``prompt_tokens_details`` wrapper). Without the + fallback in the OpenAI extractor, Mistral customers see + ``cache_read_tokens=0`` even when the inference cache hit. + """ + body = json.dumps( + { + "model": "mistral-large-latest", + "choices": [{"finish_reason": "stop"}], + "usage": { + "prompt_tokens": 57, + "completion_tokens": 18, + "total_tokens": 75, + "num_cached_tokens": 41, + }, + } + ).encode() + out = _openai_extractor(body, 200) + assert out is not None + assert out["cache_read_tokens"] == 41 + + def test_openai_v1_streaming_final_chunk(): """OpenAI v1.0+ streaming responses only carry `usage` in the LAST SSE chunk. We feed the full accumulated buffer (multiple SSE chunks @@ -142,6 +165,38 @@ def test_anthropic_error_returns_none(): assert _anthropic_extractor(body, 429) is None +def test_anthropic_extended_thinking_tokens(): + """Anthropic 4.5+ extended-thinking surfaces + ``output_tokens_details.thinking_tokens`` so callers can split + reasoning from visible output. Without the read, every + thinking-mode call is invisible in the reasoning dashboard. + """ + body = json.dumps( + { + "id": "msg_01", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "Final answer."}], + "usage": { + "input_tokens": 1200, + "output_tokens": 215, + "output_tokens_details": {"thinking_tokens": 80}, + "cache_read_input_tokens": 3500, + "cache_creation_input_tokens": 800, + }, + } + ).encode() + out = _anthropic_extractor(body, 200) + assert out is not None + assert out["reasoning_tokens"] == 80 + assert out["completion_tokens"] == 215 + # total stays as input + output (Anthropic bills thinking + # tokens at the output rate upstream, so the + # input+output sum already includes them). + assert out["total_tokens"] == 1415 + assert out["cache_read_tokens"] == 3500 + assert out["cache_write_tokens"] == 800 + + # --------------------------------------------------------------------------- # Google Gemini (Generative Language API) # --------------------------------------------------------------------------- @@ -187,19 +242,53 @@ def test_gemini_no_usage_returns_none(): assert _gemini_extractor(body, 200) is None +def test_gemini_2_5_thinking_tokens(): + """Gemini 2.5+ "thinking" mode surfaces ``thoughtsTokenCount`` in + ``usageMetadata``. Without the read the dashboard can't tell + reasoning tokens from visible output for thinking-mode calls. + """ + body = json.dumps( + { + "modelVersion": "gemini-2.5-pro", + "candidates": [{"content": {"parts": [{"text": "answer"}]}}], + "usageMetadata": { + "promptTokenCount": 1200, + "candidatesTokenCount": 47, + "thoughtsTokenCount": 30, + "totalTokenCount": 1277, + }, + } + ).encode() + out = _gemini_extractor(body, 200) + assert out is not None + assert out["reasoning_tokens"] == 30 + # totalTokenCount is authoritative when present; we don't + # double-count the reasoning tokens into total (Gemini already + # includes them in candidatesTokenCount). + assert out["total_tokens"] == 1277 + + # --------------------------------------------------------------------------- # Cohere # --------------------------------------------------------------------------- def test_cohere_v2_response(): + """Cohere v2 canonical response — actual API shape has + ``usage.tokens`` as a nested object with input/output token + counts, not a flat integer. Older fixtures had the integer + shape, which masked the v2 nested-object reality. + """ body = json.dumps( { "model": "command-r-plus", "usage": { "input_tokens": 18, "output_tokens": 4, - "tokens": 22, + "tokens": { + "input_tokens": 18, + "output_tokens": 4, + }, }, } ).encode() @@ -209,6 +298,10 @@ def test_cohere_v2_response(): assert out["completion_tokens"] == 4 assert out["total_tokens"] == 22 assert out["model"] == "command-r-plus" + # Nested usage.tokens is NOT the total value; v2 carries + # the int fields in both shapes (top-level + nested) and + # the top-level wins as the source-of-truth. + assert out["cache_read_tokens"] == 0 def test_cohere_v1_legacy_prompt_completion_keys(): @@ -227,6 +320,90 @@ def test_cohere_v1_legacy_prompt_completion_keys(): assert out["total_tokens"] == 7 +def test_cohere_v2_message_tool_calls_path(): + """Cohere v2 nests ``tool_calls`` under ``message.tool_calls`` (not + at top level). v1 still used top-level; both must work. Without + the v2 path fix, ``tool_names`` is always empty for v2 callers + and the backend's loop detection misses every Cohere tool use. + """ + body = json.dumps( + { + "id": "c14c80c3-18eb-4519-9460-6c92edd8cfb4", + "model": "command-r-plus", + "finish_reason": "COMPLETE", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "Let me check that."}], + "tool_calls": [ + { + "id": "search_docs_dkf0akqdazjb", + "type": "function", + "function": { + "name": "search_docs", + "arguments": '{"query":"tool use","top_k":3}', + }, + } + ], + }, + "usage": { + "input_tokens": 71, + "output_tokens": 418, + "tokens": { + "input_tokens": 71, + "output_tokens": 418, + }, + }, + } + ).encode() + out = _cohere_extractor(body, 200) + assert out is not None + assert out["tool_names"] == ["search_docs"], ( + "v2 message.tool_calls must be picked up; " + f"got {out['tool_names']!r}" + ) + # UPPERCASE finish_reason is normalized via _FINISH_REASON_MAP + assert out["finish_reason"] == "stop" + + +def test_cohere_v2_cached_tokens(): + """Cohere v2 exposes ``usage.tokens.cached_tokens`` for inference + cache hits. Previously always read as 0. + """ + body = json.dumps( + { + "model": "command-r-plus", + "usage": { + "input_tokens": 71, + "output_tokens": 18, + "tokens": { + "input_tokens": 71, + "output_tokens": 18, + "cached_tokens": 40, + }, + }, + } + ).encode() + out = _cohere_extractor(body, 200) + assert out is not None + assert out["cache_read_tokens"] == 40 + + +def test_cohere_v1_top_level_tool_calls_fallback(): + """v1 still surfaces tool_calls at the top level (no ``message`` + wrapper). v1 callers keep working alongside the new v2 path. + """ + body = json.dumps( + { + "model": "command", + "tool_calls": [{"name": "legacy_tool"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 2}, + } + ).encode() + out = _cohere_extractor(body, 200) + assert out is not None + assert out["tool_names"] == ["legacy_tool"] + + # --------------------------------------------------------------------------- # AWS Bedrock # --------------------------------------------------------------------------- @@ -269,6 +446,81 @@ def test_bedrock_error_returns_none(): assert _bedrock_extractor(body, 403) is None +def test_bedrock_mistral_finish_reason_via_choices(): + """Mistral-on-Bedrock / OpenAI-compat carries + ``choices[0].finish_reason`` (not ``stopReason``). Without the + extra read, the backend loses the finish signal on every + Mistral-on-Bedrock call and dashboard aggregations can never + distinguish `stop` from `length`. + + Bedrock wraps the underlying model response in + ``InvokeModelResponse`` with ``output`` as a base64-encoded body + for streaming; the unwrapped shape is what we test here because + the SDK receives the parsed body. Token fields in this shape + are camelCase (AWS-style) at the top level of ``usage`` even + for OpenAI-compat models. + """ + body = json.dumps( + { + "id": "cmpl-bedrock-1", + "model": "mistral.mistral-large-2407-v1:0", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "abc", + "type": "function", + "function": {"name": "search"}, + } + ], + }, + } + ], + "usage": { + "inputTokens": 100, + "outputTokens": 20, + "totalTokens": 120, + }, + } + ).encode() + out = _bedrock_extractor(body, 200) + assert out is not None + assert out["finish_reason"] == "tool_calls" + assert out["tool_names"] == ["search"] + + +def test_bedrock_llama_finish_reason_via_top_level(): + """Llama-on-Bedrock carries ``stop_reason`` at the top level + (snake_case). Make sure we pick that up. Note Llama's token + fields are top-level ``prompt_token_count`` / + ``generation_token_count`` (not under ``usage``), so the + ``usage`` discriminator still matches via the + ``inputTokens``/``outputTokens`` top-level fallback path. + """ + body = json.dumps( + { + "id": "bedrock-llama-1", + "stop_reason": "stop", + "inputTokens": 100, + "outputTokens": 20, + "output": { + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + } + }, + } + ).encode() + out = _bedrock_extractor(body, 200) + assert out is not None + assert out["finish_reason"] == "stop" + + # --------------------------------------------------------------------------- # _match_extractor table # ---------------------------------------------------------------------------