1"""OpenAI-compatible endpoints. Requests are validated, the model is loaded on demand, then the2request is proxied to the local worker. Usage/timings are recorded from the worker's response."""34from __future__ import annotations56import asyncio7import json8import logging9import time10from typing import Any1112import httpx13from fastapi import APIRouter, Depends, Request14from fastapi.responses import JSONResponse, StreamingResponse1516from ..auth import Principal, require_inference17from ..errors import APIError, GenerationTimeout, WorkerCrashed18from ..events import bus19from ..manager import ModelManager2021log = logging.getLogger("llm_api.openai")22router = APIRouter()232425def _manager(request: Request) -> ModelManager:26 return request.app.state.manager272829async def _parse_body(request: Request) -> dict:30 settings = request.app.state.settings31 raw = await request.body()32 if len(raw) > settings.max_body_bytes:33 raise APIError("Request body too large.", status_code=413, code="BODY_TOO_LARGE")34 try:35 body = json.loads(raw or b"{}")36 except json.JSONDecodeError:37 raise APIError("Request body is not valid JSON.")38 if not isinstance(body, dict):39 raise APIError("Request body must be a JSON object.")40 return body414243def _validate_sampling(body: dict) -> None:44 def num(k, lo, hi):45 v = body.get(k)46 if v is None:47 return48 if not isinstance(v, (int, float)) or v < lo or v > hi:49 raise APIError(f"{k} must be a number between {lo} and {hi}.", param=k)50 num("temperature", 0, 2)51 num("top_p", 0, 1)52 num("presence_penalty", -2, 2)53 num("frequency_penalty", -2, 2)54 mt = body.get("max_tokens", body.get("max_completion_tokens"))55 if mt is not None and (not isinstance(mt, int) or mt < 1 or mt > 200000):56 raise APIError("max_tokens must be a positive integer.", param="max_tokens")57 n = body.get("n")58 if n not in (None, 1):59 raise APIError("Only n=1 is supported.", param="n")60 stop = body.get("stop")61 if stop is not None and not isinstance(stop, (str, list)):62 raise APIError("stop must be a string or an array of strings.", param="stop")63 if isinstance(stop, list) and len(stop) > 8:64 raise APIError("At most 8 stop sequences are supported.", param="stop")656667async def _pick_model(request: Request, body: dict, endpoint: str) -> tuple[str, dict]:68 """Resolve model name (alias/auto) -> registry row."""69 manager = _manager(request)70 name = body.get("model")71 if not name or not isinstance(name, str):72 # default model setting73 name = await request.app.state.db.get_setting("default_model")74 if not name:75 cur = manager.current_model()76 if cur:77 name = cur.model["id"]78 if not name:79 raise APIError("model is required.", param="model")80 if name == "auto":81 from ..routing import choose_auto82 name = await choose_auto(request.app.state, body, endpoint)83 model = await manager.registry.resolve(name)84 if not model:85 from ..errors import ModelNotFound86 raise ModelNotFound(f"The model '{name}' does not exist. Use GET /v1/models to list available models.")87 return name, model888990async def _record(app_state, *, model_id: str | None, requested: str, endpoint: str, principal: Principal,91 stream: bool, usage: dict | None, timings: dict | None, status: int, error_code: str | None,92 started: float, load_wait_ms: float, prompt: Any = None, completion: str | None = None) -> None:93 db = app_state.db94 settings = app_state.settings95 log_prompts = bool(await db.get_setting("log_prompts", settings.log_prompts))96 usage = usage or {}97 timings = timings or {}98 tps = timings.get("generation_tps") or timings.get("predicted_per_second")99 ttft = timings.get("ttft_ms") or timings.get("prompt_ms")100 total_ms = (time.time() - started) * 1000101 ctoks = usage.get("completion_tokens") or timings.get("predicted_n")102 ptoks = usage.get("prompt_tokens") or timings.get("prompt_n")103 await db.execute(104 "INSERT INTO inference_requests(created_at, model_id, requested_model, endpoint, api_key_id, stream, prompt_tokens, "105 "completion_tokens, ttft_ms, total_ms, tps, load_wait_ms, status, error_code, prompt, completion) "106 "VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",107 (started, model_id, requested, endpoint, principal.id if principal.kind == "api_key" else None, int(stream),108 ptoks, ctoks, ttft, round(total_ms, 1), tps, round(load_wait_ms, 1), status, error_code,109 json.dumps(prompt, ensure_ascii=False)[:20000] if (log_prompts and prompt is not None) else None,110 (completion or "")[:20000] if log_prompts else None))111 if model_id and status < 400:112 await app_state.manager.registry.record_usage(model_id, int(ctoks or 0), tps, ttft)113 app_state.manager.stats["tokens"] += int(ctoks or 0)114 bus.publish("request", {"model_id": model_id, "endpoint": endpoint, "status": status, "tps": tps, "ttft_ms": ttft,115 "completion_tokens": ctoks, "prompt_tokens": ptoks, "total_ms": round(total_ms), "stream": stream})116117118def _extract_timings(data: dict) -> dict:119 t = data.get("timings") or {}120 if "predicted_per_second" in t and "generation_tps" not in t: # llama.cpp121 t = {"generation_tps": t.get("predicted_per_second"), "prompt_tps": t.get("prompt_per_second"),122 "ttft_ms": t.get("prompt_ms"), "prompt_ms": t.get("prompt_ms"), "generation_ms": t.get("predicted_ms"),123 "predicted_n": t.get("predicted_n"), "prompt_n": t.get("prompt_n"), "cached_tokens": t.get("cache_n")}124 return t125126127async def _proxy(request: Request, endpoint: str, body: dict, principal: Principal, *, kind: str):128 """Common path for chat/completions/embeddings/rerank."""129 app_state = request.app.state130 manager: ModelManager = app_state.manager131 started = time.time()132 requested, model = await _pick_model(request, body, endpoint)133 if kind == "chat" or kind == "completion":134 if model.get("embedding") or model.get("reranker"):135 raise APIError(f"Model '{model['id']}' is an {'embedding' if model.get('embedding') else 'reranking'} model and cannot generate text.",136 code="WRONG_MODEL_TYPE")137 if kind == "embeddings" and not model.get("embedding"):138 # Allow causal LMs served by MLX to embed (last-token pooling) only if explicitly tagged; otherwise reject139 if model["runtime"] != "mlx":140 raise APIError(f"Model '{model['id']}' is not an embedding model.", code="WRONG_MODEL_TYPE")141 if kind == "rerank" and not model.get("reranker"):142 raise APIError(f"Model '{model['id']}' is not a reranking model.", code="WRONG_MODEL_TYPE")143 stream = bool(body.get("stream")) and kind in ("chat", "completion")144 t_load = time.time()145 try:146 lm = await manager.ensure_loaded(model["id"], reason=f"{endpoint}")147 except APIError as e:148 await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,149 stream=stream, usage=None, timings=None, status=e.status_code, error_code=e.code, started=started,150 load_wait_ms=(time.time() - t_load) * 1000)151 raise152 load_wait_ms = (time.time() - t_load) * 1000153 body = dict(body)154 body["model"] = model["id"]155 if stream and kind == "chat":156 so = dict(body.get("stream_options") or {})157 so["include_usage"] = True158 body["stream_options"] = so159 url = f"{lm.handle.base_url}{endpoint}"160 timeout = httpx.Timeout(connect=10.0, read=float(app_state.settings.generation_timeout_seconds), write=60.0, pool=10.0)161 prompt_for_log = body.get("messages") or body.get("prompt") or body.get("input")162163 async with manager.use(lm):164 if not stream:165 try:166 r = await manager.client.post(url, json=body, timeout=timeout)167 except httpx.ReadTimeout:168 raise GenerationTimeout("The model did not finish generating in time.")169 except httpx.HTTPError as e:170 if not lm.handle.alive():171 raise WorkerCrashed(f"The inference worker for '{model['id']}' crashed during the request.")172 raise APIError(f"Worker connection error: {e}", status_code=502, code="WORKER_UNREACHABLE", error_type="runtime_error")173 try:174 data = r.json()175 except Exception:176 data = {"error": {"message": r.text[:500], "type": "runtime_error", "code": "WORKER_BAD_RESPONSE"}}177 if r.status_code >= 400:178 err = data.get("error") if isinstance(data, dict) else None179 code = (err or {}).get("code") if isinstance(err, dict) else None180 msg = (err or {}).get("message") if isinstance(err, dict) else str(err)181 await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,182 stream=False, usage=None, timings=None, status=r.status_code, error_code=str(code or "WORKER_ERROR"),183 started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log)184 # normalise llama.cpp error shape185 return JSONResponse(status_code=r.status_code, content={"error": {"message": msg or "worker error",186 "type": (err or {}).get("type", "runtime_error") if isinstance(err, dict) else "runtime_error",187 "code": code or "WORKER_ERROR"}})188 data["model"] = requested if requested != "auto" else model["id"]189 timings = _extract_timings(data)190 if timings:191 data["timings"] = timings192 completion = None193 if kind == "chat":194 completion = ((data.get("choices") or [{}])[0].get("message") or {}).get("content")195 elif kind == "completion":196 completion = (data.get("choices") or [{}])[0].get("text")197 await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,198 stream=False, usage=data.get("usage"), timings=timings, status=200, error_code=None,199 started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log, completion=completion)200 return JSONResponse(content=data)201202 # ---- streaming ----------------------------------------------------203 req = manager.client.build_request("POST", url, json=body, timeout=timeout)204 try:205 resp = await manager.client.send(req, stream=True)206 except httpx.HTTPError as e:207 raise APIError(f"Worker connection error: {e}", status_code=502, code="WORKER_UNREACHABLE", error_type="runtime_error")208 if resp.status_code >= 400:209 raw = await resp.aread()210 await resp.aclose()211 try:212 data = json.loads(raw)213 except Exception:214 data = {"error": {"message": raw.decode(errors="replace")[:500], "type": "runtime_error", "code": "WORKER_ERROR"}}215 err = data.get("error") or {}216 await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,217 stream=True, usage=None, timings=None, status=resp.status_code, error_code=str(err.get("code") or "WORKER_ERROR"),218 started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log)219 return JSONResponse(status_code=resp.status_code, content={"error": {"message": err.get("message", "worker error"),220 "type": err.get("type", "runtime_error"),221 "code": err.get("code", "WORKER_ERROR")}})222223 lm.in_flight += 1 # held until the generator finishes224225 async def gen():226 usage = None227 timings: dict = {}228 completion_parts: list[str] = []229 status = 200230 error_code = None231 try:232 async for line in resp.aiter_lines():233 if not line:234 continue235 if line.startswith("data: "):236 payload = line[6:]237 if payload.strip() == "[DONE]":238 yield b"data: [DONE]\n\n"239 continue240 try:241 obj = json.loads(payload)242 except json.JSONDecodeError:243 yield (line + "\n\n").encode()244 continue245 if "error" in obj and "choices" not in obj:246 status = 500247 error_code = (obj["error"] or {}).get("code", "GENERATION_FAILED")248 yield (f"data: {json.dumps(obj)}\n\n").encode()249 continue250 obj["model"] = requested if requested != "auto" else model["id"]251 if obj.get("usage"):252 usage = obj["usage"]253 if obj.get("timings"):254 timings = _extract_timings(obj)255 obj["timings"] = timings256 for ch in obj.get("choices") or []:257 d = ch.get("delta") or {}258 if d.get("content"):259 completion_parts.append(d["content"])260 if ch.get("text"):261 completion_parts.append(ch["text"])262 yield (f"data: {json.dumps(obj, ensure_ascii=False)}\n\n").encode()263 else:264 yield (line + "\n").encode()265 except (httpx.ReadTimeout, asyncio.TimeoutError):266 status, error_code = 504, "GENERATION_TIMEOUT"267 yield (f"data: {json.dumps({'error': {'message': 'generation timed out', 'type': 'runtime_error', 'code': 'GENERATION_TIMEOUT'}})}\n\n").encode()268 except httpx.HTTPError:269 status, error_code = 502, "WORKER_CRASHED" if not lm.handle.alive() else "WORKER_STREAM_ERROR"270 yield (f"data: {json.dumps({'error': {'message': 'worker stream interrupted', 'type': 'runtime_error', 'code': error_code}})}\n\n").encode()271 finally:272 await resp.aclose()273 lm.in_flight = max(0, lm.in_flight - 1)274 lm.last_used = time.time()275 try:276 await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,277 stream=True, usage=usage, timings=timings, status=status, error_code=error_code,278 started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log,279 completion="".join(completion_parts))280 except Exception:281 log.exception("record failed")282283 return StreamingResponse(gen(), media_type="text/event-stream",284 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"})285286287# ---------------------------------------------------------------------------288289290@router.get("/v1/models")291async def list_models(request: Request, principal: Principal = Depends(require_inference)):292 manager = _manager(request)293 models = await manager.registry.list_models()294 aliases = await manager.registry.aliases()295 out = []296 for m in models:297 if not m.get("enabled"):298 continue299 st = manager.status_of(m["id"])300 out.append({301 "id": m["id"], "object": "model", "created": int(m["created_at"]), "owned_by": m.get("provider") or "local",302 "root": m["id"], "parent": None,303 # llm-api extensions (ignored by SDKs)304 "name": m["name"], "family": m["family"], "runtime": m["runtime"], "quantization": m["quantization"],305 "parameter_count": m["parameter_count"], "estimated_ram_gb": m["estimated_ram_gb"],306 "context": m["recommended_context"], "max_context": m["max_context"], "task": m["task"],307 "compatibility": m["compatibility_status"], "status": st, "loaded": st == "ready",308 "capabilities": {"vision": m["vision"], "embedding": m["embedding"], "reranking": m["reranker"],309 "tools": m["tools"], "thinking": m["thinking"]},310 })311 for alias, mid in aliases.items():312 out.append({"id": alias, "object": "model", "created": 0, "owned_by": "alias", "root": mid, "parent": mid,313 "alias_of": mid})314 return {"object": "list", "data": out}315316317@router.get("/v1/models/{model_id:path}")318async def get_model(model_id: str, request: Request, principal: Principal = Depends(require_inference)):319 manager = _manager(request)320 m = await manager.registry.resolve(model_id)321 if not m:322 from ..errors import ModelNotFound323 raise ModelNotFound(f"The model '{model_id}' does not exist.")324 return {"id": m["id"], "object": "model", "created": int(m["created_at"]), "owned_by": m.get("provider") or "local",325 "status": manager.status_of(m["id"])}326327328@router.post("/v1/chat/completions")329async def chat_completions(request: Request, principal: Principal = Depends(require_inference)):330 body = await _parse_body(request)331 msgs = body.get("messages")332 if not isinstance(msgs, list) or not msgs:333 raise APIError("messages must be a non-empty array.", param="messages")334 for m in msgs:335 if not isinstance(m, dict) or "role" not in m:336 raise APIError("Each message needs a role.", param="messages")337 _validate_sampling(body)338 return await _proxy(request, "/v1/chat/completions", body, principal, kind="chat")339340341@router.post("/v1/completions")342async def completions(request: Request, principal: Principal = Depends(require_inference)):343 body = await _parse_body(request)344 if "prompt" not in body:345 raise APIError("prompt is required.", param="prompt")346 _validate_sampling(body)347 return await _proxy(request, "/v1/completions", body, principal, kind="completion")348349350@router.post("/v1/embeddings")351async def embeddings(request: Request, principal: Principal = Depends(require_inference)):352 body = await _parse_body(request)353 if "input" not in body:354 raise APIError("input is required.", param="input")355 inp = body["input"]356 n = len(inp) if isinstance(inp, list) else 1357 if n > 256:358 raise APIError("At most 256 inputs per request.", param="input")359 return await _proxy(request, "/v1/embeddings", body, principal, kind="embeddings")360361362@router.post("/v1/rerank")363async def rerank(request: Request, principal: Principal = Depends(require_inference)):364 body = await _parse_body(request)365 if not body.get("query") or not isinstance(body.get("documents"), list):366 raise APIError("query and documents are required.")367 if len(body["documents"]) > 200:368 raise APIError("At most 200 documents per request.", param="documents")369 return await _proxy(request, "/v1/rerank", body, principal, kind="rerank")370