diff --git a/.forgejo/workflows/opencode.yml b/.forgejo/workflows/opencode.yml index 4fef9e1..6d126d3 100644 --- a/.forgejo/workflows/opencode.yml +++ b/.forgejo/workflows/opencode.yml @@ -11,7 +11,7 @@ jobs: opencode: if: | contains(github.event.comment.body, '/oc') || - contains(github.event.review.body, '/oc') + contains(github.event.comment.body, '/opencode') runs-on: docker-amd64 container: image: node:lts-bookworm @@ -54,7 +54,9 @@ jobs: uses: ./.opencode-action with: nomyo_api_key: ${{ secrets.NOMYO_API_KEY }} - model: nomyo/unsloth/Qwen3.6-35B-A3B-MTP-GGUF:Q4_K_XL + model: nomyo/unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_M forgejo_api_url: https://bitfreedom.net/code/ forgejo_token: ${{ secrets.FORGEJO_TOKEN }} forgejo_push_token: ${{ secrets.FORGEJO_PUSH_TOKEN }} + + diff --git a/.gitignore b/.gitignore index 7cd8431..8cfc932 100644 --- a/.gitignore +++ b/.gitignore @@ -66,4 +66,7 @@ config.yaml # SQLite *.db* +# aissurance compliance plugin — encrypted offline retry buffer +compliance-buffer/ + *settings.json \ No newline at end of file diff --git a/README.md b/README.md index b2cf971..ac277dd 100644 --- a/README.md +++ b/README.md @@ -132,82 +132,6 @@ This way the Ollama backend servers are utilized more efficient than by simply u NOMYO Router also supports OpenAI API compatible v1 backend servers. -## OpenAI Responses API - -In addition to Chat Completions, NOMYO Router exposes the OpenAI **Responses API**: - -``` -POST /v1/responses # create a response (stream or non-stream) -GET /v1/responses/{id} # retrieve a stored response -DELETE /v1/responses/{id} # delete a stored response -POST /v1/responses/{id}/cancel # cancel a background response -``` - -It works transparently across **all** backends. When the routed model lives on a native -Responses backend (external OpenAI) the request is forwarded as-is; for Ollama and llama-server the -router translates Responses ⇄ Chat Completions in both directions (request, response, and streaming -typed SSE events), so clients get a consistent `/v1/responses` surface regardless of backend. - -### Conversation state (`store` / `previous_response_id`) - -The router **owns conversation state itself** (persisted in its SQLite DB) rather than delegating to -the upstream provider, so `store` and `previous_response_id` behave identically on every backend. -On a follow-up request the router rehydrates the prior turns from its DB and expands them into the -conversation; outbound native calls always send `store=false`. Trade-off: this forgoes OpenAI's -server-side reasoning-state reuse in exchange for uniform, backend-agnostic chaining. - -### Background mode - -`background:true` (which requires `store:true`) returns immediately with `{"status":"queued"}`; the -request runs server-side and the client polls `GET /v1/responses/{id}` until the status reaches a -terminal state (`completed` / `failed` / `cancelled`). `POST /v1/responses/{id}/cancel` aborts it. - -Limitations: streaming reconnect-resume via `starting_after` is not yet implemented. In a -multi-worker/replica deployment polling works via the shared DB, but `cancel` only reaches the -running task in the worker that started it (other workers just mark the stored row cancelled). A -background task interrupted by a server restart is reconciled to `failed` on the next startup. - -## Anthropic Messages API - -NOMYO Router also exposes the Anthropic **Messages API**: - -``` -POST /v1/messages # create a message (stream or non-stream) -POST /v1/messages/count_tokens # count input tokens for a request -``` - -It works transparently across **all** backends. For Ollama / llama-server / llama-swap the router -translates Messages ⇄ Chat Completions in both directions (request, response, and streaming typed -SSE events — `message_start` → `content_block_*` → `message_delta` → `message_stop`), so clients get -a consistent `/v1/messages` surface regardless of backend. The API is stateless — there is no store, -background mode, or conversation persistence. - -### Native Anthropic upstream - -Configure real Anthropic endpoints under the `anthropic_endpoints` config key (base URL **without** -a `/v1` suffix). Requests routed to a model advertised by such an endpoint are **forwarded verbatim** -over the Anthropic wire format — the router injects the endpoint's `api_keys` entry as the `x-api-key` -header and pins `anthropic-version`, passing through the client's `anthropic-beta`. Their advertised -models are treated as always-loaded, like external OpenAI endpoints. - -```yaml -anthropic_endpoints: - - https://api.anthropic.com -api_keys: - "https://api.anthropic.com": "${ANTHROPIC_API_KEY}" -``` - -### Thinking - -An inbound `thinking` block is mapped to the backend's `reasoning_effort` (budget → `low`/`medium`/ -`high`); a backend that streams `reasoning_content` is surfaced back as Anthropic `thinking` content -blocks / `thinking_delta` events. On native endpoints, `thinking` passes through untouched. - -### Caching - -Set `nomyo: {"cache": true}` on the request body to consult the router's semantic LLM cache; a hit is -reflected via `usage.cache_read_input_tokens` (input tokens served from cache rather than re-processed). - ## Semantic LLM Cache NOMYO Router includes an optional semantic cache that serves repeated or semantically similar LLM requests from cache — no endpoint round-trip, no token cost, response in <10 ms. @@ -248,7 +172,7 @@ Each request is keyed on `model + system_prompt` (exact) combined with a weighte ### Cached routes -`/api/chat` · `/api/generate` · `/v1/chat/completions` · `/v1/completions` · `/v1/responses` +`/api/chat` · `/api/generate` · `/v1/chat/completions` · `/v1/completions` ### Cache management @@ -257,6 +181,24 @@ curl http://localhost:12434/api/cache/stats # hit rate, counters, config curl -X POST http://localhost:12434/api/cache/invalidate # clear all entries ``` +## Compliance Plugin (aissurance.eu) + +Optional in-process plugin that periodically sends signed AI-infrastructure +evidence — model/endpoint inventory and aggregate telemetry, **never** prompt or +completion content — to [aissurance.eu](https://www.aissurance.eu) for EU AI Act +compliance. All traffic is outbound; the Router keeps proxying if the upstream is +down (evidence is encrypted and buffered for later). + +```yaml +compliance: + enabled: true + server_url: "https://www.aissurance.eu/api/v1/discovery/receive" + api_key: "${AISSURANCE_KEY}" # tenant secret, separate from the router API key + polling_interval: 300 +``` + +See the **[Compliance Guide](doc/compliance.md)** for the full reference. + ## Supplying the router API key If you set `nomyo-router-api-key` in `config.yaml` (or `NOMYO_ROUTER_API_KEY` env), every request to NOMYO Router must include the key: diff --git a/api/management.py b/api/management.py index 0e9ecc2..ac1f356 100644 --- a/api/management.py +++ b/api/management.py @@ -27,7 +27,7 @@ from state import ( _affinity_lock, ) from sse import subscribe, unsubscribe -from backends.normalize import _normalize_llama_model_name, is_llama_server, llama_endpoints +from backends.normalize import _normalize_llama_model_name from backends.probe import _endpoint_health @@ -127,6 +127,7 @@ async def affinity_stats(request: Request): now = time.monotonic() entries: list[dict] = [] + llama_eps = set(config.llama_server_endpoints) async with _affinity_lock: for fp, (ep, mdl, expires_at) in list(_affinity_map.items()): remaining = expires_at - now @@ -135,7 +136,7 @@ async def affinity_stats(request: Request): continue # Mirror the normalisation used by /api/ps_details so the dashboard # can join affinity entries to PS rows by (endpoint, model). - display_model = _normalize_llama_model_name(mdl) if is_llama_server(ep) else mdl + display_model = _normalize_llama_model_name(mdl) if ep in llama_eps else mdl entries.append({ "endpoint": ep, "model": display_model, @@ -174,12 +175,9 @@ async def config_proxy(request: Request): ollama_results = await asyncio.gather(*[check(ep) for ep in config.endpoints]) llama_results = [] - # llama-server and llama-swap render identically in the dashboard ("llama" rows), - # so health-check both and merge them into one list. - llama_eps = llama_endpoints(config) - if llama_eps: + if config.llama_server_endpoints: llama_results = await asyncio.gather( - *[check(ep) for ep in llama_eps] + *[check(ep) for ep in config.llama_server_endpoints] ) return { @@ -229,7 +227,7 @@ async def health_proxy(request: Request): # purposes. Probing /api/version alone would miss the case where the # Ollama process is up but /api/ps is failing — see issue #83. all_endpoints = list(config.endpoints) - llama_eps_extra = [ep for ep in llama_endpoints(config) if ep not in config.endpoints] + llama_eps_extra = [ep for ep in config.llama_server_endpoints if ep not in config.endpoints] all_endpoints += llama_eps_extra probe_results = await asyncio.gather( diff --git a/api/messages.py b/api/messages.py deleted file mode 100644 index 8b113c3..0000000 --- a/api/messages.py +++ /dev/null @@ -1,330 +0,0 @@ -"""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 {} - # 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, - "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)}) diff --git a/api/ollama.py b/api/ollama.py index 7585a5b..afba243 100644 --- a/api/ollama.py +++ b/api/ollama.py @@ -13,7 +13,6 @@ import asyncio import re import time from typing import Optional -from urllib.parse import quote import aiohttp import ollama @@ -41,12 +40,9 @@ from backends.health import ( from backends.normalize import ( dedupe_on_keys, is_openai_compatible, - is_llama_server, - llama_endpoints, _normalize_llama_model_name, _extract_llama_quant, ) -from backends.control import unload_model from backends.probe import fetch from backends.sessions import _make_openai_client, get_ollama_client, get_probe_session from requests.chat import _make_moe_requests @@ -167,38 +163,31 @@ async def proxy(request: Request): _affinity_key = _conversation_fingerprint(model, None, prompt) endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key) - # _guarded_stream's finally releases the reservation once we hand off; until - # then any failure during request building / client construction (including - # CancelledError on client disconnect) must release it or the counter leaks. - try: - use_openai = is_openai_compatible(endpoint) - if use_openai: - if ":latest" in model: - model = model.split(":latest") - model = model[0] - params = { - "prompt": prompt, - "model": model, - } + use_openai = is_openai_compatible(endpoint) + if use_openai: + if ":latest" in model: + model = model.split(":latest") + model = model[0] + params = { + "prompt": prompt, + "model": model, + } - optional_params = { - "stream": stream, - "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, - "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, - "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, - "seed": options.get("seed") if options and "seed" in options else None, - "stop": options.get("stop") if options and "stop" in options else None, - "top_p": options.get("top_p") if options and "top_p" in options else None, - "temperature": options.get("temperature") if options and "temperature" in options else None, - "suffix": suffix, - } - params.update({k: v for k, v in optional_params.items() if v is not None}) - oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) - else: - client = get_ollama_client(endpoint) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise + optional_params = { + "stream": stream, + "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, + "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, + "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, + "seed": options.get("seed") if options and "seed" in options else None, + "stop": options.get("stop") if options and "stop" in options else None, + "top_p": options.get("top_p") if options and "top_p" in options else None, + "temperature": options.get("temperature") if options and "temperature" in options else None, + "suffix": suffix, + } + params.update({k: v for k, v in optional_params.items() if v is not None}) + oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) + else: + client = get_ollama_client(endpoint) # 4. Async generator body (error handling + cleanup handled by _guarded_stream) async def stream_generate_response(): @@ -343,70 +332,59 @@ async def chat_proxy(request: Request): opt = False _affinity_key = _conversation_fingerprint(model, messages, None) endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key) - # Releasing the reservation is owned by _guarded_stream's finally once we hand - # off to the streaming generator. Until then, any failure during request - # building / client construction (including CancelledError on client - # disconnect) must release it here or the usage counter leaks. - try: - use_openai = is_openai_compatible(endpoint) - if use_openai: - if ":latest" in model: - model = model.split(":latest") - model = model[0] - if messages: - if any("images" in m for m in messages): - messages = await asyncio.to_thread(transform_images_to_data_urls, messages) - messages = transform_tool_calls_to_openai(messages) - messages = _strip_assistant_prefill(messages) - params = { - "messages": messages, - "model": model, - } - optional_params = { - "tools": tools, - "stream": stream, - "stream_options": {"include_usage": True} if stream else None, - "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, - "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, - "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, - "seed": options.get("seed") if options and "seed" in options else None, - "stop": options.get("stop") if options and "stop" in options else None, - "top_p": options.get("top_p") if options and "top_p" in options else None, - "temperature": options.get("temperature") if options and "temperature" in options else None, - "logprobs": logprobs if logprobs is not None else (options.get("logprobs") if options and "logprobs" in options else None), - "top_logprobs": top_logprobs if top_logprobs is not None else (options.get("top_logprobs") if options and "top_logprobs" in options else None), - "response_format": {"type": "json_schema", "json_schema": _format} if _format is not None else None - } - params.update({k: v for k, v in optional_params.items() if v is not None}) - oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) - else: - client = get_ollama_client(endpoint) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise + use_openai = is_openai_compatible(endpoint) + if use_openai: + if ":latest" in model: + model = model.split(":latest") + model = model[0] + if messages: + if any("images" in m for m in messages): + messages = await asyncio.to_thread(transform_images_to_data_urls, messages) + messages = transform_tool_calls_to_openai(messages) + messages = _strip_assistant_prefill(messages) + params = { + "messages": messages, + "model": model, + } + optional_params = { + "tools": tools, + "stream": stream, + "stream_options": {"include_usage": True} if stream else None, + "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, + "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, + "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, + "seed": options.get("seed") if options and "seed" in options else None, + "stop": options.get("stop") if options and "stop" in options else None, + "top_p": options.get("top_p") if options and "top_p" in options else None, + "temperature": options.get("temperature") if options and "temperature" in options else None, + "logprobs": logprobs if logprobs is not None else (options.get("logprobs") if options and "logprobs" in options else None), + "top_logprobs": top_logprobs if top_logprobs is not None else (options.get("top_logprobs") if options and "top_logprobs" in options else None), + "response_format": {"type": "json_schema", "json_schema": _format} if _format is not None else None + } + params.update({k: v for k, v in optional_params.items() if v is not None}) + oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) + else: + client = get_ollama_client(endpoint) # For OpenAI endpoints: make the API call in handler scope # (try/except inside async generators is unreliable with Starlette's streaming) start_ts = None async_gen = None if use_openai: start_ts = time.perf_counter() + # Proactive trim: only for small-ctx models we've already seen run out of space + _lookup_model = _normalize_llama_model_name(model) if endpoint in config.llama_server_endpoints else model + _known_nctx = _endpoint_nctx.get((endpoint, _lookup_model)) + if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT: + _pre_target = int((_known_nctx - _known_nctx // 4) / 1.2) + _pre_est = _count_message_tokens(params.get("messages", [])) + if _pre_est > _pre_target: + _pre_msgs = params.get("messages", []) + _pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target) + _dropped = len(_pre_msgs) - len(_pre_trimmed) + print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True) + params = {**params, "messages": _pre_trimmed} try: - # Proactive trim: only for small-ctx models we've already seen run out of space - _lookup_model = _normalize_llama_model_name(model) if is_llama_server(endpoint) else model - _known_nctx = _endpoint_nctx.get((endpoint, _lookup_model)) - if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT: - _pre_target = int((_known_nctx - _known_nctx // 4) / 1.2) - _pre_est = _count_message_tokens(params.get("messages", [])) - if _pre_est > _pre_target: - _pre_msgs = params.get("messages", []) - _pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target) - _dropped = len(_pre_msgs) - len(_pre_trimmed) - print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True) - params = {**params, "messages": _pre_trimmed} async_gen = await oclient.chat.completions.create(**params) - except asyncio.CancelledError: - await decrement_usage(endpoint, tracking_model) - raise except Exception as e: _e_str = str(e) print(f"[chat_proxy] caught {type(e).__name__}: {_e_str[:200]}") @@ -613,21 +591,14 @@ async def _handle_embedding_request( # 2. Endpoint logic endpoint, tracking_model = await choose_endpoint(model) - # _guarded_stream's finally releases the reservation once we hand off; until - # then any failure during client construction (including CancelledError on - # client disconnect) must release it or the counter leaks. - try: - use_openai = is_openai_compatible(endpoint) - if use_openai: - if ":latest" in model: - model = model.split(":latest") - model = model[0] - client = _make_openai_client(endpoint, api_key=config.api_keys.get(endpoint, "no-key")) - else: - client = get_ollama_client(endpoint) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise + use_openai = is_openai_compatible(endpoint) + if use_openai: + if ":latest" in model: + model = model.split(":latest") + model = model[0] + client = _make_openai_client(endpoint, api_key=config.api_keys.get(endpoint, "no-key")) + else: + client = get_ollama_client(endpoint) # 3. Async generator body (error handling + cleanup handled by _guarded_stream) async def stream_embedding_response(): @@ -964,8 +935,8 @@ async def tags_proxy(request: Request): # 1. Query all endpoints for models tasks = [fetch.endpoint_details(ep, "/api/tags", "models", skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" not in ep] tasks += [fetch.endpoint_details(ep, "/models", "data", config.api_keys[ep], skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" in ep] - # Also query llama-server / llama-swap endpoints not already covered by config.endpoints - llama_eps_for_tags = [ep for ep in llama_endpoints(config) if ep not in config.endpoints] + # Also query llama-server endpoints not already covered by config.endpoints + llama_eps_for_tags = [ep for ep in config.llama_server_endpoints if ep not in config.endpoints] tasks += [fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) for ep in llama_eps_for_tags] all_models = await asyncio.gather(*tasks) @@ -989,79 +960,27 @@ async def tags_proxy(request: Request): ) -async def _fetch_llama_swap_running(endpoint: str) -> list[dict]: - """Return the list of ready (`state == "ready"`) workers from a llama-swap - endpoint's `/running` route. llama-swap omits the per-model `status` field on - `/v1/models`, so running workers must be read here instead. - """ - config = get_config() - base_url = endpoint.rstrip("/").removesuffix("/v1") - return await fetch.endpoint_details( - base_url, "/running", "running", config.api_keys.get(endpoint), - skip_error_cache=True, timeout=8, - ) - - -# Match the context size in a llama-swap worker's `cmd` string, e.g. -# "llama-server --port 5818 -hf ... --ctx-size 131072 ...". llama.cpp accepts -# both --ctx-size and the short -c alias. -_CTX_SIZE_CMD_RE = re.compile(r"(?:--ctx-size|-c)[=\s]+(\d+)") - - -def _ctx_size_from_cmd(cmd: str) -> int | None: - """Extract n_ctx from a llama-swap worker `cmd` string, or None if absent.""" - if not cmd: - return None - m = _CTX_SIZE_CMD_RE.search(cmd) - return int(m.group(1)) if m else None - - -async def _fetch_llama_swap_nctx(endpoint: str, model_id: str) -> int | None: - """Fallback when a worker's `cmd` lacks --ctx-size: ask the underlying - llama-server via llama-swap's /upstream//props route (plain /props?model= - is not routed by llama-swap and 404s). Returns n_ctx or None on any failure. - """ - config = get_config() - base_url = endpoint.rstrip("/").removesuffix("/v1") - props_url = f"{base_url}/upstream/{quote(model_id, safe='')}/props" - headers = None - api_key = config.api_keys.get(endpoint) - if api_key: - headers = {"Authorization": f"Bearer {api_key}"} - try: - client: aiohttp.ClientSession = get_probe_session(endpoint) - async with client.get(props_url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp: - if resp.status == 200: - data = await resp.json() - return data.get("default_generation_settings", {}).get("n_ctx") - except Exception as e: - print(f"[ps_details] Failed to fetch props from {props_url}: {e}") - return None - - @router.get("/api/ps") async def ps_proxy(request: Request): """ - Proxy a ps request to all Ollama, llama-server and llama-swap endpoints and reply a unique list of all running models. + Proxy a ps request to all Ollama and llama-server endpoints and reply a unique list of all running models. For Ollama endpoints: queries /api/ps For llama-server endpoints: queries /v1/models with status.value == "loaded" - For llama-swap endpoints: queries /running (state == "ready") """ config = get_config() # 1. Query Ollama endpoints for running models via /api/ps ollama_tasks = [fetch.endpoint_details(ep, "/api/ps", "models", skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" not in ep] # 2. Query llama-server endpoints for loaded models via /v1/models + # Also query endpoints from llama_server_endpoints that may not be in config.endpoints + all_llama_endpoints = set(config.llama_server_endpoints) | set(ep for ep in config.endpoints if ep in config.llama_server_endpoints) llama_tasks = [ fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) - for ep in config.llama_server_endpoints + for ep in all_llama_endpoints ] - # 3. Query llama-swap endpoints for running workers via /running - swap_tasks = [_fetch_llama_swap_running(ep) for ep in config.llama_swap_endpoints] ollama_loaded = await asyncio.gather(*ollama_tasks) if ollama_tasks else [] llama_loaded = await asyncio.gather(*llama_tasks) if llama_tasks else [] - swap_running = await asyncio.gather(*swap_tasks) if swap_tasks else [] models = {'models': []} # Add Ollama models (if any) @@ -1084,21 +1003,6 @@ async def ps_proxy(request: Request): "status": item.get("status"), "details": {"quantization_level": quant} if quant else {} }) - # Add llama-swap running workers (already filtered on state == "ready") - if swap_running: - for runlist in swap_running: - for item in runlist: - if item.get("state") != "ready": - continue - raw_id = item.get("model", "") - normalized = _normalize_llama_model_name(raw_id) - quant = _extract_llama_quant(raw_id) - models['models'].append({ - "name": normalized, - "id": normalized, - "digest": "", - "details": {"quantization_level": quant} if quant else {} - }) # 3. Return a JSONResponse with deduplicated currently deployed models # Deduplicate on 'name' rather than 'digest': llama-server models always @@ -1197,7 +1101,16 @@ async def ps_details_proxy(request: Request): is_generation = "temperature" in dgs if is_sleeping: - await unload_model(endpoint, model_id) + unload_url = f"{base_url}/models/unload" + try: + async with client.post( + unload_url, + json={"model": model_id}, + headers=headers, + ) as unload_resp: + print(f"[ps_details] Unloaded sleeping model {model_id} from {endpoint}: {unload_resp.status}") + except Exception as ue: + print(f"[ps_details] Failed to unload sleeping model {model_id} from {endpoint}: {ue}") return n_ctx, is_sleeping, is_generation except Exception as e: @@ -1218,55 +1131,4 @@ async def ps_details_proxy(request: Request): if not is_sleeping: models.append(model_dict) - # Add llama-swap running workers (read from /running; no status/props/auto-unload — - # llama-swap omits the status field on /v1/models and manages its own TTL eviction). - if config.llama_swap_endpoints: - swap_running = await asyncio.gather( - *[_fetch_llama_swap_running(ep) for ep in config.llama_swap_endpoints] - ) - swap_nctx_fallbacks: list[tuple[str, str, dict]] = [] - for endpoint, runlist in zip(config.llama_swap_endpoints, swap_running): - for item in runlist: - if not isinstance(item, dict) or item.get("state") != "ready": - continue - raw_id = item.get("model", "") - if not raw_id: - continue - normalized = _normalize_llama_model_name(raw_id) - quant = _extract_llama_quant(raw_id) - swap_model = { - "name": normalized, - "id": normalized, - "original_name": raw_id, - "digest": "", - "details": {"quantization_level": quant} if quant else {}, - "endpoint": endpoint, - "state": item.get("state"), - "ttl": item.get("ttl"), - "proxy": item.get("proxy"), - } - # llama-swap omits n_ctx from /running, but the worker's launch - # command carries --ctx-size, so parse it from there (no extra - # request). Workers whose cmd lacks the flag fall back to an - # /upstream//props probe below. - n_ctx = _ctx_size_from_cmd(item.get("cmd", "")) - if n_ctx is not None: - swap_model["context_length"] = n_ctx - if 0 < n_ctx <= _CTX_TRIM_SMALL_LIMIT: - _endpoint_nctx[(endpoint, normalized)] = n_ctx - else: - swap_nctx_fallbacks.append((endpoint, raw_id, swap_model)) - models.append(swap_model) - - # Resolve ctx for workers whose cmd lacked --ctx-size via /upstream props. - if swap_nctx_fallbacks: - fallback_results = await asyncio.gather( - *[_fetch_llama_swap_nctx(ep, rid) for ep, rid, _ in swap_nctx_fallbacks] - ) - for (ep, _rid, swap_model), n_ctx in zip(swap_nctx_fallbacks, fallback_results): - if n_ctx is not None: - swap_model["context_length"] = n_ctx - if 0 < n_ctx <= _CTX_TRIM_SMALL_LIMIT: - _endpoint_nctx[(ep, swap_model["id"])] = n_ctx - return JSONResponse(content={"models": models}, status_code=200) diff --git a/api/openai.py b/api/openai.py index 9b64fe9..4662f50 100644 --- a/api/openai.py +++ b/api/openai.py @@ -34,8 +34,6 @@ from backends.normalize import ( ep2base, is_ext_openai_endpoint, is_openai_compatible, - is_llama_server, - llama_endpoints, _normalize_llama_model_name, ) from backends.probe import fetch @@ -48,95 +46,6 @@ from routing import choose_endpoint, decrement_usage router = APIRouter() -async def create_chat_with_retries(oclient, send_params, endpoint, model, tracking_model): - """Call ``chat.completions.create`` with the router's resilience retries. - - Encapsulates the recovery ladder shared by the chat-completions handler and - the translated ``/v1/responses`` path: - - * ``does not support tools`` → retry without ``tools`` - * llama-server context exhaustion → sliding-window message trim, with a - second retry that also strips ``tools``/``tool_choice`` - * backend connection failure → mark (endpoint, model) unhealthy so the next - request reroutes, then re-raise - * ``image input is not supported`` → strip images and retry - - The caller owns the usage reservation taken by ``choose_endpoint``: this - function never decrements it. On unrecoverable failure the exception is - re-raised so the caller's guard releases the slot exactly once. Returns the - established async generator / response. - """ - config = get_config() - try: - async_gen = await oclient.chat.completions.create(**send_params) - except Exception as e: - _e_str = str(e) - _is_ctx_err = "exceed_context_size_error" in _e_str or "exceeds the available context size" in _e_str - print(f"[ochat] caught={type(e).__name__} ctx={_is_ctx_err} msg={_e_str[:120]}", flush=True) - if "does not support tools" in _e_str: - # Model doesn't support tools — retry without them - print(f"[ochat] retry: no tools", flush=True) - params_without_tools = {k: v for k, v in send_params.items() if k != "tools"} - async_gen = await oclient.chat.completions.create(**params_without_tools) - elif _is_ctx_err: - # Backend context limit hit — apply sliding-window trim (context-shift at message level) - err_body = getattr(e, "body", {}) or {} - err_detail = err_body.get("error", {}) if isinstance(err_body, dict) else {} - n_ctx_limit = err_detail.get("n_ctx", 0) - actual_tokens = err_detail.get("n_prompt_tokens", 0) - # Fallback: parse from string if body parsing yielded nothing (SDK may not parse llama-server errors) - if not n_ctx_limit: - import re as _re - _m = _re.search(r"'n_ctx':\s*(\d+)", _e_str) - if _m: - n_ctx_limit = int(_m.group(1)) - _m = _re.search(r"'n_prompt_tokens':\s*(\d+)", _e_str) - if _m: - actual_tokens = int(_m.group(1)) - print(f"[ctx-trim] n_ctx={n_ctx_limit} actual={actual_tokens}", flush=True) - if not n_ctx_limit: - raise - if n_ctx_limit <= _CTX_TRIM_SMALL_LIMIT: - _endpoint_nctx[(endpoint, model)] = n_ctx_limit - - msgs_to_trim = send_params.get("messages", []) - try: - cal_target = _calibrated_trim_target(msgs_to_trim, n_ctx_limit, actual_tokens) - trimmed_messages = _trim_messages_for_context(msgs_to_trim, n_ctx_limit, target_tokens=cal_target) - except Exception as _helper_exc: - print(f"[ctx-trim] helper crash: {type(_helper_exc).__name__}: {str(_helper_exc)[:100]}", flush=True) - raise - dropped = len(msgs_to_trim) - len(trimmed_messages) - print(f"[ctx-trim] target={cal_target} dropped={dropped} remaining={len(trimmed_messages)} retrying-1", flush=True) - try: - async_gen = await oclient.chat.completions.create(**{**send_params, "messages": trimmed_messages}) - print(f"[ctx-trim] retry-1 ok", flush=True) - except Exception as e2: - _e2_str = str(e2) - if "exceed_context_size_error" in _e2_str or "exceeds the available context size" in _e2_str: - # Still too large — tool definitions likely consuming too many tokens, strip them too - print(f"[ctx-trim] retry-1 still exceeded, stripping tools retrying-2", flush=True) - params_no_tools = {k: v for k, v in send_params.items() if k not in ("tools", "tool_choice")} - async_gen = await oclient.chat.completions.create(**{**params_no_tools, "messages": trimmed_messages}) - print(f"[ctx-trim] retry-2 ok", flush=True) - else: - raise - elif _is_backend_connection_error(e): - # Upstream connection failed (e.g. llama-server in router mode - # whose delegated worker died). Mark (endpoint, model) so the - # next request reroutes; the client will retry this one. - print(f"[ochat] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True) - await _mark_backend_unhealthy(endpoint, model, _e_str) - raise - elif "image input is not supported" in _e_str: - # Model doesn't support images — strip and retry - print(f"[openai_chat_completions_proxy] Model {model} doesn't support images, retrying with text-only messages") - async_gen = await oclient.chat.completions.create(**{**send_params, "messages": _strip_images_from_messages(send_params.get("messages", []))}) - else: - raise - return async_gen - - @router.post("/v1/embeddings") async def openai_embedding_proxy(request: Request): """ @@ -180,14 +89,13 @@ async def openai_embedding_proxy(request: Request): # 2. Endpoint logic endpoint, tracking_model = await choose_endpoint(model) - # The finally below releases the reservation for every exit — success, error, - # or CancelledError — so client construction is kept inside the guarded block. + if is_openai_compatible(endpoint): + api_key = config.api_keys.get(endpoint, "no-key") + else: + api_key = "ollama" + oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=api_key) + try: - if is_openai_compatible(endpoint): - api_key = config.api_keys.get(endpoint, "no-key") - else: - api_key = "ollama" - oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=api_key) async_gen = await oclient.embeddings.create(input=doc, model=model) result = async_gen.model_dump() for item in result.get("data", []): @@ -336,30 +244,106 @@ async def openai_chat_completions_proxy(request: Request): # Make the API call in handler scope — try/except inside async generators is unreliable # with Starlette's streaming machinery, so we resolve errors here before the generator starts. - # The reservation taken by choose_endpoint is released by stream_ochat_response's finally - # once we hand off; until then, any failure here (including CancelledError on client - # disconnect during a cold model load) must release it or the counter leaks. + send_params = params + if not is_ext_openai_endpoint(endpoint): + resolved_msgs = await _normalize_images_in_messages(params.get("messages", [])) + send_params = {**params, "messages": resolved_msgs} + # Proactive trim: only for small-ctx models we've already seen run out of space + _lookup_model = _normalize_llama_model_name(model) if endpoint in config.llama_server_endpoints else model + _known_nctx = _endpoint_nctx.get((endpoint, _lookup_model)) + if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT: + _pre_target = int(((_known_nctx - _known_nctx // 4)) / 1.2) + _pre_est = _count_message_tokens(send_params.get("messages", [])) + if _pre_est > _pre_target: + _pre_msgs = send_params.get("messages", []) + _pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target) + _dropped = len(_pre_msgs) - len(_pre_trimmed) + print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True) + send_params = {**send_params, "messages": _pre_trimmed} try: - send_params = params - if not is_ext_openai_endpoint(endpoint): - resolved_msgs = await _normalize_images_in_messages(params.get("messages", [])) - send_params = {**params, "messages": resolved_msgs} - # Proactive trim: only for small-ctx models we've already seen run out of space - _lookup_model = _normalize_llama_model_name(model) if is_llama_server(endpoint) else model - _known_nctx = _endpoint_nctx.get((endpoint, _lookup_model)) - if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT: - _pre_target = int(((_known_nctx - _known_nctx // 4)) / 1.2) - _pre_est = _count_message_tokens(send_params.get("messages", [])) - if _pre_est > _pre_target: - _pre_msgs = send_params.get("messages", []) - _pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target) - _dropped = len(_pre_msgs) - len(_pre_trimmed) - print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True) - send_params = {**send_params, "messages": _pre_trimmed} - async_gen = await create_chat_with_retries(oclient, send_params, endpoint, model, tracking_model) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise + async_gen = await oclient.chat.completions.create(**send_params) + except Exception as e: + _e_str = str(e) + _is_ctx_err = "exceed_context_size_error" in _e_str or "exceeds the available context size" in _e_str + print(f"[ochat] caught={type(e).__name__} ctx={_is_ctx_err} msg={_e_str[:120]}", flush=True) + if "does not support tools" in _e_str: + # Model doesn't support tools — retry without them + print(f"[ochat] retry: no tools", flush=True) + try: + params_without_tools = {k: v for k, v in send_params.items() if k != "tools"} + async_gen = await oclient.chat.completions.create(**params_without_tools) + except Exception: + await decrement_usage(endpoint, tracking_model) + raise + elif _is_ctx_err: + # Backend context limit hit — apply sliding-window trim (context-shift at message level) + err_body = getattr(e, "body", {}) or {} + err_detail = err_body.get("error", {}) if isinstance(err_body, dict) else {} + n_ctx_limit = err_detail.get("n_ctx", 0) + actual_tokens = err_detail.get("n_prompt_tokens", 0) + # Fallback: parse from string if body parsing yielded nothing (SDK may not parse llama-server errors) + if not n_ctx_limit: + import re as _re + _m = _re.search(r"'n_ctx':\s*(\d+)", _e_str) + if _m: + n_ctx_limit = int(_m.group(1)) + _m = _re.search(r"'n_prompt_tokens':\s*(\d+)", _e_str) + if _m: + actual_tokens = int(_m.group(1)) + print(f"[ctx-trim] n_ctx={n_ctx_limit} actual={actual_tokens}", flush=True) + if not n_ctx_limit: + await decrement_usage(endpoint, tracking_model) + raise + if n_ctx_limit <= _CTX_TRIM_SMALL_LIMIT: + _endpoint_nctx[(endpoint, model)] = n_ctx_limit + + msgs_to_trim = send_params.get("messages", []) + try: + cal_target = _calibrated_trim_target(msgs_to_trim, n_ctx_limit, actual_tokens) + trimmed_messages = _trim_messages_for_context(msgs_to_trim, n_ctx_limit, target_tokens=cal_target) + except Exception as _helper_exc: + print(f"[ctx-trim] helper crash: {type(_helper_exc).__name__}: {str(_helper_exc)[:100]}", flush=True) + await decrement_usage(endpoint, tracking_model) + raise + dropped = len(msgs_to_trim) - len(trimmed_messages) + print(f"[ctx-trim] target={cal_target} dropped={dropped} remaining={len(trimmed_messages)} retrying-1", flush=True) + try: + async_gen = await oclient.chat.completions.create(**{**send_params, "messages": trimmed_messages}) + print(f"[ctx-trim] retry-1 ok", flush=True) + except Exception as e2: + _e2_str = str(e2) + if "exceed_context_size_error" in _e2_str or "exceeds the available context size" in _e2_str: + # Still too large — tool definitions likely consuming too many tokens, strip them too + print(f"[ctx-trim] retry-1 still exceeded, stripping tools retrying-2", flush=True) + params_no_tools = {k: v for k, v in send_params.items() if k not in ("tools", "tool_choice")} + try: + async_gen = await oclient.chat.completions.create(**{**params_no_tools, "messages": trimmed_messages}) + print(f"[ctx-trim] retry-2 ok", flush=True) + except Exception: + await decrement_usage(endpoint, tracking_model) + raise + else: + await decrement_usage(endpoint, tracking_model) + raise + elif _is_backend_connection_error(e): + # Upstream connection failed (e.g. llama-server in router mode + # whose delegated worker died). Mark (endpoint, model) so the + # next request reroutes; the client will retry this one. + print(f"[ochat] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True) + await _mark_backend_unhealthy(endpoint, model, _e_str) + await decrement_usage(endpoint, tracking_model) + raise + elif "image input is not supported" in _e_str: + # Model doesn't support images — strip and retry + print(f"[openai_chat_completions_proxy] Model {model} doesn't support images, retrying with text-only messages") + try: + async_gen = await oclient.chat.completions.create(**{**send_params, "messages": _strip_images_from_messages(send_params.get("messages", []))}) + except Exception: + await decrement_usage(endpoint, tracking_model) + raise + else: + await decrement_usage(endpoint, tracking_model) + raise # 4. Async generator — only streams the already-established async_gen async def stream_ochat_response(): @@ -540,17 +524,12 @@ async def openai_completions_proxy(request: Request): # 2. Endpoint logic _affinity_key = _conversation_fingerprint(model, None, prompt) endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key) + oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) # 3. Async generator that streams completions data and decrements the counter - # Make the API call in handler scope (try/except inside async generators is unreliable). - # The reservation is released by stream_ocompletions_response's finally once we hand off; - # until then any failure here — including CancelledError on client disconnect — releases it. + # Make the API call in handler scope (try/except inside async generators is unreliable) try: - oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) async_gen = await oclient.completions.create(**params) - except asyncio.CancelledError: - await decrement_usage(endpoint, tracking_model) - raise except Exception as e: if _is_backend_connection_error(e): print(f"[ocompl] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True) @@ -658,23 +637,17 @@ async def openai_models_proxy(request: Request): ollama_tasks = [fetch.endpoint_details(ep, "/api/tags", "models", skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" not in ep] # 2. Query external OpenAI endpoints (Groq, OpenAI, etc.) via /models ext_openai_tasks = [fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) for ep in config.endpoints if is_ext_openai_endpoint(ep)] - # 3. Query llama-server / llama-swap endpoints for advertised models via /v1/models - # Also query endpoints that may not be in config.endpoints - all_llama_endpoints = llama_endpoints(config) + # 3. Query llama-server endpoints for loaded models via /v1/models + # Also query endpoints from llama_server_endpoints that may not be in config.endpoints + all_llama_endpoints = set(config.llama_server_endpoints) | set(ep for ep in config.endpoints if ep in config.llama_server_endpoints) llama_tasks = [ fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) for ep in all_llama_endpoints ] - # 4. Query native Anthropic endpoints via /v1/models (auth headers picked by endpoint type) - anthropic_tasks = [ - fetch.endpoint_details(ep, "/v1/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) - for ep in config.anthropic_endpoints - ] ollama_models = await asyncio.gather(*ollama_tasks) if ollama_tasks else [] ext_openai_models = await asyncio.gather(*ext_openai_tasks) if ext_openai_tasks else [] llama_models = await asyncio.gather(*llama_tasks) if llama_tasks else [] - anthropic_models = await asyncio.gather(*anthropic_tasks) if anthropic_tasks else [] models = {'data': []} @@ -708,16 +681,6 @@ async def openai_models_proxy(request: Request): model['name'] = model['id'] models['data'].append(model) - # Add native Anthropic models (if any) - if anthropic_models: - for modellist in anthropic_models: - for model in modellist: - if not "id" in model.keys(): - model['id'] = model.get('name', model.get('id', '')) - else: - model['name'] = model['id'] - models['data'].append(model) - # 2. Return a JSONResponse with a deduplicated list of unique models for inference return JSONResponse( content={"data": dedupe_on_keys(models['data'], ['name'])}, @@ -789,38 +752,36 @@ async def rerank_proxy(request: Request): ), ) - # The finally below releases the reservation for every exit (success, error, - # or CancelledError), so request building and session lookup stay inside it. - try: - if ":latest" in model: - model = model.split(":latest")[0] + if ":latest" in model: + model = model.split(":latest")[0] - # Build upstream rerank request body – forward only recognised fields - upstream_payload: dict = {"model": model, "query": query, "documents": documents} - for optional_key in ("top_n", "return_documents", "max_tokens_per_doc"): - if optional_key in payload: - upstream_payload[optional_key] = payload[optional_key] + # Build upstream rerank request body – forward only recognised fields + upstream_payload: dict = {"model": model, "query": query, "documents": documents} + for optional_key in ("top_n", "return_documents", "max_tokens_per_doc"): + if optional_key in payload: + upstream_payload[optional_key] = payload[optional_key] - # Determine upstream URL: - # llama-server / llama-swap expose /v1/rerank (base already contains /v1) - # External OpenAI endpoints expose /rerank under their /v1 base - if is_llama_server(endpoint): - # llama-server / llama-swap: endpoint may or may not already contain /v1 - if "/v1" in endpoint: - rerank_url = f"{endpoint}/rerank" - else: - rerank_url = f"{endpoint}/v1/rerank" + # Determine upstream URL: + # llama-server exposes /v1/rerank (base already contains /v1 for llama_server_endpoints) + # External OpenAI endpoints expose /rerank under their /v1 base + if endpoint in config.llama_server_endpoints: + # llama-server: endpoint may or may not already contain /v1 + if "/v1" in endpoint: + rerank_url = f"{endpoint}/rerank" else: - # External OpenAI-compatible: ep2base gives us the /v1 base - rerank_url = f"{ep2base(endpoint)}/rerank" + rerank_url = f"{endpoint}/v1/rerank" + else: + # External OpenAI-compatible: ep2base gives us the /v1 base + rerank_url = f"{ep2base(endpoint)}/rerank" - api_key = config.api_keys.get(endpoint, "no-key") - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } + api_key = config.api_keys.get(endpoint, "no-key") + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } - client: aiohttp.ClientSession = get_session(endpoint) + client: aiohttp.ClientSession = get_session(endpoint) + try: async with client.post(rerank_url, json=upstream_payload, headers=headers) as resp: response_bytes = await resp.read() if resp.status >= 400: @@ -841,82 +802,3 @@ async def rerank_proxy(request: Request): return JSONResponse(content=data) finally: await decrement_usage(endpoint, tracking_model) - - -async def _resolve_llama_swap_endpoint(model_id: str) -> str | None: - """Pick the llama-swap endpoint that serves ``model_id``. - - Prefers an endpoint that already has the worker running; falls back to any - that advertises the model. Returns None if none do. - """ - config = get_config() - swap_eps = config.llama_swap_endpoints - if not swap_eps: - return None - - advertised = await asyncio.gather( - *[fetch.available_models(ep, config.api_keys.get(ep)) for ep in swap_eps] - ) - candidates = [ep for ep, models in zip(swap_eps, advertised) if model_id in models] - if not candidates: - return None - if len(candidates) == 1: - return candidates[0] - - loaded = await asyncio.gather(*[fetch.loaded_models(ep) for ep in candidates]) - for ep, lm in zip(candidates, loaded): - if model_id in lm: - return ep - return candidates[0] - - -@router.api_route("/upstream/{model_id}/{path:path}", methods=["GET", "POST"]) -async def llama_swap_upstream(model_id: str, path: str, request: Request): - """Bypass llama-swap and reach a model's underlying llama-server worker directly - via llama-swap's ``/upstream/:model_id`` route. - - Lets clients use llama-server features that llama-swap itself does not forward - (e.g. token-array prompts), while still letting the router pick the backend that - actually hosts the model. ``/upstream`` is a root route, so the ``/v1`` suffix is - stripped from the configured endpoint. - """ - config = get_config() - endpoint = await _resolve_llama_swap_endpoint(model_id) - if endpoint is None: - raise HTTPException( - status_code=404, - detail=f"No configured llama-swap endpoint serves model '{model_id}'.", - ) - - base_url = endpoint.rstrip("/").removesuffix("/v1") - url = f"{base_url}/upstream/{model_id}/{path}" - if request.url.query: - url = f"{url}?{request.url.query}" - - headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} - content_type = request.headers.get("content-type") - if content_type: - headers["Content-Type"] = content_type - api_key = config.api_keys.get(endpoint) - if api_key is not None: - headers["Authorization"] = "Bearer " + api_key - - body = await request.body() - client: aiohttp.ClientSession = get_session(endpoint) - try: - resp = await client.request(request.method, url, data=body or None, headers=headers) - except Exception as e: - raise HTTPException(status_code=502, detail=f"Upstream request to {url} failed: {e}") - - async def _iter(): - try: - async for chunk in resp.content.iter_any(): - yield chunk - finally: - resp.release() - - return StreamingResponse( - _iter(), - status_code=resp.status, - media_type=resp.headers.get("Content-Type"), - ) diff --git a/api/responses.py b/api/responses.py deleted file mode 100644 index 5627c98..0000000 --- a/api/responses.py +++ /dev/null @@ -1,420 +0,0 @@ -"""OpenAI **Responses API** routes (``/v1/responses`` and its retrieve / delete / -cancel companions). - -The router speaks Chat Completions to its backends, so this layer: - - * **native** (external OpenAI): forwards via ``oclient.responses.create`` and - streams the SDK's typed events straight back, rewriting the response ``id`` to - a router-owned ``resp_`` id so chaining stays router-managed. - * **translated** (Ollama / llama-server): converts the request to chat, reuses - the resilient ``create_chat_with_retries`` ladder, and re-emits the result as - Responses typed SSE events (``requests/responses.py``). - -State (``store`` / ``previous_response_id``) and background-task status live in the -router's SQLite DB (``db.py``); the router mints and owns every response id. -""" -import asyncio -import secrets -import time - -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 db import get_db -from fingerprint import _conversation_fingerprint -from state import token_queue, default_headers -from backends.normalize import is_ext_openai_endpoint -from backends.sessions import _make_openai_client -from routing import choose_endpoint, decrement_usage -from api.openai import create_chat_with_retries -from requests.responses import ( - ChatToResponsesStream, - build_response_object, - chat_message_to_output_items, - messages_to_responses_input, - responses_input_to_messages, - responses_object_to_sse, - tools_responses_to_chat, - usage_chat_to_responses, -) - -router = APIRouter() - -# In-memory handles for background tasks so /cancel can reach a running task in -# this worker. Cross-worker cancel falls back to marking the DB row cancelled. -_background_tasks: dict[str, asyncio.Task] = {} - - -# --------------------------------------------------------------------------- -# small helpers -# --------------------------------------------------------------------------- -def _usage_tokens(usage): - """Return ``(prompt, completion)`` tokens from a chat- or responses-shaped usage.""" - if not usage: - return 0, 0 - if "input_tokens" in usage: - return usage.get("input_tokens", 0) or 0, usage.get("output_tokens", 0) or 0 - return usage.get("prompt_tokens", 0) or 0, usage.get("completion_tokens", 0) or 0 - - -def _text_format_to_response_format(text): - """Map Responses ``text.format`` → Chat Completions ``response_format`` (best effort).""" - if not isinstance(text, dict): - return None - fmt = text.get("format") - if not isinstance(fmt, dict): - return None - ftype = fmt.get("type") - if ftype == "json_object": - return {"type": "json_object"} - if ftype == "json_schema": - return {"type": "json_schema", "json_schema": { - k: fmt[k] for k in ("name", "schema", "strict", "description") if k in fmt - }} - return None - - -def _native_usage_from_response(data): - return data.get("usage") - - -async def _resolve_history_messages(previous_response_id): - """Rebuild prior-turn chat messages from the stored response chain.""" - if not previous_response_id: - return [] - db = get_db() - chain = await db.get_response_chain(previous_response_id) - messages = [] - for turn in chain: - # Each turn stored the chat messages that produced it + its output items. - for m in turn.get("input_messages") or []: - messages.append(m) - for item in turn.get("output_items") or []: - if item.get("type") == "message": - text = "".join( - p.get("text", "") for p in item.get("content") or [] - if p.get("type") == "output_text" - ) - if text: - messages.append({"role": "assistant", "content": text}) - elif item.get("type") == "function_call": - messages.append({ - "role": "assistant", "content": None, - "tool_calls": [{"id": item.get("call_id"), "type": "function", - "function": {"name": item.get("name"), - "arguments": item.get("arguments", "")}}], - }) - return messages - - -class _NativeStream: - """Re-emit an SDK Responses event stream, rewriting the response id and - capturing the final output/usage for storage.""" - - def __init__(self, response_id): - self.response_id = response_id - self.output_items = [] - self.usage = None - - async def events(self, sdk_gen): - async for event in sdk_gen: - data = event.model_dump() if hasattr(event, "model_dump") else event - etype = data.get("type", "") - resp = data.get("response") - if isinstance(resp, dict) and resp.get("id"): - resp["id"] = self.response_id - if etype in ("response.completed", "response.incomplete", "response.failed") \ - and isinstance(resp, dict): - self.output_items = resp.get("output", []) or [] - self.usage = resp.get("usage") - yield f"event: {etype}\ndata: {orjson.dumps(data).decode('utf-8')}\n\n".encode("utf-8") - - -# --------------------------------------------------------------------------- -# backend execution (non-streaming, used by background + non-stream sync) -# --------------------------------------------------------------------------- -async def _run_to_completion(*, native, oclient, endpoint, model, tracking_model, - send_params, native_params): - """Drive the backend to completion (no client streaming). - - Returns ``(output_items, usage)`` where usage is responses-shaped. The caller - owns the usage reservation and must release it (this function and - ``create_chat_with_retries`` never decrement).""" - if native: - resp_obj = await oclient.responses.create(stream=False, **native_params) - data = resp_obj.model_dump() - return data.get("output", []) or [], data.get("usage") - async_gen = await create_chat_with_retries(oclient, {**send_params, "stream": False}, - endpoint, model, tracking_model) - message = async_gen.choices[0].message.model_dump() if async_gen.choices else {} - output_items = chat_message_to_output_items(message) - usage = usage_chat_to_responses( - async_gen.usage.model_dump() if async_gen.usage is not None else None - ) - return output_items, usage - - -# --------------------------------------------------------------------------- -# POST /v1/responses -# --------------------------------------------------------------------------- -@router.post("/v1/responses") -async def openai_responses_proxy(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") - input_data = payload.get("input") - instructions = payload.get("instructions") - stream = bool(payload.get("stream")) - store = payload.get("store", True) - background = bool(payload.get("background")) - previous_response_id = payload.get("previous_response_id") - tools = payload.get("tools") - metadata = payload.get("metadata") or {} - _cache_enabled = payload.get("nomyo", {}).get("cache", False) - - if not model: - raise HTTPException(status_code=400, detail="Missing required field 'model'") - if input_data is None: - raise HTTPException(status_code=400, detail="Missing required field 'input'") - if background and not store: - raise HTTPException(status_code=400, detail="background mode requires store=true") - - if ":latest" in model: - model = model.split(":latest")[0] - - # Resolve conversation: prior turns (from store) + this turn's input. - history = await _resolve_history_messages(previous_response_id) - messages = history + responses_input_to_messages(input_data, instructions) - - response_id = f"resp_{secrets.token_hex(24)}" - created_at = int(time.time()) - - # Cache lookup (foreground only) — before endpoint selection. - _cache = get_llm_cache() - if _cache is not None and _cache_enabled and not background: - cached = await _cache.get_chat("openai_responses", model, messages) - if cached is not None: - resp_obj = orjson.loads(cached) - resp_obj["id"] = response_id - if stream: - async def _served_cached(): - yield responses_object_to_sse(resp_obj) - return StreamingResponse(_served_cached(), media_type="text/event-stream") - return JSONResponse(content=resp_obj) - - # Endpoint selection (reserves a slot — must be released exactly once). The - # release is owned by the per-branch finally (_bg_run / _stream / the - # non-streaming try) once we hand off; any failure during client/param - # construction (including CancelledError on client disconnect) must release - # it here or the usage counter leaks. - _affinity_key = _conversation_fingerprint(model, messages, None) - endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key) - try: - oclient = _make_openai_client(endpoint, default_headers=default_headers, - api_key=config.api_keys.get(endpoint, "no-key")) - native = is_ext_openai_endpoint(endpoint) - - # Build backend params for both shapes. - send_params = {"messages": messages, "model": model} - _opt = { - "temperature": payload.get("temperature"), - "top_p": payload.get("top_p"), - "max_tokens": payload.get("max_output_tokens"), - "tools": tools_responses_to_chat(tools), - "tool_choice": payload.get("tool_choice"), - "response_format": _text_format_to_response_format(payload.get("text")), - } - send_params.update({k: v for k, v in _opt.items() if v is not None}) - - native_instructions, native_input = messages_to_responses_input(messages) - native_params = {"model": model, "input": native_input, "store": False} - _nopt = { - "instructions": native_instructions, - "temperature": payload.get("temperature"), - "top_p": payload.get("top_p"), - "max_output_tokens": payload.get("max_output_tokens"), - "tools": tools, - "tool_choice": payload.get("tool_choice"), - "text": payload.get("text"), - "reasoning": payload.get("reasoning"), - } - native_params.update({k: v for k, v in _nopt.items() if v is not None}) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise - - async def _persist(status, output_items=None, usage=None, error=None, insert=False): - if not store: - return - db = get_db() - if insert: - await db.store_response( - response_id, previous_response_id=previous_response_id, model=model, - status=status, created_at=created_at, input_messages=messages, - output_items=output_items, usage=usage, instructions=instructions, error=error) - else: - await db.update_response_status(response_id, status, output_items=output_items, - usage=usage, error=error) - - async def _track(usage): - prompt_tok, comp_tok = _usage_tokens(usage) - if prompt_tok or comp_tok: - await token_queue.put((endpoint, tracking_model, prompt_tok, comp_tok)) - - async def _cache_store(output_items, usage): - if _cache is None or not _cache_enabled or not output_items: - return - obj = build_response_object(response_id=response_id, model=model, - output_items=output_items, usage=usage, - created_at=created_at, - previous_response_id=previous_response_id, - instructions=instructions, metadata=metadata) - try: - await _cache.set_chat("openai_responses", model, messages, orjson.dumps(obj)) - except Exception as _ce: - print(f"[cache] set_chat (openai_responses) failed: {_ce}") - - # ---- background: run detached, return queued immediately -------------- - if background: - # Once the task is created, _bg_run's finally owns the release. Guard the - # pre-task setup so a failure there (queued persist, task creation, or a - # client disconnect) still releases the reservation. - try: - await _persist("queued", insert=True) - - async def _bg_run(): - try: - await get_db().update_response_status(response_id, "in_progress") - output_items, usage = await _run_to_completion( - native=native, oclient=oclient, endpoint=endpoint, model=model, - tracking_model=tracking_model, send_params=send_params, - native_params=native_params) - await _track(usage) - await _persist("completed", output_items=output_items, usage=usage) - await _cache_store(output_items, usage) - except asyncio.CancelledError: - await get_db().update_response_status(response_id, "cancelled") - raise - except Exception as e: - await get_db().update_response_status( - response_id, "failed", - error={"message": str(e)[:500], "type": type(e).__name__}) - finally: - await decrement_usage(endpoint, tracking_model) - _background_tasks.pop(response_id, None) - - task = asyncio.create_task(_bg_run()) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise - _background_tasks[response_id] = task - queued = build_response_object(response_id=response_id, model=model, output_items=[], - status="queued", created_at=created_at, - previous_response_id=previous_response_id, - instructions=instructions, metadata=metadata) - return JSONResponse(content=queued, status_code=200) - - # ---- streaming sync ---------------------------------------------------- - if stream: - # _stream's finally owns the release once iteration starts. Establishing - # the source can fail (or be cancelled) before that — release here, since - # create_chat_with_retries no longer self-decrements. - try: - if native: - source = await oclient.responses.create(stream=True, **native_params) - translator = _NativeStream(response_id) - else: - source = await create_chat_with_retries( - oclient, {**send_params, "stream": True, - "stream_options": {"include_usage": True}}, - endpoint, model, tracking_model) - translator = ChatToResponsesStream( - response_id, model, created_at=created_at, - previous_response_id=previous_response_id, instructions=instructions, - metadata=metadata) - except BaseException: - await decrement_usage(endpoint, tracking_model) - raise - - async def _stream(): - await _persist("in_progress", insert=True) - try: - async for sse in translator.events(source): - yield sse - await _track(translator.usage) - await _persist("completed", output_items=translator.output_items, - usage=translator.usage) - await _cache_store(translator.output_items, translator.usage) - finally: - await decrement_usage(endpoint, tracking_model) - - return StreamingResponse(_stream(), media_type="text/event-stream") - - # ---- non-streaming sync ------------------------------------------------ - try: - output_items, usage = await _run_to_completion( - native=native, oclient=oclient, endpoint=endpoint, model=model, - tracking_model=tracking_model, send_params=send_params, - native_params=native_params) - await _track(usage) - await _persist("completed", output_items=output_items, usage=usage, insert=True) - await _cache_store(output_items, usage) - finally: - await decrement_usage(endpoint, tracking_model) - - resp_obj = build_response_object( - response_id=response_id, model=model, output_items=output_items, usage=usage, - created_at=created_at, previous_response_id=previous_response_id, - instructions=instructions, metadata=metadata) - return JSONResponse(content=resp_obj) - - -# --------------------------------------------------------------------------- -# GET / DELETE / cancel -# --------------------------------------------------------------------------- -def _stored_to_response_object(row): - return build_response_object( - response_id=row["response_id"], model=row.get("model"), - output_items=row.get("output_items") or [], usage=row.get("usage"), - status=row.get("status") or "completed", created_at=row.get("created_at"), - previous_response_id=row.get("previous_response_id"), - instructions=row.get("instructions"), error=row.get("error")) - - -@router.get("/v1/responses/{response_id}") -async def get_response(response_id: str): - row = await get_db().get_response(response_id) - if row is None: - raise HTTPException(status_code=404, detail=f"Response '{response_id}' not found") - return JSONResponse(content=_stored_to_response_object(row)) - - -@router.delete("/v1/responses/{response_id}") -async def delete_response(response_id: str): - deleted = await get_db().delete_response(response_id) - if not deleted: - raise HTTPException(status_code=404, detail=f"Response '{response_id}' not found") - return JSONResponse(content={"id": response_id, "object": "response.deleted", "deleted": True}) - - -@router.post("/v1/responses/{response_id}/cancel") -async def cancel_response(response_id: str): - row = await get_db().get_response(response_id) - if row is None: - raise HTTPException(status_code=404, detail=f"Response '{response_id}' not found") - # Cancel the running task if it lives in this worker; otherwise just mark the - # DB row so a polling client sees a terminal state (cross-worker limitation). - task = _background_tasks.get(response_id) - if task is not None and not task.done(): - task.cancel() - elif row.get("status") in ("queued", "in_progress"): - await get_db().update_response_status(response_id, "cancelled") - row = await get_db().get_response(response_id) - return JSONResponse(content=_stored_to_response_object(row)) diff --git a/backends/control.py b/backends/control.py deleted file mode 100644 index fda5fe3..0000000 --- a/backends/control.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Backend control operations (model unload). - -llama-server and llama-swap evict a resident model through different routes: - * llama-server → ``POST {base}/models/unload`` with body ``{"model": id}`` - * llama-swap → ``POST {base}/api/models/unload/{id}`` (path parameter) - -``unload_model`` dispatches on the configured backend type so callers don't -have to know which one they are talking to. Both routes live at the endpoint -root, so any ``/v1`` suffix is stripped first. -""" -from typing import Optional - -import aiohttp - -from config import get_config -from state import default_headers -from backends.sessions import get_probe_session -from backends.normalize import is_llama_swap -from backends.health import _format_connection_issue - - -async def unload_model(endpoint: str, model_id: str) -> bool: - """Ask ``endpoint`` to unload ``model_id``. Returns True on a 2xx response. - - ``model_id`` must be the backend's native model identifier (the raw HF id - for llama-server / llama-swap), not the router-normalized display name. - """ - cfg = get_config() - base_url = endpoint.rstrip("/").removesuffix("/v1") - headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} - api_key: Optional[str] = cfg.api_keys.get(endpoint) - if api_key is not None: - headers["Authorization"] = "Bearer " + api_key - - if is_llama_swap(endpoint): - url = f"{base_url}/api/models/unload/{model_id}" - json_body = None - else: - url = f"{base_url}/models/unload" - json_body = {"model": model_id} - - client: aiohttp.ClientSession = get_probe_session(endpoint) - try: - async with client.post(url, json=json_body, headers=headers) as resp: - ok = resp.status < 400 - print(f"[unload_model] {model_id} on {endpoint}: {resp.status}") - return ok - except Exception as e: - print(f"[unload_model] {_format_connection_issue(url, e)}") - return False diff --git a/backends/normalize.py b/backends/normalize.py index eef57cc..6603f9d 100644 --- a/backends/normalize.py +++ b/backends/normalize.py @@ -50,56 +50,27 @@ def dedupe_on_keys(dicts, key_fields): return out -def is_llama_swap(endpoint: str) -> bool: - """True if the endpoint is a configured llama-swap front.""" - return endpoint in get_config().llama_swap_endpoints - - -def is_llama_server(endpoint: str) -> bool: - """True for a llama.cpp llama-server OR a llama-swap front. - - Both speak the same OpenAI-compatible surface, so the router treats them - identically everywhere except loaded-model detection and model unload. - """ - cfg = get_config() - return endpoint in cfg.llama_server_endpoints or endpoint in cfg.llama_swap_endpoints - - -def llama_endpoints(cfg) -> list: - """Combined, de-duplicated llama-server + llama-swap endpoints (order preserved).""" - return list(dict.fromkeys([*cfg.llama_server_endpoints, *cfg.llama_swap_endpoints])) - - -def is_anthropic_endpoint(endpoint: str) -> bool: - """True if the endpoint is a configured native Anthropic Messages-API backend. - - These speak the Anthropic wire format (``x-api-key`` / ``anthropic-version`` - headers, ``/v1/messages``), not OpenAI Chat Completions, so requests routed to - them are forwarded verbatim rather than translated. - """ - return endpoint in get_config().anthropic_endpoints - - def is_ext_openai_endpoint(endpoint: str) -> bool: """ - Determine if an endpoint is an external OpenAI-compatible endpoint (not Ollama, llama-server or llama-swap). + Determine if an endpoint is an external OpenAI-compatible endpoint (not Ollama or llama-server). Returns True for: - External services like OpenAI.com, Groq, etc. Returns False for: - Ollama endpoints (without /v1, or with /v1 but default port 11434) - - llama-server / llama-swap endpoints (explicitly configured) + - llama-server endpoints (explicitly configured in llama_server_endpoints) """ - # Check if it's a llama-server / llama-swap endpoint (has /v1 and is in a configured list) - if is_llama_server(endpoint): + cfg = get_config() + # Check if it's a llama-server endpoint (has /v1 and is in the configured list) + if endpoint in cfg.llama_server_endpoints: return False if "/v1" not in endpoint: return False base_endpoint = endpoint.replace('/v1', '') - if base_endpoint in get_config().endpoints: + if base_endpoint in cfg.endpoints: return False # It's Ollama's /v1 # Check for default Ollama port @@ -112,9 +83,9 @@ def is_ext_openai_endpoint(endpoint: str) -> bool: def is_openai_compatible(endpoint: str) -> bool: """ Return True if the endpoint speaks the OpenAI API (not native Ollama). - This includes external OpenAI endpoints AND llama-server / llama-swap endpoints. + This includes external OpenAI endpoints AND llama-server endpoints. """ - return "/v1" in endpoint or is_llama_server(endpoint) + return "/v1" in endpoint or endpoint in get_config().llama_server_endpoints def get_tracking_model(endpoint: str, model: str) -> str: @@ -127,12 +98,12 @@ def get_tracking_model(endpoint: str, model: str) -> str: This ensures consistent model naming across all routes for usage tracking. """ - # External OpenAI / native Anthropic endpoints are not shown in PS, keep as-is - if is_ext_openai_endpoint(endpoint) or is_anthropic_endpoint(endpoint): + # External OpenAI endpoints are not shown in PS, keep as-is + if is_ext_openai_endpoint(endpoint): return model - # llama-server / llama-swap endpoints use normalized names in PS - if is_llama_server(endpoint): + # llama-server endpoints use normalized names in PS + if endpoint in get_config().llama_server_endpoints: return _normalize_llama_model_name(model) # Ollama endpoints: append ":latest" if no version suffix diff --git a/backends/probe.py b/backends/probe.py index 0d10da8..3ce089f 100644 --- a/backends/probe.py +++ b/backends/probe.py @@ -46,32 +46,7 @@ from backends.health import ( _format_connection_issue, _is_llama_model_loaded, ) -from backends.normalize import ( - is_ext_openai_endpoint, - is_openai_compatible, - is_llama_server, - is_llama_swap, - is_anthropic_endpoint, -) - -# Anthropic Messages API version pinned on every native probe/proxy request. -ANTHROPIC_VERSION = "2023-06-01" - - -def _auth_headers(endpoint: str, api_key: Optional[str]) -> dict: - """Build outbound auth headers for a probe, honoring the backend's scheme. - - Native Anthropic endpoints authenticate with ``x-api-key`` + a pinned - ``anthropic-version`` header; everything else uses OpenAI-style Bearer auth. - """ - headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} - if is_anthropic_endpoint(endpoint): - if api_key is not None: - headers["x-api-key"] = api_key - headers["anthropic-version"] = ANTHROPIC_VERSION - elif api_key is not None: - headers["Authorization"] = "Bearer " + api_key - return headers +from backends.normalize import is_ext_openai_endpoint, is_openai_compatible class fetch: @@ -81,16 +56,15 @@ class fetch: This is called by available_models() after checking caches and in-flight requests. """ cfg = get_config() - headers = _auth_headers(endpoint, api_key) + headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} + if api_key is not None: + headers["Authorization"] = "Bearer " + api_key ep_base = endpoint.rstrip("/") - if is_anthropic_endpoint(endpoint): + if endpoint in cfg.llama_server_endpoints and "/v1" not in endpoint: endpoint_url = f"{ep_base}/v1/models" key = "data" - elif is_llama_server(endpoint) and "/v1" not in endpoint: - endpoint_url = f"{ep_base}/v1/models" - key = "data" - elif "/v1" in endpoint or is_llama_server(endpoint): + elif "/v1" in endpoint or endpoint in cfg.llama_server_endpoints: endpoint_url = f"{ep_base}/models" key = "data" else: @@ -220,38 +194,6 @@ class fetch: client: aiohttp.ClientSession = get_probe_session(endpoint) cfg = get_config() - # llama-swap: loaded/running workers are reported at /running (state == "ready"), - # NOT via a status field on /v1/models (which it omits). /running is a root route, - # so strip any /v1 suffix from the configured endpoint. - if is_llama_swap(endpoint): - base_url = endpoint.rstrip("/").removesuffix("/v1") - headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} - api_key = cfg.api_keys.get(endpoint) - if api_key is not None: - headers["Authorization"] = "Bearer " + api_key - try: - async with client.get(f"{base_url}/running", headers=headers) as resp: - await _ensure_success(resp) - data = await resp.json() - - models = { - item.get("model") - for item in data.get("running", []) - if item.get("model") and item.get("state") == "ready" - } - - async with _loaded_models_cache_lock: - _loaded_models_cache[endpoint] = (models, time.time()) - async with _loaded_error_cache_lock: - _loaded_error_cache.pop(endpoint, None) - return models - except Exception as e: - message = _format_connection_issue(f"{base_url}/running", e) - print(f"[fetch.loaded_models] {message}") - async with _loaded_error_cache_lock: - _loaded_error_cache[endpoint] = time.time() - return set() - # Check if this is a llama-server endpoint if endpoint in cfg.llama_server_endpoints: # Query /v1/models for llama-server. Send the configured key as a @@ -346,18 +288,8 @@ class fetch: Uses request coalescing to prevent cache stampede and stale-while-revalidate to serve requests immediately even when cache is stale (refreshing in background). """ - if is_ext_openai_endpoint(endpoint) or is_anthropic_endpoint(endpoint): - # External OpenAI-compatible backends (vLLM, OpenAI, Groq, …) and - # native Anthropic endpoints keep every advertised model permanently - # resident — there is no /api/ps-style "loaded" subset to probe. - # Report the advertised set as the loaded set so choose_endpoint's - # "loaded & free" preference - # treats them on par with Ollama/llama-server backends that have the - # model hot, instead of relegating them to the free-slot fallback and - # never using them when an Ollama box advertises the same model. - # See issue #128. (llama-server / llama-swap are NOT ext-openai — they - # keep their real resident-subset detection below.) - return await fetch.available_models(endpoint, get_config().api_keys.get(endpoint)) + if is_ext_openai_endpoint(endpoint): + return set() # Check loaded models cache with lock protection async with _loaded_models_cache_lock: @@ -423,7 +355,9 @@ class fetch: if _is_fresh(_available_error_cache[endpoint], 300): return [] - headers = _auth_headers(endpoint, api_key) + headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} + if api_key is not None: + headers["Authorization"] = "Bearer " + api_key request_url = f"{endpoint.rstrip('/')}/{route.lstrip('/')}" client: aiohttp.ClientSession = get_probe_session(endpoint) @@ -459,7 +393,9 @@ async def _raw_probe( (unlike `fetch.endpoint_details`, which returns [] on either). Returns `(ok, payload_or_error_message)`. """ - headers = _auth_headers(ep, api_key) + headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")} + if api_key is not None: + headers["Authorization"] = "Bearer " + api_key url = f"{ep.rstrip('/')}/{route.lstrip('/')}" req_kwargs = {} if timeout is not None: @@ -482,14 +418,6 @@ async def _endpoint_health(ep: str, *, timeout: Optional[float] = None) -> dict: path (issue #83) is reported as `error` rather than `ok`. OpenAI-compatible endpoints use a single `/models` probe. """ - if is_anthropic_endpoint(ep): - ok, payload = await _raw_probe( - ep, "/v1/models", get_config().api_keys.get(ep), timeout=timeout, - ) - if ok: - return {"status": "ok", "version": "latest"} - return {"status": "error", "detail": str(payload)} - if is_openai_compatible(ep): ok, payload = await _raw_probe( ep, "/models", get_config().api_keys.get(ep), timeout=timeout, diff --git a/compliance/__init__.py b/compliance/__init__.py new file mode 100644 index 0000000..5a03b05 --- /dev/null +++ b/compliance/__init__.py @@ -0,0 +1,22 @@ +"""aissurance — the aissurance.eu compliance plugin. + +An in-process module that periodically snapshots the Router's AI-infrastructure +landscape and sends structured, HMAC-signed evidence to aissurance.eu (the EU +AI Act compliance platform). It only reads infrastructure metadata the Router +already maintains for routing — never prompt/completion content. + +Public entry points: + * ``AissurancePlugin`` — owns the scheduler task; ``start()`` / ``stop()`` are + called from router.py's startup/shutdown events. +""" +__all__ = ["AissurancePlugin"] + + +def __getattr__(name): + # Lazy export so that ``config.py`` can import ``compliance.settings`` without + # eagerly loading the plugin chain (plugin → scheduler → collector → config), + # which would form an import cycle. router.py triggers this after startup. + if name == "AissurancePlugin": + from compliance.plugin import AissurancePlugin + return AissurancePlugin + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/compliance/buffer.py b/compliance/buffer.py new file mode 100644 index 0000000..20ac8be --- /dev/null +++ b/compliance/buffer.py @@ -0,0 +1,94 @@ +"""Persistent, encrypted retry buffer for offline periods. + +When the upstream is unreachable, signed payloads are written to ``buffer_dir`` +as individual ``{timestamp}-{router_id}.json.enc`` files, encrypted with a key +derived from the tenant API key (see compliance.crypto). On reconnection they +are replayed oldest-first; each is deleted only after a ``202 Accepted``. + +The buffer is bounded (``max_buffer_payloads``); the oldest file is dropped +when the limit is exceeded. +""" +import json +import re +from pathlib import Path +from typing import Optional + +from compliance import crypto +from compliance.settings import ComplianceSettings + +_SAFE = re.compile(r"[^A-Za-z0-9_.-]") + + +def _safe(s: str) -> str: + return _SAFE.sub("_", s) + + +class BufferStore: + def __init__(self, cfg: ComplianceSettings): + self._cfg = cfg + self._dir = Path(cfg.buffer_dir) + + def _ensure_dir(self) -> None: + self._dir.mkdir(parents=True, exist_ok=True) + + def _files(self) -> list[Path]: + if not self._dir.exists(): + return [] + # Filenames are ``{epoch_ms}-{router_id}.json.enc``; lexical sort on the + # zero-padded epoch prefix is chronological. + return sorted(self._dir.glob("*.json.enc")) + + def count(self) -> int: + return len(self._files()) + + def save(self, payload: dict, timestamp_ms: int) -> Optional[Path]: + """Encrypt and persist ``payload``. Returns the path, or None on failure.""" + if not self._cfg.api_key: + return None + self._ensure_dir() + self._enforce_limit(reserve=1) + name = f"{timestamp_ms:020d}-{_safe(self._cfg.router_id)}.json.enc" + path = self._dir / name + token = crypto.encrypt(json.dumps(payload).encode("utf-8"), self._cfg.api_key) + path.write_bytes(token) + return path + + def load(self, path: Path) -> Optional[dict]: + """Decrypt and parse a buffer file. Returns None (and deletes) on corruption.""" + if not self._cfg.api_key: + return None + try: + token = path.read_bytes() + except OSError: + return None + plaintext = crypto.decrypt(token, self._cfg.api_key) + if plaintext is None: + # Wrong key or corrupt file — unrecoverable, drop it. + self.delete(path) + return None + try: + return json.loads(plaintext) + except json.JSONDecodeError: + self.delete(path) + return None + + def pending(self) -> list[Path]: + """Buffered files, oldest first.""" + return self._files() + + def delete(self, path: Path) -> None: + try: + path.unlink() + except OSError: + pass + + def clear(self) -> None: + for p in self._files(): + self.delete(p) + + def _enforce_limit(self, reserve: int = 0) -> None: + """Drop oldest files until under (max - reserve).""" + limit = max(self._cfg.max_buffer_payloads - reserve, 0) + files = self._files() + while len(files) > limit: + self.delete(files.pop(0)) diff --git a/compliance/collector.py b/compliance/collector.py new file mode 100644 index 0000000..305fd0d --- /dev/null +++ b/compliance/collector.py @@ -0,0 +1,219 @@ +"""Snapshot collector — turns live Router state into the discovery data model. + +Reads only structures the Router already maintains for routing: + * ``state.usage_counts`` — live per-endpoint/model active connections + * ``config.endpoints`` / ``llama_server_endpoints`` / ``api_keys`` / limits + * ``backends.probe`` model discovery + endpoint health + * ``cache.LLMCache.stats()`` — global cache hit rate + * the token time-series DB — 24h token totals and per-model ``last_used`` + +Fields the Router does not currently instrument (per-model latency percentiles, +request/error counts, per-model cache hit rate, token averages) are emitted as +``None`` rather than fabricated. See the plugin docs / spec §4 for the rationale. +""" +import time +from datetime import datetime, timezone +from typing import Optional + +from config import Config +from state import usage_counts +from cache import get_llm_cache +from db import get_db +from routing import get_max_connections +from backends.probe import fetch, _endpoint_health +from backends.normalize import ( + is_ext_openai_endpoint, + _normalize_llama_model_name, + _extract_llama_quant, +) + +_DAY_SECONDS = 86400 + + +def _iso_now() -> str: + return datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _iso_from_unix(ts: Optional[int]) -> Optional[str]: + if not ts: + return None + return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _endpoint_type(cfg: Config, ep: str) -> str: + """Best-effort classification from config membership + URL heuristics.""" + low = ep.lower() + if "googleapis" in low or "gemini" in low: + return "gemini" + if "anthropic" in low: + return "anthropic" + if "cohere" in low: + return "cohere" + if "cerebras" in low: + return "cerebras" + if "inceptionlabs" in low or "inception" in low: + return "inception labs" + if ep in cfg.llama_server_endpoints: + # llama_server_endpoints covers both llama.cpp and vLLM (both OpenAI- + # compatible). vLLM commonly serves on :8000; treat the rest as llama.cpp. + return "vllm" if ":8000" in ep else "llama_cpp" + if is_ext_openai_endpoint(ep): + return "openai" + return "ollama" + + +def _auth_method(cfg: Config, ep: str) -> str: + return "bearer" if cfg.api_keys.get(ep) else "none" + + +def _status_label(health: dict) -> str: + # _endpoint_health only distinguishes ok/error; map onto the spec's enum. + return "healthy" if health.get("status") == "ok" else "unhealthy" + + +async def _collect_endpoints(cfg: Config) -> tuple[list[dict], int]: + """Build the endpoint registry and the global active-connection count.""" + seen: list[str] = [] + for ep in list(cfg.endpoints) + list(cfg.llama_server_endpoints): + if ep not in seen: + seen.append(ep) + + endpoints: list[dict] = [] + total_active = 0 + for ep in seen: + concurrent = sum(usage_counts.get(ep, {}).values()) + total_active += concurrent + try: + health = await _endpoint_health(ep, timeout=5) + except Exception: + health = {"status": "error"} + endpoints.append({ + "url": ep, + "type": _endpoint_type(cfg, ep), + "status": _status_label(health), + "concurrent_connections": concurrent, + "max_concurrent_connections": get_max_connections(ep), + "tls_enabled": ep.lower().startswith("https://"), + "auth_method": _auth_method(cfg, ep), + }) + return endpoints, total_active + + +def _ollama_detail_index(details: list[dict]) -> dict: + """Index Ollama /api/tags entries by model name for size/quant lookup.""" + index: dict[str, dict] = {} + for item in details or []: + name = item.get("name") or item.get("model") + if name: + index[name] = item + return index + + +async def _collect_models(cfg: Config, last_used: dict) -> list[dict]: + models: list[dict] = [] + for ep in list(cfg.endpoints) + list(cfg.llama_server_endpoints): + is_llama = ep in cfg.llama_server_endpoints + try: + available = await fetch.available_models(ep, cfg.api_keys.get(ep)) + except Exception: + available = set() + if not available: + continue + + # Ollama exposes weight size + quantization via /api/tags details. + detail_index: dict = {} + if not is_llama and not is_ext_openai_endpoint(ep): + try: + details = await fetch.endpoint_details(ep, "/api/tags", "models") + detail_index = _ollama_detail_index(details) + except Exception: + detail_index = {} + + for raw_name in sorted(available): + if is_llama: + name = _normalize_llama_model_name(raw_name) + quant = _extract_llama_quant(raw_name) or "none" + version = raw_name.split(":", 1)[1] if ":" in raw_name else "latest" + size_gb = None + else: + name = raw_name + version = raw_name.split(":", 1)[1] if ":" in raw_name else "latest" + meta = detail_index.get(raw_name, {}) + det = meta.get("details", {}) if isinstance(meta, dict) else {} + quant = det.get("quantization_level") or "none" + size_bytes = meta.get("size") if isinstance(meta, dict) else None + size_gb = round(size_bytes / 1e9, 3) if size_bytes else None + + models.append({ + "name": name, + "version": version, + "quantization": quant, + "size_gb": size_gb, + "endpoint": ep, + "last_used": _iso_from_unix(last_used.get((ep, raw_name)) or last_used.get((ep, name))), + # Un-instrumented today — emitted as null rather than fabricated. + "request_count_24h": None, + "avg_latency_ms": None, + "p99_latency_ms": None, + "error_count_24h": None, + "cache_hit_rate": None, + "avg_input_tokens": None, + "avg_output_tokens": None, + }) + return models + + +async def _collect_telemetry(cfg: Config, active_connections: int, started_at: float) -> dict: + cache = get_llm_cache() + if cache is not None: + cstats = cache.stats() + total_cache_hits = cstats.get("hits", 0) + cache_hit_rate = cstats.get("hit_rate", 0.0) + else: + total_cache_hits = 0 + cache_hit_rate = 0.0 + + cutoff = int(time.time()) - _DAY_SECONDS + try: + tokens_in, tokens_out = await get_db().get_token_totals_since(cutoff) + except Exception: + tokens_in, tokens_out = 0, 0 + + uptime_hours = round((time.monotonic() - started_at) / 3600.0, 3) + + return { + # Request/error totals are un-instrumented today (null, not fabricated). + "total_requests_24h": None, + "total_cache_hits": total_cache_hits, + "cache_hit_rate": cache_hit_rate, + "error_rate": None, + # Router tracks tokens, not raw bytes; surfaced under the byte fields as + # the closest available signal (24h windowed token counts). + "total_bytes_in_24h": tokens_in, + "total_bytes_out_24h": tokens_out, + "active_connections": active_connections, + "uptime_hours": uptime_hours, + } + + +async def collect_snapshot(cfg: Config, started_at: float) -> dict: + """Assemble the full discovery snapshot body (without signature). + + ``started_at`` is a ``time.monotonic()`` reading taken at plugin start, + used to derive ``uptime_hours``. + """ + try: + last_used = await get_db().get_last_used_map() + except Exception: + last_used = {} + + endpoints, active_connections = await _collect_endpoints(cfg) + models = await _collect_models(cfg, last_used) + telemetry = await _collect_telemetry(cfg, active_connections, started_at) + + return { + "timestamp": _iso_now(), + "models": models, + "endpoints": endpoints, + "telemetry": telemetry, + } diff --git a/compliance/crypto.py b/compliance/crypto.py new file mode 100644 index 0000000..84cfb79 --- /dev/null +++ b/compliance/crypto.py @@ -0,0 +1,45 @@ +"""Buffer encryption for offline payloads. + +Buffered discovery payloads contain only infrastructure metadata (model names, +endpoint URLs, latency/error aggregates) — never prompt or completion content. +They are still encrypted at rest with a key derived from the tenant's +``AISSURANCE_KEY`` via HKDF-SHA256, so a leaked buffer file is useless without +the tenant secret. + +The on-disk format is Fernet (AES-128-CBC + HMAC-SHA256), which authenticates +the ciphertext and embeds a timestamp. +""" +import base64 +from typing import Optional + +from cryptography.fernet import Fernet, InvalidToken +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.hkdf import HKDF + +# Domain-separation label so this derived key can never collide with a key +# derived from the same secret for a different purpose. +_HKDF_INFO = b"aissurance-buffer-v1" + + +def _fernet(api_key: str) -> Fernet: + """Derive a Fernet instance from the tenant API key.""" + derived = HKDF( + algorithm=hashes.SHA256(), + length=32, + salt=None, + info=_HKDF_INFO, + ).derive(api_key.encode("utf-8")) + return Fernet(base64.urlsafe_b64encode(derived)) + + +def encrypt(plaintext: bytes, api_key: str) -> bytes: + """Encrypt ``plaintext`` for at-rest buffer storage.""" + return _fernet(api_key).encrypt(plaintext) + + +def decrypt(token: bytes, api_key: str) -> Optional[bytes]: + """Decrypt a buffer file. Returns None if the key is wrong or data is corrupt.""" + try: + return _fernet(api_key).decrypt(token) + except (InvalidToken, ValueError): + return None diff --git a/compliance/log.py b/compliance/log.py new file mode 100644 index 0000000..5a5fc35 --- /dev/null +++ b/compliance/log.py @@ -0,0 +1,15 @@ +"""Structured JSON logging to stderr (spec §12.2).""" +import json +import sys +from datetime import datetime, timezone + + +def log(event: str, level: str = "info", **fields) -> None: + record = { + "timestamp": datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + "level": level, + "module": "aissurance", + "event": event, + **fields, + } + print(json.dumps(record), file=sys.stderr, flush=True) diff --git a/compliance/payload.py b/compliance/payload.py new file mode 100644 index 0000000..267a468 --- /dev/null +++ b/compliance/payload.py @@ -0,0 +1,81 @@ +"""Payload assembly and HMAC-SHA256 signing. + +The signature covers every payload field *except* ``hmac_signature`` itself, +over a canonical ``json.dumps(..., sort_keys=True)`` serialization (spec §6.2), +so any tampering with the body invalidates it. +""" +import hashlib +import hmac +import json +from typing import Optional + +from compliance.settings import ComplianceSettings + +# Router release this plugin ships with. Sent as ``router_version`` and in the +# ``X-Router-Version`` header. +ROUTER_VERSION = "0.9" + + +def _signing_key(cfg: ComplianceSettings) -> str: + """HMAC key: the per-tenant override if set, else the API key (spec §6.2).""" + return cfg.hmac_secret or cfg.api_key or "" + + +def canonical_message(body: dict) -> str: + """Canonical serialization signed by both client and server.""" + fields = {k: v for k, v in body.items() if k != "hmac_signature"} + return json.dumps(fields, sort_keys=True, separators=(",", ":")) + + +def sign(body: dict, cfg: ComplianceSettings) -> str: + """Return the ``sha256=`` signature for ``body``.""" + digest = hmac.new( + _signing_key(cfg).encode("utf-8"), + canonical_message(body).encode("utf-8"), + hashlib.sha256, + ).hexdigest() + return f"sha256={digest}" + + +def build_payload(cfg: ComplianceSettings, snapshot: dict) -> dict: + """Combine config identity + collector snapshot into a signed payload.""" + body = { + "tenant_id": cfg.tenant_id, + "router_id": cfg.router_id, + "router_version": ROUTER_VERSION, + "timestamp": snapshot["timestamp"], + "models": snapshot["models"], + "endpoints": snapshot["endpoints"], + "telemetry": snapshot["telemetry"], + } + body["hmac_signature"] = sign(body, cfg) + return body + + +def build_config_report(cfg: ComplianceSettings) -> dict: + """The body posted to ``/config/report`` (spec §5.3).""" + return { + "polling_interval": cfg.polling_interval, + "batch_size": cfg.batch_size, + "server_url": cfg.server_url, + "router_version": ROUTER_VERSION, + } + + +def split_into_batches(payload: dict, batch_size: int, cfg: ComplianceSettings) -> list[dict]: + """Split a payload whose ``models`` list exceeds ``batch_size`` into several + signed payloads sharing the same timestamp/identity (used on HTTP 413). + + Each batch is re-signed because its ``models`` slice differs. + """ + models = payload.get("models", []) + if batch_size <= 0 or len(models) <= batch_size: + return [payload] + + batches: list[dict] = [] + for i in range(0, len(models), batch_size): + chunk = {k: v for k, v in payload.items() if k != "hmac_signature"} + chunk["models"] = models[i : i + batch_size] + chunk["hmac_signature"] = sign(chunk, cfg) + batches.append(chunk) + return batches diff --git a/compliance/plugin.py b/compliance/plugin.py new file mode 100644 index 0000000..2ca0bf9 --- /dev/null +++ b/compliance/plugin.py @@ -0,0 +1,129 @@ +"""Plugin lifecycle: owns the periodic tasks and shared transport/buffer state. + +Instantiated and driven by router.py's startup/shutdown events. The plugin is +non-critical to the proxy: any failure here is logged and contained, never +propagated into request routing. +""" +import asyncio +import time +from typing import Optional + +from compliance import scheduler +from compliance.buffer import BufferStore +from compliance.log import log +from compliance.settings import ComplianceSettings +from compliance.transport import Transport + + +class AissurancePlugin: + def __init__(self, settings: ComplianceSettings): + self._settings = settings + self.started_at = time.monotonic() + self.stopping = False + self.sent_total = 0 + + self.transport: Optional[Transport] = None + self.buffer = BufferStore(settings) + + self._stop_event = asyncio.Event() + self._tasks: list[asyncio.Task] = [] + self._auth_failed = False + + # -- accessors used by the scheduler ------------------------------------- + @property + def cfg(self) -> ComplianceSettings: + return self._settings + + @property + def cfg_obj(self): + """Live Router Config, so the collector sees current endpoints.""" + from config import get_config + return get_config() + + async def sleep(self, seconds: float) -> None: + """Sleep up to ``seconds``, returning early if the plugin is stopping.""" + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + async def handle_auth_failure(self) -> None: + """401/403: stop retrying, clear the buffer, alert, and go dormant.""" + if self._auth_failed: + return + self._auth_failed = True + self.buffer.clear() + log("auth_failed", level="error", router_id=self._settings.router_id, + message="AISSURANCE_KEY rejected — compliance reporting halted") + self.stopping = True + self._stop_event.set() + + # -- lifecycle ----------------------------------------------------------- + async def start(self) -> None: + errors = self._settings.validation_errors() + if errors: + log("plugin_disabled", level="error", reasons=errors) + return + + self.transport = Transport(self._settings) + self.started_at = time.monotonic() + + if not self._settings.tenant_id: + log("tenant_id_missing", level="warning", + message="No AISSURANCE_TENANT_ID / compliance.tenant_id; payload tenant_id will be null") + + # Startup handshake: connectivity check + config report (best effort). + try: + await scheduler.run_health_cycle(self) + await scheduler.run_config_report(self) + except Exception as exc: # never let handshake failure block startup + log("handshake_failed", level="warning", error=str(exc)) + + self._tasks = [ + asyncio.create_task(self._poll_loop()), + asyncio.create_task(self._health_loop()), + ] + log("plugin_started", router_id=self._settings.router_id, + polling_interval=self._settings.polling_interval, + buffered=self.buffer.count()) + + async def _poll_loop(self) -> None: + # Emit an initial snapshot immediately so evidence appears without + # waiting a full interval, then settle into the configured cadence. + while not self.stopping: + try: + await scheduler.run_discovery_cycle(self) + except Exception as exc: # contain — plugin is non-critical + log("cycle_error", level="error", error=str(exc)) + if self.stopping: + break + await self.sleep(self._settings.polling_interval) + + async def _health_loop(self) -> None: + while not self.stopping: + await self.sleep(self._settings.health_interval) + if self.stopping: + break + try: + await scheduler.run_health_cycle(self) + except Exception as exc: + log("health_error", level="warning", error=str(exc)) + + async def stop(self) -> None: + self.stopping = True + self._stop_event.set() + for t in self._tasks: + t.cancel() + try: + await t + except (asyncio.CancelledError, Exception): + pass + self._tasks = [] + if self.transport is not None: + try: + await self.transport.aclose() + except Exception: + pass + self.transport = None + log("plugin_stopped", router_id=self._settings.router_id, + buffered=self.buffer.count(), successful=self.sent_total) diff --git a/compliance/scheduler.py b/compliance/scheduler.py new file mode 100644 index 0000000..a17ca35 --- /dev/null +++ b/compliance/scheduler.py @@ -0,0 +1,132 @@ +"""Snapshot lifecycle: collect → sign → send → buffer-on-failure, plus replay. + +These coroutines operate on an ``AissurancePlugin`` instance (duck-typed to +avoid an import cycle) and contain the per-cycle decision logic. The periodic +timing loops live in plugin.py. +""" +import asyncio +import time +from typing import TYPE_CHECKING + +from compliance import payload as payload_mod +from compliance.collector import collect_snapshot +from compliance.log import log +from compliance.transport import Outcome, SendResult + +if TYPE_CHECKING: # pragma: no cover + from compliance.plugin import AissurancePlugin + +# Snapshot collection budget (spec §9.2). Exceeding it skips the cycle. +_SNAPSHOT_TIMEOUT = 10.0 + + +async def _send_once(plugin: "AissurancePlugin", body: dict) -> SendResult: + result = await plugin.transport.send_payload(body) + if result.outcome is Outcome.TOO_LARGE: + # 413 — split by batch_size and send each chunk; worst chunk outcome wins. + batches = payload_mod.split_into_batches(body, plugin.cfg.batch_size, plugin.cfg) + if len(batches) > 1: + worst = SendResult(Outcome.OK, 202) + for chunk in batches: + r = await plugin.transport.send_payload(chunk) + if r.outcome is not Outcome.OK: + worst = r + return worst + return result + + +async def send_with_retry(plugin: "AissurancePlugin", body: dict) -> Outcome: + """Send one payload, retrying transient failures with exponential backoff. + + Returns the final outcome. The caller decides whether to buffer. + """ + attempts = plugin.cfg.max_retry_attempts + base = plugin.cfg.retry_backoff_base + for attempt in range(max(attempts, 1)): + result = await _send_once(plugin, body) + if result.outcome is Outcome.OK: + plugin.sent_total += 1 + log("payload_sent", tenant_id=plugin.cfg.tenant_id, router_id=plugin.cfg.router_id, + models_count=len(body.get("models", [])), response_code=result.status_code, + latency_ms=result.latency_ms) + return Outcome.OK + if result.outcome in (Outcome.AUTH_FAILED, Outcome.RATE_LIMITED): + # Non-retryable within this cycle. + return result.outcome + # RETRY (network/5xx): back off unless this was the last attempt. + if attempt < attempts - 1 and not plugin.stopping: + log("payload_retry", level="warning", attempt=attempt + 1, + response_code=result.status_code) + await plugin.sleep(min(base ** attempt, 60)) + return Outcome.RETRY + + +async def replay_buffer(plugin: "AissurancePlugin") -> None: + """Replay buffered payloads oldest-first. Stops at the first non-OK send.""" + for path in plugin.buffer.pending(): + if plugin.stopping: + return + body = plugin.buffer.load(path) + if body is None: + continue # corrupt file already dropped by load() + result = await plugin.transport.send_payload(body) + if result.outcome is Outcome.OK: + plugin.buffer.delete(path) + plugin.sent_total += 1 + log("payload_sent", router_id=plugin.cfg.router_id, response_code=result.status_code, + latency_ms=result.latency_ms, replayed=True) + elif result.outcome is Outcome.AUTH_FAILED: + await plugin.handle_auth_failure() + return + else: + # Still offline / rate limited — keep the rest buffered for next time. + return + + +async def run_discovery_cycle(plugin: "AissurancePlugin") -> None: + """One scheduled discovery snapshot.""" + try: + snapshot = await asyncio.wait_for( + collect_snapshot(plugin.cfg_obj, plugin.started_at), timeout=_SNAPSHOT_TIMEOUT + ) + except asyncio.TimeoutError: + log("snapshot_too_slow", level="warning", timeout_s=_SNAPSHOT_TIMEOUT) + return + except Exception as exc: # serialization / state read error — skip this cycle + log("snapshot_failed", level="error", error=str(exc)) + return + + body = payload_mod.build_payload(plugin.cfg, snapshot) + outcome = await send_with_retry(plugin, body) + + if outcome is Outcome.OK: + await replay_buffer(plugin) + elif outcome is Outcome.AUTH_FAILED: + await plugin.handle_auth_failure() + elif outcome is Outcome.RATE_LIMITED: + # Server already rate-limited us; drop this snapshot (next is as good). + log("payload_rate_limited", level="warning") + else: # RETRY exhausted — persist for later. + path = plugin.buffer.save(body, int(time.time() * 1000)) + log("payload_buffered", level="warning", buffered=plugin.buffer.count(), + path=str(path) if path else None) + + +async def run_config_report(plugin: "AissurancePlugin") -> None: + """POST current config to /config/report (handshake / on change).""" + result = await plugin.transport.send_config_report( + payload_mod.build_config_report(plugin.cfg) + ) + if result.outcome is Outcome.AUTH_FAILED: + await plugin.handle_auth_failure() + return + suggested = (result.body or {}).get("server_suggested_polling_interval") + log("config_reported", response_code=result.status_code, + server_suggested_polling_interval=suggested) + + +async def run_health_cycle(plugin: "AissurancePlugin") -> None: + """Periodic health/keepalive + config-sync probe.""" + result = await plugin.transport.send_health() + if result.outcome is Outcome.AUTH_FAILED: + await plugin.handle_auth_failure() diff --git a/compliance/settings.py b/compliance/settings.py new file mode 100644 index 0000000..812aee3 --- /dev/null +++ b/compliance/settings.py @@ -0,0 +1,98 @@ +"""Configuration model for the aissurance compliance plugin. + +Populated from the ``compliance`` block of the Router's ``config.yaml`` (see +config.Config). ``${VAR}`` references in the YAML are already expanded by +``Config._expand_env_refs`` before this model is constructed; on top of that, +the aissurance-specific environment variables (``AISSURANCE_KEY`` etc.) are +honoured as defaults so the plugin works with env-only configuration too. + +This module imports nothing from the Router, so config.py can import it +without creating a cycle. +""" +import os +import socket +from typing import Optional + +from pydantic import BaseModel, Field, field_validator + + +def _env_or_none(name: str) -> Optional[str]: + val = os.getenv(name) + return val if val else None + + +class ComplianceSettings(BaseModel): + """The ``compliance:`` config block.""" + + enabled: bool = False + + # Full URL of the discovery receive endpoint. Health and config-report URLs + # are derived from this by replacing the trailing ``/receive`` path. + server_url: str = "https://www.aissurance.eu/api/v1/discovery/receive" + + # Tenant-provisioned secret (AISSURANCE_KEY). Used both as the Bearer token + # and as the HMAC signing key unless ``hmac_secret`` overrides the latter. + api_key: Optional[str] = Field(default_factory=lambda: _env_or_none("AISSURANCE_KEY")) + # Optional per-tenant HMAC secret override (AISSURANCE_HMAC_SECRET). + hmac_secret: Optional[str] = Field(default_factory=lambda: _env_or_none("AISSURANCE_HMAC_SECRET")) + # Tenant UUID, if not embedded in the key (AISSURANCE_TENANT_ID). + tenant_id: Optional[str] = Field(default_factory=lambda: _env_or_none("AISSURANCE_TENANT_ID")) + + # Unique identifier for this Router instance. Multiple instances per tenant + # are allowed, so this defaults to the hostname rather than the tenant id. + router_id: str = Field(default_factory=socket.gethostname) + + polling_interval: int = 300 # seconds between discovery snapshots + health_interval: int = 3600 # seconds between health/config-sync calls + batch_size: int = 50 # max models per payload + max_retry_attempts: int = 10 # in-cycle send retries before buffering + retry_backoff_base: int = 2 # exponential backoff base (seconds) + + buffer_dir: str = "./compliance-buffer" # persistent encrypted retry storage + max_buffer_payloads: int = 100 # oldest dropped first when full + + # TLS verification. AISSURANCE_VERIFY_TLS=0 disables it (dev only). + verify_tls: bool = Field( + default_factory=lambda: os.getenv("AISSURANCE_VERIFY_TLS", "1") != "0" + ) + + @field_validator("api_key", "hmac_secret", "tenant_id", mode="before") + @classmethod + def _empty_to_none(cls, v): + # config.yaml ``${VAR}`` expansion yields "" for unset vars; treat as None + # so the env-var default_factory can take over. + if v == "": + return None + return v + + @property + def discovery_base_url(self) -> str: + """Base discovery URL, derived by stripping a trailing ``/receive``.""" + url = self.server_url.rstrip("/") + if url.endswith("/receive"): + url = url[: -len("/receive")] + return url + + @property + def receive_url(self) -> str: + return self.server_url + + @property + def health_url(self) -> str: + tid = self.tenant_id or "unknown" + return f"{self.discovery_base_url}/health/{tid}" + + @property + def config_report_url(self) -> str: + return f"{self.discovery_base_url}/config/report" + + def validation_errors(self) -> list[str]: + """Return human-readable reasons the plugin cannot run, or [] if OK.""" + errors: list[str] = [] + if not self.api_key: + errors.append("AISSURANCE_KEY / compliance.api_key is not set") + if not self.server_url.lower().startswith("https://"): + errors.append(f"server_url must be HTTPS (got {self.server_url!r})") + if self.polling_interval < 60: + errors.append("polling_interval must be >= 60 seconds") + return errors diff --git a/compliance/transport.py b/compliance/transport.py new file mode 100644 index 0000000..bb2d39a --- /dev/null +++ b/compliance/transport.py @@ -0,0 +1,89 @@ +"""Outbound HTTP transport to aissurance.eu. + +All communication is Router-initiated (outbound only). This module performs the +three POSTs in the protocol — discovery ``/receive``, ``/health/{tenant_id}``, +and ``/config/report`` — and classifies the HTTP response into an action the +scheduler can act on (spec §5, §12.1). It owns a single long-lived +``httpx.AsyncClient``. +""" +import enum +import time +from dataclasses import dataclass +from typing import Optional + +import httpx + +from compliance.payload import ROUTER_VERSION +from compliance.settings import ComplianceSettings + + +class Outcome(enum.Enum): + OK = "ok" # 202 — stored + AUTH_FAILED = "auth" # 401/403 — stop, alert, clear buffer + RATE_LIMITED = "rate" # 429 — back off + TOO_LARGE = "too_large" # 413 — split + retry + RETRY = "retry" # network error / 5xx — buffer + backoff + + +@dataclass +class SendResult: + outcome: Outcome + status_code: Optional[int] = None + latency_ms: int = 0 + body: Optional[dict] = None + + +def _classify(status: int) -> Outcome: + if status == 202 or status == 200: + return Outcome.OK + if status in (401, 403): + return Outcome.AUTH_FAILED + if status == 429: + return Outcome.RATE_LIMITED + if status == 413: + return Outcome.TOO_LARGE + return Outcome.RETRY + + +class Transport: + def __init__(self, cfg: ComplianceSettings): + self._cfg = cfg + self._client = httpx.AsyncClient( + timeout=httpx.Timeout(10.0, connect=5.0), + verify=cfg.verify_tls, + ) + + async def aclose(self) -> None: + await self._client.aclose() + + def _headers(self) -> dict: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {self._cfg.api_key}", + "X-Router-Version": ROUTER_VERSION, + "X-Router-Id": self._cfg.router_id, + } + + async def _post(self, url: str, body: dict) -> SendResult: + start = time.monotonic() + try: + resp = await self._client.post(url, json=body, headers=self._headers()) + except httpx.HTTPError: + return SendResult(Outcome.RETRY, None, int((time.monotonic() - start) * 1000)) + latency_ms = int((time.monotonic() - start) * 1000) + parsed: Optional[dict] = None + try: + parsed = resp.json() + except ValueError: + parsed = None + return SendResult(_classify(resp.status_code), resp.status_code, latency_ms, parsed) + + async def send_payload(self, payload: dict) -> SendResult: + return await self._post(self._cfg.receive_url, payload) + + async def send_health(self) -> SendResult: + # Health is a keepalive/config-sync probe; an empty body is sufficient. + return await self._post(self._cfg.health_url, {"router_id": self._cfg.router_id}) + + async def send_config_report(self, body: dict) -> SendResult: + return await self._post(self._cfg.config_report_url, body) diff --git a/config.py b/config.py index ea98ee5..7db9c11 100644 --- a/config.py +++ b/config.py @@ -13,6 +13,8 @@ import yaml from pydantic import Field from pydantic_settings import BaseSettings +from compliance.settings import ComplianceSettings + class Config(BaseSettings): # List of Ollama endpoints @@ -23,16 +25,6 @@ class Config(BaseSettings): ) # List of llama-server endpoints (OpenAI-compatible with /v1/models status info) llama_server_endpoints: List[str] = Field(default_factory=list) - # List of llama-swap endpoints (OpenAI-compatible front for multiple llama-server - # workers). Same surface as llama_server_endpoints, but loaded models are read from - # /running (not /v1/models status) and unload uses POST /api/models/unload/:model_id. - llama_swap_endpoints: List[str] = Field(default_factory=list) - # List of native Anthropic Messages-API endpoints (e.g. https://api.anthropic.com). - # Configure the base URL WITHOUT a /v1 suffix; the router appends /v1/models and - # /v1/messages itself. Requests routed here are forwarded verbatim (no Messages⇄Chat - # translation); the endpoint's api_keys entry is sent as the x-api-key header. Their - # advertised models are treated as always-loaded, like external OpenAI endpoints. - anthropic_endpoints: List[str] = Field(default_factory=list) # Max concurrent connections per endpoint‑model pair, see OLLAMA_NUM_PARALLEL max_concurrent_connections: int = 1 # Per-endpoint overrides: {endpoint_url: {max_concurrent_connections: N}} @@ -72,6 +64,9 @@ class Config(BaseSettings): # 0.3 = 30% history context signal, 70% question signal cache_history_weight: float = Field(default=0.3) + # aissurance.eu compliance plugin (`compliance:` block). See compliance/. + compliance: ComplianceSettings = Field(default_factory=ComplianceSettings) + class Config: # YAML loading is handled manually via Config.from_yaml(); env vars use this prefix. env_prefix = "NOMYO_ROUTER_" diff --git a/config.yaml b/config.yaml index 3e26597..998810f 100644 --- a/config.yaml +++ b/config.yaml @@ -6,23 +6,7 @@ endpoints: - https://api.openai.com/v1 llama_server_endpoints: - - http://192.168.0.51:8889/v1 - -# llama-swap endpoints (OpenAI-compatible front for multiple llama-server workers). -# Same surface as llama_server_endpoints, but the router reads loaded/running workers -# from /running (state == "ready") instead of a /v1/models status field, and unloads via -# POST /api/models/unload/:model_id. The router also exposes /upstream/:model_id/ -# to bypass llama-swap and reach a model's underlying llama-server worker directly. -llama_swap_endpoints: - - http://192.168.0.52:8890/v1 - -# Native Anthropic Messages-API endpoints (optional). Configure the base URL WITHOUT a -# /v1 suffix; the router appends /v1/models and /v1/messages itself. Requests routed to a -# model advertised here are forwarded verbatim (no Messages⇄Chat translation), with the -# matching api_keys entry sent as the x-api-key header. Advertised models are treated as -# always-loaded, like external OpenAI endpoints. -# anthropic_endpoints: -# - https://api.anthropic.com + - http://192.168.0.50:8889/v1 # Maximum concurrent connections *per endpoint‑model pair* (equals to OLLAMA_NUM_PARALLEL) # This is the global default; individual endpoints can override it via endpoint_config below. @@ -73,8 +57,7 @@ api_keys: "http://192.168.0.51:11434": "ollama" "http://192.168.0.52:11434": "ollama" "https://api.openai.com/v1": "${OPENAI_KEY}" - "http://192.168.0.51:8889/v1": "llama" - "http://192.168.0.52:8889/v1": "llama-swap" + "http://192.168.0.50:8889/v1": "llama" # ------------------------------------------------------------- # Semantic LLM Cache (optional — disabled by default) @@ -108,4 +91,25 @@ api_keys: # Weight of the BM25-weighted chat-history embedding vs last-user-message embedding. # 0.3 = 30% history context signal, 70% question signal. # Only relevant when cache_similarity < 1.0. -# cache_history_weight: 0.3 \ No newline at end of file +# cache_history_weight: 0.3 + +# ------------------------------------------------------------- +# aissurance.eu compliance plugin (optional — disabled by default) +# Periodically snapshots the AI-infrastructure landscape (model/endpoint +# inventory + aggregate telemetry — never prompt/completion content) and +# sends HMAC-signed evidence to aissurance.eu. All traffic is outbound. +# See compliance/ for the implementation. +# ------------------------------------------------------------- +# compliance: +# enabled: true +# server_url: "https://www.aissurance.eu/api/v1/discovery/receive" +# api_key: "${AISSURANCE_KEY}" # tenant secret (distinct from nomyo-router-api-key) +# tenant_id: "${AISSURANCE_TENANT_ID}" # optional if embedded in the key +# router_id: "router-prod-001" # defaults to the machine hostname +# polling_interval: 300 # seconds between discovery snapshots +# health_interval: 3600 # seconds between health/keepalive calls +# batch_size: 50 # max models per payload (413 → auto-split) +# max_retry_attempts: 10 # in-cycle send retries before buffering +# retry_backoff_base: 2 # exponential backoff base (seconds) +# buffer_dir: "./compliance-buffer" # encrypted offline retry storage +# max_buffer_payloads: 100 # oldest dropped first when full \ No newline at end of file diff --git a/db.py b/db.py index 03a5393..051fed0 100644 --- a/db.py +++ b/db.py @@ -1,4 +1,4 @@ -import aiosqlite, asyncio, orjson +import aiosqlite, asyncio from typing import Optional from pathlib import Path from datetime import datetime, timezone @@ -75,24 +75,6 @@ class TokenDatabase: ''') await db.execute('CREATE INDEX IF NOT EXISTS idx_token_time_series_timestamp ON token_time_series(timestamp)') await db.execute('CREATE INDEX IF NOT EXISTS idx_token_time_series_model_ts ON token_time_series(model, timestamp)') - # Responses API state — the router owns conversation state for the - # /v1/responses family (store / previous_response_id) and tracks - # background-task status here so polling survives across workers. - await db.execute(''' - CREATE TABLE IF NOT EXISTS stored_responses ( - response_id TEXT PRIMARY KEY, - previous_response_id TEXT, - model TEXT, - status TEXT, - created_at INTEGER, - input_messages TEXT, - output_items TEXT, - usage TEXT, - instructions TEXT, - error TEXT - ) - ''') - await db.execute('CREATE INDEX IF NOT EXISTS idx_stored_responses_prev ON stored_responses(previous_response_id)') await db.commit() async def update_token_counts(self, endpoint: str, model: str, input_tokens: int, output_tokens: int): @@ -273,6 +255,34 @@ class TokenDatabase: } return None + async def get_last_used_map(self) -> dict: + """Return {(endpoint, model): latest_unix_timestamp} from the time series. + + Read-only aggregate used by the compliance collector to derive each + model's ``last_used``. Computed in SQL in a single pass. + """ + db = await self._get_connection() + async with self._operation_lock: + async with db.execute(''' + SELECT endpoint, model, MAX(timestamp) + FROM token_time_series + GROUP BY endpoint, model + ''') as cursor: + rows = await cursor.fetchall() + return {(row[0], row[1]): row[2] for row in rows} + + async def get_token_totals_since(self, cutoff_ts: int) -> tuple[int, int]: + """Return (input_tokens, output_tokens) summed over entries at/after cutoff_ts.""" + db = await self._get_connection() + async with self._operation_lock: + async with db.execute(''' + SELECT COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0) + FROM token_time_series + WHERE timestamp >= ? + ''', (cutoff_ts,)) as cursor: + row = await cursor.fetchone() + return (int(row[0]), int(row[1])) if row else (0, 0) + async def aggregate_time_series_older_than(self, days: int, trim_old: bool = False) -> int: """ Aggregate time_series entries older than 'days' days into daily aggregates by @@ -337,158 +347,3 @@ class TokenDatabase: await db.commit() return aggregated_count - - # ----------------------------------------------------------------- - # Responses API state (store / previous_response_id / background) - # ----------------------------------------------------------------- - @staticmethod - def _row_to_response(row) -> dict: - """Map a stored_responses row to a plain dict, decoding JSON columns.""" - def _loads(val): - if val is None: - return None - try: - return orjson.loads(val) - except (orjson.JSONDecodeError, TypeError): - return None - return { - 'response_id': row[0], - 'previous_response_id': row[1], - 'model': row[2], - 'status': row[3], - 'created_at': row[4], - 'input_messages': _loads(row[5]), - 'output_items': _loads(row[6]), - 'usage': _loads(row[7]), - 'instructions': row[8], - 'error': _loads(row[9]), - } - - async def store_response( - self, - response_id: str, - *, - previous_response_id: Optional[str], - model: str, - status: str, - created_at: int, - input_messages: list, - output_items: Optional[list] = None, - usage: Optional[dict] = None, - instructions: Optional[str] = None, - error: Optional[dict] = None, - ): - """Insert or replace a stored Responses-API response row.""" - db = await self._get_connection() - async with self._operation_lock: - await db.execute(''' - INSERT INTO stored_responses - (response_id, previous_response_id, model, status, created_at, - input_messages, output_items, usage, instructions, error) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT (response_id) DO UPDATE SET - previous_response_id = excluded.previous_response_id, - model = excluded.model, - status = excluded.status, - created_at = excluded.created_at, - input_messages = excluded.input_messages, - output_items = excluded.output_items, - usage = excluded.usage, - instructions = excluded.instructions, - error = excluded.error - ''', ( - response_id, previous_response_id, model, status, created_at, - orjson.dumps(input_messages).decode("utf-8"), - orjson.dumps(output_items).decode("utf-8") if output_items is not None else None, - orjson.dumps(usage).decode("utf-8") if usage is not None else None, - instructions, - orjson.dumps(error).decode("utf-8") if error is not None else None, - )) - await db.commit() - - async def update_response_status( - self, - response_id: str, - status: str, - *, - output_items: Optional[list] = None, - usage: Optional[dict] = None, - error: Optional[dict] = None, - ): - """Update the status (and optionally output/usage/error) of a stored response.""" - db = await self._get_connection() - async with self._operation_lock: - await db.execute(''' - UPDATE stored_responses - SET status = ?, - output_items = COALESCE(?, output_items), - usage = COALESCE(?, usage), - error = COALESCE(?, error) - WHERE response_id = ? - ''', ( - status, - orjson.dumps(output_items).decode("utf-8") if output_items is not None else None, - orjson.dumps(usage).decode("utf-8") if usage is not None else None, - orjson.dumps(error).decode("utf-8") if error is not None else None, - response_id, - )) - await db.commit() - - async def get_response(self, response_id: str) -> Optional[dict]: - """Return a stored response as a dict, or None if not found.""" - db = await self._get_connection() - async with self._operation_lock: - async with db.execute(''' - SELECT response_id, previous_response_id, model, status, created_at, - input_messages, output_items, usage, instructions, error - FROM stored_responses WHERE response_id = ? - ''', (response_id,)) as cursor: - row = await cursor.fetchone() - return self._row_to_response(row) if row is not None else None - - async def delete_response(self, response_id: str) -> bool: - """Delete a stored response. Returns True if a row was removed.""" - db = await self._get_connection() - async with self._operation_lock: - cursor = await db.execute( - 'DELETE FROM stored_responses WHERE response_id = ?', (response_id,) - ) - await db.commit() - return cursor.rowcount > 0 - - async def get_response_chain(self, response_id: str, max_turns: int = 50) -> list: - """Walk previous_response_id back to the root, returned oldest-first. - - Bounded to ``max_turns`` so a pathological chain cannot stall a request. - Missing links terminate the walk gracefully. - """ - chain: list = [] - seen: set = set() - current = response_id - while current and current not in seen and len(chain) < max_turns: - seen.add(current) - resp = await self.get_response(current) - if resp is None: - break - chain.append(resp) - current = resp.get('previous_response_id') - chain.reverse() - return chain - - async def fail_orphaned_responses(self) -> int: - """Mark non-terminal responses as failed (called on startup). - - A background task lives in a worker's event loop; a process restart loses - it while the DB row stays ``queued``/``in_progress`` forever. Reconcile - those to ``failed`` so polling clients get a terminal state. - """ - db = await self._get_connection() - async with self._operation_lock: - cursor = await db.execute(''' - UPDATE stored_responses - SET status = 'failed', - error = ? - WHERE status IN ('queued', 'in_progress') - ''', (orjson.dumps({"message": "Response interrupted by server restart", "type": "server_error"}).decode("utf-8"),)) - await db.commit() - return cursor.rowcount diff --git a/doc/README.md b/doc/README.md index 3909089..9fde189 100644 --- a/doc/README.md +++ b/doc/README.md @@ -11,6 +11,7 @@ doc/ ├── usage.md # API usage examples ├── deployment.md # Deployment scenarios ├── monitoring.md # Monitoring and troubleshooting +├── compliance.md # aissurance.eu compliance plugin └── examples/ # Example configurations and scripts ├── docker-compose.yml ├── sample-config.yaml @@ -59,6 +60,7 @@ uvicorn router:app --host 0.0.0.0 --port 12434 - **Real-time Monitoring**: Server-Sent Events for live usage updates - **OpenAI Compatibility**: Full OpenAI API compatibility layer - **MOE System**: Multiple Opinions Ensemble for improved responses with smaller models +- **Compliance Plugin**: Optional EU AI Act evidence reporting to aissurance.eu ## Documentation Guides @@ -82,6 +84,10 @@ Step-by-step deployment guides for bare metal, Docker, Kubernetes, and productio Monitoring endpoints, troubleshooting guides, performance tuning, and best practices for maintaining your router. +### [Compliance](compliance.md) + +The optional aissurance.eu compliance plugin: what it collects (and what it never does), configuration, payload structure, security model, and troubleshooting. + ## Examples The [examples](examples/) directory contains ready-to-use configuration files: diff --git a/doc/compliance.md b/doc/compliance.md new file mode 100644 index 0000000..fba735b --- /dev/null +++ b/doc/compliance.md @@ -0,0 +1,342 @@ +# Compliance Plugin (aissurance.eu) + +The compliance plugin turns the NOMYO Router from a pure traffic proxy into an +**AI compliance agent**. When enabled, it periodically snapshots your +AI-infrastructure landscape — model inventory, endpoint registry, and aggregate +telemetry — and sends signed, structured evidence to +[aissurance.eu](https://www.aissurance.eu), the EU AI Act compliance platform. + +The plugin runs **in-process** inside the Router; it is not a separate service. +It only reads infrastructure metadata the Router already maintains for routing. +**No prompt or completion content is ever read, stored, or transmitted.** + +> **Non-critical by design.** Any failure inside the plugin is logged and +> contained — it never affects request routing. If the upstream is unreachable, +> the Router keeps proxying and buffers evidence for later. + +--- + +## How it works + +``` + Router state (models, endpoints, usage, cache, token DB) + │ read-only snapshot, every polling_interval + ▼ + Snapshot Collector → Payload Builder → HMAC-SHA256 Signer + │ + ▼ + HTTPS POST (outbound only) ──────────────► aissurance.eu + │ 202 Accepted → done /api/v1/discovery/receive + │ network error / 5xx → encrypt + buffer to disk, retry next cycle + ▼ + Encrypted offline buffer (replayed oldest-first on reconnect) +``` + +**All communication is outbound from the Router.** aissurance.eu never +initiates contact, so the Router can run behind NAT, firewalls, or strict egress +policies. + +--- + +## What is and isn't collected + +This is the most important part to understand before enabling the plugin. + +### Collected (infrastructure metadata only) + +- **Model inventory** — name, version, quantization, weight size, serving endpoint +- **Endpoint registry** — URL, type, health status, concurrency, TLS, auth method +- **Aggregate telemetry** — active connections, uptime, global cache hit rate, + 24h token totals + +### Never collected + +- Individual request payloads (prompt / completion content) +- User identity or auth tokens from proxied requests +- Conversation content passing through the proxy +- End-user client IP addresses + +### Emitted as `null` (not yet instrumented) + +The Router does not currently measure the following, so the plugin emits `null` +rather than fabricating values. They are candidates for a future instrumentation +phase: + +| Field | Where | +|---|---| +| `request_count_24h`, `error_count_24h` | per model | +| `avg_latency_ms`, `p99_latency_ms` | per model | +| `cache_hit_rate`, `avg_input_tokens`, `avg_output_tokens` | per model | +| `total_requests_24h`, `error_rate` | telemetry | + +> **Note on byte fields.** The Router tracks **tokens**, not raw bytes. +> `total_bytes_in_24h` / `total_bytes_out_24h` therefore carry 24h **token** +> counts as the closest available signal. + +--- + +## Quick start + +1. Obtain your `AISSURANCE_KEY` (and optionally `AISSURANCE_TENANT_ID`) from + aissurance.eu during tenant signup. This is **distinct** from the Router's + own `nomyo-router-api-key` — see [Dual-key model](#dual-key-model). + +2. Export the secret (never commit it): + + ```bash + export AISSURANCE_KEY="aiss_live_…" + export AISSURANCE_TENANT_ID="550e8400-e29b-41d4-a716-446655440000" # optional + ``` + +3. Add the `compliance` block to `config.yaml`: + + ```yaml + compliance: + enabled: true + server_url: "https://www.aissurance.eu/api/v1/discovery/receive" + api_key: "${AISSURANCE_KEY}" + tenant_id: "${AISSURANCE_TENANT_ID}" + polling_interval: 300 + ``` + +4. Restart the Router. On startup you should see structured logs: + + ```json + {"event":"plugin_started","router_id":"router-prod-001","polling_interval":300,"buffered":0} + ``` + +--- + +## Configuration options + +All options live under the `compliance:` block in `config.yaml`. + +| Option | Type | Default | Description | +|---|---|---|---| +| `enabled` | `bool` | `false` | Master switch. All other options are ignored when `false`. | +| `server_url` | `str` | `https://www.aissurance.eu/api/v1/discovery/receive` | Discovery receive URL. Must be HTTPS. Health and config-report URLs are derived from it. | +| `api_key` | `str` | `${AISSURANCE_KEY}` | Tenant secret. Used as the Bearer token and as the HMAC signing key. | +| `tenant_id` | `str` | `${AISSURANCE_TENANT_ID}` | Tenant UUID. Optional if embedded in the key. | +| `router_id` | `str` | machine hostname | Unique identifier for this Router instance. Multiple instances per tenant are allowed. | +| `polling_interval` | `int` | `300` | Seconds between discovery snapshots. Minimum `60`. | +| `health_interval` | `int` | `3600` | Seconds between health / keepalive calls. | +| `batch_size` | `int` | `50` | Max models per payload. On `413` the payload is auto-split into batches. | +| `max_retry_attempts` | `int` | `10` | In-cycle send retries before the payload is buffered to disk. | +| `retry_backoff_base` | `int` | `2` | Exponential backoff base in seconds (`base ** attempt`, capped at 60s). | +| `buffer_dir` | `str` | `./compliance-buffer` | Directory for encrypted offline retry files. | +| `max_buffer_payloads` | `int` | `100` | Buffer cap. Oldest file is dropped first when exceeded. | + +### `polling_interval` + +How often a snapshot is collected and sent. The server enforces a maximum of one +payload per `polling_interval` per tenant; sending faster yields `429`. Values +below `60` are rejected at startup. + +### `batch_size` + +Caps how many models a single payload carries. If the upstream replies `413 +Payload Too Large`, the plugin splits the model list into `batch_size`-sized +chunks, re-signs each chunk, and sends them individually. + +### `buffer_dir` and `max_buffer_payloads` + +When the upstream is unreachable, each signed payload is encrypted and written to +`buffer_dir`. The directory is bounded by `max_buffer_payloads`; once full, the +oldest file is dropped before a new one is written. Buffered payloads are +replayed oldest-first on the next successful cycle. Add `buffer_dir` to your +backup exclusions — files are encrypted but transient. + +--- + +## Environment variables + +| Variable | Required | Description | +|---|---|---| +| `AISSURANCE_KEY` | Yes | Tenant secret issued at signup. Referenced by `api_key: "${AISSURANCE_KEY}"`. | +| `AISSURANCE_HMAC_SECRET` | No | Per-tenant HMAC secret override. When set, it signs payloads instead of `AISSURANCE_KEY`. | +| `AISSURANCE_TENANT_ID` | No | Tenant UUID, if not embedded in the key. | +| `AISSURANCE_VERIFY_TLS` | No | Set to `"0"` to allow self-signed certs (development only — logged as a warning). | + +Values referenced as `${VAR}` in `config.yaml` are expanded at load time. If a +referenced variable is unset, the corresponding environment-variable default +above takes over. + +--- + +## Outbound endpoints + +All three are Router-initiated POSTs. URLs are derived from `server_url` by +replacing the trailing `/receive` path. + +| Endpoint | When called | Purpose | +|---|---|---| +| `POST …/discovery/receive` | every `polling_interval` | Send the signed discovery payload. | +| `POST …/discovery/health/{tenant_id}` | on startup + every `health_interval` | Connectivity check / keepalive / config sync. | +| `POST …/discovery/config/report` | on startup handshake | Report the active compliance config for acknowledgment. | + +Every request carries: + +``` +Authorization: Bearer ${AISSURANCE_KEY} +X-Router-Version: +X-Router-Id: +Content-Type: application/json +``` + +### Response handling + +| Status | Meaning | Plugin action | +|---|---|---| +| `202` / `200` | Accepted | Mark sent; replay any buffered payloads. | +| `401` / `403` | Invalid key / blacklisted | **Stop reporting**, clear the buffer, log `auth_failed`. | +| `429` | Rate limited | Drop this snapshot (the next one is just as fresh). | +| `413` | Too large | Split by `batch_size`, re-sign, resend. | +| `5xx` / network error | Server / connectivity failure | Retry with backoff, then encrypt + buffer to disk. | + +--- + +## Payload structure + +```json +{ + "tenant_id": "550e8400-e29b-41d4-a716-446655440000", + "router_id": "router-prod-001", + "router_version": "0.9", + "timestamp": "2026-06-10T10:30:00Z", + "models": [ + { + "name": "mistral-7b-instruct", + "version": "latest", + "quantization": "Q4_K_M", + "size_gb": 4.1, + "endpoint": "http://192.168.0.50:11434", + "last_used": "2026-06-10T10:28:14Z", + "request_count_24h": null, + "avg_latency_ms": null, + "p99_latency_ms": null, + "error_count_24h": null, + "cache_hit_rate": null, + "avg_input_tokens": null, + "avg_output_tokens": null + } + ], + "endpoints": [ + { + "url": "http://192.168.0.50:11434", + "type": "ollama", + "status": "healthy", + "concurrent_connections": 2, + "max_concurrent_connections": 4, + "tls_enabled": false, + "auth_method": "none" + } + ], + "telemetry": { + "total_requests_24h": null, + "total_cache_hits": 1547, + "cache_hit_rate": 0.634, + "error_rate": null, + "total_bytes_in_24h": 9120344, + "total_bytes_out_24h": 14882910, + "active_connections": 2, + "uptime_hours": 72.4 + }, + "hmac_signature": "sha256=abcdef…" +} +``` + +`endpoint.type` is best-effort: `ollama`, `llama_cpp`, `vllm`, `openai`, +`gemini`, `anthropic`, `cohere`, `cerebras`, `inception labs`, or `other`. +`endpoint.status` maps the Router's health probe onto `healthy` / `unhealthy`. + +--- + +## Security + +### Dual-key model + +| Credential | Protects | Managed by | +|---|---|---| +| `nomyo-router-api-key` | The Router's own API / dashboard | Router operator | +| `AISSURANCE_KEY` | The upstream reporting channel | aissurance.eu (issued at signup) | + +These are independent. See [Router API key usage](configuration.md#using-the-router-api-key) +for the former. + +### Payload signing + +Every payload is HMAC-SHA256 signed. The signature covers all fields **except** +`hmac_signature`, over a canonical `sort_keys` serialization, so any tampering +with the body invalidates it. The signing key is `AISSURANCE_HMAC_SECRET` if set, +otherwise `AISSURANCE_KEY`. + +### Buffer encryption + +Offline buffer files are encrypted at rest with a key derived from your +`AISSURANCE_KEY` via HKDF-SHA256 (Fernet / AES-128-CBC + HMAC). A leaked buffer +file is useless without the tenant secret. Corrupt or wrong-key files are dropped +on read. + +### Transport + +TLS is required (`server_url` must be HTTPS). `AISSURANCE_VERIFY_TLS=0` disables +certificate verification for development only and is logged as a warning. + +--- + +## Observability + +The plugin writes structured JSON logs to stderr (one object per line): + +```json +{"timestamp":"2026-06-10T10:30:00Z","level":"info","module":"aissurance","event":"payload_sent","tenant_id":"abc123","router_id":"router-prod-001","models_count":12,"response_code":202,"latency_ms":245} +``` + +Event types: + +| Event | Meaning | +|---|---| +| `plugin_started` / `plugin_stopped` | Lifecycle. `plugin_stopped` reports buffered + successful counts. | +| `plugin_disabled` | Startup validation failed (e.g. missing key, non-HTTPS URL). Reasons included. | +| `payload_sent` | A discovery payload (or a replayed buffer file) was accepted. | +| `payload_buffered` | Upstream unreachable; payload encrypted to disk. | +| `payload_retry` | A transient failure is being retried with backoff. | +| `payload_rate_limited` | Server returned `429`; snapshot dropped. | +| `auth_failed` | Key rejected (`401`/`403`); reporting halted, buffer cleared. **Operator action required.** | +| `config_reported` | Config handshake completed; any server suggestion included. | +| `snapshot_too_slow` | Collection exceeded its 10s budget; cycle skipped. | +| `tenant_id_missing` | No tenant id resolved; payloads will carry `tenant_id: null`. | + +--- + +## Lifecycle + +**Startup** — validate config (HTTPS URL, key present, `polling_interval ≥ 60`); +on failure the plugin logs `plugin_disabled` and stays dormant without affecting +the Router. On success it performs a health + config-report handshake, then +starts the snapshot and health loops. + +**Runtime** — the first snapshot is emitted immediately so evidence appears +without waiting a full interval, then the plugin settles into `polling_interval`. +Snapshot collection runs under a 10s budget and reads from shared Router state +without blocking request routing. + +**Shutdown** — the loops stop, the HTTP client closes, and a `plugin_stopped` +log records how many payloads were buffered vs. sent. The Router continues and +completes its own shutdown normally. + +> **Config changes require a restart.** Like the rest of the Router config, the +> `compliance` block is read at startup. Hot-reload of the compliance settings is +> not yet supported. + +--- + +## Troubleshooting + +| Symptom | Likely cause | Fix | +|---|---|---| +| `plugin_disabled` at startup | Missing key, non-HTTPS `server_url`, or `polling_interval < 60` | Check the `reasons` array in the log line. | +| `auth_failed`, reporting stops | `AISSURANCE_KEY` invalid/expired/blacklisted | Re-provision the key on aissurance.eu and restart. | +| Buffer files accumulating | Upstream unreachable | Confirm egress to aissurance.eu over HTTPS; check firewall/NAT. | +| `tenant_id` is `null` in payloads | No `AISSURANCE_TENANT_ID` / `compliance.tenant_id` | Set it, unless your key embeds the tenant. | +| `snapshot_too_slow` warnings | Slow/unhealthy backends delaying health probes | Investigate endpoint health; the cycle self-recovers next tick. | diff --git a/doc/configuration.md b/doc/configuration.md index e067207..ef9ebfd 100644 --- a/doc/configuration.md +++ b/doc/configuration.md @@ -78,37 +78,6 @@ endpoints: - OpenAI-compatible endpoints use `/v1` prefix - The router automatically detects endpoint type based on URL pattern -### `llama_server_endpoints` - -**Type**: `list[str]` (optional) - -**Default**: `[]` - -**Description**: List of [llama.cpp `llama-server`](https://github.com/ggml-org/llama.cpp) endpoints (OpenAI-compatible, configured with the `/v1` suffix). The router reads each backend's loaded models from `/v1/models` (entries with `status == "loaded"`) and unloads idle models via `POST /models/unload`. - -```yaml -llama_server_endpoints: - - http://192.168.0.50:8889/v1 -``` - -### `llama_swap_endpoints` - -**Type**: `list[str]` (optional) - -**Default**: `[]` - -**Description**: List of [llama-swap](https://github.com/mostlygeek/llama-swap) endpoints (OpenAI-compatible, configured with the `/v1` suffix). llama-swap fronts multiple `llama-server` workers behind one address. It is treated like `llama_server_endpoints` for routing, model discovery, and reranking, but differs in two ways the router handles automatically: - -- **Loaded-model detection** — llama-swap's `/v1/models` omits the per-model `status` field, so running workers are read from `GET /running` (entries with `state == "ready"`). -- **Model unload** — done via `POST /api/models/unload/:model_id` (path parameter), not the `llama-server` body form. - -The router also exposes a passthrough route, `GET|POST /upstream/:model_id/`, which forwards directly to a model's underlying `llama-server` worker (via llama-swap's `/upstream`), letting clients use `llama-server` features that llama-swap does not forward (e.g. token-array prompts). - -```yaml -llama_swap_endpoints: - - http://192.168.0.50:8890/v1 -``` - ### `max_concurrent_connections` **Type**: `int` @@ -563,6 +532,28 @@ docker build -t nomyo-router . docker build --build-arg SEMANTIC_CACHE=true -t nomyo-router:semantic . ``` +## Compliance Plugin (aissurance.eu) + +The optional `compliance` block enables the aissurance.eu compliance plugin, +which periodically sends signed AI-infrastructure evidence (model/endpoint +inventory + aggregate telemetry — never prompt or completion content) to +aissurance.eu for EU AI Act compliance. + +```yaml +compliance: + enabled: true + server_url: "https://www.aissurance.eu/api/v1/discovery/receive" + api_key: "${AISSURANCE_KEY}" # tenant secret (distinct from nomyo-router-api-key) + tenant_id: "${AISSURANCE_TENANT_ID}" # optional if embedded in the key + polling_interval: 300 # seconds between discovery snapshots +``` + +The `AISSURANCE_KEY` is separate from `router_api_key` — one secures the proxy, +the other secures the upstream reporting channel. All traffic is outbound. + +See the **[Compliance Guide](compliance.md)** for the full option reference, +payload structure, security model, and troubleshooting. + ## Configuration Validation The router validates the configuration at startup: diff --git a/doc/examples/sample-config.yaml b/doc/examples/sample-config.yaml index a1081af..f252ec7 100644 --- a/doc/examples/sample-config.yaml +++ b/doc/examples/sample-config.yaml @@ -63,4 +63,24 @@ api_keys: # Weight of the BM25-weighted chat-history embedding vs last-user-message embedding. # 0.3 = 30% history context signal, 70% question signal. # Only relevant when cache_similarity < 1.0. -# cache_history_weight: 0.3 \ No newline at end of file +# cache_history_weight: 0.3 + +# ------------------------------------------------------------- +# aissurance.eu compliance plugin (optional — disabled by default) +# Sends signed AI-infrastructure evidence (model/endpoint inventory + +# aggregate telemetry — never prompt/completion content) to aissurance.eu. +# All traffic is outbound. See doc/compliance.md for the full reference. +# ------------------------------------------------------------- +# compliance: +# enabled: true +# server_url: "https://www.aissurance.eu/api/v1/discovery/receive" +# api_key: "${AISSURANCE_KEY}" # tenant secret (separate from nomyo-router-api-key) +# tenant_id: "${AISSURANCE_TENANT_ID}" # optional if embedded in the key +# router_id: "router-prod-001" # defaults to the machine hostname +# polling_interval: 300 # seconds between discovery snapshots +# health_interval: 3600 # seconds between health/keepalive calls +# batch_size: 50 # max models per payload (413 → auto-split) +# max_retry_attempts: 10 # in-cycle send retries before buffering +# retry_backoff_base: 2 # exponential backoff base (seconds) +# buffer_dir: "./compliance-buffer" # encrypted offline retry storage +# max_buffer_payloads: 100 # oldest dropped first when full \ No newline at end of file diff --git a/requests/anthropic.py b/requests/anthropic.py deleted file mode 100644 index b6550d4..0000000 --- a/requests/anthropic.py +++ /dev/null @@ -1,558 +0,0 @@ -"""Translation between the Anthropic **Messages API** and **Chat Completions**. - -The router speaks Chat Completions to its local backends (Ollama, llama-server, -llama-swap). To expose ``/v1/messages`` transparently on top of that, this module -converts in both directions: - - * request: Anthropic ``system`` / ``messages`` / ``tools`` → chat ``messages`` / ``tools`` - * response: chat ``choices[0].message`` → Anthropic ``content`` blocks - * stream: chat completion deltas → Anthropic typed SSE events - -Pure functions / a stream-translator class — no I/O, mirroring ``requests/responses.py``. -The native passthrough path (configured ``anthropic_endpoints``) does not use this -module; it forwards the Anthropic wire format straight through. - -Reasoning: an inbound ``thinking`` block maps to the backend's ``reasoning_effort``; -a backend that streams ``reasoning_content`` is surfaced back as Anthropic -``thinking`` content blocks / ``thinking_delta`` events. -""" -import secrets - -import orjson - - -# --------------------------------------------------------------------------- -# Request direction: Anthropic Messages → Chat Completions -# --------------------------------------------------------------------------- -def _system_to_text(system): - """Flatten Anthropic ``system`` (string or text blocks) to a plain string.""" - if system is None: - return None - if isinstance(system, str): - return system - if isinstance(system, list): - parts = [] - for b in system: - if isinstance(b, dict) and b.get("type") == "text": - parts.append(b.get("text", "")) - elif isinstance(b, str): - parts.append(b) - return "\n\n".join(p for p in parts if p) or None - return None - - -def _image_block_to_chat(block): - """Convert an Anthropic ``image`` block to an OpenAI ``image_url`` part (or None).""" - src = block.get("source") or {} - stype = src.get("type") - if stype == "base64": - media = src.get("media_type", "image/png") - data = src.get("data", "") - return {"type": "image_url", "image_url": {"url": f"data:{media};base64,{data}"}} - if stype == "url" and src.get("url"): - return {"type": "image_url", "image_url": {"url": src["url"]}} - return None - - -def _tool_result_content_to_str(content): - """Flatten an Anthropic ``tool_result`` content (string or blocks) to a string.""" - if content is None: - return "" - if isinstance(content, str): - return content - if isinstance(content, list): - parts = [] - for b in content: - if isinstance(b, dict): - if b.get("type") == "text": - parts.append(b.get("text", "")) - elif b.get("type") == "image": - parts.append("[image]") - else: - parts.append(orjson.dumps(b).decode("utf-8")) - else: - parts.append(str(b)) - return "\n".join(parts) - return orjson.dumps(content).decode("utf-8") - - -def _user_content_to_chat(content): - """Split an Anthropic user ``content`` into ``(chat_content, tool_messages)``. - - ``tool_result`` blocks become standalone OpenAI ``role:"tool"`` messages; the - remaining text/image parts collapse to a chat ``content`` (string or list). - """ - tool_messages = [] - if content is None or isinstance(content, str): - return content, tool_messages - parts = [] - for b in content if isinstance(content, list) else []: - if not isinstance(b, dict): - parts.append({"type": "text", "text": str(b)}) - continue - btype = b.get("type") - if btype == "text": - parts.append({"type": "text", "text": b.get("text", "")}) - elif btype == "image": - img = _image_block_to_chat(b) - if img: - parts.append(img) - elif btype == "tool_result": - tool_messages.append({ - "role": "tool", - "tool_call_id": b.get("tool_use_id"), - "content": _tool_result_content_to_str(b.get("content")), - }) - # document / other blocks: no chat equivalent → skip - if not parts: - chat_content = None - elif len(parts) == 1 and parts[0].get("type") == "text": - chat_content = parts[0]["text"] - else: - chat_content = parts - return chat_content, tool_messages - - -def _assistant_content_to_chat(content): - """Convert an Anthropic assistant ``content`` to ``(text_or_parts, tool_calls)``.""" - if content is None or isinstance(content, str): - return content, None - text_parts = [] - tool_calls = [] - for b in content if isinstance(content, list) else []: - if not isinstance(b, dict): - continue - btype = b.get("type") - if btype == "text": - text_parts.append(b.get("text", "")) - elif btype == "tool_use": - tool_calls.append({ - "id": b.get("id"), - "type": "function", - "function": { - "name": b.get("name"), - "arguments": orjson.dumps(b.get("input") or {}).decode("utf-8"), - }, - }) - # thinking blocks in history have no chat equivalent → drop - text = "".join(text_parts) if text_parts else None - return text, (tool_calls or None) - - -def anthropic_messages_to_chat(system, messages): - """Build a Chat Completions ``messages`` list from Anthropic ``system`` + ``messages``.""" - chat = [] - sys_text = _system_to_text(system) - if sys_text: - chat.append({"role": "system", "content": sys_text}) - for m in messages or []: - role = m.get("role") - content = m.get("content") - if role == "user": - chat_content, tool_messages = _user_content_to_chat(content) - # Anthropic packs tool results into a user turn; OpenAI wants each as a - # separate tool message *before* any trailing user text. - chat.extend(tool_messages) - if chat_content is not None and chat_content != []: - chat.append({"role": "user", "content": chat_content}) - elif role == "assistant": - text, tool_calls = _assistant_content_to_chat(content) - msg = {"role": "assistant", "content": text} - if tool_calls: - msg["tool_calls"] = tool_calls - chat.append(msg) - else: - chat.append({"role": role or "user", "content": content}) - return chat - - -def tools_anthropic_to_chat(tools): - """Map Anthropic tool definitions → Chat Completions function tools.""" - if not tools: - return None - out = [] - for t in tools: - if not isinstance(t, dict): - continue - # Server tools (web_search, …) carry a ``type`` but no ``input_schema`` — - # local backends can't run them, so skip. - if "input_schema" not in t and t.get("type") not in (None, "custom"): - continue - fn = {"name": t.get("name")} - if t.get("description"): - fn["description"] = t["description"] - fn["parameters"] = t.get("input_schema") or {"type": "object", "properties": {}} - out.append({"type": "function", "function": fn}) - return out or None - - -def tool_choice_anthropic_to_chat(tool_choice): - """Map Anthropic ``tool_choice`` → Chat Completions ``tool_choice``.""" - if not isinstance(tool_choice, dict): - return None - ttype = tool_choice.get("type") - if ttype == "auto": - return "auto" - if ttype == "any": - return "required" - if ttype == "none": - return "none" - if ttype == "tool" and tool_choice.get("name"): - return {"type": "function", "function": {"name": tool_choice["name"]}} - return None - - -def _thinking_to_reasoning_effort(thinking): - """Map an Anthropic ``thinking`` config to an OpenAI ``reasoning_effort`` level.""" - if not isinstance(thinking, dict): - return None - if thinking.get("type") == "disabled": - return None - budget = thinking.get("budget_tokens") - if not isinstance(budget, int): - return "medium" # adaptive / enabled without an explicit budget - if budget < 2048: - return "low" - if budget < 8192: - return "medium" - return "high" - - -def anthropic_to_chat_send_params(payload, chat_messages, model): - """Assemble the Chat Completions request body from an Anthropic payload.""" - send = {"messages": chat_messages, "model": model} - if payload.get("max_tokens") is not None: - send["max_tokens"] = payload["max_tokens"] - opt = { - "temperature": payload.get("temperature"), - "top_p": payload.get("top_p"), - "stop": payload.get("stop_sequences"), - "tools": tools_anthropic_to_chat(payload.get("tools")), - "tool_choice": tool_choice_anthropic_to_chat(payload.get("tool_choice")), - "reasoning_effort": _thinking_to_reasoning_effort(payload.get("thinking")), - } - send.update({k: v for k, v in opt.items() if v is not None}) - return send - - -# --------------------------------------------------------------------------- -# Response direction: Chat Completions → Anthropic Messages -# --------------------------------------------------------------------------- -def new_message_id(): - return f"msg_{secrets.token_hex(24)}" - - -_STOP_REASON_MAP = { - "stop": "end_turn", - "length": "max_tokens", - "tool_calls": "tool_use", - "function_call": "tool_use", - "content_filter": "end_turn", -} - - -def finish_reason_to_stop_reason(finish_reason, has_tool_use=False): - if has_tool_use: - return "tool_use" - 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. - - ``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": max(prompt - cache_read, 0), - "output_tokens": completion, - "cache_creation_input_tokens": cache_creation_tokens, - "cache_read_input_tokens": cache_read, - } - - -def chat_message_to_content_blocks(message): - """Convert an assistant chat message (dict) into Anthropic content blocks. - - Ordering follows the Anthropic convention: thinking → text → tool_use. - """ - blocks = [] - reasoning = message.get("reasoning_content") or message.get("reasoning") - if reasoning: - blocks.append({"type": "thinking", "thinking": reasoning}) - content = message.get("content") - if content: - blocks.append({"type": "text", "text": content}) - for tc in message.get("tool_calls") or []: - fn = tc.get("function", {}) - try: - args = orjson.loads(fn.get("arguments") or "{}") - except (orjson.JSONDecodeError, TypeError): - args = {} - blocks.append({ - "type": "tool_use", - "id": tc.get("id") or f"toolu_{secrets.token_hex(12)}", - "name": fn.get("name"), - "input": args, - }) - return blocks - - -def build_message_object(*, message_id, model, content_blocks, stop_reason, - usage, stop_sequence=None): - """Assemble a full ``type:"message"`` body for a non-streaming reply.""" - return { - "id": message_id, - "type": "message", - "role": "assistant", - "model": model, - "content": content_blocks or [], - "stop_reason": stop_reason, - "stop_sequence": stop_sequence, - "usage": usage, - } - - -# --------------------------------------------------------------------------- -# Cache-hit replay: a finished message object → Anthropic SSE stream -# --------------------------------------------------------------------------- -def _sse(etype, payload): - body = {"type": etype, **payload} - return f"event: {etype}\ndata: {orjson.dumps(body).decode('utf-8')}\n\n".encode("utf-8") - - -def message_object_to_sse(msg): - """Render a *finished* message object as a valid Anthropic SSE event stream. - - Used to serve cache hits to streaming clients without a backend call. - """ - out = [] - start = {**msg, "content": [], "stop_reason": None, "stop_sequence": None, - "usage": {**(msg.get("usage") or {}), "output_tokens": 0}} - out.append(_sse("message_start", {"message": start})) - for i, block in enumerate(msg.get("content") or []): - btype = block.get("type") - if btype == "text": - out.append(_sse("content_block_start", { - "index": i, "content_block": {"type": "text", "text": ""}})) - out.append(_sse("content_block_delta", { - "index": i, "delta": {"type": "text_delta", "text": block.get("text", "")}})) - elif btype == "thinking": - out.append(_sse("content_block_start", { - "index": i, "content_block": {"type": "thinking", "thinking": ""}})) - out.append(_sse("content_block_delta", { - "index": i, - "delta": {"type": "thinking_delta", "thinking": block.get("thinking", "")}})) - elif btype == "tool_use": - out.append(_sse("content_block_start", { - "index": i, - "content_block": {"type": "tool_use", "id": block.get("id"), - "name": block.get("name"), "input": {}}})) - out.append(_sse("content_block_delta", { - "index": i, - "delta": {"type": "input_json_delta", - "partial_json": orjson.dumps(block.get("input") or {}).decode("utf-8")}})) - out.append(_sse("content_block_stop", {"index": i})) - out.append(_sse("message_delta", { - "delta": {"stop_reason": msg.get("stop_reason"), - "stop_sequence": msg.get("stop_sequence")}, - "usage": {"output_tokens": (msg.get("usage") or {}).get("output_tokens", 0)}})) - out.append(_sse("message_stop", {})) - return b"".join(out) - - -# --------------------------------------------------------------------------- -# Streaming direction: Chat Completions deltas → Anthropic typed SSE events -# --------------------------------------------------------------------------- -class ChatToMessagesStream: - """Translate a Chat Completions streaming generator into Anthropic events. - - Usage:: - - translator = ChatToMessagesStream(message_id, model) - async for sse_bytes in translator.events(chat_async_gen): - yield sse_bytes - # translator.content_blocks / usage / stop_reason now populated for storage - - Emits ``message_start`` → (``content_block_start`` → ``*_delta``* → - ``content_block_stop``)* → ``message_delta`` (stop_reason + usage) → - ``message_stop``. A single monotonically increasing ``index`` is assigned across - thinking / text / tool_use blocks. - """ - - def __init__(self, message_id, model, cache_read_tokens=0, cache_creation_tokens=0): - self.message_id = message_id - self.model = model - self.cache_read_tokens = cache_read_tokens - self.cache_creation_tokens = cache_creation_tokens - self.usage = None - self.stop_reason = "end_turn" - self.content_blocks = [] - - def _open_block(self, index, content_block): - return _sse("content_block_start", {"index": index, "content_block": content_block}) - - def _delta(self, index, delta): - return _sse("content_block_delta", {"index": index, "delta": delta}) - - def _stop(self, index): - return _sse("content_block_stop", {"index": index}) - - async def events(self, async_gen): - yield _sse("message_start", {"message": { - "id": self.message_id, "type": "message", "role": "assistant", - "model": self.model, "content": [], "stop_reason": None, - "stop_sequence": None, - "usage": {"input_tokens": 0, "output_tokens": 0, - "cache_creation_input_tokens": self.cache_creation_tokens, - "cache_read_input_tokens": self.cache_read_tokens}, - }}) - - next_index = 0 - # open block state: ("thinking"|"text", index) or None - open_kind = None - open_index = None - thinking_text = [] - text_parts = [] - finish_reason = None - # tool call state: chat tool-call index -> {index, id, name, args} - tc_state = {} - - def close_open(): - nonlocal open_kind, open_index - if open_kind is not None: - ev = self._stop(open_index) - open_kind, open_index = None, None - return ev - return None - - async for chunk in async_gen: - usage = getattr(chunk, "usage", None) - if usage is not None: - 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: - continue - choice = choices[0] - if getattr(choice, "finish_reason", None): - finish_reason = choice.finish_reason - delta = choice.delta - - reasoning = getattr(delta, "reasoning_content", None) or getattr(delta, "reasoning", None) - if reasoning: - if open_kind != "thinking": - ev = close_open() - if ev: - yield ev - open_kind, open_index = "thinking", next_index - next_index += 1 - yield self._open_block(open_index, {"type": "thinking", "thinking": ""}) - thinking_text.append(reasoning) - yield self._delta(open_index, {"type": "thinking_delta", "thinking": reasoning}) - - content_piece = getattr(delta, "content", None) - if content_piece: - if open_kind != "text": - ev = close_open() - if ev: - yield ev - open_kind, open_index = "text", next_index - next_index += 1 - yield self._open_block(open_index, {"type": "text", "text": ""}) - text_parts.append(content_piece) - yield self._delta(open_index, {"type": "text_delta", "text": content_piece}) - - for tc in getattr(delta, "tool_calls", None) or []: - idx = tc.index - fn = getattr(tc, "function", None) - if idx not in tc_state: - ev = close_open() - if ev: - yield ev - block_index = next_index - next_index += 1 - state = { - "index": block_index, - "id": getattr(tc, "id", None) or f"toolu_{secrets.token_hex(12)}", - "name": (fn.name if fn else None), - "args": "", - } - tc_state[idx] = state - open_kind, open_index = "tool", block_index - yield self._open_block(block_index, { - "type": "tool_use", "id": state["id"], "name": state["name"], "input": {}}) - else: - state = tc_state[idx] - if getattr(tc, "id", None): - state["id"] = tc.id - if fn and fn.name: - state["name"] = fn.name - if fn and fn.arguments: - state["args"] += fn.arguments - yield self._delta(state["index"], { - "type": "input_json_delta", "partial_json": fn.arguments}) - - ev = close_open() - if ev: - yield ev - - # Assemble final content blocks (thinking → text → tool_use) for storage. - if thinking_text: - self.content_blocks.append({"type": "thinking", "thinking": "".join(thinking_text)}) - if text_parts: - self.content_blocks.append({"type": "text", "text": "".join(text_parts)}) - for idx in sorted(tc_state.keys()): - state = tc_state[idx] - try: - parsed = orjson.loads(state["args"]) if state["args"] else {} - except (orjson.JSONDecodeError, TypeError): - parsed = {} - self.content_blocks.append({ - "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)) - 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": { - "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/requests/chat.py b/requests/chat.py index 168e713..5bf1572 100644 --- a/requests/chat.py +++ b/requests/chat.py @@ -44,40 +44,38 @@ async def _make_chat_request(model: str, messages: list, tools=None, stream: boo """ config = get_config() endpoint, tracking_model = await choose_endpoint(model) # selects and atomically reserves - # The finally below releases the reservation on every exit — success, error, - # or CancelledError — so request building and client construction stay inside it. - try: - use_openai = is_openai_compatible(endpoint) - if use_openai: - if ":latest" in model: - model = model.split(":latest")[0] - if messages: - if any("images" in m for m in messages): - messages = await asyncio.to_thread(transform_images_to_data_urls, messages) - messages = transform_tool_calls_to_openai(messages) - messages = _strip_assistant_prefill(messages) - params = { - "messages": messages, - "model": model, - } - optional_params = { - "tools": tools, - "stream": stream, - "stream_options": {"include_usage": True} if stream else None, - "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, - "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, - "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, - "seed": options.get("seed") if options and "seed" in options else None, - "stop": options.get("stop") if options and "stop" in options else None, - "top_p": options.get("top_p") if options and "top_p" in options else None, - "temperature": options.get("temperature") if options and "temperature" in options else None, - "response_format": {"type": "json_schema", "json_schema": format} if format is not None else None - } - params.update({k: v for k, v in optional_params.items() if v is not None}) - oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) - else: - client = ollama.AsyncClient(host=endpoint) + use_openai = is_openai_compatible(endpoint) + if use_openai: + if ":latest" in model: + model = model.split(":latest")[0] + if messages: + if any("images" in m for m in messages): + messages = await asyncio.to_thread(transform_images_to_data_urls, messages) + messages = transform_tool_calls_to_openai(messages) + messages = _strip_assistant_prefill(messages) + params = { + "messages": messages, + "model": model, + } + optional_params = { + "tools": tools, + "stream": stream, + "stream_options": {"include_usage": True} if stream else None, + "max_tokens": options.get("num_predict") if options and "num_predict" in options else None, + "frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None, + "presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None, + "seed": options.get("seed") if options and "seed" in options else None, + "stop": options.get("stop") if options and "stop" in options else None, + "top_p": options.get("top_p") if options and "top_p" in options else None, + "temperature": options.get("temperature") if options and "temperature" in options else None, + "response_format": {"type": "json_schema", "json_schema": format} if format is not None else None + } + params.update({k: v for k, v in optional_params.items() if v is not None}) + oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key")) + else: + client = ollama.AsyncClient(host=endpoint) + try: if use_openai: start_ts = time.perf_counter() try: diff --git a/requests/responses.py b/requests/responses.py deleted file mode 100644 index 907dc42..0000000 --- a/requests/responses.py +++ /dev/null @@ -1,492 +0,0 @@ -"""Translation between the OpenAI **Responses API** and **Chat Completions**. - -The router speaks Chat Completions to every backend (Ollama, llama-server, -external OpenAI). To expose ``/v1/responses`` transparently on top of that, this -module converts in both directions: - - * request: Responses ``input`` / ``instructions`` / ``tools`` → chat ``messages`` / ``tools`` - * response: chat ``choices[0].message`` → Responses ``output`` items - * stream: chat completion deltas → Responses typed SSE events - -Pure functions / a stream-translator class — no I/O, mirroring the style of -``requests/messages.py``. The native passthrough path (external OpenAI) does not -use this module; it forwards the SDK's Responses objects directly. -""" -import secrets -import time - -import orjson - -from requests.messages import _accumulate_openai_tc_delta - - -# --------------------------------------------------------------------------- -# Request direction: Responses → Chat Completions -# --------------------------------------------------------------------------- -def _responses_content_to_chat(content): - """Convert a Responses message ``content`` into Chat Completions content. - - Collapses a single text part to a plain string (what most backends expect); - keeps a multimodal list otherwise. - """ - if content is None or isinstance(content, str): - return content - if not isinstance(content, list): - return str(content) - parts = [] - for p in content: - if not isinstance(p, dict): - parts.append({"type": "text", "text": str(p)}) - continue - ptype = p.get("type") - if ptype in ("input_text", "output_text", "text"): - parts.append({"type": "text", "text": p.get("text", "")}) - elif ptype in ("input_image", "image_url"): - url = p.get("image_url") - if isinstance(url, dict): - url = url.get("url") - if url: - parts.append({"type": "image_url", "image_url": {"url": url}}) - # input_file / refusal / reasoning parts have no chat equivalent → skip - if len(parts) == 1 and parts[0].get("type") == "text": - return parts[0]["text"] - return parts - - -def _input_item_to_message(item): - """Convert a single Responses ``input`` item to a chat message (or None).""" - if isinstance(item, str): - return {"role": "user", "content": item} - if not isinstance(item, dict): - return None - - itype = item.get("type") - - if itype == "function_call": - return { - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": item.get("call_id") or item.get("id"), - "type": "function", - "function": { - "name": item.get("name"), - "arguments": item.get("arguments", ""), - }, - }], - } - - if itype == "function_call_output": - output = item.get("output", "") - if not isinstance(output, str): - output = orjson.dumps(output).decode("utf-8") - return { - "role": "tool", - "tool_call_id": item.get("call_id") or item.get("id"), - "content": output, - } - - if itype in ("reasoning",): - # No Chat Completions equivalent — drop. - return None - - # "message" item or a bare {role, content} chat-style item - role = item.get("role") - if role is None: - return None - return {"role": role, "content": _responses_content_to_chat(item.get("content"))} - - -def responses_input_to_messages(input_data, instructions=None): - """Build a Chat Completions ``messages`` list from Responses ``input``. - - ``instructions`` becomes a leading system message; a string ``input`` becomes - a single user message; a list ``input`` is mapped item-by-item. - """ - messages = [] - if instructions: - messages.append({"role": "system", "content": instructions}) - if input_data is None: - return messages - if isinstance(input_data, str): - messages.append({"role": "user", "content": input_data}) - return messages - if isinstance(input_data, list): - for item in input_data: - msg = _input_item_to_message(item) - if msg is not None: - messages.append(msg) - return messages - - -def _chat_content_to_responses_parts(content, assistant=False): - """Convert chat message content → Responses content parts.""" - text_type = "output_text" if assistant else "input_text" - if content is None: - return [] - if isinstance(content, str): - return [{"type": text_type, "text": content}] - parts = [] - for p in content if isinstance(content, list) else []: - if not isinstance(p, dict): - parts.append({"type": text_type, "text": str(p)}) - elif p.get("type") == "text": - parts.append({"type": text_type, "text": p.get("text", "")}) - elif p.get("type") == "image_url": - url = (p.get("image_url") or {}).get("url") - if url: - parts.append({"type": "input_image", "image_url": url}) - return parts - - -def messages_to_responses_input(messages): - """Convert chat messages → ``(instructions, Responses input items)``. - - Used for the native passthrough path: history that the router has resolved in - chat-message space is re-expressed as Responses ``input``. Leading/standalone - system messages are merged into ``instructions``. - """ - instructions_parts = [] - items = [] - for m in messages: - role = m.get("role") - if role == "system": - c = m.get("content") - instructions_parts.append(c if isinstance(c, str) else orjson.dumps(c).decode("utf-8")) - continue - if role == "tool": - out = m.get("content") - if not isinstance(out, str): - out = orjson.dumps(out).decode("utf-8") - items.append({"type": "function_call_output", - "call_id": m.get("tool_call_id"), "output": out}) - continue - if role == "assistant" and m.get("tool_calls"): - for tc in m["tool_calls"]: - fn = tc.get("function", {}) - items.append({"type": "function_call", "call_id": tc.get("id"), - "name": fn.get("name"), "arguments": fn.get("arguments", "")}) - if m.get("content"): - items.append({"role": "assistant", - "content": _chat_content_to_responses_parts(m["content"], assistant=True)}) - continue - items.append({"role": role, - "content": _chat_content_to_responses_parts(m.get("content"), - assistant=(role == "assistant"))}) - instructions = "\n\n".join(p for p in instructions_parts if p) or None - return instructions, items - - -def responses_object_to_sse(resp): - """Render a *finished* Responses object as a valid SSE event stream. - - Used to serve cache/store hits to streaming clients without a backend call. - """ - seq = [-1] - - def ev(etype, payload): - seq[0] += 1 - body = {"type": etype, "sequence_number": seq[0], **payload} - return f"event: {etype}\ndata: {orjson.dumps(body).decode('utf-8')}\n\n".encode("utf-8") - - parts_out = [] - in_progress = {**resp, "status": "in_progress", "output": [], "output_text": ""} - parts_out.append(ev("response.created", {"response": in_progress})) - parts_out.append(ev("response.in_progress", {"response": in_progress})) - for oi, item in enumerate(resp.get("output", [])): - parts_out.append(ev("response.output_item.added", - {"output_index": oi, "item": {**item, "status": "in_progress"}})) - if item.get("type") == "message": - for ci, part in enumerate(item.get("content", [])): - if part.get("type") == "output_text": - iid = item.get("id") - parts_out.append(ev("response.content_part.added", { - "item_id": iid, "output_index": oi, "content_index": ci, - "part": {"type": "output_text", "text": "", "annotations": []}})) - parts_out.append(ev("response.output_text.delta", { - "item_id": iid, "output_index": oi, "content_index": ci, - "delta": part.get("text", "")})) - parts_out.append(ev("response.output_text.done", { - "item_id": iid, "output_index": oi, "content_index": ci, - "text": part.get("text", "")})) - parts_out.append(ev("response.content_part.done", { - "item_id": iid, "output_index": oi, "content_index": ci, "part": part})) - parts_out.append(ev("response.output_item.done", {"output_index": oi, "item": item})) - parts_out.append(ev("response.completed", {"response": resp})) - return b"".join(parts_out) - - -def tools_responses_to_chat(tools): - """Map Responses tool definitions (flattened) → Chat Completions (nested).""" - if not tools: - return None - out = [] - for t in tools: - if isinstance(t, dict) and t.get("type") == "function" and "function" not in t: - fn = {k: t[k] for k in ("name", "description", "parameters", "strict") if k in t} - out.append({"type": "function", "function": fn}) - else: - out.append(t) - return out - - -# --------------------------------------------------------------------------- -# Response direction: Chat Completions → Responses -# --------------------------------------------------------------------------- -def _new_id(prefix): - return f"{prefix}_{secrets.token_hex(16)}" - - -def chat_message_to_output_items(message): - """Convert an assistant chat message (dict) into Responses output items.""" - items = [] - content = message.get("content") - if content: - items.append({ - "type": "message", - "id": _new_id("msg"), - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": content, "annotations": []}], - }) - for tc in message.get("tool_calls") or []: - fn = tc.get("function", {}) - items.append({ - "type": "function_call", - "id": _new_id("fc"), - "call_id": tc.get("id"), - "name": fn.get("name"), - "arguments": fn.get("arguments", ""), - "status": "completed", - }) - return items - - -def usage_chat_to_responses(usage): - """Map chat usage ``{prompt_tokens, completion_tokens}`` → Responses usage.""" - if not usage: - return None - prompt = usage.get("prompt_tokens") or 0 - completion = usage.get("completion_tokens") or 0 - return { - "input_tokens": prompt, - "output_tokens": completion, - "total_tokens": usage.get("total_tokens") or (prompt + completion), - } - - -def output_items_to_text(output_items): - """Concatenate the ``output_text`` parts of all message items.""" - chunks = [] - for item in output_items or []: - if item.get("type") != "message": - continue - for part in item.get("content") or []: - if part.get("type") == "output_text": - chunks.append(part.get("text", "")) - return "".join(chunks) - - -def build_response_object( - *, - response_id, - model, - output_items=None, - usage=None, - status="completed", - created_at=None, - previous_response_id=None, - instructions=None, - error=None, - metadata=None, -): - """Assemble a full ``object:"response"`` body for a non-streaming reply.""" - output_items = output_items or [] - return { - "id": response_id, - "object": "response", - "created_at": created_at or int(time.time()), - "status": status, - "model": model, - "output": output_items, - "output_text": output_items_to_text(output_items), - "instructions": instructions, - "previous_response_id": previous_response_id, - "usage": usage_chat_to_responses(usage) if usage and "input_tokens" not in usage else usage, - "error": error, - "metadata": metadata or {}, - } - - -# --------------------------------------------------------------------------- -# Streaming direction: Chat Completions deltas → Responses typed SSE events -# --------------------------------------------------------------------------- -class ChatToResponsesStream: - """Translate a Chat Completions streaming generator into Responses events. - - Usage:: - - translator = ChatToResponsesStream(response_id, model, created_at) - async for sse_bytes in translator.events(chat_async_gen): - yield sse_bytes - # translator.output_items / translator.usage now populated for storage - - Emits the ordered event family - ``response.created`` → ``response.in_progress`` → - (``response.output_item.added`` → ``response.content_part.added`` → - ``response.output_text.delta``* → ``response.output_text.done`` → - ``response.content_part.done`` → ``response.output_item.done``) and/or - function-call item events → ``response.completed`` (carrying usage). - """ - - def __init__(self, response_id, model, created_at=None, - previous_response_id=None, instructions=None, metadata=None): - self.response_id = response_id - self.model = model - self.created_at = created_at or int(time.time()) - self.previous_response_id = previous_response_id - self.instructions = instructions - self.metadata = metadata or {} - self.seq = -1 - self.output_items = [] - self.usage = None - - def _snapshot(self, status, output=None): - return build_response_object( - response_id=self.response_id, - model=self.model, - output_items=output if output is not None else [], - usage=self.usage, - status=status, - created_at=self.created_at, - previous_response_id=self.previous_response_id, - instructions=self.instructions, - metadata=self.metadata, - ) - - def _event(self, etype, payload): - self.seq += 1 - body = {"type": etype, "sequence_number": self.seq, **payload} - return f"event: {etype}\ndata: {orjson.dumps(body).decode('utf-8')}\n\n".encode("utf-8") - - async def events(self, async_gen): - yield self._event("response.created", {"response": self._snapshot("in_progress")}) - yield self._event("response.in_progress", {"response": self._snapshot("in_progress")}) - - next_oi = 0 - # text message state - msg_item_id = None - msg_oi = None - text_parts = [] - # function-call state, keyed by chat tool_call index - tc_state = {} # idx -> {oi, item_id, call_id, name, args} - - async for chunk in async_gen: - usage = getattr(chunk, "usage", None) - if usage is not None: - self.usage = { - "prompt_tokens": usage.prompt_tokens or 0, - "completion_tokens": usage.completion_tokens or 0, - } - choices = getattr(chunk, "choices", None) - if not choices: - continue - delta = choices[0].delta - - content_piece = getattr(delta, "content", None) - if content_piece: - if msg_item_id is None: - msg_item_id = _new_id("msg") - msg_oi = next_oi - next_oi += 1 - item = { - "id": msg_item_id, "type": "message", "status": "in_progress", - "role": "assistant", "content": [], - } - yield self._event("response.output_item.added", - {"output_index": msg_oi, "item": item}) - yield self._event("response.content_part.added", { - "item_id": msg_item_id, "output_index": msg_oi, "content_index": 0, - "part": {"type": "output_text", "text": "", "annotations": []}, - }) - text_parts.append(content_piece) - yield self._event("response.output_text.delta", { - "item_id": msg_item_id, "output_index": msg_oi, "content_index": 0, - "delta": content_piece, - }) - - for tc in getattr(delta, "tool_calls", None) or []: - idx = tc.index - fn = getattr(tc, "function", None) - if idx not in tc_state: - item_id = _new_id("fc") - state = { - "oi": next_oi, "item_id": item_id, - "call_id": getattr(tc, "id", None) or _new_id("call"), - "name": (fn.name if fn else None), "args": "", - } - next_oi += 1 - tc_state[idx] = state - yield self._event("response.output_item.added", { - "output_index": state["oi"], - "item": { - "id": item_id, "type": "function_call", "status": "in_progress", - "call_id": state["call_id"], "name": state["name"], "arguments": "", - }, - }) - else: - state = tc_state[idx] - if getattr(tc, "id", None): - state["call_id"] = tc.id - if fn and fn.name: - state["name"] = fn.name - if fn and fn.arguments: - state["args"] += fn.arguments - yield self._event("response.function_call_arguments.delta", { - "item_id": state["item_id"], "output_index": state["oi"], - "delta": fn.arguments, - }) - - # finalize message item - if msg_item_id is not None: - full_text = "".join(text_parts) - yield self._event("response.output_text.done", { - "item_id": msg_item_id, "output_index": msg_oi, "content_index": 0, - "text": full_text, - }) - done_part = {"type": "output_text", "text": full_text, "annotations": []} - yield self._event("response.content_part.done", { - "item_id": msg_item_id, "output_index": msg_oi, "content_index": 0, - "part": done_part, - }) - msg_item = { - "id": msg_item_id, "type": "message", "status": "completed", - "role": "assistant", "content": [done_part], - } - yield self._event("response.output_item.done", - {"output_index": msg_oi, "item": msg_item}) - - # finalize function-call items (in output-index order) - tc_items = {} - for idx, state in tc_state.items(): - yield self._event("response.function_call_arguments.done", { - "item_id": state["item_id"], "output_index": state["oi"], - "arguments": state["args"], - }) - fc_item = { - "id": state["item_id"], "type": "function_call", "status": "completed", - "call_id": state["call_id"], "name": state["name"], "arguments": state["args"], - } - tc_items[state["oi"]] = fc_item - yield self._event("response.output_item.done", - {"output_index": state["oi"], "item": fc_item}) - - # assemble final output items ordered by output index - ordered = [] - if msg_item_id is not None: - ordered.append((msg_oi, msg_item)) - ordered.extend(tc_items.items()) - self.output_items = [item for _, item in sorted(ordered, key=lambda kv: kv[0])] - - yield self._event("response.completed", - {"response": self._snapshot("completed", self.output_items)}) diff --git a/requirements.txt b/requirements.txt index b11bf62..0d61e91 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,44 +1,45 @@ -aiohappyeyeballs==2.7.1 -aiohttp==3.14.1 +aiohappyeyeballs==2.6.1 +aiohttp==3.14.0 aiosignal==1.4.0 annotated-types==0.7.0 -anyio==4.14.2 +anyio==4.13.0 async-timeout==5.0.1 attrs==26.1.0 -certifi==2026.6.17 -click==8.4.2 +certifi==2026.4.22 +click==8.4.0 +cryptography==48.0.1 distro==1.9.0 exceptiongroup==1.3.1 -fastapi==0.139.0 +fastapi==0.136.1 fastapi-sse==1.1.1 frozenlist==1.8.0 h11==0.16.0 httpcore==1.0.9 httpx==0.28.1 -idna==3.18 -jiter==0.16.0 +idna==3.15 +jiter==0.14.0 multidict==6.7.1 ollama==0.6.2 -openai==2.45.0 +openai==2.37.0 orjson>=3.11.5 numpy>=1.26 -pillow==12.3.0 +pillow==12.2.0 propcache==0.5.2 pydantic==2.13.4 -pydantic-settings==2.14.2 +pydantic-settings==2.14.1 pydantic_core==2.46.4 python-dotenv==1.2.2 PyYAML==6.0.3 sniffio==1.3.1 -starlette>=1.0.1 +starlette==0.52.1 truststore==0.10.4 tiktoken==0.13.0 -tqdm==4.68.4 +tqdm==4.67.3 typing-inspection==0.4.2 -typing_extensions==4.16.0 -uvicorn==0.49.0 +typing_extensions==4.15.0 +uvicorn==0.47.0 uvloop -yarl==1.24.2 +yarl==1.23.0 aiosqlite # Semantic LLM cache — base install (exact-match mode, no heavy ML deps) # For semantic mode use the :semantic Docker image tag (adds sentence-transformers + torch) diff --git a/router.py b/router.py index d6bb41f..bbd6b15 100644 --- a/router.py +++ b/router.py @@ -68,6 +68,10 @@ from state import ( token_worker_task: asyncio.Task | None = None flush_task: asyncio.Task | None = None +# aissurance compliance plugin instance (created on startup if enabled). +from compliance import AissurancePlugin +compliance_plugin: "AissurancePlugin | None" = None + from config import Config, _config_path_from_env from ollama._types import TokenLogprob, Logprob @@ -231,7 +235,6 @@ from backends.health import ( from backends.normalize import ( is_ext_openai_endpoint, is_openai_compatible, - llama_endpoints, get_tracking_model, ) @@ -291,10 +294,6 @@ from api.management import router as management_router app.include_router(management_router) from api.openai import router as openai_router app.include_router(openai_router) -from api.responses import router as responses_router -app.include_router(responses_router) -from api.messages import router as messages_router -app.include_router(messages_router) from api.ollama import router as ollama_router app.include_router(ollama_router) @@ -313,7 +312,6 @@ async def startup_event() -> None: f"Loaded configuration from {config_path}:\n" f" endpoints={config.endpoints},\n" f" llama_server_endpoints={config.llama_server_endpoints},\n" - f" llama_swap_endpoints={config.llama_swap_endpoints},\n" f" max_concurrent_connections={config.max_concurrent_connections},\n" f" endpoint_config={config.endpoint_config},\n" f" priority_routing={config.priority_routing}" @@ -328,13 +326,6 @@ async def startup_event() -> None: db = TokenDatabase(config.db_path) await db.init_db() - # Reconcile Responses-API background tasks lost across a restart: their - # in-memory asyncio task is gone but the DB row may still read queued / - # in_progress, so mark those failed to give polling clients a terminal state. - _orphaned = await db.fail_orphaned_responses() - if _orphaned: - print(f"[startup] Marked {_orphaned} orphaned background response(s) as failed.") - # Load existing token counts from database async for count_entry in db.load_token_counts(): endpoint = count_entry['endpoint'] @@ -377,15 +368,8 @@ async def startup_event() -> None: if is_ext_openai_endpoint(ep): app_state["httpx_clients"][ep] = httpx.AsyncClient(timeout=30.0) - # Native Anthropic Messages-API endpoints are forwarded over httpx. Use a - # long read timeout so streamed completions aren't cut short. Closed on - # shutdown by the shared httpx_clients cleanup. - for ep in config.anthropic_endpoints: - app_state["httpx_clients"][ep] = httpx.AsyncClient( - timeout=httpx.Timeout(300.0, connect=15.0)) - # Create per-endpoint Unix socket sessions for .sock endpoints - for ep in llama_endpoints(config): + for ep in config.llama_server_endpoints: if _is_unix_socket_endpoint(ep): sock_path = _get_socket_path(ep) sock_connector = aiohttp.UnixConnector(path=sock_path) @@ -402,7 +386,7 @@ async def startup_event() -> None: # client (/api/chat, /api/generate) and the OpenAI client (/v1/* routes), # so warm both; OpenAI-compatible endpoints only need the OpenAI client. _warm_endpoints = config.endpoints + [ - ep for ep in llama_endpoints(config) if ep not in config.endpoints + ep for ep in config.llama_server_endpoints if ep not in config.endpoints ] for ep in _warm_endpoints: try: @@ -419,10 +403,29 @@ async def startup_event() -> None: flush_task = asyncio.create_task(flush_buffer()) await init_llm_cache(config) + # Start the aissurance compliance plugin (non-critical; guarded by config). + global compliance_plugin + if config.compliance.enabled: + try: + compliance_plugin = AissurancePlugin(config.compliance) + await compliance_plugin.start() + except Exception as e: + print(f"[startup] aissurance compliance plugin failed to start: {e}") + compliance_plugin = None + @app.on_event("shutdown") async def shutdown_event() -> None: await close_all_sse_queues() + # Stop the compliance plugin first (best-effort buffer flush on its own). + global compliance_plugin + if compliance_plugin is not None: + try: + await compliance_plugin.stop() + except Exception as e: + print(f"[shutdown] Error stopping aissurance compliance plugin: {e}") + compliance_plugin = None + # Stop background tasks first so they stop touching the DB before we close it. for t in (token_worker_task, flush_task): if t is not None: diff --git a/routing.py b/routing.py index df2da84..059cf9a 100644 --- a/routing.py +++ b/routing.py @@ -32,9 +32,6 @@ from backends.health import _is_fresh from backends.normalize import ( is_ext_openai_endpoint, is_openai_compatible, - is_llama_server, - is_anthropic_endpoint, - llama_endpoints, get_tracking_model, ) from backends.probe import fetch @@ -96,24 +93,13 @@ async def choose_endpoint(model: str, reserve: bool = True, """ config = get_config() # 1️⃣ Gather advertised‑model sets for all endpoints concurrently - # Include config.endpoints plus any llama-server / llama-swap endpoints - llama_eps_extra = [ep for ep in llama_endpoints(config) if ep not in config.endpoints] - anthropic_eps_extra = [ep for ep in config.anthropic_endpoints if ep not in config.endpoints] - all_endpoints = config.endpoints + llama_eps_extra + anthropic_eps_extra + # Include both config.endpoints and config.llama_server_endpoints + llama_eps_extra = [ep for ep in config.llama_server_endpoints if ep not in config.endpoints] + all_endpoints = config.endpoints + llama_eps_extra - # Build the probe tasks in the SAME order as ``all_endpoints`` so the - # gathered results stay aligned for the ``zip(all_endpoints, advertised_sets)`` - # below. (A previous partition into non-OpenAI / OpenAI groups reordered the - # tasks relative to ``all_endpoints``, pairing each endpoint with another - # endpoint's advertised models — so an OpenAI-compatible backend listed - # before an Ollama one would inherit the Ollama model set and get a request - # for a model it cannot serve, 404ing. See issue #128.) - def _advertised_task(ep: str): - if is_openai_compatible(ep) or is_anthropic_endpoint(ep): - return fetch.available_models(ep, config.api_keys.get(ep)) - return fetch.available_models(ep) - - tag_tasks = [_advertised_task(ep) for ep in all_endpoints] + tag_tasks = [fetch.available_models(ep) for ep in config.endpoints if not is_openai_compatible(ep)] + tag_tasks += [fetch.available_models(ep, config.api_keys.get(ep)) for ep in config.endpoints if is_openai_compatible(ep)] + tag_tasks += [fetch.available_models(ep, config.api_keys.get(ep)) for ep in llama_eps_extra] advertised_sets = await asyncio.gather(*tag_tasks) # 2️⃣ Filter endpoints that advertise the requested model @@ -128,7 +114,7 @@ async def choose_endpoint(model: str, reserve: bool = True, model_without_latest = model.split(":latest")[0] candidate_endpoints = [ ep for ep, models in zip(all_endpoints, advertised_sets) - if model_without_latest in models and (is_ext_openai_endpoint(ep) or is_llama_server(ep)) + if model_without_latest in models and (is_ext_openai_endpoint(ep) or ep in config.llama_server_endpoints) ] if not candidate_endpoints: # Only add :latest suffix if model doesn't already have a version suffix @@ -216,43 +202,6 @@ async def choose_endpoint(model: str, reserve: bool = True, def utilization_ratio(ep: str) -> float: return tracking_usage(ep) / get_max_connections(ep) - def total_load(ep: str) -> int: - """Sum of in-flight requests across *all* models on the endpoint.""" - return sum(usage_counts.get(ep, {}).values()) - - # How many models each candidate currently has *resident* (from the - # /api/ps probe). With infinite keep-alive a model stays loaded long - # after its in-flight count drops to zero, so this is the signal that - # spreads *distinct* models across backends. - ep_loaded_counts = { - ep: len(models) for ep, models in zip(candidate_endpoints, loaded_sets) - } - - def loaded_count(ep: str) -> int: - return ep_loaded_counts.get(ep, 0) - - def pick_least_loaded(eps: list[str]) -> str: - """Pick the least-committed endpoint, breaking ties at random. - - Ordering key is ``(total_load, loaded_count)``: - - * ``total_load`` (in-flight requests across *all* models) keeps a - request off a backend already busy with a *different* model — - otherwise the per-model count reads zero everywhere and the - ranking is discarded (cold model B landing on the box serving A). - * ``loaded_count`` (number of *resident* models) then spreads - distinct models across backends. Two different cold models (27b, - 35b) requested back-to-back must not pile onto the same box: once - 27b is resident there, that box has loaded_count 1 while the idle - backends have 0, so the next cold model prefers an empty backend - even though every backend reports zero in-flight load. - - ``random.choice`` only breaks genuine ties on both keys, so a single - idle cluster still distributes the very first cold model evenly.""" - best = min((total_load(ep), loaded_count(ep)) for ep in eps) - tied = [ep for ep in eps if (total_load(ep), loaded_count(ep)) == best] - return random.choice(tied) - # Priority map: position in all_endpoints list (lower = higher priority) ep_priority = {ep: i for i, ep in enumerate(all_endpoints)} @@ -286,11 +235,15 @@ async def choose_endpoint(model: str, reserve: bool = True, loaded_and_free.sort(key=utilization_ratio) selected = loaded_and_free[0] else: - # All endpoints here already have the model loaded, so there - # is no model-switching cost to optimise for. Pick the least - # *total*-loaded one (tie broken at random) so we steer away - # from a backend busy serving other models. - selected = pick_least_loaded(loaded_and_free) + # Sort ascending for load balancing — all endpoints here already have the + # model loaded, so there is no model-switching cost to optimise for. + loaded_and_free.sort(key=tracking_usage) + # When all candidates are equally idle, randomise to avoid always picking + # the first entry in a stable sort. + if all(tracking_usage(ep) == 0 for ep in loaded_and_free): + selected = random.choice(loaded_and_free) + else: + selected = loaded_and_free[0] else: # 4️⃣ Endpoints among the candidates that simply have a free slot endpoints_with_free_slot = [ @@ -304,10 +257,14 @@ async def choose_endpoint(model: str, reserve: bool = True, endpoints_with_free_slot.sort(key=utilization_ratio) selected = endpoints_with_free_slot[0] else: - # Prefer the endpoint with the lowest *total* load so the - # cold-start cost lands on genuinely idle hardware rather - # than a backend already busy with a different model. - selected = pick_least_loaded(endpoints_with_free_slot) + # Sort by total endpoint load (ascending) to prefer idle endpoints. + endpoints_with_free_slot.sort( + key=lambda ep: sum(usage_counts.get(ep, {}).values()) + ) + if all(tracking_usage(ep) == 0 for ep in endpoints_with_free_slot): + selected = random.choice(endpoints_with_free_slot) + else: + selected = endpoints_with_free_slot[0] else: # 5️⃣ All candidate endpoints are saturated – pick the least-busy one (will queue) if config.priority_routing: @@ -323,23 +280,15 @@ async def choose_endpoint(model: str, reserve: bool = True, if reserve: usage_counts[selected][tracking_model] += 1 snapshot = _capture_snapshot() - # The slot is now reserved. Any failure (including CancelledError on client - # disconnect) between here and `return` would otherwise leak it — the caller - # never receives (endpoint, tracking_model) and so can never decrement it. - try: - if snapshot is not None: - await _distribute_snapshot(snapshot) - # Record / refresh affinity *after* releasing usage_lock. - if reserve and config.conversation_affinity and affinity_key: - expires_at = time.monotonic() + config.conversation_affinity_ttl - async with _affinity_lock: - _affinity_map[affinity_key] = (selected, model, expires_at) - if len(_affinity_map) > _AFFINITY_MAX_ENTRIES: - now = time.monotonic() - for k in [k for k, v in _affinity_map.items() if v[2] < now]: - _affinity_map.pop(k, None) - except BaseException: - if reserve: - await decrement_usage(selected, tracking_model) - raise + if snapshot is not None: + await _distribute_snapshot(snapshot) + # Record / refresh affinity *after* releasing usage_lock. + if reserve and config.conversation_affinity and affinity_key: + expires_at = time.monotonic() + config.conversation_affinity_ttl + async with _affinity_lock: + _affinity_map[affinity_key] = (selected, model, expires_at) + if len(_affinity_map) > _AFFINITY_MAX_ENTRIES: + now = time.monotonic() + for k in [k for k, v in _affinity_map.items() if v[2] < now]: + _affinity_map.pop(k, None) return selected, tracking_model diff --git a/test/config_test.yaml b/test/config_test.yaml index ed96542..30f2fa3 100644 --- a/test/config_test.yaml +++ b/test/config_test.yaml @@ -4,14 +4,10 @@ endpoints: llama_server_endpoints: - http://192.168.0.51:12434/v1 -llama_swap_endpoints: - - http://192.168.0.51:12435/v1 - max_concurrent_connections: 2 api_keys: "http://192.168.0.51:12434": "ollama" "http://192.168.0.51:12434/v1": "llama" - "http://192.168.0.51:12435/v1": "llama-swap" cache_enabled: false diff --git a/test/conftest.py b/test/conftest.py index da7dacf..c5142da 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -57,7 +57,6 @@ def mock_config(): cfg = MagicMock() cfg.endpoints = [TEST_OLLAMA] cfg.llama_server_endpoints = [TEST_LLAMA] - cfg.llama_swap_endpoints = [] cfg.api_keys = {TEST_OLLAMA: "ollama", TEST_LLAMA: "llama"} cfg.max_concurrent_connections = 2 cfg.router_api_key = None @@ -71,7 +70,6 @@ def mock_config_no_llama(): cfg = MagicMock() cfg.endpoints = [TEST_OLLAMA] cfg.llama_server_endpoints = [] - cfg.llama_swap_endpoints = [] cfg.api_keys = {TEST_OLLAMA: "ollama"} cfg.max_concurrent_connections = 2 cfg.router_api_key = None @@ -85,7 +83,6 @@ def mock_config_with_key(): cfg = MagicMock() cfg.endpoints = [TEST_OLLAMA] cfg.llama_server_endpoints = [] - cfg.llama_swap_endpoints = [] cfg.api_keys = {} cfg.max_concurrent_connections = 2 cfg.router_api_key = "test-secret-key" diff --git a/test/test_choose_endpoint.py b/test/test_choose_endpoint.py index 2bc4985..ece609a 100644 --- a/test/test_choose_endpoint.py +++ b/test/test_choose_endpoint.py @@ -12,11 +12,10 @@ EP3 = "http://ep3:11434" LLAMA_EP = "http://llama:8080/v1" -def _make_cfg(endpoints, llama_eps=None, swap_eps=None, max_conn=2, endpoint_config=None, priority_routing=False): +def _make_cfg(endpoints, llama_eps=None, max_conn=2, endpoint_config=None, priority_routing=False): cfg = MagicMock() cfg.endpoints = endpoints cfg.llama_server_endpoints = llama_eps or [] - cfg.llama_swap_endpoints = swap_eps or [] cfg.api_keys = {} cfg.max_concurrent_connections = max_conn cfg.endpoint_config = endpoint_config or {} @@ -47,27 +46,6 @@ class TestChooseEndpointBasic: assert ep == EP1 assert tracking == "llama3.2:latest" - async def test_llama_swap_endpoint_is_a_candidate(self): - swap_ep = "http://swap:8080/v1" - cfg = _make_cfg([EP1], swap_eps=[swap_ep]) - - async def available(ep, *_): - # Only the llama-swap backend advertises this model - return {"org/model:Q4_K_M"} if ep == swap_ep else set() - - async def loaded(ep): - return {"org/model:Q4_K_M"} if ep == swap_ep else set() - - with ( - patch.object(router, "config", cfg), - patch.object(router.fetch, "available_models", side_effect=available), - patch.object(router.fetch, "loaded_models", side_effect=loaded), - ): - ep, tracking = await router.choose_endpoint("org/model:Q4_K_M") - assert ep == swap_ep - # llama-swap models are tracked under their normalized name - assert tracking == "model" - async def test_raises_when_no_endpoint_has_model(self): cfg = _make_cfg([EP1, EP2]) with ( @@ -107,64 +85,6 @@ class TestChooseEndpointBasic: ep, _ = await router.choose_endpoint("llama3.2:latest") assert ep in (EP1, EP2) - async def test_cold_model_avoids_backend_busy_with_other_model(self): - # Regression: heterogeneous cluster. A cold model B (loaded nowhere) - # must not be routed to a backend already serving a *different* model - # while other backends sit idle. The step-4 idle check used to look at - # per-model usage (zero everywhere for B) and discard the total-load - # ranking, so B could land on the busy backend at random. - cfg = _make_cfg([EP1, EP2, EP3], max_conn=4) - - async def available(ep, *_): - return {"model-a:latest", "model-b:latest"} - - # EP3 is busy with model A; EP1 and EP2 are completely idle. Model B - # is loaded nowhere. - router.usage_counts[EP3]["model-a:latest"] = 1 - - with ( - patch.object(router, "config", cfg), - patch.object(router.fetch, "available_models", side_effect=available), - patch.object(router.fetch, "loaded_models", AsyncMock(return_value=set())), - ): - # Run repeatedly: the busy backend must be excluded every time, - # the idle two share the load at random. - for _ in range(50): - ep, _ = await router.choose_endpoint("model-b:latest", reserve=False) - assert ep in (EP1, EP2) - assert ep != EP3 - - async def test_two_cold_models_spread_across_backends(self): - # Regression: 3 backends all advertise all models. Two *different* - # cold models requested back-to-back must land on *different* - # backends. Once model-a is resident on the chosen backend (infinite - # keep-alive), its in-flight count drops back to 0 — so only the - # resident-model count distinguishes the backends. Without it, the - # second cold model would randomly re-collide on the busy backend. - cfg = _make_cfg([EP1, EP2, EP3], max_conn=4) - - async def available(ep, *_): - return {"model-a:latest", "model-b:latest"} - - # model-a finished loading on EP1 and stays resident; its request has - # completed so EP1 has zero in-flight load, same as EP2/EP3. - loaded = {EP1: {"model-a:latest"}, EP2: set(), EP3: set()} - - async def loaded_models(ep): - return loaded[ep] - - with ( - patch.object(router, "config", cfg), - patch.object(router.fetch, "available_models", side_effect=available), - patch.object(router.fetch, "loaded_models", side_effect=loaded_models), - ): - # A cold model-b must avoid EP1 (which already holds model-a) and - # go to one of the empty backends, every time. - for _ in range(50): - ep, _ = await router.choose_endpoint("model-b:latest", reserve=False) - assert ep in (EP2, EP3) - assert ep != EP1 - async def test_saturated_picks_least_busy(self): cfg = _make_cfg([EP1, EP2]) cfg.max_concurrent_connections = 1 @@ -279,57 +199,6 @@ class TestChooseEndpointModelNaming: assert ep == EP1 -class TestChooseEndpointAdvertisedAlignment: - """Regression for issue #128 — advertised-model sets must stay aligned with - their endpoints regardless of config ordering. A previous implementation - partitioned the probe tasks (non-OpenAI first, then OpenAI-compatible) while - the endpoint list kept config order, so an OpenAI-compatible endpoint listed - *before* an Ollama one inherited the Ollama model set and received requests - for models it does not serve (404).""" - - OPENAI_EP = "http://vllm:8000/v1" - - async def test_openai_endpoint_before_ollama_does_not_misroute(self): - # OpenAI-compatible endpoint deliberately listed FIRST in config. - cfg = _make_cfg([self.OPENAI_EP, EP1]) - - async def available(ep, *_): - if ep == self.OPENAI_EP: - return {"vllm-only-model"} - return {"qwen2.5-coder:1.5b-base"} # only the Ollama box advertises this - - async def loaded(ep): - return {"qwen2.5-coder:1.5b-base"} if ep == EP1 else set() - - with ( - patch.object(router, "config", cfg), - patch.object(router.fetch, "available_models", side_effect=available), - patch.object(router.fetch, "loaded_models", side_effect=loaded), - ): - ep, _ = await router.choose_endpoint("qwen2.5-coder:1.5b-base") - # Must route to the endpoint that actually advertises the model. - assert ep == EP1 - - async def test_openai_endpoint_before_ollama_still_routes_its_own_model(self): - cfg = _make_cfg([self.OPENAI_EP, EP1]) - - async def available(ep, *_): - if ep == self.OPENAI_EP: - return {"vllm-only-model"} - return {"qwen2.5-coder:1.5b-base"} - - async def loaded(ep): - return {"qwen2.5-coder:1.5b-base"} if ep == EP1 else set() - - with ( - patch.object(router, "config", cfg), - patch.object(router.fetch, "available_models", side_effect=available), - patch.object(router.fetch, "loaded_models", side_effect=loaded), - ): - ep, _ = await router.choose_endpoint("vllm-only-model") - assert ep == self.OPENAI_EP - - class TestChooseEndpointLoadBalancing: async def test_random_selection_among_idle(self): cfg = _make_cfg([EP1, EP2, EP3]) diff --git a/test/test_compliance.py b/test/test_compliance.py new file mode 100644 index 0000000..754a345 --- /dev/null +++ b/test/test_compliance.py @@ -0,0 +1,230 @@ +"""Unit tests for the aissurance compliance plugin. + +Covers the pieces that don't require a live upstream or backend: + * config/settings resolution and URL derivation + * HMAC payload signing (determinism + tamper detection) + * buffer encrypt/decrypt round-trip + bound enforcement + * transport status-code → outcome classification + * collector field mapping against a mocked Router state +""" +import json +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from compliance.settings import ComplianceSettings +from compliance import payload as payload_mod +from compliance import crypto +from compliance.buffer import BufferStore +from compliance.transport import Outcome, _classify + + +# ── settings ──────────────────────────────────────────────────────────────── + +def test_url_derivation(): + s = ComplianceSettings( + server_url="https://www.aissurance.eu/api/v1/discovery/receive", + tenant_id="tid-123", + ) + assert s.receive_url.endswith("/discovery/receive") + assert s.health_url == "https://www.aissurance.eu/api/v1/discovery/health/tid-123" + assert s.config_report_url == "https://www.aissurance.eu/api/v1/discovery/config/report" + + +def test_health_url_without_tenant(): + s = ComplianceSettings(tenant_id=None) + assert s.health_url.endswith("/health/unknown") + + +def test_validation_errors(): + assert ComplianceSettings(api_key=None).validation_errors() # missing key + assert ComplianceSettings(api_key="k", server_url="http://x").validation_errors() # not https + assert ComplianceSettings(api_key="k", polling_interval=5).validation_errors() # too fast + assert ComplianceSettings( + api_key="k", server_url="https://x/receive", polling_interval=300 + ).validation_errors() == [] + + +def test_empty_string_env_becomes_none(): + # config.yaml ${VAR} expansion yields "" for unset vars. + s = ComplianceSettings(api_key="", tenant_id="") + assert s.api_key is None + assert s.tenant_id is None + + +# ── payload signing ───────────────────────────────────────────────────────── + +def _cfg(**kw): + base = dict(api_key="secret-key", tenant_id="t1", router_id="r1", + server_url="https://x/api/v1/discovery/receive") + base.update(kw) + return ComplianceSettings(**base) + + +def test_signature_deterministic_and_key_order_independent(): + cfg = _cfg() + snap = {"timestamp": "2026-01-01T00:00:00Z", "models": [{"b": 1, "a": 2}], + "endpoints": [], "telemetry": {"x": 1}} + p1 = payload_mod.build_payload(cfg, snap) + p2 = payload_mod.build_payload(cfg, snap) + assert p1["hmac_signature"] == p2["hmac_signature"] + assert p1["hmac_signature"].startswith("sha256=") + + +def test_signature_detects_tampering(): + cfg = _cfg() + snap = {"timestamp": "t", "models": [], "endpoints": [], "telemetry": {}} + p = payload_mod.build_payload(cfg, snap) + sig = p.pop("hmac_signature") + assert payload_mod.sign(p, cfg) == sig # untouched verifies + p["router_id"] = "evil" + assert payload_mod.sign(p, cfg) != sig # tampered does not + + +def test_hmac_secret_overrides_api_key(): + snap = {"timestamp": "t", "models": [], "endpoints": [], "telemetry": {}} + a = payload_mod.build_payload(_cfg(), snap) + b = payload_mod.build_payload(_cfg(hmac_secret="other"), snap) + assert a["hmac_signature"] != b["hmac_signature"] + + +def test_split_into_batches_resigns(): + cfg = _cfg() + snap = {"timestamp": "t", "models": [{"name": f"m{i}"} for i in range(5)], + "endpoints": [], "telemetry": {}} + p = payload_mod.build_payload(cfg, snap) + batches = payload_mod.split_into_batches(p, batch_size=2, cfg=cfg) + assert len(batches) == 3 + assert [len(b["models"]) for b in batches] == [2, 2, 1] + for b in batches: + sig = b.pop("hmac_signature") + assert payload_mod.sign(b, cfg) == sig + + +# ── buffer crypto + store ─────────────────────────────────────────────────── + +def test_crypto_roundtrip(): + ct = crypto.encrypt(b"hello", "key-1") + assert ct != b"hello" + assert crypto.decrypt(ct, "key-1") == b"hello" + assert crypto.decrypt(ct, "wrong-key") is None # wrong key -> None, no raise + + +def test_buffer_save_load_roundtrip(tmp_path): + cfg = _cfg(buffer_dir=str(tmp_path / "buf")) + store = BufferStore(cfg) + body = {"router_id": "r1", "models": [{"name": "x"}]} + path = store.save(body, timestamp_ms=1000) + assert path is not None and path.exists() + # raw file must not contain plaintext + assert b"router_id" not in path.read_bytes() + assert store.count() == 1 + assert store.load(path) == body + + +def test_buffer_replay_order_and_bound(tmp_path): + cfg = _cfg(buffer_dir=str(tmp_path / "buf"), max_buffer_payloads=3) + store = BufferStore(cfg) + for ts in (5000, 1000, 3000): + store.save({"ts": ts}, timestamp_ms=ts) + pending = store.pending() + loaded_ts = [store.load(p)["ts"] for p in pending] + assert loaded_ts == [1000, 3000, 5000] # oldest-first + + # exceeding the bound drops the oldest + for ts in (7000, 9000): + store.save({"ts": ts}, timestamp_ms=ts) + remaining = [store.load(p)["ts"] for p in store.pending()] + assert 1000 not in remaining + assert len(remaining) == 3 + + +def test_buffer_drops_corrupt_file(tmp_path): + cfg = _cfg(buffer_dir=str(tmp_path / "buf")) + store = BufferStore(cfg) + p = store.save({"a": 1}, timestamp_ms=1) + p.write_bytes(b"not-a-valid-fernet-token") + assert store.load(p) is None + assert not p.exists() # corrupt file removed + + +# ── transport classification ──────────────────────────────────────────────── + +@pytest.mark.parametrize("status,expected", [ + (202, Outcome.OK), + (200, Outcome.OK), + (401, Outcome.AUTH_FAILED), + (403, Outcome.AUTH_FAILED), + (429, Outcome.RATE_LIMITED), + (413, Outcome.TOO_LARGE), + (500, Outcome.RETRY), + (502, Outcome.RETRY), +]) +def test_status_classification(status, expected): + assert _classify(status) == expected + + +# ── collector field mapping ───────────────────────────────────────────────── + +async def test_collector_maps_endpoints_and_models(): + from compliance import collector + + cfg = SimpleNamespace( + endpoints=["http://ollama:11434"], + llama_server_endpoints=["https://llama:8000/v1"], + api_keys={"https://llama:8000/v1": "k"}, + endpoint_config={}, + max_concurrent_connections=4, + ) + + async def fake_health(ep, timeout=5): + return {"status": "ok"} + + async def fake_available(ep, api_key=None): + if "ollama" in ep: + return {"mistral:7b"} + return {"org/gpt-oss-20b:Q4_K_M"} + + async def fake_details(ep, route, detail): + return [{"name": "mistral:7b", "size": 4_100_000_000, + "details": {"quantization_level": "Q4_0"}}] + + fake_db = SimpleNamespace( + get_last_used_map=lambda: _async_return({}), + get_token_totals_since=lambda ts: _async_return((10, 20)), + ) + + with patch.object(collector, "_endpoint_health", fake_health), \ + patch.object(collector.fetch, "available_models", fake_available), \ + patch.object(collector.fetch, "endpoint_details", fake_details), \ + patch.object(collector, "get_db", lambda: fake_db), \ + patch.object(collector, "get_llm_cache", lambda: None), \ + patch.dict(collector.usage_counts, {"http://ollama:11434": {"mistral:7b": 2}}, clear=True), \ + patch.object(collector, "get_max_connections", lambda ep: 4): + snap = await collector.collect_snapshot(cfg, started_at=0.0) + + eps = {e["url"]: e for e in snap["endpoints"]} + assert eps["http://ollama:11434"]["type"] == "ollama" + assert eps["http://ollama:11434"]["tls_enabled"] is False + assert eps["http://ollama:11434"]["concurrent_connections"] == 2 + assert eps["https://llama:8000/v1"]["type"] == "vllm" + assert eps["https://llama:8000/v1"]["tls_enabled"] is True + assert eps["https://llama:8000/v1"]["auth_method"] == "bearer" + + models = {m["name"]: m for m in snap["models"]} + assert models["mistral:7b"]["quantization"] == "Q4_0" + assert models["mistral:7b"]["size_gb"] == pytest.approx(4.1, rel=1e-3) + assert models["gpt-oss-20b"]["quantization"] == "Q4_K_M" + # un-instrumented fields are null, not fabricated + assert models["mistral:7b"]["avg_latency_ms"] is None + assert models["mistral:7b"]["request_count_24h"] is None + + tel = snap["telemetry"] + assert tel["active_connections"] == 2 + assert tel["total_bytes_in_24h"] == 10 + assert tel["total_requests_24h"] is None + + +async def _async_return(value): + return value diff --git a/test/test_fetch.py b/test/test_fetch.py index ef7e30e..76121e1 100644 --- a/test/test_fetch.py +++ b/test/test_fetch.py @@ -20,11 +20,10 @@ MOCK_OLLAMA_EP = "http://mock-ollama:11434" MOCK_LLAMA_EP = "http://mock-llama:8080/v1" -def _make_cfg(ollama_eps=None, llama_eps=None, swap_eps=None, api_keys=None): +def _make_cfg(ollama_eps=None, llama_eps=None, api_keys=None): cfg = MagicMock() cfg.endpoints = ollama_eps or [MOCK_OLLAMA_EP] cfg.llama_server_endpoints = llama_eps or [MOCK_LLAMA_EP] - cfg.llama_swap_endpoints = swap_eps or [] cfg.api_keys = api_keys or {} cfg.max_concurrent_connections = 2 cfg.router_api_key = None @@ -229,30 +228,6 @@ class TestFetchLoadedModels: models = await router.fetch.loaded_models(MOCK_LLAMA_EP) assert "always-on-model" in models - async def test_llama_swap_reads_running_state_ready(self): - # llama-swap omits the /v1/models status field, so loaded workers come - # from /running (a root route — the /v1 suffix must be stripped). - swap_ep = "http://mock-swap:8080/v1" - cfg = _make_cfg(llama_eps=[], swap_eps=[swap_ep]) - with patch.object(router, "config", cfg), mock_probe() as m: - m.add_get( - "http://mock-swap:8080/running", - payload={"running": [ - {"model": "org/ready-model:Q4_K_M", "state": "ready"}, - {"model": "org/starting-model:Q8_0", "state": "starting"}, - ]}, - ) - models = await router.fetch.loaded_models(swap_ep) - assert models == {"org/ready-model:Q4_K_M"} - - async def test_llama_swap_records_error_on_failure(self): - swap_ep = "http://mock-swap:8080/v1" - cfg = _make_cfg(llama_eps=[], swap_eps=[swap_ep]) - with patch.object(router, "config", cfg), mock_probe() as m: - m.add_get("http://mock-swap:8080/running", status=502, payload={}) - await router.fetch.loaded_models(swap_ep) - assert swap_ep in router._loaded_error_cache - async def test_returns_empty_on_error(self): cfg = _make_cfg(ollama_eps=[MOCK_OLLAMA_EP], llama_eps=[]) with patch.object(router, "config", cfg), mock_probe() as m: @@ -260,20 +235,12 @@ class TestFetchLoadedModels: models = await router.fetch.loaded_models(MOCK_OLLAMA_EP) assert models == set() - async def test_ext_openai_reports_advertised_as_loaded(self): - # Issue #128: external OpenAI-compatible backends (vLLM, OpenAI, …) keep - # every advertised model permanently resident, so loaded_models mirrors - # the advertised set rather than returning empty (which made them lose - # choose_endpoint's "loaded & free" preference to Ollama boxes). + async def test_ext_openai_always_empty(self): ext_ep = "https://api.openai.com/v1" cfg = _make_cfg(ollama_eps=[ext_ep], llama_eps=[]) - with patch.object(router, "config", cfg), mock_probe() as m: - m.add_get( - f"{ext_ep}/models", - payload={"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini"}]}, - ) + with patch.object(router, "config", cfg): models = await router.fetch.loaded_models(ext_ep) - assert models == {"gpt-4o", "gpt-4o-mini"} + assert models == set() async def test_caches_result(self): cfg = _make_cfg(ollama_eps=[MOCK_OLLAMA_EP], llama_eps=[]) diff --git a/test/test_llama_swap.py b/test/test_llama_swap.py deleted file mode 100644 index 74197d4..0000000 --- a/test/test_llama_swap.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Tests for llama-swap specific behavior: unload dispatch + /upstream resolution.""" -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -import router -import backends.control as control -import api.openai as openai_api -import api.ollama as ollama_api - -SWAP_EP = "http://swap:8080/v1" -SERVER_EP = "http://server:8080/v1" - - -def _cfg(*, server=None, swap=None, api_keys=None): - cfg = MagicMock() - cfg.endpoints = [] - cfg.llama_server_endpoints = server or [] - cfg.llama_swap_endpoints = swap or [] - cfg.api_keys = api_keys or {} - return cfg - - -class _RecordingSession: - """Captures the most recent ``post`` call and returns a 200 response.""" - - def __init__(self, status=200): - self.calls = [] - self._status = status - - def post(self, url, **kwargs): - self.calls.append((url, kwargs)) - resp = MagicMock() - resp.status = self._status - - class _Ctx: - async def __aenter__(self_): - return resp - - async def __aexit__(self_, *exc): - return False - - return _Ctx() - - -class TestUnloadDispatch: - async def test_llama_swap_uses_path_param(self): - sess = _RecordingSession() - cfg = _cfg(swap=[SWAP_EP]) - with ( - patch.object(router, "config", cfg), - patch.object(control, "get_probe_session", lambda ep: sess), - ): - ok = await control.unload_model(SWAP_EP, "org/model:Q4_K_M") - assert ok is True - url, kwargs = sess.calls[0] - # /v1 stripped, model id is a path param, no JSON body - assert url == "http://swap:8080/api/models/unload/org/model:Q4_K_M" - assert kwargs.get("json") is None - - async def test_llama_server_uses_body(self): - sess = _RecordingSession() - cfg = _cfg(server=[SERVER_EP]) - with ( - patch.object(router, "config", cfg), - patch.object(control, "get_probe_session", lambda ep: sess), - ): - ok = await control.unload_model(SERVER_EP, "org/model:Q4_K_M") - assert ok is True - url, kwargs = sess.calls[0] - assert url == "http://server:8080/models/unload" - assert kwargs.get("json") == {"model": "org/model:Q4_K_M"} - - async def test_unload_failure_returns_false(self): - sess = _RecordingSession(status=500) - cfg = _cfg(swap=[SWAP_EP]) - with ( - patch.object(router, "config", cfg), - patch.object(control, "get_probe_session", lambda ep: sess), - ): - ok = await control.unload_model(SWAP_EP, "m") - assert ok is False - - -class TestUpstreamResolution: - async def test_resolves_endpoint_that_advertises_model(self): - cfg = _cfg(swap=[SWAP_EP]) - with ( - patch.object(openai_api, "get_config", lambda: cfg), - patch.object(openai_api.fetch, "available_models", - AsyncMock(return_value={"org/model:Q4_K_M"})), - ): - ep = await openai_api._resolve_llama_swap_endpoint("org/model:Q4_K_M") - assert ep == SWAP_EP - - async def test_returns_none_when_unserved(self): - cfg = _cfg(swap=[SWAP_EP]) - with ( - patch.object(openai_api, "get_config", lambda: cfg), - patch.object(openai_api.fetch, "available_models", - AsyncMock(return_value=set())), - ): - ep = await openai_api._resolve_llama_swap_endpoint("missing") - assert ep is None - - async def test_returns_none_without_swap_endpoints(self): - cfg = _cfg(swap=[]) - with patch.object(openai_api, "get_config", lambda: cfg): - ep = await openai_api._resolve_llama_swap_endpoint("any") - assert ep is None - - -class TestCtxSizeFromCmd: - """ctx-size parsing from a /running worker's launch `cmd` string.""" - - def test_parses_long_flag(self): - cmd = ("llama-server --port 5818\n -hf unsloth/gpt-oss-20b-GGUF:F16\n" - " --ctx-size 131072\n --temp 1.0\n") - assert ollama_api._ctx_size_from_cmd(cmd) == 131072 - - def test_parses_short_flag(self): - assert ollama_api._ctx_size_from_cmd("llama-server -c 8192 --port 1") == 8192 - - def test_parses_equals_form(self): - assert ollama_api._ctx_size_from_cmd("llama-server --ctx-size=4096") == 4096 - - def test_returns_none_when_absent(self): - assert ollama_api._ctx_size_from_cmd("llama-server --port 5818") is None - - def test_returns_none_for_empty(self): - assert ollama_api._ctx_size_from_cmd("") is None diff --git a/test/test_messages.py b/test/test_messages.py deleted file mode 100644 index 5c8700e..0000000 --- a/test/test_messages.py +++ /dev/null @@ -1,477 +0,0 @@ -"""Tests for the Anthropic Messages API support (api/messages.py + requests/anthropic.py). - -Covers the pure translation layer, the translated (Ollama-style) and native -(configured-anthropic-endpoint) backend paths, thinking→reasoning mapping, -streaming event shape, token counting, and the nomyo-cache reflection. -""" -from contextlib import ExitStack, contextmanager -from types import SimpleNamespace as NS -from unittest.mock import AsyncMock, MagicMock, patch - -import orjson -import pytest - -import router -from api import messages as api_messages -from requests import anthropic as at - - -# ────────────────────────────────────────────────────────────────────────────── -# Pure translation unit tests (no app / no I/O) -# ────────────────────────────────────────────────────────────────────────────── - -class TestRequestTranslation: - def test_system_string(self): - chat = at.anthropic_messages_to_chat("be brief", [{"role": "user", "content": "hi"}]) - assert chat[0] == {"role": "system", "content": "be brief"} - assert chat[1] == {"role": "user", "content": "hi"} - - def test_system_blocks_join(self): - chat = at.anthropic_messages_to_chat( - [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], - [{"role": "user", "content": "hi"}]) - assert chat[0] == {"role": "system", "content": "a\n\nb"} - - def test_image_base64_to_data_url(self): - chat = at.anthropic_messages_to_chat(None, [{ - "role": "user", "content": [ - {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAA"}}, - {"type": "text", "text": "what?"}, - ]}]) - parts = chat[0]["content"] - assert parts[0] == {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}} - assert parts[1] == {"type": "text", "text": "what?"} - - def test_tool_use_and_result_roundtrip(self): - chat = at.anthropic_messages_to_chat(None, [ - {"role": "assistant", "content": [ - {"type": "text", "text": "calling"}, - {"type": "tool_use", "id": "toolu_1", "name": "get", "input": {"x": 1}}]}, - {"role": "user", "content": [ - {"type": "tool_result", "tool_use_id": "toolu_1", "content": "42"}]}, - ]) - assert chat[0]["role"] == "assistant" - assert chat[0]["content"] == "calling" - assert chat[0]["tool_calls"][0]["id"] == "toolu_1" - assert chat[0]["tool_calls"][0]["function"]["arguments"] == '{"x":1}' - assert chat[1] == {"role": "tool", "tool_call_id": "toolu_1", "content": "42"} - - def test_tool_result_blocks_flattened(self): - chat = at.anthropic_messages_to_chat(None, [{ - "role": "user", "content": [ - {"type": "tool_result", "tool_use_id": "t1", - "content": [{"type": "text", "text": "line"}]}]}]) - assert chat[0] == {"role": "tool", "tool_call_id": "t1", "content": "line"} - - def test_tools_and_choice(self): - sp = at.anthropic_to_chat_send_params( - {"max_tokens": 10, "tools": [{"name": "get", "description": "d", - "input_schema": {"type": "object"}}], - "tool_choice": {"type": "tool", "name": "get"}}, - [], "m") - assert sp["tools"] == [{"type": "function", "function": { - "name": "get", "description": "d", "parameters": {"type": "object"}}}] - assert sp["tool_choice"] == {"type": "function", "function": {"name": "get"}} - - def test_tool_choice_variants(self): - assert at.tool_choice_anthropic_to_chat({"type": "auto"}) == "auto" - assert at.tool_choice_anthropic_to_chat({"type": "any"}) == "required" - assert at.tool_choice_anthropic_to_chat({"type": "none"}) == "none" - - def test_server_tools_dropped(self): - # web_search has a type but no input_schema → not runnable locally - assert at.tools_anthropic_to_chat([{"type": "web_search_20260209", "name": "web_search"}]) is None - - def test_stop_sequences_and_sampling(self): - sp = at.anthropic_to_chat_send_params( - {"max_tokens": 5, "stop_sequences": ["X"], "temperature": 0.2, "top_p": 0.9}, [], "m") - assert sp["stop"] == ["X"] - assert sp["temperature"] == 0.2 - assert sp["top_p"] == 0.9 - - def test_thinking_maps_to_reasoning_effort(self): - assert at.anthropic_to_chat_send_params( - {"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 1000}}, [], "m" - )["reasoning_effort"] == "low" - assert at.anthropic_to_chat_send_params( - {"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 4096}}, [], "m" - )["reasoning_effort"] == "medium" - assert at.anthropic_to_chat_send_params( - {"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 20000}}, [], "m" - )["reasoning_effort"] == "high" - assert "reasoning_effort" not in at.anthropic_to_chat_send_params( - {"max_tokens": 5, "thinking": {"type": "disabled"}}, [], "m") - - -class TestResponseTranslation: - def test_content_blocks_order(self): - blocks = at.chat_message_to_content_blocks({ - "role": "assistant", "reasoning_content": "hmm", "content": "answer", - "tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": '{"a":1}'}}]}) - assert [b["type"] for b in blocks] == ["thinking", "text", "tool_use"] - assert blocks[0]["thinking"] == "hmm" - assert blocks[1]["text"] == "answer" - assert blocks[2] == {"type": "tool_use", "id": "c1", "name": "f", "input": {"a": 1}} - - def test_stop_reason_mapping(self): - assert at.finish_reason_to_stop_reason("stop") == "end_turn" - assert at.finish_reason_to_stop_reason("length") == "max_tokens" - 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_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} - - -# ────────────────────────────────────────────────────────────────────────────── -# Streaming translator (pure) -# ────────────────────────────────────────────────────────────────────────────── - -def _chunk(content=None, tool_calls=None, reasoning=None, finish_reason=None): - delta = NS(content=content, tool_calls=tool_calls) - if reasoning is not None: - delta.reasoning_content = reasoning - return NS(choices=[NS(delta=delta, finish_reason=finish_reason)], usage=None) - - -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): - frames = [] - async for sse in translator.events(gen): - frames.append(sse.decode()) - return _parse_sse("".join(frames)) - - -def _parse_sse(text): - out = [] - for frame in text.strip().split("\n\n"): - if not frame.strip(): - continue - etype = data = None - for line in frame.splitlines(): - if line.startswith("event: "): - etype = line[len("event: "):] - elif line.startswith("data: "): - data = orjson.loads(line[len("data: "):]) - out.append((etype, data)) - return out - - -class TestStreamTranslator: - async def test_text_stream(self): - async def gen(): - yield _chunk(content="Hel") - yield _chunk(content="lo", finish_reason="stop") - yield _usage_chunk(3, 5) - tr = at.ChatToMessagesStream("msg_1", "m") - events = await _collect(tr, gen()) - types = [e[0] for e in events] - assert types[0] == "message_start" - assert types[-1] == "message_stop" - assert "content_block_start" in types and "content_block_stop" in types - text = "".join(d["delta"]["text"] for t, d in events - if t == "content_block_delta" and d["delta"]["type"] == "text_delta") - assert text == "Hello" - md = [d for t, d in events if t == "message_delta"][0] - assert md["delta"]["stop_reason"] == "end_turn" - 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 ") - yield _chunk(reasoning="more") - yield _chunk(content="answer", finish_reason="stop") - tr = at.ChatToMessagesStream("msg_1", "m") - events = await _collect(tr, gen()) - # thinking block (index 0) then text block (index 1) - starts = [(d["index"], d["content_block"]["type"]) - for t, d in events if t == "content_block_start"] - assert starts == [(0, "thinking"), (1, "text")] - think = "".join(d["delta"]["thinking"] for t, d in events - if t == "content_block_delta" and d["delta"]["type"] == "thinking_delta") - assert think == "think more" - assert tr.content_blocks[0] == {"type": "thinking", "thinking": "think more"} - - async def test_tool_call_stream(self): - tc0 = NS(index=0, id="call_1", function=NS(name="lookup", arguments='{"q":')) - tc1 = NS(index=0, id=None, function=NS(name=None, arguments='"hi"}')) - - async def gen(): - yield _chunk(tool_calls=[tc0]) - yield _chunk(tool_calls=[tc1], finish_reason="tool_calls") - yield _usage_chunk(4, 2) - tr = at.ChatToMessagesStream("msg_1", "m") - events = await _collect(tr, gen()) - assert any(t == "content_block_start" and d["content_block"]["type"] == "tool_use" - for t, d in events) - partial = "".join(d["delta"]["partial_json"] for t, d in events - if t == "content_block_delta" and d["delta"]["type"] == "input_json_delta") - assert partial == '{"q":"hi"}' - md = [d for t, d in events if t == "message_delta"][0] - assert md["delta"]["stop_reason"] == "tool_use" - assert tr.content_blocks[-1] == {"type": "tool_use", "id": "call_1", - "name": "lookup", "input": {"q": "hi"}} - - -# ────────────────────────────────────────────────────────────────────────────── -# Cache-hit replay -# ────────────────────────────────────────────────────────────────────────────── - -class TestCacheReplay: - def test_message_object_to_sse_roundtrip(self): - msg = at.build_message_object( - message_id="msg_1", model="m", - content_blocks=[{"type": "text", "text": "hi"}], - stop_reason="end_turn", - usage=at.usage_chat_to_anthropic({"prompt_tokens": 2, "completion_tokens": 1})) - events = _parse_sse(at.message_object_to_sse(msg).decode()) - types = [e[0] for e in events] - assert types[0] == "message_start" - assert types[-1] == "message_stop" - text = "".join(d["delta"]["text"] for t, d in events - if t == "content_block_delta" and d["delta"]["type"] == "text_delta") - assert text == "hi" - - -# ────────────────────────────────────────────────────────────────────────────── -# Route-level tests -# ────────────────────────────────────────────────────────────────────────────── - -@contextmanager -def _enter(*cms): - with ExitStack() as stack: - for cm in cms: - stack.enter_context(cm) - yield - - -def _fake_completion(content="hello", usage=(3, 5), reasoning=None, tool_calls=None, - finish_reason="stop", cached=None): - md = {"role": "assistant", "content": content} - if reasoning is not None: - md["reasoning_content"] = reasoning - if tool_calls is not None: - 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 = usage_dump - return NS(choices=[NS(message=msg, finish_reason=finish_reason)], usage=usage_obj) - - -def _patch_backend(native=False, endpoint="http://ollama:11434", cache=None): - return ( - patch.object(api_messages, "choose_endpoint", - AsyncMock(return_value=(endpoint, "test-model:latest"))), - patch.object(api_messages, "decrement_usage", AsyncMock()), - patch.object(api_messages, "is_anthropic_endpoint", return_value=native), - patch.object(api_messages, "_make_openai_client", return_value=MagicMock()), - patch.object(api_messages, "get_llm_cache", return_value=cache), - ) - - -class TestTranslatedRoute: - async def test_nonstream(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_messages, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("hello world")))): - resp = await client.post("/v1/messages", - json={"model": "test-model", "max_tokens": 100, - "messages": [{"role": "user", "content": "hi"}]}) - assert resp.status_code == 200 - body = resp.json() - assert body["type"] == "message" - assert body["role"] == "assistant" - assert body["content"] == [{"type": "text", "text": "hello world"}] - assert body["stop_reason"] == "end_turn" - 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", - json={"model": "m", "messages": [{"role": "user", "content": "hi"}]}) - assert resp.status_code == 400 - - async def test_nonstream_tool_use(self, client): - tc = [{"id": "c1", "function": {"name": "get", "arguments": '{"a":1}'}}] - with _enter(*_patch_backend(native=False), - patch.object(api_messages, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion( - content=None, tool_calls=tc, finish_reason="tool_calls")))): - resp = await client.post("/v1/messages", - json={"model": "m", "max_tokens": 50, - "messages": [{"role": "user", "content": "call it"}]}) - body = resp.json() - assert body["stop_reason"] == "tool_use" - assert body["content"][0] == {"type": "tool_use", "id": "c1", "name": "get", "input": {"a": 1}} - - async def test_stream_event_sequence(self, client): - async def _text(): - yield _chunk(content="Hi", finish_reason="stop") - yield _usage_chunk(3, 2) - with _enter(*_patch_backend(native=False), - patch.object(api_messages, "create_chat_with_retries", - AsyncMock(return_value=_text()))): - resp = await client.post("/v1/messages", - json={"model": "m", "max_tokens": 50, "stream": True, - "messages": [{"role": "user", "content": "hi"}]}) - assert resp.headers["content-type"].startswith("text/event-stream") - events = _parse_sse(resp.content.decode()) - types = [e[0] for e in events] - assert types[0] == "message_start" and types[-1] == "message_stop" - text = "".join(d["delta"]["text"] for t, d in events - if t == "content_block_delta" and d["delta"]["type"] == "text_delta") - assert text == "Hi" - - async def test_thinking_passed_as_reasoning_effort(self, client): - captured = {} - - async def _spy(oclient, send_params, endpoint, model, tracking_model): - captured.update(send_params) - return _fake_completion("ok") - - with _enter(*_patch_backend(native=False), - patch.object(api_messages, "create_chat_with_retries", _spy)): - await client.post("/v1/messages", - json={"model": "m", "max_tokens": 50, - "thinking": {"type": "enabled", "budget_tokens": 20000}, - "messages": [{"role": "user", "content": "hi"}]}) - assert captured["reasoning_effort"] == "high" - - -class TestCacheRoute: - async def test_hit_reports_cache_read(self, client): - stored = at.build_message_object( - message_id="msg_old", model="m", - content_blocks=[{"type": "text", "text": "cached"}], stop_reason="end_turn", - usage=at.usage_chat_to_anthropic({"prompt_tokens": 9, "completion_tokens": 4})) - fake_cache = MagicMock() - fake_cache.get_chat = AsyncMock(return_value=orjson.dumps(stored)) - with _enter(*_patch_backend(native=False, cache=fake_cache)): - resp = await client.post("/v1/messages", - json={"model": "m", "max_tokens": 50, "nomyo": {"cache": True}, - "messages": [{"role": "user", "content": "hi"}]}) - body = resp.json() - assert body["content"] == [{"type": "text", "text": "cached"}] - assert body["usage"]["cache_read_input_tokens"] == 9 - assert body["usage"]["input_tokens"] == 0 - assert body["id"].startswith("msg_") and body["id"] != "msg_old" - - -class TestNativeRoute: - def _fake_client(self, *, post_return=None, stream_frames=None, status=200): - client = MagicMock() - if post_return is not None: - resp = MagicMock() - resp.status_code = status - resp.json.return_value = post_return - client.post = AsyncMock(return_value=resp) - if stream_frames is not None: - class _Stream: - async def __aenter__(self_): - return self_ - async def __aexit__(self_, *a): - return False - async def aiter_bytes(self_): - for f in stream_frames: - yield f - client.stream = MagicMock(return_value=_Stream()) - return client - - async def test_nonstream_passthrough(self, client): - upstream = {"id": "msg_upstream", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "native hi"}], - "stop_reason": "end_turn", - "usage": {"input_tokens": 2, "output_tokens": 3}} - fake = self._fake_client(post_return=upstream) - with _enter(*_patch_backend(native=True, endpoint="https://api.anthropic.com"), - patch.object(api_messages, "_anthropic_http_client", return_value=fake)): - resp = await client.post("/v1/messages", - json={"model": "claude-x", "max_tokens": 50, - "messages": [{"role": "user", "content": "hi"}]}) - assert resp.status_code == 200 - assert resp.json()["content"][0]["text"] == "native hi" - # request forwarded verbatim with stream disabled, nomyo stripped - sent = fake.post.call_args.kwargs["json"] - assert sent["stream"] is False and "nomyo" not in sent - headers = fake.post.call_args.kwargs["headers"] - assert headers["x-api-key"] and headers["anthropic-version"] - - async def test_stream_passthrough(self, client): - frames = [ - b'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":5}}}\n\n', - b'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":7}}\n\n', - b'event: message_stop\ndata: {"type":"message_stop"}\n\n', - ] - fake = self._fake_client(stream_frames=frames) - tracked = [] - with _enter(*_patch_backend(native=True, endpoint="https://api.anthropic.com"), - patch.object(api_messages, "_anthropic_http_client", return_value=fake), - patch.object(api_messages, "_track", - AsyncMock(side_effect=lambda *a: tracked.append(a)))): - resp = await client.post("/v1/messages", - json={"model": "claude-x", "max_tokens": 50, "stream": True, - "messages": [{"role": "user", "content": "hi"}]}) - body = resp.content.decode() - assert "message_start" in body and "message_stop" in body - # usage parsed out of the proxied stream for token tracking - assert tracked and tracked[0][2] == 5 and tracked[0][3] == 7 - - -class TestCountTokens: - async def test_local_estimate(self, client): - with _enter(patch.object(api_messages, "choose_endpoint", - AsyncMock(return_value=("http://ollama:11434", "m"))), - patch.object(api_messages, "is_anthropic_endpoint", return_value=False)): - resp = await client.post("/v1/messages/count_tokens", - json={"model": "m", - "messages": [{"role": "user", "content": "count me"}]}) - assert resp.status_code == 200 - assert isinstance(resp.json()["input_tokens"], int) - assert resp.json()["input_tokens"] > 0 diff --git a/test/test_responses.py b/test/test_responses.py deleted file mode 100644 index 73d3ff2..0000000 --- a/test/test_responses.py +++ /dev/null @@ -1,460 +0,0 @@ -"""Tests for the OpenAI Responses API support (api/responses.py + requests/responses.py). - -Covers the pure translation layer, the translated (Ollama-style) and native -(external-OpenAI) backend paths, conversation storage / chaining, background mode, -and the retrieve / delete / cancel routes. -""" -import asyncio -from contextlib import ExitStack, contextmanager -from types import SimpleNamespace as NS -from unittest.mock import AsyncMock, MagicMock, patch - -import orjson -import pytest - -import router -from api import responses as api_responses -from requests import responses as rt - - -# ────────────────────────────────────────────────────────────────────────────── -# Pure translation unit tests (no app / no I/O) -# ────────────────────────────────────────────────────────────────────────────── - -class TestTranslationInputToMessages: - def test_string_input(self): - msgs = rt.responses_input_to_messages("hello") - assert msgs == [{"role": "user", "content": "hello"}] - - def test_instructions_become_system(self): - msgs = rt.responses_input_to_messages("hi", instructions="be brief") - assert msgs[0] == {"role": "system", "content": "be brief"} - assert msgs[1] == {"role": "user", "content": "hi"} - - def test_item_list_text_and_image(self): - items = [{ - "type": "message", "role": "user", - "content": [ - {"type": "input_text", "text": "describe"}, - {"type": "input_image", "image_url": "http://x/y.png"}, - ], - }] - msgs = rt.responses_input_to_messages(items) - assert msgs[0]["role"] == "user" - assert msgs[0]["content"] == [ - {"type": "text", "text": "describe"}, - {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, - ] - - def test_single_text_part_collapses_to_string(self): - items = [{"type": "message", "role": "user", - "content": [{"type": "input_text", "text": "yo"}]}] - assert rt.responses_input_to_messages(items)[0]["content"] == "yo" - - def test_function_call_roundtrip(self): - items = [ - {"type": "function_call", "call_id": "c1", "name": "get", "arguments": "{\"x\":1}"}, - {"type": "function_call_output", "call_id": "c1", "output": "42"}, - ] - msgs = rt.responses_input_to_messages(items) - assert msgs[0]["role"] == "assistant" - assert msgs[0]["tool_calls"][0]["id"] == "c1" - assert msgs[0]["tool_calls"][0]["function"]["name"] == "get" - assert msgs[1] == {"role": "tool", "tool_call_id": "c1", "content": "42"} - - -class TestTranslationResponseDirection: - def test_chat_message_to_output_items_text(self): - items = rt.chat_message_to_output_items({"role": "assistant", "content": "hi there"}) - assert len(items) == 1 - assert items[0]["type"] == "message" - assert items[0]["content"][0] == {"type": "output_text", "text": "hi there", "annotations": []} - - def test_chat_message_to_output_items_tool_call(self): - items = rt.chat_message_to_output_items({ - "role": "assistant", "content": None, - "tool_calls": [{"id": "c9", "function": {"name": "f", "arguments": "{}"}}], - }) - assert items[0]["type"] == "function_call" - assert items[0]["call_id"] == "c9" - assert items[0]["name"] == "f" - - def test_usage_mapping(self): - u = rt.usage_chat_to_responses({"prompt_tokens": 7, "completion_tokens": 3}) - assert u == {"input_tokens": 7, "output_tokens": 3, "total_tokens": 10} - - def test_build_response_object_output_text(self): - items = rt.chat_message_to_output_items({"role": "assistant", "content": "abc"}) - obj = rt.build_response_object(response_id="resp_1", model="m", output_items=items) - assert obj["object"] == "response" - assert obj["output_text"] == "abc" - assert obj["status"] == "completed" - - def test_tools_responses_to_chat(self): - tools = [{"type": "function", "name": "f", "description": "d", "parameters": {"type": "object"}}] - chat_tools = rt.tools_responses_to_chat(tools) - assert chat_tools == [{"type": "function", - "function": {"name": "f", "description": "d", - "parameters": {"type": "object"}}}] - - def test_messages_to_responses_input(self): - instr, items = rt.messages_to_responses_input([ - {"role": "system", "content": "sys"}, - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "yo"}, - ]) - assert instr == "sys" - assert items[0] == {"role": "user", "content": [{"type": "input_text", "text": "hi"}]} - assert items[1] == {"role": "assistant", "content": [{"type": "output_text", "text": "yo"}]} - - -# ────────────────────────────────────────────────────────────────────────────── -# Fakes for backend generators -# ────────────────────────────────────────────────────────────────────────────── - -def _fake_completion(content="hello world", usage=(3, 5)): - msg = MagicMock() - msg.model_dump.return_value = {"role": "assistant", "content": content} - usage_obj = MagicMock() - usage_obj.model_dump.return_value = { - "prompt_tokens": usage[0], "completion_tokens": usage[1], "total_tokens": sum(usage)} - return NS(choices=[NS(message=msg)], usage=usage_obj) - - -def _chunk(content=None, tool_calls=None): - return NS(choices=[NS(delta=NS(content=content, tool_calls=tool_calls), - finish_reason=None)], usage=None) - - -def _usage_chunk(p, c): - return NS(choices=[], usage=NS(prompt_tokens=p, completion_tokens=c)) - - -def _text_chunks(): - async def _gen(): - yield _chunk(content="Hel") - yield _chunk(content="lo") - yield _usage_chunk(3, 5) - return _gen() - - -def _toolcall_chunks(): - tc0 = NS(index=0, id="call_1", function=NS(name="lookup", arguments='{"q":')) - tc1 = NS(index=0, id=None, function=NS(name=None, arguments='"hi"}')) - - async def _gen(): - yield _chunk(tool_calls=[tc0]) - yield _chunk(tool_calls=[tc1]) - yield _usage_chunk(4, 2) - return _gen() - - -class _FakeEvent: - def __init__(self, data): - self._data = data - - def model_dump(self): - return self._data - - -def _native_event_stream(): - async def _gen(): - yield _FakeEvent({"type": "response.created", - "response": {"id": "resp_openai", "status": "in_progress", "output": []}}) - yield _FakeEvent({"type": "response.output_text.delta", - "item_id": "msg_1", "output_index": 0, "delta": "hi"}) - yield _FakeEvent({"type": "response.completed", "response": { - "id": "resp_openai", "status": "completed", - "output": [{"type": "message", "role": "assistant", - "content": [{"type": "output_text", "text": "hi"}]}], - "usage": {"input_tokens": 2, "output_tokens": 1, "total_tokens": 3}}}) - return _gen() - - -def _sse_events(text): - """Split an SSE body into a list of (event_type, data_dict).""" - out = [] - for frame in text.strip().split("\n\n"): - if not frame.strip(): - continue - etype = data = None - for line in frame.splitlines(): - if line.startswith("event: "): - etype = line[len("event: "):] - elif line.startswith("data: "): - data = orjson.loads(line[len("data: "):]) - out.append((etype, data)) - return out - - -@contextmanager -def _enter(*cms): - """Enter a variable number of context managers (works with *unpacked tuples).""" - with ExitStack() as stack: - for cm in cms: - stack.enter_context(cm) - yield - - -def _patch_backend(native=False, endpoint="http://ollama:11434"): - """Context managers patching endpoint selection + client construction.""" - return ( - patch.object(api_responses, "choose_endpoint", - AsyncMock(return_value=(endpoint, "test-model:latest"))), - patch.object(api_responses, "decrement_usage", AsyncMock()), - patch.object(api_responses, "is_ext_openai_endpoint", return_value=native), - patch.object(api_responses, "_make_openai_client", return_value=MagicMock()), - patch.object(api_responses, "get_llm_cache", return_value=None), - ) - - -# ────────────────────────────────────────────────────────────────────────────── -# Translated path (Ollama-style backend) -# ────────────────────────────────────────────────────────────────────────────── - -class TestTranslatedPath: - async def test_nonstream(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("hello world")))): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", "store": False}) - assert resp.status_code == 200 - body = resp.json() - assert body["object"] == "response" - assert body["output_text"] == "hello world" - assert body["usage"] == {"input_tokens": 3, "output_tokens": 5, "total_tokens": 8} - assert body["id"].startswith("resp_") - - async def test_stream_event_sequence(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_text_chunks()))): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", - "stream": True, "store": False}) - assert resp.status_code == 200 - assert resp.headers["content-type"].startswith("text/event-stream") - events = _sse_events(resp.content.decode()) - types = [e[0] for e in events] - assert types[0] == "response.created" - assert "response.output_text.delta" in types - assert types[-1] == "response.completed" - # concatenated deltas reconstruct the content - deltas = "".join(d["delta"] for t, d in events if t == "response.output_text.delta") - assert deltas == "Hello" - # completed event carries usage - completed = [d for t, d in events if t == "response.completed"][0] - assert completed["response"]["usage"]["input_tokens"] == 3 - - async def test_stream_tool_calls(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_toolcall_chunks()))): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "lookup hi", - "stream": True, "store": False}) - events = _sse_events(resp.content.decode()) - types = [e[0] for e in events] - assert "response.function_call_arguments.delta" in types - assert "response.function_call_arguments.done" in types - args = "".join(d["delta"] for t, d in events - if t == "response.function_call_arguments.delta") - assert args == '{"q":"hi"}' - completed = [d for t, d in events if t == "response.completed"][0] - fc = [i for i in completed["response"]["output"] if i["type"] == "function_call"][0] - assert fc["name"] == "lookup" - assert fc["arguments"] == '{"q":"hi"}' - - -# ────────────────────────────────────────────────────────────────────────────── -# Native path (external OpenAI backend) -# ────────────────────────────────────────────────────────────────────────────── - -class TestNativePath: - async def test_nonstream_passthrough_rewrites_id(self, client): - oclient = MagicMock() - resp_obj = MagicMock() - resp_obj.model_dump.return_value = { - "id": "resp_openai", "status": "completed", - "output": [{"type": "message", "role": "assistant", - "content": [{"type": "output_text", "text": "native hi"}]}], - "usage": {"input_tokens": 2, "output_tokens": 3, "total_tokens": 5}} - oclient.responses.create = AsyncMock(return_value=resp_obj) - with (patch.object(api_responses, "choose_endpoint", - AsyncMock(return_value=("https://api.openai.com/v1", "gpt"))), - patch.object(api_responses, "decrement_usage", AsyncMock()), - patch.object(api_responses, "is_ext_openai_endpoint", return_value=True), - patch.object(api_responses, "_make_openai_client", return_value=oclient), - patch.object(api_responses, "get_llm_cache", return_value=None)): - resp = await client.post("/v1/responses", - json={"model": "gpt", "input": "hi", "store": False}) - body = resp.json() - assert body["output_text"] == "native hi" - assert body["id"].startswith("resp_") and body["id"] != "resp_openai" - # native call must not delegate state upstream - assert oclient.responses.create.call_args.kwargs["store"] is False - - async def test_stream_passthrough(self, client): - oclient = MagicMock() - oclient.responses.create = AsyncMock(return_value=_native_event_stream()) - with (patch.object(api_responses, "choose_endpoint", - AsyncMock(return_value=("https://api.openai.com/v1", "gpt"))), - patch.object(api_responses, "decrement_usage", AsyncMock()), - patch.object(api_responses, "is_ext_openai_endpoint", return_value=True), - patch.object(api_responses, "_make_openai_client", return_value=oclient), - patch.object(api_responses, "get_llm_cache", return_value=None)): - resp = await client.post("/v1/responses", - json={"model": "gpt", "input": "hi", - "stream": True, "store": False}) - events = _sse_events(resp.content.decode()) - # the completed event's response id is rewritten to the router id - completed = [d for t, d in events if t == "response.completed"][0] - assert completed["response"]["id"].startswith("resp_") - assert completed["response"]["id"] != "resp_openai" - - -# ────────────────────────────────────────────────────────────────────────────── -# Storage + chaining + retrieve/delete -# ────────────────────────────────────────────────────────────────────────────── - -class TestStorageAndChaining: - async def test_store_and_retrieve(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("remembered")))): - created = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", "store": True}) - rid = created.json()["id"] - got = await client.get(f"/v1/responses/{rid}") - assert got.status_code == 200 - assert got.json()["output_text"] == "remembered" - - async def test_previous_response_id_rehydrates_history(self, client): - # First turn - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("turn-one")))): - first = await client.post("/v1/responses", - json={"model": "test-model", "input": "first?", "store": True}) - rid = first.json()["id"] - - # Second turn references the first — capture the messages sent to the backend - capture = AsyncMock(return_value=_fake_completion("turn-two")) - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", capture)): - await client.post("/v1/responses", - json={"model": "test-model", "input": "second?", - "previous_response_id": rid, "store": True}) - sent_messages = capture.call_args.args[1]["messages"] - contents = [m.get("content") for m in sent_messages] - assert "first?" in contents # prior user turn replayed - assert "turn-one" in contents # prior assistant turn replayed - assert "second?" in contents # current turn appended - - async def test_delete(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("bye")))): - created = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", "store": True}) - rid = created.json()["id"] - deleted = await client.delete(f"/v1/responses/{rid}") - assert deleted.status_code == 200 - assert deleted.json()["deleted"] is True - assert (await client.get(f"/v1/responses/{rid}")).status_code == 404 - - async def test_retrieve_missing_404(self, client): - assert (await client.get("/v1/responses/resp_missing")).status_code == 404 - - -# ────────────────────────────────────────────────────────────────────────────── -# Background mode -# ────────────────────────────────────────────────────────────────────────────── - -class TestBackgroundMode: - async def test_background_requires_store(self, client): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", - "background": True, "store": False}) - assert resp.status_code == 400 - - async def test_background_lifecycle(self, client): - with _enter(*_patch_backend(native=False), - patch.object(api_responses, "create_chat_with_retries", - AsyncMock(return_value=_fake_completion("bg-done")))): - created = await client.post("/v1/responses", - json={"model": "test-model", "input": "hi", - "background": True, "store": True}) - assert created.status_code == 200 - assert created.json()["status"] == "queued" - rid = created.json()["id"] - # poll until terminal - status = None - for _ in range(100): - await asyncio.sleep(0.01) - got = await client.get(f"/v1/responses/{rid}") - status = got.json()["status"] - if status in ("completed", "failed", "cancelled"): - break - assert status == "completed" - assert got.json()["output_text"] == "bg-done" - - async def test_fail_orphaned_responses(self, client): - db = router.db - await db.store_response("resp_orphan", previous_response_id=None, model="m", - status="in_progress", created_at=0, input_messages=[]) - n = await db.fail_orphaned_responses() - assert n >= 1 - row = await db.get_response("resp_orphan") - assert row["status"] == "failed" - - -# ────────────────────────────────────────────────────────────────────────────── -# Cache parity -# ────────────────────────────────────────────────────────────────────────────── - -class _FakeCache: - def __init__(self, response_bytes): - self._resp = response_bytes - self.calls = [] - - async def get_chat(self, route, model, messages): - self.calls.append((route, model, messages)) - return self._resp - - -class TestCacheParity: - async def test_cache_hit_served_as_response(self, client): - cached = orjson.dumps(rt.build_response_object( - response_id="resp_cached", model="test-model", - output_items=rt.chat_message_to_output_items( - {"role": "assistant", "content": "from-cache"}))) - fake = _FakeCache(cached) - with (patch.object(api_responses, "get_llm_cache", return_value=fake), - patch.object(api_responses, "choose_endpoint", - AsyncMock(side_effect=AssertionError("backend must not be reached")))): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "ping", - "store": False, "nomyo": {"cache": True}}) - assert resp.status_code == 200 - assert resp.json()["output_text"] == "from-cache" - assert fake.calls and fake.calls[0][0] == "openai_responses" - - async def test_cache_hit_served_as_sse(self, client): - cached = orjson.dumps(rt.build_response_object( - response_id="resp_cached", model="test-model", - output_items=rt.chat_message_to_output_items( - {"role": "assistant", "content": "from-cache"}))) - fake = _FakeCache(cached) - with (patch.object(api_responses, "get_llm_cache", return_value=fake), - patch.object(api_responses, "choose_endpoint", - AsyncMock(side_effect=AssertionError("backend must not be reached")))): - resp = await client.post("/v1/responses", - json={"model": "test-model", "input": "ping", - "stream": True, "store": False, - "nomyo": {"cache": True}}) - assert resp.headers["content-type"].startswith("text/event-stream") - events = _sse_events(resp.content.decode()) - deltas = "".join(d["delta"] for t, d in events if t == "response.output_text.delta") - assert deltas == "from-cache" diff --git a/test/test_unit_helpers.py b/test/test_unit_helpers.py index def7082..d38eb37 100644 --- a/test/test_unit_helpers.py +++ b/test/test_unit_helpers.py @@ -277,49 +277,3 @@ class TestGetTrackingModel: with patch.object(router, "config", cfg): result = router.get_tracking_model(ep, "unsloth/model:Q8_0") assert result == "model" - - -class TestLlamaSwapClassification: - def _cfg(self, *, server=None, swap=None): - cfg = MagicMock() - cfg.endpoints = [] - cfg.llama_server_endpoints = server or [] - cfg.llama_swap_endpoints = swap or [] - return cfg - - def test_is_llama_swap_only_for_swap_list(self): - from backends.normalize import is_llama_swap - swap_ep = "http://host:8890/v1" - server_ep = "http://host:8889/v1" - cfg = self._cfg(server=[server_ep], swap=[swap_ep]) - with patch.object(router, "config", cfg): - assert is_llama_swap(swap_ep) is True - assert is_llama_swap(server_ep) is False - - def test_is_llama_server_covers_both(self): - from backends.normalize import is_llama_server - swap_ep = "http://host:8890/v1" - server_ep = "http://host:8889/v1" - cfg = self._cfg(server=[server_ep], swap=[swap_ep]) - with patch.object(router, "config", cfg): - assert is_llama_server(swap_ep) is True - assert is_llama_server(server_ep) is True - assert is_llama_server("http://host:11434") is False - - def test_swap_is_openai_compatible_not_ext(self): - swap_ep = "http://host:8890/v1" - cfg = self._cfg(swap=[swap_ep]) - with patch.object(router, "config", cfg): - assert router.is_openai_compatible(swap_ep) is True - assert router.is_ext_openai_endpoint(swap_ep) is False - - def test_swap_tracking_model_normalized(self): - swap_ep = "http://host:8890/v1" - cfg = self._cfg(swap=[swap_ep]) - with patch.object(router, "config", cfg): - assert router.get_tracking_model(swap_ep, "unsloth/model:Q8_0") == "model" - - def test_llama_endpoints_dedupes_and_orders(self): - from backends.normalize import llama_endpoints - cfg = self._cfg(server=["a", "b"], swap=["b", "c"]) - assert llama_endpoints(cfg) == ["a", "b", "c"]