Compare commits

..

2 commits

Author SHA1 Message Date
2d9bbd06b5 doc/compliance.md aktualisiert 2026-06-10 13:25:17 +02:00
392034f839
feat: aissurance plugin initial commit 2026-06-10 13:22:44 +02:00
43 changed files with 2002 additions and 4231 deletions

View file

@ -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 }}

3
.gitignore vendored
View file

@ -66,4 +66,7 @@ config.yaml
# SQLite
*.db*
# aissurance compliance plugin — encrypted offline retry buffer
compliance-buffer/
*settings.json

View file

@ -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:

View file

@ -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(

View file

@ -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)})

View file

@ -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/<model>/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/<model>/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)

View file

@ -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"),
)

View file

@ -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))

View file

@ -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

View file

@ -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

View file

@ -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,

22
compliance/__init__.py Normal file
View file

@ -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}")

94
compliance/buffer.py Normal file
View file

@ -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))

219
compliance/collector.py Normal file
View file

@ -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,
}

45
compliance/crypto.py Normal file
View file

@ -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

15
compliance/log.py Normal file
View file

@ -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)

81
compliance/payload.py Normal file
View file

@ -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=<hex>`` 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

129
compliance/plugin.py Normal file
View file

@ -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)

132
compliance/scheduler.py Normal file
View file

@ -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()

98
compliance/settings.py Normal file
View file

@ -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

89
compliance/transport.py Normal file
View file

@ -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)

View file

@ -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 endpointmodel 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_"

View file

@ -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/<path>
# 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 endpointmodel 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
# 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

203
db.py
View file

@ -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

View file

@ -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:

342
doc/compliance.md Normal file
View file

@ -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: <router version>
X-Router-Id: <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. |

View file

@ -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/<path>`, 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:

View file

@ -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
# 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

View file

@ -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", {})

View file

@ -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:

View file

@ -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)})

View file

@ -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)

View file

@ -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:

View file

@ -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 advertisedmodel 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

View file

@ -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

View file

@ -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"

View file

@ -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])

230
test/test_compliance.py Normal file
View file

@ -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

View file

@ -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=[])

View file

@ -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

View file

@ -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, thinkingreasoning 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

View file

@ -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"

View file

@ -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"]