329 lines
14 KiB
Python
329 lines
14 KiB
Python
"""Anthropic **Messages API** routes (``/v1/messages`` and ``/v1/messages/count_tokens``).
|
|
|
|
The router speaks Chat Completions to its local backends, so this layer:
|
|
|
|
* **native** (configured ``anthropic_endpoints``): forwards the Anthropic request
|
|
verbatim over httpx with ``x-api-key`` / ``anthropic-version`` headers and streams
|
|
the upstream SSE straight back.
|
|
* **translated** (Ollama / llama-server / llama-swap): converts the request to chat,
|
|
reuses the resilient ``create_chat_with_retries`` ladder, and re-emits the result as
|
|
Anthropic typed SSE events (``requests/anthropic.py``).
|
|
|
|
The Messages API is stateless, so — unlike ``/v1/responses`` — there is no store,
|
|
background mode, or DB. An optional ``nomyo.cache`` extension field reflects hits back
|
|
through the router's semantic LLM cache (a hit is reported via ``usage.cache_read_input_tokens``).
|
|
"""
|
|
import codecs
|
|
|
|
import httpx
|
|
import orjson
|
|
from fastapi import APIRouter, HTTPException, Request
|
|
from starlette.responses import JSONResponse, StreamingResponse
|
|
|
|
from cache import get_llm_cache
|
|
from config import get_config
|
|
from context_window import _count_message_tokens
|
|
from fingerprint import _conversation_fingerprint
|
|
from state import app_state, token_queue, default_headers
|
|
from backends.normalize import is_anthropic_endpoint
|
|
from backends.probe import ANTHROPIC_VERSION
|
|
from backends.sessions import _make_openai_client
|
|
from routing import choose_endpoint, decrement_usage
|
|
from api.openai import create_chat_with_retries
|
|
from requests.anthropic import (
|
|
ChatToMessagesStream,
|
|
anthropic_messages_to_chat,
|
|
anthropic_to_chat_send_params,
|
|
build_message_object,
|
|
chat_message_to_content_blocks,
|
|
finish_reason_to_stop_reason,
|
|
message_object_to_sse,
|
|
new_message_id,
|
|
usage_chat_to_anthropic,
|
|
)
|
|
|
|
router = APIRouter()
|
|
|
|
CACHE_ROUTE = "anthropic_messages"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# helpers
|
|
# ---------------------------------------------------------------------------
|
|
def _anthropic_http_client(endpoint: str) -> httpx.AsyncClient:
|
|
"""Return the warmed httpx client for a native Anthropic endpoint.
|
|
|
|
Startup pre-creates one per configured endpoint; fall back to an on-demand
|
|
client (cached in app_state) for tests that skip the lifespan startup.
|
|
"""
|
|
client = app_state["httpx_clients"].get(endpoint)
|
|
if client is None:
|
|
client = httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=15.0))
|
|
app_state["httpx_clients"][endpoint] = client
|
|
return client
|
|
|
|
|
|
def _native_headers(request: Request, api_key: str) -> dict:
|
|
"""Build outbound headers for a native Anthropic forward.
|
|
|
|
Injects the router's stored key as ``x-api-key`` and pins ``anthropic-version``,
|
|
passing through the client's ``anthropic-beta`` / ``anthropic-version`` when present.
|
|
"""
|
|
headers = {
|
|
"content-type": "application/json",
|
|
"x-api-key": api_key,
|
|
"anthropic-version": request.headers.get("anthropic-version", ANTHROPIC_VERSION),
|
|
}
|
|
beta = request.headers.get("anthropic-beta")
|
|
if beta:
|
|
headers["anthropic-beta"] = beta
|
|
return headers
|
|
|
|
|
|
async def _track(endpoint, tracking_model, prompt_tok, comp_tok):
|
|
if prompt_tok or comp_tok:
|
|
await token_queue.put((endpoint, tracking_model, prompt_tok, comp_tok))
|
|
|
|
|
|
def _serve_cache_hit(cached: bytes, message_id: str, stream: bool):
|
|
"""Serve a stored message object as a cache hit (input tokens → cache_read)."""
|
|
obj = orjson.loads(cached)
|
|
obj["id"] = message_id
|
|
u = obj.get("usage") or {}
|
|
read = u.get("input_tokens", 0) or 0
|
|
obj["usage"] = {
|
|
**u,
|
|
"input_tokens": 0,
|
|
"cache_read_input_tokens": read,
|
|
"cache_creation_input_tokens": 0,
|
|
}
|
|
if stream:
|
|
async def _served():
|
|
yield message_object_to_sse(obj)
|
|
return StreamingResponse(_served(), media_type="text/event-stream")
|
|
return JSONResponse(content=obj)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# POST /v1/messages
|
|
# ---------------------------------------------------------------------------
|
|
@router.post("/v1/messages")
|
|
async def anthropic_messages_proxy(request: Request):
|
|
config = get_config()
|
|
raw_body = await request.body()
|
|
try:
|
|
payload = orjson.loads(raw_body.decode("utf-8"))
|
|
except orjson.JSONDecodeError as e:
|
|
raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}") from e
|
|
|
|
model = payload.get("model")
|
|
messages = payload.get("messages")
|
|
system = payload.get("system")
|
|
stream = bool(payload.get("stream"))
|
|
_cache_enabled = payload.get("nomyo", {}).get("cache", False)
|
|
|
|
if not model:
|
|
raise HTTPException(status_code=400, detail="Missing required field 'model'")
|
|
if messages is None:
|
|
raise HTTPException(status_code=400, detail="Missing required field 'messages'")
|
|
if payload.get("max_tokens") is None:
|
|
raise HTTPException(status_code=400, detail="Missing required field 'max_tokens'")
|
|
|
|
if ":latest" in model:
|
|
model = model.split(":latest")[0]
|
|
|
|
chat_messages = anthropic_messages_to_chat(system, messages)
|
|
message_id = new_message_id()
|
|
|
|
# Cache lookup (foreground) — before endpoint selection, keyed on the chat form.
|
|
_cache = get_llm_cache()
|
|
if _cache is not None and _cache_enabled:
|
|
cached = await _cache.get_chat(CACHE_ROUTE, model, chat_messages)
|
|
if cached is not None:
|
|
return _serve_cache_hit(cached, message_id, stream)
|
|
|
|
async def _cache_store(obj):
|
|
if _cache is None or not _cache_enabled or not obj.get("content"):
|
|
return
|
|
try:
|
|
await _cache.set_chat(CACHE_ROUTE, model, chat_messages, orjson.dumps(obj))
|
|
except Exception as _ce:
|
|
print(f"[cache] set_chat ({CACHE_ROUTE}) failed: {_ce}")
|
|
|
|
# Endpoint selection reserves a slot — released exactly once per branch.
|
|
_affinity_key = _conversation_fingerprint(model, chat_messages, None)
|
|
endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key)
|
|
try:
|
|
native = is_anthropic_endpoint(endpoint)
|
|
if not native:
|
|
oclient = _make_openai_client(endpoint, default_headers=default_headers,
|
|
api_key=config.api_keys.get(endpoint, "no-key"))
|
|
send_params = anthropic_to_chat_send_params(payload, chat_messages, model)
|
|
except BaseException:
|
|
await decrement_usage(endpoint, tracking_model)
|
|
raise
|
|
|
|
# ---- native passthrough -----------------------------------------------
|
|
if native:
|
|
return await _handle_native(
|
|
request, payload, endpoint, tracking_model, stream,
|
|
api_key=config.api_keys.get(endpoint, "no-key"), cache_store=_cache_store)
|
|
|
|
# ---- translated streaming ---------------------------------------------
|
|
if stream:
|
|
try:
|
|
source = await create_chat_with_retries(
|
|
oclient, {**send_params, "stream": True,
|
|
"stream_options": {"include_usage": True}},
|
|
endpoint, model, tracking_model)
|
|
except BaseException:
|
|
await decrement_usage(endpoint, tracking_model)
|
|
raise
|
|
translator = ChatToMessagesStream(message_id, model)
|
|
|
|
async def _stream():
|
|
try:
|
|
async for sse in translator.events(source):
|
|
yield sse
|
|
prompt = (translator.usage or {}).get("prompt_tokens", 0)
|
|
comp = (translator.usage or {}).get("completion_tokens", 0)
|
|
await _track(endpoint, tracking_model, prompt, comp)
|
|
obj = build_message_object(
|
|
message_id=message_id, model=model,
|
|
content_blocks=translator.content_blocks,
|
|
stop_reason=translator.stop_reason,
|
|
usage=usage_chat_to_anthropic(translator.usage))
|
|
await _cache_store(obj)
|
|
finally:
|
|
await decrement_usage(endpoint, tracking_model)
|
|
|
|
return StreamingResponse(_stream(), media_type="text/event-stream")
|
|
|
|
# ---- translated non-streaming -----------------------------------------
|
|
try:
|
|
result = await create_chat_with_retries(
|
|
oclient, {**send_params, "stream": False}, endpoint, model, tracking_model)
|
|
message = result.choices[0].message.model_dump() if result.choices else {}
|
|
usage = result.usage.model_dump() if result.usage is not None else None
|
|
content_blocks = chat_message_to_content_blocks(message)
|
|
finish_reason = getattr(result.choices[0], "finish_reason", None) if result.choices else None
|
|
has_tool_use = any(b.get("type") == "tool_use" for b in content_blocks)
|
|
stop_reason = finish_reason_to_stop_reason(finish_reason, has_tool_use=has_tool_use)
|
|
await _track(endpoint, tracking_model,
|
|
(usage or {}).get("prompt_tokens", 0),
|
|
(usage or {}).get("completion_tokens", 0))
|
|
finally:
|
|
await decrement_usage(endpoint, tracking_model)
|
|
|
|
obj = build_message_object(
|
|
message_id=message_id, model=model, content_blocks=content_blocks,
|
|
stop_reason=stop_reason, usage=usage_chat_to_anthropic(usage))
|
|
await _cache_store(obj)
|
|
return JSONResponse(content=obj)
|
|
|
|
|
|
async def _handle_native(request, payload, endpoint, tracking_model, stream,
|
|
*, api_key, cache_store):
|
|
"""Forward an Anthropic request verbatim to a native upstream."""
|
|
client = _anthropic_http_client(endpoint)
|
|
headers = _native_headers(request, api_key)
|
|
forward = {k: v for k, v in payload.items() if k != "nomyo"}
|
|
url = f"{endpoint.rstrip('/')}/v1/messages"
|
|
|
|
if not stream:
|
|
forward["stream"] = False
|
|
try:
|
|
resp = await client.post(url, headers=headers, json=forward)
|
|
except BaseException:
|
|
await decrement_usage(endpoint, tracking_model)
|
|
raise
|
|
try:
|
|
data = resp.json()
|
|
except Exception:
|
|
data = {"type": "error", "error": {"message": resp.text[:500]}}
|
|
if resp.status_code == 200 and isinstance(data, dict):
|
|
u = data.get("usage") or {}
|
|
await _track(endpoint, tracking_model,
|
|
u.get("input_tokens", 0) or 0, u.get("output_tokens", 0) or 0)
|
|
await cache_store(data)
|
|
await decrement_usage(endpoint, tracking_model)
|
|
return JSONResponse(content=data, status_code=resp.status_code)
|
|
|
|
forward["stream"] = True
|
|
|
|
async def _proxy_stream():
|
|
decoder = codecs.getincrementaldecoder("utf-8")()
|
|
buffer = ""
|
|
input_tok = 0
|
|
output_tok = 0
|
|
try:
|
|
async with client.stream("POST", url, headers=headers, json=forward) as resp:
|
|
async for raw in resp.aiter_bytes():
|
|
if not raw:
|
|
continue
|
|
yield raw
|
|
# Parse a light copy to capture usage for token tracking.
|
|
buffer += decoder.decode(raw)
|
|
while "\n" in buffer:
|
|
line, buffer = buffer.split("\n", 1)
|
|
line = line.strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
payload_str = line[len("data:"):].strip()
|
|
if not payload_str or payload_str == "[DONE]":
|
|
continue
|
|
try:
|
|
evt = orjson.loads(payload_str)
|
|
except orjson.JSONDecodeError:
|
|
continue
|
|
if evt.get("type") == "message_start":
|
|
u = (evt.get("message") or {}).get("usage") or {}
|
|
input_tok = u.get("input_tokens", 0) or input_tok
|
|
elif evt.get("type") == "message_delta":
|
|
u = evt.get("usage") or {}
|
|
output_tok = u.get("output_tokens", 0) or output_tok
|
|
finally:
|
|
await _track(endpoint, tracking_model, input_tok, output_tok)
|
|
await decrement_usage(endpoint, tracking_model)
|
|
|
|
return StreamingResponse(_proxy_stream(), media_type="text/event-stream")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# POST /v1/messages/count_tokens
|
|
# ---------------------------------------------------------------------------
|
|
@router.post("/v1/messages/count_tokens")
|
|
async def anthropic_count_tokens(request: Request):
|
|
config = get_config()
|
|
try:
|
|
payload = orjson.loads((await request.body()).decode("utf-8"))
|
|
except orjson.JSONDecodeError as e:
|
|
raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}") from e
|
|
|
|
model = payload.get("model")
|
|
messages = payload.get("messages")
|
|
if not model:
|
|
raise HTTPException(status_code=400, detail="Missing required field 'model'")
|
|
if messages is None:
|
|
raise HTTPException(status_code=400, detail="Missing required field 'messages'")
|
|
if ":latest" in model:
|
|
model = model.split(":latest")[0]
|
|
|
|
chat_messages = anthropic_messages_to_chat(payload.get("system"), messages)
|
|
|
|
# No slot reservation — this is a metadata call, not a completion.
|
|
endpoint, _tracking = await choose_endpoint(model, reserve=False)
|
|
|
|
if is_anthropic_endpoint(endpoint):
|
|
client = _anthropic_http_client(endpoint)
|
|
headers = _native_headers(request, config.api_keys.get(endpoint, "no-key"))
|
|
forward = {k: v for k, v in payload.items() if k != "nomyo"}
|
|
try:
|
|
resp = await client.post(
|
|
f"{endpoint.rstrip('/')}/v1/messages/count_tokens",
|
|
headers=headers, json=forward)
|
|
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=502, detail=f"count_tokens upstream failed: {e}") from e
|
|
|
|
return JSONResponse(content={"input_tokens": _count_message_tokens(chat_messages)})
|