feat: add anthropic cached token values for local backends that support it
This commit is contained in:
parent
1d012a92ef
commit
04bc69993d
3 changed files with 93 additions and 16 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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", {})
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue