Coverage for src/lilbee/server/anthropic_api/routes.py: 100%
123 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""HTTP routes for the Anthropic-compatible ``/v1/messages`` surface."""
3from __future__ import annotations
5import asyncio
6import logging
7import uuid
8from collections.abc import AsyncGenerator
10from litestar import Request, Router, post
11from litestar.background_tasks import BackgroundTask
12from litestar.exceptions import NotAuthorizedException, ValidationException
13from litestar.response import Response, Stream
15from lilbee.app.services import get_services
16from lilbee.core.config import cfg
17from lilbee.core.config.enums import ReasoningMode
18from lilbee.retrieval.reasoning import effective_reasoning_cap
19from lilbee.server.anthropic_api.errors import anthropic_error_body, anthropic_error_type
20from lilbee.server.anthropic_api.models import (
21 _THINKING_DISABLED,
22 AnthropicEventType,
23 CountTokensRequest,
24 CountTokensResponse,
25 MessagesRequest,
26 MessagesResponse,
27)
28from lilbee.server.anthropic_api.streaming import encode_anthropic_event, encode_anthropic_sse
29from lilbee.server.anthropic_api.translate import (
30 canonical_stream_to_anthropic_events,
31 canonical_to_messages_response,
32 count_tokens_to_canonical_request,
33 messages_to_canonical_request,
34 resolve_reasoning_mode,
35)
36from lilbee.server.auth import auth_checked_in_handler, session_manager
37from lilbee.server.chat_completions_api.errors import (
38 CompletionsErrorCode,
39 classify_provider_error,
40)
41from lilbee.server.chat_dispatch.canonical import CanonicalChatRequest, images_dropped_headers
42from lilbee.server.chat_dispatch.concurrency import (
43 ChatBusyError,
44 ChatSlotGuard,
45 acquire_chat_slot_or_busy,
46)
47from lilbee.server.chat_dispatch.dispatch import (
48 count_request_tokens,
49 dispatch_chat_stream,
50 preflight_chat_request,
51 resolve_served_model,
52)
53from lilbee.server.chat_dispatch.reasoning_cap import (
54 budget_capped_chars,
55 cap_aware_chat,
56 cap_aware_chat_stream,
57)
58from lilbee.server.handlers.sse import SSE_MEDIA_TYPE
59from lilbee.server.validation_format import format_validation
61log = logging.getLogger(__name__)
63_INTERNAL_ERROR_MESSAGE = "Internal server error. Check the server logs for details."
64# Says whether the count_tokens answer was measured on the backend or estimated.
65COUNT_ACCURACY_HEADER = "X-Lilbee-Token-Count-Accuracy"
68async def _auth_before_request(request: Request) -> Response | None:
69 """Reject an unauthenticated caller before Litestar parses the body.
71 Same rationale as the completions surface: the auth answer must ride this
72 surface's own error envelope, and it must win over body validation.
73 """
74 return _auth_failure(request)
77@post("/v1/messages", status_code=200, before_request=_auth_before_request)
78@auth_checked_in_handler
79async def messages_endpoint(request: Request, data: MessagesRequest) -> Response | Stream:
80 """``/v1/messages`` (stream + non-stream + tools), Anthropic wire format.
82 ``stream: true`` switches the 200 response from JSON to the Anthropic SSE
83 event stream; the request body picks the arm, matching Anthropic's own
84 contract.
85 """
86 # Thinking is opt-in per request; the setting only presents it.
87 mode = resolve_reasoning_mode(data.thinking, default=cfg.messages_reasoning)
88 cap_chars = budget_capped_chars(effective_reasoning_cap(), _budget_tokens(data))
89 try:
90 req = messages_to_canonical_request(data, mode=mode)
91 except ValueError as exc:
92 # Wire-valid but untranslatable (image content, bare tool choice).
93 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, str(exc))
95 preflight = await _preflight_resolved_model(req)
96 if isinstance(preflight, Response):
97 return preflight
98 resolved_model = preflight
100 try:
101 await acquire_chat_slot_or_busy(get_services().provider.max_concurrent_chats())
102 except ChatBusyError:
103 return _error_response(
104 429,
105 CompletionsErrorCode.RATE_LIMIT_EXCEEDED,
106 "Backend is busy. Retry in a moment.",
107 headers={"Retry-After": "1"},
108 )
110 guard = ChatSlotGuard()
111 if req.stream:
112 # The after-send hook frees the slot when a disconnect lands before the
113 # generator's first iteration (its finally never runs in that case).
114 return Stream(
115 _gated_messages_stream(
116 req, guard, model=resolved_model, mode=mode, cap_chars=cap_chars
117 ),
118 media_type=SSE_MEDIA_TYPE,
119 background=BackgroundTask(guard.release),
120 headers=images_dropped_headers(req),
121 )
122 return await _run_non_stream(
123 req, guard, canonical_model=resolved_model, mode=mode, cap_chars=cap_chars
124 )
127@post("/v1/messages/count_tokens", status_code=200, before_request=_auth_before_request)
128@auth_checked_in_handler
129async def count_tokens_endpoint(request: Request, data: CountTokensRequest) -> Response:
130 """``/v1/messages/count_tokens``: the served model's own count for a request.
132 Clients read ``input_tokens`` for context accounting, and
133 ``X-Lilbee-Token-Count-Accuracy`` for whether that number was measured on the
134 backend (``exact``) or estimated (``estimated``). Counting never runs a tool,
135 so the only thing a request has to be is translatable: the tool-capability
136 check ``/v1/messages`` runs is not applied. The count does not take a chat
137 slot, so it never answers 429; a count that finds the model cold loads it,
138 as a chat call would.
139 """
140 mode = resolve_reasoning_mode(data.thinking, default=cfg.messages_reasoning)
141 try:
142 req = count_tokens_to_canonical_request(data, mode=mode)
143 except ValueError as exc:
144 # Wire-valid but untranslatable (image content, bare tool choice).
145 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, str(exc))
147 try:
148 canonical_model = await asyncio.to_thread(resolve_served_model, req)
149 count = await asyncio.to_thread(count_request_tokens, req, canonical_model=canonical_model)
150 except Exception as exc:
151 return _classified_error_response(exc)
152 body = CountTokensResponse(input_tokens=count.tokens)
153 return Response(
154 body.model_dump(),
155 media_type="application/json",
156 headers={COUNT_ACCURACY_HEADER: count.accuracy},
157 )
160def _budget_tokens(data: MessagesRequest) -> int | None:
161 """The thinking budget this request asks for; ``disabled`` carries none."""
162 if data.thinking is None or data.thinking.type == _THINKING_DISABLED:
163 return None
164 return data.thinking.budget_tokens
167async def _preflight_resolved_model(req: CanonicalChatRequest) -> str | Response:
168 """Validate *req* before any streaming response starts (see completions)."""
169 try:
170 return await asyncio.to_thread(preflight_chat_request, req)
171 except Exception as exc:
172 return _classified_error_response(exc)
175def _classified_error_response(exc: Exception) -> Response:
176 """The envelope for *exc*, or the generic 500 when nothing classifies it."""
177 classified = classify_provider_error(exc)
178 if classified is None:
179 return _internal_error_response()
180 return _error_response(classified.http_status, classified.code, classified.message)
183async def _run_non_stream(
184 req: CanonicalChatRequest,
185 guard: ChatSlotGuard,
186 *,
187 canonical_model: str,
188 mode: ReasoningMode = ReasoningMode.SEPARATE,
189 cap_chars: int = 0,
190) -> Response:
191 """Dispatch a non-streaming chat call, translating errors to the envelope."""
192 try:
193 resp = await asyncio.to_thread(
194 cap_aware_chat, req, canonical_model=canonical_model, cap_chars=cap_chars
195 )
196 except Exception as exc:
197 return _classified_error_response(exc)
198 finally:
199 await guard.release()
200 body: MessagesResponse = canonical_to_messages_response(
201 resp, response_id=_response_id(), mode=mode
202 )
203 return Response(
204 body.model_dump(), media_type="application/json", headers=images_dropped_headers(req)
205 )
208async def _gated_messages_stream(
209 req: CanonicalChatRequest,
210 guard: ChatSlotGuard,
211 *,
212 model: str,
213 mode: ReasoningMode = ReasoningMode.SEPARATE,
214 cap_chars: int = 0,
215) -> AsyncGenerator[bytes, None]:
216 """Drive dispatch -> translate -> SSE-encode, freeing the slot on exit.
218 A mid-stream failure surfaces as Anthropic's ``event: error`` frame; the
219 headers are already flushed at 200 by then, so the frame is the only
220 channel left.
221 """
222 response_id = _response_id()
223 try:
224 try:
225 events = cap_aware_chat_stream(
226 dispatch_chat_stream(req, canonical_model=model),
227 req,
228 canonical_model=model,
229 cap_chars=cap_chars,
230 )
231 pairs = canonical_stream_to_anthropic_events(
232 events, model=model, response_id=response_id, mode=mode
233 )
234 async for frame in encode_anthropic_sse(pairs):
235 yield frame
236 except Exception as exc:
237 classified = classify_provider_error(exc)
238 if classified is None:
239 log.exception("anthropic messages stream failed")
240 body = anthropic_error_body("api_error", _INTERNAL_ERROR_MESSAGE)
241 else:
242 body = anthropic_error_body(
243 anthropic_error_type(classified.code), classified.message
244 )
245 yield encode_anthropic_event(AnthropicEventType.ERROR, body)
246 finally:
247 await guard.release()
250def _internal_error_response() -> Response:
251 """Log and return the generic api_error 500 envelope."""
252 log.exception("anthropic messages surface failed")
253 return _error_response(500, CompletionsErrorCode.INTERNAL_ERROR, _INTERNAL_ERROR_MESSAGE)
256def _error_response(
257 status: int,
258 code: CompletionsErrorCode,
259 message: str,
260 *,
261 headers: dict[str, str] | None = None,
262) -> Response:
263 return Response(
264 anthropic_error_body(anthropic_error_type(code), message),
265 status_code=status,
266 headers=headers or {},
267 media_type="application/json",
268 )
271def _auth_failure(request: Request) -> Response | None:
272 """Return a 401 envelope if the bearer token is missing/wrong, else None.
274 Claude Code sends ``ANTHROPIC_AUTH_TOKEN`` as a bearer Authorization
275 header, which is exactly the session token check every /v1 route runs.
276 """
277 auth_header = request.headers.get("authorization", "")
278 try:
279 authorized = session_manager.validate(auth_header)
280 except NotAuthorizedException:
281 authorized = False
282 if authorized:
283 return None
284 return _error_response(401, CompletionsErrorCode.INVALID_API_KEY, "Missing or invalid API key.")
287def _validation_exception_handler(_: Request, exc: ValidationException) -> Response:
288 """Wrap Litestar's body-parse failures in the Anthropic error envelope."""
289 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, format_validation(exc))
292def _response_id() -> str:
293 """Anthropic-style ``msg_*`` id."""
294 return f"msg_{uuid.uuid4().hex[:24]}"
297anthropic_router = Router(
298 path="/",
299 route_handlers=[messages_endpoint, count_tokens_endpoint],
300 exception_handlers={ValidationException: _validation_exception_handler},
301)