diff --git a/api/messages.py b/api/messages.py index 90dd945..8b113c3 100644 --- a/api/messages.py +++ b/api/messages.py @@ -90,7 +90,8 @@ def _serve_cache_hit(cached: bytes, message_id: str, stream: bool): obj = orjson.loads(cached) obj["id"] = message_id u = obj.get("usage") or {} - read = u.get("input_tokens", 0) or 0 + # Whole prompt served from cache: fold the stored uncached + read tokens into read. + read = (u.get("input_tokens", 0) or 0) + (u.get("cache_read_input_tokens", 0) or 0) obj["usage"] = { **u, "input_tokens": 0, diff --git a/requests/anthropic.py b/requests/anthropic.py index bef32b0..b6550d4 100644 --- a/requests/anthropic.py +++ b/requests/anthropic.py @@ -257,15 +257,46 @@ def finish_reason_to_stop_reason(finish_reason, has_tool_use=False): return _STOP_REASON_MAP.get(finish_reason, "end_turn") +def _cached_prompt_tokens(usage): + """Read ``prompt_tokens_details.cached_tokens`` from a chat usage (dict or SDK obj). + + OpenAI-compatible backends with automatic prefix caching (vLLM, recent + llama-server) report the reused-prefix token count here; backends that don't + populate it yield 0. + """ + if not usage: + return 0 + details = usage.get("prompt_tokens_details") if isinstance(usage, dict) \ + else getattr(usage, "prompt_tokens_details", None) + if details is None: + return 0 + if isinstance(details, dict): + return details.get("cached_tokens") or 0 + return getattr(details, "cached_tokens", 0) or 0 + + def usage_chat_to_anthropic(usage, cache_read_tokens=0, cache_creation_tokens=0): - """Map chat usage → Anthropic usage, folding in nomyo-cache attribution.""" - prompt = (usage or {}).get("prompt_tokens") or 0 - completion = (usage or {}).get("completion_tokens") or 0 + """Map chat usage → Anthropic usage. + + ``cached_tokens`` reported by the backend (automatic prefix-cache reuse) plus any + explicit ``cache_read_tokens`` become ``cache_read_input_tokens``; ``input_tokens`` + is the remaining uncached prompt so ``input + cache_read`` equals the prompt total. + + ``cache_creation_input_tokens`` stays at ``cache_creation_tokens`` (0 by default): + local backends do automatic KV-prefix reuse with no explicit ``cache_control`` + write/breakpoint concept, so there is no honest "tokens written" figure to report — + that value is only real on the native Anthropic passthrough path. + """ + prompt = (usage or {}).get("prompt_tokens") or 0 if isinstance(usage, dict) \ + else (getattr(usage, "prompt_tokens", 0) or 0) + completion = (usage or {}).get("completion_tokens") or 0 if isinstance(usage, dict) \ + else (getattr(usage, "completion_tokens", 0) or 0) + cache_read = _cached_prompt_tokens(usage) + cache_read_tokens return { - "input_tokens": prompt, + "input_tokens": max(prompt - cache_read, 0), "output_tokens": completion, "cache_creation_input_tokens": cache_creation_tokens, - "cache_read_input_tokens": cache_read_tokens, + "cache_read_input_tokens": cache_read, } @@ -430,6 +461,7 @@ class ChatToMessagesStream: self.usage = { "prompt_tokens": getattr(usage, "prompt_tokens", 0) or 0, "completion_tokens": getattr(usage, "completion_tokens", 0) or 0, + "prompt_tokens_details": {"cached_tokens": _cached_prompt_tokens(usage)}, } choices = getattr(chunk, "choices", None) if not choices: @@ -512,8 +544,15 @@ class ChatToMessagesStream: "type": "tool_use", "id": state["id"], "name": state["name"], "input": parsed}) self.stop_reason = finish_reason_to_stop_reason(finish_reason, has_tool_use=bool(tc_state)) - out_tokens = (self.usage or {}).get("completion_tokens", 0) + final_usage = usage_chat_to_anthropic( + self.usage, cache_read_tokens=self.cache_read_tokens, + cache_creation_tokens=self.cache_creation_tokens) yield _sse("message_delta", { "delta": {"stop_reason": self.stop_reason, "stop_sequence": None}, - "usage": {"output_tokens": out_tokens}}) + "usage": { + "input_tokens": final_usage["input_tokens"], + "output_tokens": final_usage["output_tokens"], + "cache_read_input_tokens": final_usage["cache_read_input_tokens"], + "cache_creation_input_tokens": final_usage["cache_creation_input_tokens"], + }}) yield _sse("message_stop", {}) diff --git a/test/test_messages.py b/test/test_messages.py index ab081c9..5c8700e 100644 --- a/test/test_messages.py +++ b/test/test_messages.py @@ -119,10 +119,17 @@ class TestResponseTranslation: assert at.finish_reason_to_stop_reason("tool_calls") == "tool_use" assert at.finish_reason_to_stop_reason("stop", has_tool_use=True) == "tool_use" - def test_usage_mapping(self): - u = at.usage_chat_to_anthropic({"prompt_tokens": 7, "completion_tokens": 3}, - cache_read_tokens=5) + def test_usage_mapping_no_cache(self): + u = at.usage_chat_to_anthropic({"prompt_tokens": 7, "completion_tokens": 3}) assert u == {"input_tokens": 7, "output_tokens": 3, + "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0} + + def test_usage_mapping_cached_tokens_subtracted(self): + # Backend reports 5 of 7 prompt tokens served from its prefix cache. + u = at.usage_chat_to_anthropic({ + "prompt_tokens": 7, "completion_tokens": 3, + "prompt_tokens_details": {"cached_tokens": 5}}) + assert u == {"input_tokens": 2, "output_tokens": 3, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 5} @@ -137,8 +144,11 @@ def _chunk(content=None, tool_calls=None, reasoning=None, finish_reason=None): return NS(choices=[NS(delta=delta, finish_reason=finish_reason)], usage=None) -def _usage_chunk(p, c): - return NS(choices=[], usage=NS(prompt_tokens=p, completion_tokens=c)) +def _usage_chunk(p, c, cached=None): + usage = NS(prompt_tokens=p, completion_tokens=c) + if cached is not None: + usage.prompt_tokens_details = NS(cached_tokens=cached) + return NS(choices=[], usage=usage) async def _collect(translator, gen): @@ -183,6 +193,17 @@ class TestStreamTranslator: assert md["usage"]["output_tokens"] == 5 assert tr.content_blocks == [{"type": "text", "text": "Hello"}] + async def test_cached_tokens_surface_as_cache_read(self): + async def gen(): + yield _chunk(content="hi", finish_reason="stop") + yield _usage_chunk(10, 2, cached=6) + tr = at.ChatToMessagesStream("msg_1", "m") + events = await _collect(tr, gen()) + md = [d for t, d in events if t == "message_delta"][0] + assert md["usage"]["cache_read_input_tokens"] == 6 + assert md["usage"]["input_tokens"] == 4 # 10 prompt − 6 cached + assert md["usage"]["cache_creation_input_tokens"] == 0 + async def test_thinking_then_text(self): async def gen(): yield _chunk(reasoning="think ") @@ -253,7 +274,7 @@ def _enter(*cms): def _fake_completion(content="hello", usage=(3, 5), reasoning=None, tool_calls=None, - finish_reason="stop"): + finish_reason="stop", cached=None): md = {"role": "assistant", "content": content} if reasoning is not None: md["reasoning_content"] = reasoning @@ -261,9 +282,12 @@ def _fake_completion(content="hello", usage=(3, 5), reasoning=None, tool_calls=N md["tool_calls"] = tool_calls msg = MagicMock() msg.model_dump.return_value = md + usage_dump = {"prompt_tokens": usage[0], "completion_tokens": usage[1], + "total_tokens": sum(usage)} + if cached is not None: + usage_dump["prompt_tokens_details"] = {"cached_tokens": cached} usage_obj = MagicMock() - usage_obj.model_dump.return_value = { - "prompt_tokens": usage[0], "completion_tokens": usage[1], "total_tokens": sum(usage)} + usage_obj.model_dump.return_value = usage_dump return NS(choices=[NS(message=msg, finish_reason=finish_reason)], usage=usage_obj) @@ -295,6 +319,19 @@ class TestTranslatedRoute: assert body["usage"]["input_tokens"] == 3 and body["usage"]["output_tokens"] == 5 assert body["id"].startswith("msg_") + async def test_nonstream_cached_tokens(self, client): + with _enter(*_patch_backend(native=False), + patch.object(api_messages, "create_chat_with_retries", + AsyncMock(return_value=_fake_completion( + "hi", usage=(10, 4), cached=6)))): + resp = await client.post("/v1/messages", + json={"model": "test-model", "max_tokens": 100, + "messages": [{"role": "user", "content": "hi"}]}) + u = resp.json()["usage"] + assert u["cache_read_input_tokens"] == 6 + assert u["input_tokens"] == 4 + assert u["cache_creation_input_tokens"] == 0 + async def test_missing_max_tokens_400(self, client): with _enter(*_patch_backend(native=False)): resp = await client.post("/v1/messages",