Compare commits
2 commits
main
...
aissurance
| Author | SHA1 | Date | |
|---|---|---|---|
| 2d9bbd06b5 | |||
| 392034f839 |
43 changed files with 2002 additions and 4231 deletions
|
|
@ -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
3
.gitignore
vendored
|
|
@ -66,4 +66,7 @@ config.yaml
|
|||
# SQLite
|
||||
*.db*
|
||||
|
||||
# aissurance compliance plugin — encrypted offline retry buffer
|
||||
compliance-buffer/
|
||||
|
||||
*settings.json
|
||||
96
README.md
96
README.md
|
|
@ -132,82 +132,6 @@ This way the Ollama backend servers are utilized more efficient than by simply u
|
|||
|
||||
NOMYO Router also supports OpenAI API compatible v1 backend servers.
|
||||
|
||||
## OpenAI Responses API
|
||||
|
||||
In addition to Chat Completions, NOMYO Router exposes the OpenAI **Responses API**:
|
||||
|
||||
```
|
||||
POST /v1/responses # create a response (stream or non-stream)
|
||||
GET /v1/responses/{id} # retrieve a stored response
|
||||
DELETE /v1/responses/{id} # delete a stored response
|
||||
POST /v1/responses/{id}/cancel # cancel a background response
|
||||
```
|
||||
|
||||
It works transparently across **all** backends. When the routed model lives on a native
|
||||
Responses backend (external OpenAI) the request is forwarded as-is; for Ollama and llama-server the
|
||||
router translates Responses ⇄ Chat Completions in both directions (request, response, and streaming
|
||||
typed SSE events), so clients get a consistent `/v1/responses` surface regardless of backend.
|
||||
|
||||
### Conversation state (`store` / `previous_response_id`)
|
||||
|
||||
The router **owns conversation state itself** (persisted in its SQLite DB) rather than delegating to
|
||||
the upstream provider, so `store` and `previous_response_id` behave identically on every backend.
|
||||
On a follow-up request the router rehydrates the prior turns from its DB and expands them into the
|
||||
conversation; outbound native calls always send `store=false`. Trade-off: this forgoes OpenAI's
|
||||
server-side reasoning-state reuse in exchange for uniform, backend-agnostic chaining.
|
||||
|
||||
### Background mode
|
||||
|
||||
`background:true` (which requires `store:true`) returns immediately with `{"status":"queued"}`; the
|
||||
request runs server-side and the client polls `GET /v1/responses/{id}` until the status reaches a
|
||||
terminal state (`completed` / `failed` / `cancelled`). `POST /v1/responses/{id}/cancel` aborts it.
|
||||
|
||||
Limitations: streaming reconnect-resume via `starting_after` is not yet implemented. In a
|
||||
multi-worker/replica deployment polling works via the shared DB, but `cancel` only reaches the
|
||||
running task in the worker that started it (other workers just mark the stored row cancelled). A
|
||||
background task interrupted by a server restart is reconciled to `failed` on the next startup.
|
||||
|
||||
## Anthropic Messages API
|
||||
|
||||
NOMYO Router also exposes the Anthropic **Messages API**:
|
||||
|
||||
```
|
||||
POST /v1/messages # create a message (stream or non-stream)
|
||||
POST /v1/messages/count_tokens # count input tokens for a request
|
||||
```
|
||||
|
||||
It works transparently across **all** backends. For Ollama / llama-server / llama-swap the router
|
||||
translates Messages ⇄ Chat Completions in both directions (request, response, and streaming typed
|
||||
SSE events — `message_start` → `content_block_*` → `message_delta` → `message_stop`), so clients get
|
||||
a consistent `/v1/messages` surface regardless of backend. The API is stateless — there is no store,
|
||||
background mode, or conversation persistence.
|
||||
|
||||
### Native Anthropic upstream
|
||||
|
||||
Configure real Anthropic endpoints under the `anthropic_endpoints` config key (base URL **without**
|
||||
a `/v1` suffix). Requests routed to a model advertised by such an endpoint are **forwarded verbatim**
|
||||
over the Anthropic wire format — the router injects the endpoint's `api_keys` entry as the `x-api-key`
|
||||
header and pins `anthropic-version`, passing through the client's `anthropic-beta`. Their advertised
|
||||
models are treated as always-loaded, like external OpenAI endpoints.
|
||||
|
||||
```yaml
|
||||
anthropic_endpoints:
|
||||
- https://api.anthropic.com
|
||||
api_keys:
|
||||
"https://api.anthropic.com": "${ANTHROPIC_API_KEY}"
|
||||
```
|
||||
|
||||
### Thinking
|
||||
|
||||
An inbound `thinking` block is mapped to the backend's `reasoning_effort` (budget → `low`/`medium`/
|
||||
`high`); a backend that streams `reasoning_content` is surfaced back as Anthropic `thinking` content
|
||||
blocks / `thinking_delta` events. On native endpoints, `thinking` passes through untouched.
|
||||
|
||||
### Caching
|
||||
|
||||
Set `nomyo: {"cache": true}` on the request body to consult the router's semantic LLM cache; a hit is
|
||||
reflected via `usage.cache_read_input_tokens` (input tokens served from cache rather than re-processed).
|
||||
|
||||
## Semantic LLM Cache
|
||||
|
||||
NOMYO Router includes an optional semantic cache that serves repeated or semantically similar LLM requests from cache — no endpoint round-trip, no token cost, response in <10 ms.
|
||||
|
|
@ -248,7 +172,7 @@ Each request is keyed on `model + system_prompt` (exact) combined with a weighte
|
|||
|
||||
### Cached routes
|
||||
|
||||
`/api/chat` · `/api/generate` · `/v1/chat/completions` · `/v1/completions` · `/v1/responses`
|
||||
`/api/chat` · `/api/generate` · `/v1/chat/completions` · `/v1/completions`
|
||||
|
||||
### Cache management
|
||||
|
||||
|
|
@ -257,6 +181,24 @@ curl http://localhost:12434/api/cache/stats # hit rate, counters, config
|
|||
curl -X POST http://localhost:12434/api/cache/invalidate # clear all entries
|
||||
```
|
||||
|
||||
## Compliance Plugin (aissurance.eu)
|
||||
|
||||
Optional in-process plugin that periodically sends signed AI-infrastructure
|
||||
evidence — model/endpoint inventory and aggregate telemetry, **never** prompt or
|
||||
completion content — to [aissurance.eu](https://www.aissurance.eu) for EU AI Act
|
||||
compliance. All traffic is outbound; the Router keeps proxying if the upstream is
|
||||
down (evidence is encrypted and buffered for later).
|
||||
|
||||
```yaml
|
||||
compliance:
|
||||
enabled: true
|
||||
server_url: "https://www.aissurance.eu/api/v1/discovery/receive"
|
||||
api_key: "${AISSURANCE_KEY}" # tenant secret, separate from the router API key
|
||||
polling_interval: 300
|
||||
```
|
||||
|
||||
See the **[Compliance Guide](doc/compliance.md)** for the full reference.
|
||||
|
||||
## Supplying the router API key
|
||||
|
||||
If you set `nomyo-router-api-key` in `config.yaml` (or `NOMYO_ROUTER_API_KEY` env), every request to NOMYO Router must include the key:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
330
api/messages.py
330
api/messages.py
|
|
@ -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)})
|
||||
324
api/ollama.py
324
api/ollama.py
|
|
@ -13,7 +13,6 @@ import asyncio
|
|||
import re
|
||||
import time
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
import aiohttp
|
||||
import ollama
|
||||
|
|
@ -41,12 +40,9 @@ from backends.health import (
|
|||
from backends.normalize import (
|
||||
dedupe_on_keys,
|
||||
is_openai_compatible,
|
||||
is_llama_server,
|
||||
llama_endpoints,
|
||||
_normalize_llama_model_name,
|
||||
_extract_llama_quant,
|
||||
)
|
||||
from backends.control import unload_model
|
||||
from backends.probe import fetch
|
||||
from backends.sessions import _make_openai_client, get_ollama_client, get_probe_session
|
||||
from requests.chat import _make_moe_requests
|
||||
|
|
@ -167,38 +163,31 @@ async def proxy(request: Request):
|
|||
|
||||
_affinity_key = _conversation_fingerprint(model, None, prompt)
|
||||
endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key)
|
||||
# _guarded_stream's finally releases the reservation once we hand off; until
|
||||
# then any failure during request building / client construction (including
|
||||
# CancelledError on client disconnect) must release it or the counter leaks.
|
||||
try:
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
params = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
}
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
params = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
}
|
||||
|
||||
optional_params = {
|
||||
"stream": stream,
|
||||
"max_tokens": options.get("num_predict") if options and "num_predict" in options else None,
|
||||
"frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None,
|
||||
"presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None,
|
||||
"seed": options.get("seed") if options and "seed" in options else None,
|
||||
"stop": options.get("stop") if options and "stop" in options else None,
|
||||
"top_p": options.get("top_p") if options and "top_p" in options else None,
|
||||
"temperature": options.get("temperature") if options and "temperature" in options else None,
|
||||
"suffix": suffix,
|
||||
}
|
||||
params.update({k: v for k, v in optional_params.items() if v is not None})
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
except BaseException:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
optional_params = {
|
||||
"stream": stream,
|
||||
"max_tokens": options.get("num_predict") if options and "num_predict" in options else None,
|
||||
"frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None,
|
||||
"presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None,
|
||||
"seed": options.get("seed") if options and "seed" in options else None,
|
||||
"stop": options.get("stop") if options and "stop" in options else None,
|
||||
"top_p": options.get("top_p") if options and "top_p" in options else None,
|
||||
"temperature": options.get("temperature") if options and "temperature" in options else None,
|
||||
"suffix": suffix,
|
||||
}
|
||||
params.update({k: v for k, v in optional_params.items() if v is not None})
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
|
||||
# 4. Async generator body (error handling + cleanup handled by _guarded_stream)
|
||||
async def stream_generate_response():
|
||||
|
|
@ -343,70 +332,59 @@ async def chat_proxy(request: Request):
|
|||
opt = False
|
||||
_affinity_key = _conversation_fingerprint(model, messages, None)
|
||||
endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key)
|
||||
# Releasing the reservation is owned by _guarded_stream's finally once we hand
|
||||
# off to the streaming generator. Until then, any failure during request
|
||||
# building / client construction (including CancelledError on client
|
||||
# disconnect) must release it here or the usage counter leaks.
|
||||
try:
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
if messages:
|
||||
if any("images" in m for m in messages):
|
||||
messages = await asyncio.to_thread(transform_images_to_data_urls, messages)
|
||||
messages = transform_tool_calls_to_openai(messages)
|
||||
messages = _strip_assistant_prefill(messages)
|
||||
params = {
|
||||
"messages": messages,
|
||||
"model": model,
|
||||
}
|
||||
optional_params = {
|
||||
"tools": tools,
|
||||
"stream": stream,
|
||||
"stream_options": {"include_usage": True} if stream else None,
|
||||
"max_tokens": options.get("num_predict") if options and "num_predict" in options else None,
|
||||
"frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None,
|
||||
"presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None,
|
||||
"seed": options.get("seed") if options and "seed" in options else None,
|
||||
"stop": options.get("stop") if options and "stop" in options else None,
|
||||
"top_p": options.get("top_p") if options and "top_p" in options else None,
|
||||
"temperature": options.get("temperature") if options and "temperature" in options else None,
|
||||
"logprobs": logprobs if logprobs is not None else (options.get("logprobs") if options and "logprobs" in options else None),
|
||||
"top_logprobs": top_logprobs if top_logprobs is not None else (options.get("top_logprobs") if options and "top_logprobs" in options else None),
|
||||
"response_format": {"type": "json_schema", "json_schema": _format} if _format is not None else None
|
||||
}
|
||||
params.update({k: v for k, v in optional_params.items() if v is not None})
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
except BaseException:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
if messages:
|
||||
if any("images" in m for m in messages):
|
||||
messages = await asyncio.to_thread(transform_images_to_data_urls, messages)
|
||||
messages = transform_tool_calls_to_openai(messages)
|
||||
messages = _strip_assistant_prefill(messages)
|
||||
params = {
|
||||
"messages": messages,
|
||||
"model": model,
|
||||
}
|
||||
optional_params = {
|
||||
"tools": tools,
|
||||
"stream": stream,
|
||||
"stream_options": {"include_usage": True} if stream else None,
|
||||
"max_tokens": options.get("num_predict") if options and "num_predict" in options else None,
|
||||
"frequency_penalty": options.get("frequency_penalty") if options and "frequency_penalty" in options else None,
|
||||
"presence_penalty": options.get("presence_penalty") if options and "presence_penalty" in options else None,
|
||||
"seed": options.get("seed") if options and "seed" in options else None,
|
||||
"stop": options.get("stop") if options and "stop" in options else None,
|
||||
"top_p": options.get("top_p") if options and "top_p" in options else None,
|
||||
"temperature": options.get("temperature") if options and "temperature" in options else None,
|
||||
"logprobs": logprobs if logprobs is not None else (options.get("logprobs") if options and "logprobs" in options else None),
|
||||
"top_logprobs": top_logprobs if top_logprobs is not None else (options.get("top_logprobs") if options and "top_logprobs" in options else None),
|
||||
"response_format": {"type": "json_schema", "json_schema": _format} if _format is not None else None
|
||||
}
|
||||
params.update({k: v for k, v in optional_params.items() if v is not None})
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
# For OpenAI endpoints: make the API call in handler scope
|
||||
# (try/except inside async generators is unreliable with Starlette's streaming)
|
||||
start_ts = None
|
||||
async_gen = None
|
||||
if use_openai:
|
||||
start_ts = time.perf_counter()
|
||||
# Proactive trim: only for small-ctx models we've already seen run out of space
|
||||
_lookup_model = _normalize_llama_model_name(model) if endpoint in config.llama_server_endpoints else model
|
||||
_known_nctx = _endpoint_nctx.get((endpoint, _lookup_model))
|
||||
if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_pre_target = int((_known_nctx - _known_nctx // 4) / 1.2)
|
||||
_pre_est = _count_message_tokens(params.get("messages", []))
|
||||
if _pre_est > _pre_target:
|
||||
_pre_msgs = params.get("messages", [])
|
||||
_pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target)
|
||||
_dropped = len(_pre_msgs) - len(_pre_trimmed)
|
||||
print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True)
|
||||
params = {**params, "messages": _pre_trimmed}
|
||||
try:
|
||||
# Proactive trim: only for small-ctx models we've already seen run out of space
|
||||
_lookup_model = _normalize_llama_model_name(model) if is_llama_server(endpoint) else model
|
||||
_known_nctx = _endpoint_nctx.get((endpoint, _lookup_model))
|
||||
if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_pre_target = int((_known_nctx - _known_nctx // 4) / 1.2)
|
||||
_pre_est = _count_message_tokens(params.get("messages", []))
|
||||
if _pre_est > _pre_target:
|
||||
_pre_msgs = params.get("messages", [])
|
||||
_pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target)
|
||||
_dropped = len(_pre_msgs) - len(_pre_trimmed)
|
||||
print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True)
|
||||
params = {**params, "messages": _pre_trimmed}
|
||||
async_gen = await oclient.chat.completions.create(**params)
|
||||
except asyncio.CancelledError:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
except Exception as e:
|
||||
_e_str = str(e)
|
||||
print(f"[chat_proxy] caught {type(e).__name__}: {_e_str[:200]}")
|
||||
|
|
@ -613,21 +591,14 @@ async def _handle_embedding_request(
|
|||
|
||||
# 2. Endpoint logic
|
||||
endpoint, tracking_model = await choose_endpoint(model)
|
||||
# _guarded_stream's finally releases the reservation once we hand off; until
|
||||
# then any failure during client construction (including CancelledError on
|
||||
# client disconnect) must release it or the counter leaks.
|
||||
try:
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
client = _make_openai_client(endpoint, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
except BaseException:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
use_openai = is_openai_compatible(endpoint)
|
||||
if use_openai:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")
|
||||
model = model[0]
|
||||
client = _make_openai_client(endpoint, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
else:
|
||||
client = get_ollama_client(endpoint)
|
||||
|
||||
# 3. Async generator body (error handling + cleanup handled by _guarded_stream)
|
||||
async def stream_embedding_response():
|
||||
|
|
@ -964,8 +935,8 @@ async def tags_proxy(request: Request):
|
|||
# 1. Query all endpoints for models
|
||||
tasks = [fetch.endpoint_details(ep, "/api/tags", "models", skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" not in ep]
|
||||
tasks += [fetch.endpoint_details(ep, "/models", "data", config.api_keys[ep], skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" in ep]
|
||||
# Also query llama-server / llama-swap endpoints not already covered by config.endpoints
|
||||
llama_eps_for_tags = [ep for ep in llama_endpoints(config) if ep not in config.endpoints]
|
||||
# Also query llama-server endpoints not already covered by config.endpoints
|
||||
llama_eps_for_tags = [ep for ep in config.llama_server_endpoints if ep not in config.endpoints]
|
||||
tasks += [fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) for ep in llama_eps_for_tags]
|
||||
all_models = await asyncio.gather(*tasks)
|
||||
|
||||
|
|
@ -989,79 +960,27 @@ async def tags_proxy(request: Request):
|
|||
)
|
||||
|
||||
|
||||
async def _fetch_llama_swap_running(endpoint: str) -> list[dict]:
|
||||
"""Return the list of ready (`state == "ready"`) workers from a llama-swap
|
||||
endpoint's `/running` route. llama-swap omits the per-model `status` field on
|
||||
`/v1/models`, so running workers must be read here instead.
|
||||
"""
|
||||
config = get_config()
|
||||
base_url = endpoint.rstrip("/").removesuffix("/v1")
|
||||
return await fetch.endpoint_details(
|
||||
base_url, "/running", "running", config.api_keys.get(endpoint),
|
||||
skip_error_cache=True, timeout=8,
|
||||
)
|
||||
|
||||
|
||||
# Match the context size in a llama-swap worker's `cmd` string, e.g.
|
||||
# "llama-server --port 5818 -hf ... --ctx-size 131072 ...". llama.cpp accepts
|
||||
# both --ctx-size and the short -c alias.
|
||||
_CTX_SIZE_CMD_RE = re.compile(r"(?:--ctx-size|-c)[=\s]+(\d+)")
|
||||
|
||||
|
||||
def _ctx_size_from_cmd(cmd: str) -> int | None:
|
||||
"""Extract n_ctx from a llama-swap worker `cmd` string, or None if absent."""
|
||||
if not cmd:
|
||||
return None
|
||||
m = _CTX_SIZE_CMD_RE.search(cmd)
|
||||
return int(m.group(1)) if m else None
|
||||
|
||||
|
||||
async def _fetch_llama_swap_nctx(endpoint: str, model_id: str) -> int | None:
|
||||
"""Fallback when a worker's `cmd` lacks --ctx-size: ask the underlying
|
||||
llama-server via llama-swap's /upstream/<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)
|
||||
|
|
|
|||
388
api/openai.py
388
api/openai.py
|
|
@ -34,8 +34,6 @@ from backends.normalize import (
|
|||
ep2base,
|
||||
is_ext_openai_endpoint,
|
||||
is_openai_compatible,
|
||||
is_llama_server,
|
||||
llama_endpoints,
|
||||
_normalize_llama_model_name,
|
||||
)
|
||||
from backends.probe import fetch
|
||||
|
|
@ -48,95 +46,6 @@ from routing import choose_endpoint, decrement_usage
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
async def create_chat_with_retries(oclient, send_params, endpoint, model, tracking_model):
|
||||
"""Call ``chat.completions.create`` with the router's resilience retries.
|
||||
|
||||
Encapsulates the recovery ladder shared by the chat-completions handler and
|
||||
the translated ``/v1/responses`` path:
|
||||
|
||||
* ``does not support tools`` → retry without ``tools``
|
||||
* llama-server context exhaustion → sliding-window message trim, with a
|
||||
second retry that also strips ``tools``/``tool_choice``
|
||||
* backend connection failure → mark (endpoint, model) unhealthy so the next
|
||||
request reroutes, then re-raise
|
||||
* ``image input is not supported`` → strip images and retry
|
||||
|
||||
The caller owns the usage reservation taken by ``choose_endpoint``: this
|
||||
function never decrements it. On unrecoverable failure the exception is
|
||||
re-raised so the caller's guard releases the slot exactly once. Returns the
|
||||
established async generator / response.
|
||||
"""
|
||||
config = get_config()
|
||||
try:
|
||||
async_gen = await oclient.chat.completions.create(**send_params)
|
||||
except Exception as e:
|
||||
_e_str = str(e)
|
||||
_is_ctx_err = "exceed_context_size_error" in _e_str or "exceeds the available context size" in _e_str
|
||||
print(f"[ochat] caught={type(e).__name__} ctx={_is_ctx_err} msg={_e_str[:120]}", flush=True)
|
||||
if "does not support tools" in _e_str:
|
||||
# Model doesn't support tools — retry without them
|
||||
print(f"[ochat] retry: no tools", flush=True)
|
||||
params_without_tools = {k: v for k, v in send_params.items() if k != "tools"}
|
||||
async_gen = await oclient.chat.completions.create(**params_without_tools)
|
||||
elif _is_ctx_err:
|
||||
# Backend context limit hit — apply sliding-window trim (context-shift at message level)
|
||||
err_body = getattr(e, "body", {}) or {}
|
||||
err_detail = err_body.get("error", {}) if isinstance(err_body, dict) else {}
|
||||
n_ctx_limit = err_detail.get("n_ctx", 0)
|
||||
actual_tokens = err_detail.get("n_prompt_tokens", 0)
|
||||
# Fallback: parse from string if body parsing yielded nothing (SDK may not parse llama-server errors)
|
||||
if not n_ctx_limit:
|
||||
import re as _re
|
||||
_m = _re.search(r"'n_ctx':\s*(\d+)", _e_str)
|
||||
if _m:
|
||||
n_ctx_limit = int(_m.group(1))
|
||||
_m = _re.search(r"'n_prompt_tokens':\s*(\d+)", _e_str)
|
||||
if _m:
|
||||
actual_tokens = int(_m.group(1))
|
||||
print(f"[ctx-trim] n_ctx={n_ctx_limit} actual={actual_tokens}", flush=True)
|
||||
if not n_ctx_limit:
|
||||
raise
|
||||
if n_ctx_limit <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_endpoint_nctx[(endpoint, model)] = n_ctx_limit
|
||||
|
||||
msgs_to_trim = send_params.get("messages", [])
|
||||
try:
|
||||
cal_target = _calibrated_trim_target(msgs_to_trim, n_ctx_limit, actual_tokens)
|
||||
trimmed_messages = _trim_messages_for_context(msgs_to_trim, n_ctx_limit, target_tokens=cal_target)
|
||||
except Exception as _helper_exc:
|
||||
print(f"[ctx-trim] helper crash: {type(_helper_exc).__name__}: {str(_helper_exc)[:100]}", flush=True)
|
||||
raise
|
||||
dropped = len(msgs_to_trim) - len(trimmed_messages)
|
||||
print(f"[ctx-trim] target={cal_target} dropped={dropped} remaining={len(trimmed_messages)} retrying-1", flush=True)
|
||||
try:
|
||||
async_gen = await oclient.chat.completions.create(**{**send_params, "messages": trimmed_messages})
|
||||
print(f"[ctx-trim] retry-1 ok", flush=True)
|
||||
except Exception as e2:
|
||||
_e2_str = str(e2)
|
||||
if "exceed_context_size_error" in _e2_str or "exceeds the available context size" in _e2_str:
|
||||
# Still too large — tool definitions likely consuming too many tokens, strip them too
|
||||
print(f"[ctx-trim] retry-1 still exceeded, stripping tools retrying-2", flush=True)
|
||||
params_no_tools = {k: v for k, v in send_params.items() if k not in ("tools", "tool_choice")}
|
||||
async_gen = await oclient.chat.completions.create(**{**params_no_tools, "messages": trimmed_messages})
|
||||
print(f"[ctx-trim] retry-2 ok", flush=True)
|
||||
else:
|
||||
raise
|
||||
elif _is_backend_connection_error(e):
|
||||
# Upstream connection failed (e.g. llama-server in router mode
|
||||
# whose delegated worker died). Mark (endpoint, model) so the
|
||||
# next request reroutes; the client will retry this one.
|
||||
print(f"[ochat] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True)
|
||||
await _mark_backend_unhealthy(endpoint, model, _e_str)
|
||||
raise
|
||||
elif "image input is not supported" in _e_str:
|
||||
# Model doesn't support images — strip and retry
|
||||
print(f"[openai_chat_completions_proxy] Model {model} doesn't support images, retrying with text-only messages")
|
||||
async_gen = await oclient.chat.completions.create(**{**send_params, "messages": _strip_images_from_messages(send_params.get("messages", []))})
|
||||
else:
|
||||
raise
|
||||
return async_gen
|
||||
|
||||
|
||||
@router.post("/v1/embeddings")
|
||||
async def openai_embedding_proxy(request: Request):
|
||||
"""
|
||||
|
|
@ -180,14 +89,13 @@ async def openai_embedding_proxy(request: Request):
|
|||
|
||||
# 2. Endpoint logic
|
||||
endpoint, tracking_model = await choose_endpoint(model)
|
||||
# The finally below releases the reservation for every exit — success, error,
|
||||
# or CancelledError — so client construction is kept inside the guarded block.
|
||||
if is_openai_compatible(endpoint):
|
||||
api_key = config.api_keys.get(endpoint, "no-key")
|
||||
else:
|
||||
api_key = "ollama"
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=api_key)
|
||||
|
||||
try:
|
||||
if is_openai_compatible(endpoint):
|
||||
api_key = config.api_keys.get(endpoint, "no-key")
|
||||
else:
|
||||
api_key = "ollama"
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=api_key)
|
||||
async_gen = await oclient.embeddings.create(input=doc, model=model)
|
||||
result = async_gen.model_dump()
|
||||
for item in result.get("data", []):
|
||||
|
|
@ -336,30 +244,106 @@ async def openai_chat_completions_proxy(request: Request):
|
|||
|
||||
# Make the API call in handler scope — try/except inside async generators is unreliable
|
||||
# with Starlette's streaming machinery, so we resolve errors here before the generator starts.
|
||||
# The reservation taken by choose_endpoint is released by stream_ochat_response's finally
|
||||
# once we hand off; until then, any failure here (including CancelledError on client
|
||||
# disconnect during a cold model load) must release it or the counter leaks.
|
||||
send_params = params
|
||||
if not is_ext_openai_endpoint(endpoint):
|
||||
resolved_msgs = await _normalize_images_in_messages(params.get("messages", []))
|
||||
send_params = {**params, "messages": resolved_msgs}
|
||||
# Proactive trim: only for small-ctx models we've already seen run out of space
|
||||
_lookup_model = _normalize_llama_model_name(model) if endpoint in config.llama_server_endpoints else model
|
||||
_known_nctx = _endpoint_nctx.get((endpoint, _lookup_model))
|
||||
if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_pre_target = int(((_known_nctx - _known_nctx // 4)) / 1.2)
|
||||
_pre_est = _count_message_tokens(send_params.get("messages", []))
|
||||
if _pre_est > _pre_target:
|
||||
_pre_msgs = send_params.get("messages", [])
|
||||
_pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target)
|
||||
_dropped = len(_pre_msgs) - len(_pre_trimmed)
|
||||
print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True)
|
||||
send_params = {**send_params, "messages": _pre_trimmed}
|
||||
try:
|
||||
send_params = params
|
||||
if not is_ext_openai_endpoint(endpoint):
|
||||
resolved_msgs = await _normalize_images_in_messages(params.get("messages", []))
|
||||
send_params = {**params, "messages": resolved_msgs}
|
||||
# Proactive trim: only for small-ctx models we've already seen run out of space
|
||||
_lookup_model = _normalize_llama_model_name(model) if is_llama_server(endpoint) else model
|
||||
_known_nctx = _endpoint_nctx.get((endpoint, _lookup_model))
|
||||
if _known_nctx and _known_nctx <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_pre_target = int(((_known_nctx - _known_nctx // 4)) / 1.2)
|
||||
_pre_est = _count_message_tokens(send_params.get("messages", []))
|
||||
if _pre_est > _pre_target:
|
||||
_pre_msgs = send_params.get("messages", [])
|
||||
_pre_trimmed = _trim_messages_for_context(_pre_msgs, _known_nctx, target_tokens=_pre_target)
|
||||
_dropped = len(_pre_msgs) - len(_pre_trimmed)
|
||||
print(f"[ctx-pre] n_ctx={_known_nctx} est={_pre_est} target={_pre_target} dropped={_dropped}", flush=True)
|
||||
send_params = {**send_params, "messages": _pre_trimmed}
|
||||
async_gen = await create_chat_with_retries(oclient, send_params, endpoint, model, tracking_model)
|
||||
except BaseException:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
async_gen = await oclient.chat.completions.create(**send_params)
|
||||
except Exception as e:
|
||||
_e_str = str(e)
|
||||
_is_ctx_err = "exceed_context_size_error" in _e_str or "exceeds the available context size" in _e_str
|
||||
print(f"[ochat] caught={type(e).__name__} ctx={_is_ctx_err} msg={_e_str[:120]}", flush=True)
|
||||
if "does not support tools" in _e_str:
|
||||
# Model doesn't support tools — retry without them
|
||||
print(f"[ochat] retry: no tools", flush=True)
|
||||
try:
|
||||
params_without_tools = {k: v for k, v in send_params.items() if k != "tools"}
|
||||
async_gen = await oclient.chat.completions.create(**params_without_tools)
|
||||
except Exception:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
elif _is_ctx_err:
|
||||
# Backend context limit hit — apply sliding-window trim (context-shift at message level)
|
||||
err_body = getattr(e, "body", {}) or {}
|
||||
err_detail = err_body.get("error", {}) if isinstance(err_body, dict) else {}
|
||||
n_ctx_limit = err_detail.get("n_ctx", 0)
|
||||
actual_tokens = err_detail.get("n_prompt_tokens", 0)
|
||||
# Fallback: parse from string if body parsing yielded nothing (SDK may not parse llama-server errors)
|
||||
if not n_ctx_limit:
|
||||
import re as _re
|
||||
_m = _re.search(r"'n_ctx':\s*(\d+)", _e_str)
|
||||
if _m:
|
||||
n_ctx_limit = int(_m.group(1))
|
||||
_m = _re.search(r"'n_prompt_tokens':\s*(\d+)", _e_str)
|
||||
if _m:
|
||||
actual_tokens = int(_m.group(1))
|
||||
print(f"[ctx-trim] n_ctx={n_ctx_limit} actual={actual_tokens}", flush=True)
|
||||
if not n_ctx_limit:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
if n_ctx_limit <= _CTX_TRIM_SMALL_LIMIT:
|
||||
_endpoint_nctx[(endpoint, model)] = n_ctx_limit
|
||||
|
||||
msgs_to_trim = send_params.get("messages", [])
|
||||
try:
|
||||
cal_target = _calibrated_trim_target(msgs_to_trim, n_ctx_limit, actual_tokens)
|
||||
trimmed_messages = _trim_messages_for_context(msgs_to_trim, n_ctx_limit, target_tokens=cal_target)
|
||||
except Exception as _helper_exc:
|
||||
print(f"[ctx-trim] helper crash: {type(_helper_exc).__name__}: {str(_helper_exc)[:100]}", flush=True)
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
dropped = len(msgs_to_trim) - len(trimmed_messages)
|
||||
print(f"[ctx-trim] target={cal_target} dropped={dropped} remaining={len(trimmed_messages)} retrying-1", flush=True)
|
||||
try:
|
||||
async_gen = await oclient.chat.completions.create(**{**send_params, "messages": trimmed_messages})
|
||||
print(f"[ctx-trim] retry-1 ok", flush=True)
|
||||
except Exception as e2:
|
||||
_e2_str = str(e2)
|
||||
if "exceed_context_size_error" in _e2_str or "exceeds the available context size" in _e2_str:
|
||||
# Still too large — tool definitions likely consuming too many tokens, strip them too
|
||||
print(f"[ctx-trim] retry-1 still exceeded, stripping tools retrying-2", flush=True)
|
||||
params_no_tools = {k: v for k, v in send_params.items() if k not in ("tools", "tool_choice")}
|
||||
try:
|
||||
async_gen = await oclient.chat.completions.create(**{**params_no_tools, "messages": trimmed_messages})
|
||||
print(f"[ctx-trim] retry-2 ok", flush=True)
|
||||
except Exception:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
else:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
elif _is_backend_connection_error(e):
|
||||
# Upstream connection failed (e.g. llama-server in router mode
|
||||
# whose delegated worker died). Mark (endpoint, model) so the
|
||||
# next request reroutes; the client will retry this one.
|
||||
print(f"[ochat] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True)
|
||||
await _mark_backend_unhealthy(endpoint, model, _e_str)
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
elif "image input is not supported" in _e_str:
|
||||
# Model doesn't support images — strip and retry
|
||||
print(f"[openai_chat_completions_proxy] Model {model} doesn't support images, retrying with text-only messages")
|
||||
try:
|
||||
async_gen = await oclient.chat.completions.create(**{**send_params, "messages": _strip_images_from_messages(send_params.get("messages", []))})
|
||||
except Exception:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
else:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
|
||||
# 4. Async generator — only streams the already-established async_gen
|
||||
async def stream_ochat_response():
|
||||
|
|
@ -540,17 +524,12 @@ async def openai_completions_proxy(request: Request):
|
|||
# 2. Endpoint logic
|
||||
_affinity_key = _conversation_fingerprint(model, None, prompt)
|
||||
endpoint, tracking_model = await choose_endpoint(model, affinity_key=_affinity_key)
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
|
||||
# 3. Async generator that streams completions data and decrements the counter
|
||||
# Make the API call in handler scope (try/except inside async generators is unreliable).
|
||||
# The reservation is released by stream_ocompletions_response's finally once we hand off;
|
||||
# until then any failure here — including CancelledError on client disconnect — releases it.
|
||||
# Make the API call in handler scope (try/except inside async generators is unreliable)
|
||||
try:
|
||||
oclient = _make_openai_client(endpoint, default_headers=default_headers, api_key=config.api_keys.get(endpoint, "no-key"))
|
||||
async_gen = await oclient.completions.create(**params)
|
||||
except asyncio.CancelledError:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
raise
|
||||
except Exception as e:
|
||||
if _is_backend_connection_error(e):
|
||||
print(f"[ocompl] backend connection error → marking ({endpoint}, {model}) unhealthy", flush=True)
|
||||
|
|
@ -658,23 +637,17 @@ async def openai_models_proxy(request: Request):
|
|||
ollama_tasks = [fetch.endpoint_details(ep, "/api/tags", "models", skip_error_cache=True, timeout=8) for ep in config.endpoints if "/v1" not in ep]
|
||||
# 2. Query external OpenAI endpoints (Groq, OpenAI, etc.) via /models
|
||||
ext_openai_tasks = [fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8) for ep in config.endpoints if is_ext_openai_endpoint(ep)]
|
||||
# 3. Query llama-server / llama-swap endpoints for advertised models via /v1/models
|
||||
# Also query endpoints that may not be in config.endpoints
|
||||
all_llama_endpoints = llama_endpoints(config)
|
||||
# 3. Query llama-server endpoints for loaded models via /v1/models
|
||||
# Also query endpoints from llama_server_endpoints that may not be in config.endpoints
|
||||
all_llama_endpoints = set(config.llama_server_endpoints) | set(ep for ep in config.endpoints if ep in config.llama_server_endpoints)
|
||||
llama_tasks = [
|
||||
fetch.endpoint_details(ep, "/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8)
|
||||
for ep in all_llama_endpoints
|
||||
]
|
||||
# 4. Query native Anthropic endpoints via /v1/models (auth headers picked by endpoint type)
|
||||
anthropic_tasks = [
|
||||
fetch.endpoint_details(ep, "/v1/models", "data", config.api_keys.get(ep), skip_error_cache=True, timeout=8)
|
||||
for ep in config.anthropic_endpoints
|
||||
]
|
||||
|
||||
ollama_models = await asyncio.gather(*ollama_tasks) if ollama_tasks else []
|
||||
ext_openai_models = await asyncio.gather(*ext_openai_tasks) if ext_openai_tasks else []
|
||||
llama_models = await asyncio.gather(*llama_tasks) if llama_tasks else []
|
||||
anthropic_models = await asyncio.gather(*anthropic_tasks) if anthropic_tasks else []
|
||||
|
||||
models = {'data': []}
|
||||
|
||||
|
|
@ -708,16 +681,6 @@ async def openai_models_proxy(request: Request):
|
|||
model['name'] = model['id']
|
||||
models['data'].append(model)
|
||||
|
||||
# Add native Anthropic models (if any)
|
||||
if anthropic_models:
|
||||
for modellist in anthropic_models:
|
||||
for model in modellist:
|
||||
if not "id" in model.keys():
|
||||
model['id'] = model.get('name', model.get('id', ''))
|
||||
else:
|
||||
model['name'] = model['id']
|
||||
models['data'].append(model)
|
||||
|
||||
# 2. Return a JSONResponse with a deduplicated list of unique models for inference
|
||||
return JSONResponse(
|
||||
content={"data": dedupe_on_keys(models['data'], ['name'])},
|
||||
|
|
@ -789,38 +752,36 @@ async def rerank_proxy(request: Request):
|
|||
),
|
||||
)
|
||||
|
||||
# The finally below releases the reservation for every exit (success, error,
|
||||
# or CancelledError), so request building and session lookup stay inside it.
|
||||
try:
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")[0]
|
||||
if ":latest" in model:
|
||||
model = model.split(":latest")[0]
|
||||
|
||||
# Build upstream rerank request body – forward only recognised fields
|
||||
upstream_payload: dict = {"model": model, "query": query, "documents": documents}
|
||||
for optional_key in ("top_n", "return_documents", "max_tokens_per_doc"):
|
||||
if optional_key in payload:
|
||||
upstream_payload[optional_key] = payload[optional_key]
|
||||
# Build upstream rerank request body – forward only recognised fields
|
||||
upstream_payload: dict = {"model": model, "query": query, "documents": documents}
|
||||
for optional_key in ("top_n", "return_documents", "max_tokens_per_doc"):
|
||||
if optional_key in payload:
|
||||
upstream_payload[optional_key] = payload[optional_key]
|
||||
|
||||
# Determine upstream URL:
|
||||
# llama-server / llama-swap expose /v1/rerank (base already contains /v1)
|
||||
# External OpenAI endpoints expose /rerank under their /v1 base
|
||||
if is_llama_server(endpoint):
|
||||
# llama-server / llama-swap: endpoint may or may not already contain /v1
|
||||
if "/v1" in endpoint:
|
||||
rerank_url = f"{endpoint}/rerank"
|
||||
else:
|
||||
rerank_url = f"{endpoint}/v1/rerank"
|
||||
# Determine upstream URL:
|
||||
# llama-server exposes /v1/rerank (base already contains /v1 for llama_server_endpoints)
|
||||
# External OpenAI endpoints expose /rerank under their /v1 base
|
||||
if endpoint in config.llama_server_endpoints:
|
||||
# llama-server: endpoint may or may not already contain /v1
|
||||
if "/v1" in endpoint:
|
||||
rerank_url = f"{endpoint}/rerank"
|
||||
else:
|
||||
# External OpenAI-compatible: ep2base gives us the /v1 base
|
||||
rerank_url = f"{ep2base(endpoint)}/rerank"
|
||||
rerank_url = f"{endpoint}/v1/rerank"
|
||||
else:
|
||||
# External OpenAI-compatible: ep2base gives us the /v1 base
|
||||
rerank_url = f"{ep2base(endpoint)}/rerank"
|
||||
|
||||
api_key = config.api_keys.get(endpoint, "no-key")
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
api_key = config.api_keys.get(endpoint, "no-key")
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
client: aiohttp.ClientSession = get_session(endpoint)
|
||||
client: aiohttp.ClientSession = get_session(endpoint)
|
||||
try:
|
||||
async with client.post(rerank_url, json=upstream_payload, headers=headers) as resp:
|
||||
response_bytes = await resp.read()
|
||||
if resp.status >= 400:
|
||||
|
|
@ -841,82 +802,3 @@ async def rerank_proxy(request: Request):
|
|||
return JSONResponse(content=data)
|
||||
finally:
|
||||
await decrement_usage(endpoint, tracking_model)
|
||||
|
||||
|
||||
async def _resolve_llama_swap_endpoint(model_id: str) -> str | None:
|
||||
"""Pick the llama-swap endpoint that serves ``model_id``.
|
||||
|
||||
Prefers an endpoint that already has the worker running; falls back to any
|
||||
that advertises the model. Returns None if none do.
|
||||
"""
|
||||
config = get_config()
|
||||
swap_eps = config.llama_swap_endpoints
|
||||
if not swap_eps:
|
||||
return None
|
||||
|
||||
advertised = await asyncio.gather(
|
||||
*[fetch.available_models(ep, config.api_keys.get(ep)) for ep in swap_eps]
|
||||
)
|
||||
candidates = [ep for ep, models in zip(swap_eps, advertised) if model_id in models]
|
||||
if not candidates:
|
||||
return None
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
|
||||
loaded = await asyncio.gather(*[fetch.loaded_models(ep) for ep in candidates])
|
||||
for ep, lm in zip(candidates, loaded):
|
||||
if model_id in lm:
|
||||
return ep
|
||||
return candidates[0]
|
||||
|
||||
|
||||
@router.api_route("/upstream/{model_id}/{path:path}", methods=["GET", "POST"])
|
||||
async def llama_swap_upstream(model_id: str, path: str, request: Request):
|
||||
"""Bypass llama-swap and reach a model's underlying llama-server worker directly
|
||||
via llama-swap's ``/upstream/:model_id`` route.
|
||||
|
||||
Lets clients use llama-server features that llama-swap itself does not forward
|
||||
(e.g. token-array prompts), while still letting the router pick the backend that
|
||||
actually hosts the model. ``/upstream`` is a root route, so the ``/v1`` suffix is
|
||||
stripped from the configured endpoint.
|
||||
"""
|
||||
config = get_config()
|
||||
endpoint = await _resolve_llama_swap_endpoint(model_id)
|
||||
if endpoint is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"No configured llama-swap endpoint serves model '{model_id}'.",
|
||||
)
|
||||
|
||||
base_url = endpoint.rstrip("/").removesuffix("/v1")
|
||||
url = f"{base_url}/upstream/{model_id}/{path}"
|
||||
if request.url.query:
|
||||
url = f"{url}?{request.url.query}"
|
||||
|
||||
headers = {"Referer": default_headers.get("HTTP-Referer", "https://nomyo.ai")}
|
||||
content_type = request.headers.get("content-type")
|
||||
if content_type:
|
||||
headers["Content-Type"] = content_type
|
||||
api_key = config.api_keys.get(endpoint)
|
||||
if api_key is not None:
|
||||
headers["Authorization"] = "Bearer " + api_key
|
||||
|
||||
body = await request.body()
|
||||
client: aiohttp.ClientSession = get_session(endpoint)
|
||||
try:
|
||||
resp = await client.request(request.method, url, data=body or None, headers=headers)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=502, detail=f"Upstream request to {url} failed: {e}")
|
||||
|
||||
async def _iter():
|
||||
try:
|
||||
async for chunk in resp.content.iter_any():
|
||||
yield chunk
|
||||
finally:
|
||||
resp.release()
|
||||
|
||||
return StreamingResponse(
|
||||
_iter(),
|
||||
status_code=resp.status,
|
||||
media_type=resp.headers.get("Content-Type"),
|
||||
)
|
||||
|
|
|
|||
420
api/responses.py
420
api/responses.py
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
22
compliance/__init__.py
Normal 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
94
compliance/buffer.py
Normal 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
219
compliance/collector.py
Normal 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
45
compliance/crypto.py
Normal 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
15
compliance/log.py
Normal 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
81
compliance/payload.py
Normal 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
129
compliance/plugin.py
Normal 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
132
compliance/scheduler.py
Normal 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
98
compliance/settings.py
Normal 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
89
compliance/transport.py
Normal 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)
|
||||
15
config.py
15
config.py
|
|
@ -13,6 +13,8 @@ import yaml
|
|||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
from compliance.settings import ComplianceSettings
|
||||
|
||||
|
||||
class Config(BaseSettings):
|
||||
# List of Ollama endpoints
|
||||
|
|
@ -23,16 +25,6 @@ class Config(BaseSettings):
|
|||
)
|
||||
# List of llama-server endpoints (OpenAI-compatible with /v1/models status info)
|
||||
llama_server_endpoints: List[str] = Field(default_factory=list)
|
||||
# List of llama-swap endpoints (OpenAI-compatible front for multiple llama-server
|
||||
# workers). Same surface as llama_server_endpoints, but loaded models are read from
|
||||
# /running (not /v1/models status) and unload uses POST /api/models/unload/:model_id.
|
||||
llama_swap_endpoints: List[str] = Field(default_factory=list)
|
||||
# List of native Anthropic Messages-API endpoints (e.g. https://api.anthropic.com).
|
||||
# Configure the base URL WITHOUT a /v1 suffix; the router appends /v1/models and
|
||||
# /v1/messages itself. Requests routed here are forwarded verbatim (no Messages⇄Chat
|
||||
# translation); the endpoint's api_keys entry is sent as the x-api-key header. Their
|
||||
# advertised models are treated as always-loaded, like external OpenAI endpoints.
|
||||
anthropic_endpoints: List[str] = Field(default_factory=list)
|
||||
# Max concurrent connections per endpoint‑model pair, see OLLAMA_NUM_PARALLEL
|
||||
max_concurrent_connections: int = 1
|
||||
# Per-endpoint overrides: {endpoint_url: {max_concurrent_connections: N}}
|
||||
|
|
@ -72,6 +64,9 @@ class Config(BaseSettings):
|
|||
# 0.3 = 30% history context signal, 70% question signal
|
||||
cache_history_weight: float = Field(default=0.3)
|
||||
|
||||
# aissurance.eu compliance plugin (`compliance:` block). See compliance/.
|
||||
compliance: ComplianceSettings = Field(default_factory=ComplianceSettings)
|
||||
|
||||
class Config:
|
||||
# YAML loading is handled manually via Config.from_yaml(); env vars use this prefix.
|
||||
env_prefix = "NOMYO_ROUTER_"
|
||||
|
|
|
|||
44
config.yaml
44
config.yaml
|
|
@ -6,23 +6,7 @@ endpoints:
|
|||
- https://api.openai.com/v1
|
||||
|
||||
llama_server_endpoints:
|
||||
- http://192.168.0.51:8889/v1
|
||||
|
||||
# llama-swap endpoints (OpenAI-compatible front for multiple llama-server workers).
|
||||
# Same surface as llama_server_endpoints, but the router reads loaded/running workers
|
||||
# from /running (state == "ready") instead of a /v1/models status field, and unloads via
|
||||
# POST /api/models/unload/:model_id. The router also exposes /upstream/:model_id/<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 endpoint‑model pair* (equals to OLLAMA_NUM_PARALLEL)
|
||||
# This is the global default; individual endpoints can override it via endpoint_config below.
|
||||
|
|
@ -73,8 +57,7 @@ api_keys:
|
|||
"http://192.168.0.51:11434": "ollama"
|
||||
"http://192.168.0.52:11434": "ollama"
|
||||
"https://api.openai.com/v1": "${OPENAI_KEY}"
|
||||
"http://192.168.0.51:8889/v1": "llama"
|
||||
"http://192.168.0.52:8889/v1": "llama-swap"
|
||||
"http://192.168.0.50:8889/v1": "llama"
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# Semantic LLM Cache (optional — disabled by default)
|
||||
|
|
@ -108,4 +91,25 @@ api_keys:
|
|||
# Weight of the BM25-weighted chat-history embedding vs last-user-message embedding.
|
||||
# 0.3 = 30% history context signal, 70% question signal.
|
||||
# Only relevant when cache_similarity < 1.0.
|
||||
# cache_history_weight: 0.3
|
||||
# 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
203
db.py
|
|
@ -1,4 +1,4 @@
|
|||
import aiosqlite, asyncio, orjson
|
||||
import aiosqlite, asyncio
|
||||
from typing import Optional
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -75,24 +75,6 @@ class TokenDatabase:
|
|||
''')
|
||||
await db.execute('CREATE INDEX IF NOT EXISTS idx_token_time_series_timestamp ON token_time_series(timestamp)')
|
||||
await db.execute('CREATE INDEX IF NOT EXISTS idx_token_time_series_model_ts ON token_time_series(model, timestamp)')
|
||||
# Responses API state — the router owns conversation state for the
|
||||
# /v1/responses family (store / previous_response_id) and tracks
|
||||
# background-task status here so polling survives across workers.
|
||||
await db.execute('''
|
||||
CREATE TABLE IF NOT EXISTS stored_responses (
|
||||
response_id TEXT PRIMARY KEY,
|
||||
previous_response_id TEXT,
|
||||
model TEXT,
|
||||
status TEXT,
|
||||
created_at INTEGER,
|
||||
input_messages TEXT,
|
||||
output_items TEXT,
|
||||
usage TEXT,
|
||||
instructions TEXT,
|
||||
error TEXT
|
||||
)
|
||||
''')
|
||||
await db.execute('CREATE INDEX IF NOT EXISTS idx_stored_responses_prev ON stored_responses(previous_response_id)')
|
||||
await db.commit()
|
||||
|
||||
async def update_token_counts(self, endpoint: str, model: str, input_tokens: int, output_tokens: int):
|
||||
|
|
@ -273,6 +255,34 @@ class TokenDatabase:
|
|||
}
|
||||
return None
|
||||
|
||||
async def get_last_used_map(self) -> dict:
|
||||
"""Return {(endpoint, model): latest_unix_timestamp} from the time series.
|
||||
|
||||
Read-only aggregate used by the compliance collector to derive each
|
||||
model's ``last_used``. Computed in SQL in a single pass.
|
||||
"""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
async with db.execute('''
|
||||
SELECT endpoint, model, MAX(timestamp)
|
||||
FROM token_time_series
|
||||
GROUP BY endpoint, model
|
||||
''') as cursor:
|
||||
rows = await cursor.fetchall()
|
||||
return {(row[0], row[1]): row[2] for row in rows}
|
||||
|
||||
async def get_token_totals_since(self, cutoff_ts: int) -> tuple[int, int]:
|
||||
"""Return (input_tokens, output_tokens) summed over entries at/after cutoff_ts."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
async with db.execute('''
|
||||
SELECT COALESCE(SUM(input_tokens), 0), COALESCE(SUM(output_tokens), 0)
|
||||
FROM token_time_series
|
||||
WHERE timestamp >= ?
|
||||
''', (cutoff_ts,)) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return (int(row[0]), int(row[1])) if row else (0, 0)
|
||||
|
||||
async def aggregate_time_series_older_than(self, days: int, trim_old: bool = False) -> int:
|
||||
"""
|
||||
Aggregate time_series entries older than 'days' days into daily aggregates by
|
||||
|
|
@ -337,158 +347,3 @@ class TokenDatabase:
|
|||
await db.commit()
|
||||
|
||||
return aggregated_count
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# Responses API state (store / previous_response_id / background)
|
||||
# -----------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _row_to_response(row) -> dict:
|
||||
"""Map a stored_responses row to a plain dict, decoding JSON columns."""
|
||||
def _loads(val):
|
||||
if val is None:
|
||||
return None
|
||||
try:
|
||||
return orjson.loads(val)
|
||||
except (orjson.JSONDecodeError, TypeError):
|
||||
return None
|
||||
return {
|
||||
'response_id': row[0],
|
||||
'previous_response_id': row[1],
|
||||
'model': row[2],
|
||||
'status': row[3],
|
||||
'created_at': row[4],
|
||||
'input_messages': _loads(row[5]),
|
||||
'output_items': _loads(row[6]),
|
||||
'usage': _loads(row[7]),
|
||||
'instructions': row[8],
|
||||
'error': _loads(row[9]),
|
||||
}
|
||||
|
||||
async def store_response(
|
||||
self,
|
||||
response_id: str,
|
||||
*,
|
||||
previous_response_id: Optional[str],
|
||||
model: str,
|
||||
status: str,
|
||||
created_at: int,
|
||||
input_messages: list,
|
||||
output_items: Optional[list] = None,
|
||||
usage: Optional[dict] = None,
|
||||
instructions: Optional[str] = None,
|
||||
error: Optional[dict] = None,
|
||||
):
|
||||
"""Insert or replace a stored Responses-API response row."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
await db.execute('''
|
||||
INSERT INTO stored_responses
|
||||
(response_id, previous_response_id, model, status, created_at,
|
||||
input_messages, output_items, usage, instructions, error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT (response_id) DO UPDATE SET
|
||||
previous_response_id = excluded.previous_response_id,
|
||||
model = excluded.model,
|
||||
status = excluded.status,
|
||||
created_at = excluded.created_at,
|
||||
input_messages = excluded.input_messages,
|
||||
output_items = excluded.output_items,
|
||||
usage = excluded.usage,
|
||||
instructions = excluded.instructions,
|
||||
error = excluded.error
|
||||
''', (
|
||||
response_id, previous_response_id, model, status, created_at,
|
||||
orjson.dumps(input_messages).decode("utf-8"),
|
||||
orjson.dumps(output_items).decode("utf-8") if output_items is not None else None,
|
||||
orjson.dumps(usage).decode("utf-8") if usage is not None else None,
|
||||
instructions,
|
||||
orjson.dumps(error).decode("utf-8") if error is not None else None,
|
||||
))
|
||||
await db.commit()
|
||||
|
||||
async def update_response_status(
|
||||
self,
|
||||
response_id: str,
|
||||
status: str,
|
||||
*,
|
||||
output_items: Optional[list] = None,
|
||||
usage: Optional[dict] = None,
|
||||
error: Optional[dict] = None,
|
||||
):
|
||||
"""Update the status (and optionally output/usage/error) of a stored response."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
await db.execute('''
|
||||
UPDATE stored_responses
|
||||
SET status = ?,
|
||||
output_items = COALESCE(?, output_items),
|
||||
usage = COALESCE(?, usage),
|
||||
error = COALESCE(?, error)
|
||||
WHERE response_id = ?
|
||||
''', (
|
||||
status,
|
||||
orjson.dumps(output_items).decode("utf-8") if output_items is not None else None,
|
||||
orjson.dumps(usage).decode("utf-8") if usage is not None else None,
|
||||
orjson.dumps(error).decode("utf-8") if error is not None else None,
|
||||
response_id,
|
||||
))
|
||||
await db.commit()
|
||||
|
||||
async def get_response(self, response_id: str) -> Optional[dict]:
|
||||
"""Return a stored response as a dict, or None if not found."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
async with db.execute('''
|
||||
SELECT response_id, previous_response_id, model, status, created_at,
|
||||
input_messages, output_items, usage, instructions, error
|
||||
FROM stored_responses WHERE response_id = ?
|
||||
''', (response_id,)) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return self._row_to_response(row) if row is not None else None
|
||||
|
||||
async def delete_response(self, response_id: str) -> bool:
|
||||
"""Delete a stored response. Returns True if a row was removed."""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
cursor = await db.execute(
|
||||
'DELETE FROM stored_responses WHERE response_id = ?', (response_id,)
|
||||
)
|
||||
await db.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def get_response_chain(self, response_id: str, max_turns: int = 50) -> list:
|
||||
"""Walk previous_response_id back to the root, returned oldest-first.
|
||||
|
||||
Bounded to ``max_turns`` so a pathological chain cannot stall a request.
|
||||
Missing links terminate the walk gracefully.
|
||||
"""
|
||||
chain: list = []
|
||||
seen: set = set()
|
||||
current = response_id
|
||||
while current and current not in seen and len(chain) < max_turns:
|
||||
seen.add(current)
|
||||
resp = await self.get_response(current)
|
||||
if resp is None:
|
||||
break
|
||||
chain.append(resp)
|
||||
current = resp.get('previous_response_id')
|
||||
chain.reverse()
|
||||
return chain
|
||||
|
||||
async def fail_orphaned_responses(self) -> int:
|
||||
"""Mark non-terminal responses as failed (called on startup).
|
||||
|
||||
A background task lives in a worker's event loop; a process restart loses
|
||||
it while the DB row stays ``queued``/``in_progress`` forever. Reconcile
|
||||
those to ``failed`` so polling clients get a terminal state.
|
||||
"""
|
||||
db = await self._get_connection()
|
||||
async with self._operation_lock:
|
||||
cursor = await db.execute('''
|
||||
UPDATE stored_responses
|
||||
SET status = 'failed',
|
||||
error = ?
|
||||
WHERE status IN ('queued', 'in_progress')
|
||||
''', (orjson.dumps({"message": "Response interrupted by server restart", "type": "server_error"}).decode("utf-8"),))
|
||||
await db.commit()
|
||||
return cursor.rowcount
|
||||
|
|
|
|||
|
|
@ -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
342
doc/compliance.md
Normal 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. |
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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", {})
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)})
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
47
router.py
47
router.py
|
|
@ -68,6 +68,10 @@ from state import (
|
|||
token_worker_task: asyncio.Task | None = None
|
||||
flush_task: asyncio.Task | None = None
|
||||
|
||||
# aissurance compliance plugin instance (created on startup if enabled).
|
||||
from compliance import AissurancePlugin
|
||||
compliance_plugin: "AissurancePlugin | None" = None
|
||||
|
||||
from config import Config, _config_path_from_env
|
||||
|
||||
from ollama._types import TokenLogprob, Logprob
|
||||
|
|
@ -231,7 +235,6 @@ from backends.health import (
|
|||
from backends.normalize import (
|
||||
is_ext_openai_endpoint,
|
||||
is_openai_compatible,
|
||||
llama_endpoints,
|
||||
get_tracking_model,
|
||||
)
|
||||
|
||||
|
|
@ -291,10 +294,6 @@ from api.management import router as management_router
|
|||
app.include_router(management_router)
|
||||
from api.openai import router as openai_router
|
||||
app.include_router(openai_router)
|
||||
from api.responses import router as responses_router
|
||||
app.include_router(responses_router)
|
||||
from api.messages import router as messages_router
|
||||
app.include_router(messages_router)
|
||||
from api.ollama import router as ollama_router
|
||||
app.include_router(ollama_router)
|
||||
|
||||
|
|
@ -313,7 +312,6 @@ async def startup_event() -> None:
|
|||
f"Loaded configuration from {config_path}:\n"
|
||||
f" endpoints={config.endpoints},\n"
|
||||
f" llama_server_endpoints={config.llama_server_endpoints},\n"
|
||||
f" llama_swap_endpoints={config.llama_swap_endpoints},\n"
|
||||
f" max_concurrent_connections={config.max_concurrent_connections},\n"
|
||||
f" endpoint_config={config.endpoint_config},\n"
|
||||
f" priority_routing={config.priority_routing}"
|
||||
|
|
@ -328,13 +326,6 @@ async def startup_event() -> None:
|
|||
db = TokenDatabase(config.db_path)
|
||||
await db.init_db()
|
||||
|
||||
# Reconcile Responses-API background tasks lost across a restart: their
|
||||
# in-memory asyncio task is gone but the DB row may still read queued /
|
||||
# in_progress, so mark those failed to give polling clients a terminal state.
|
||||
_orphaned = await db.fail_orphaned_responses()
|
||||
if _orphaned:
|
||||
print(f"[startup] Marked {_orphaned} orphaned background response(s) as failed.")
|
||||
|
||||
# Load existing token counts from database
|
||||
async for count_entry in db.load_token_counts():
|
||||
endpoint = count_entry['endpoint']
|
||||
|
|
@ -377,15 +368,8 @@ async def startup_event() -> None:
|
|||
if is_ext_openai_endpoint(ep):
|
||||
app_state["httpx_clients"][ep] = httpx.AsyncClient(timeout=30.0)
|
||||
|
||||
# Native Anthropic Messages-API endpoints are forwarded over httpx. Use a
|
||||
# long read timeout so streamed completions aren't cut short. Closed on
|
||||
# shutdown by the shared httpx_clients cleanup.
|
||||
for ep in config.anthropic_endpoints:
|
||||
app_state["httpx_clients"][ep] = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(300.0, connect=15.0))
|
||||
|
||||
# Create per-endpoint Unix socket sessions for .sock endpoints
|
||||
for ep in llama_endpoints(config):
|
||||
for ep in config.llama_server_endpoints:
|
||||
if _is_unix_socket_endpoint(ep):
|
||||
sock_path = _get_socket_path(ep)
|
||||
sock_connector = aiohttp.UnixConnector(path=sock_path)
|
||||
|
|
@ -402,7 +386,7 @@ async def startup_event() -> None:
|
|||
# client (/api/chat, /api/generate) and the OpenAI client (/v1/* routes),
|
||||
# so warm both; OpenAI-compatible endpoints only need the OpenAI client.
|
||||
_warm_endpoints = config.endpoints + [
|
||||
ep for ep in llama_endpoints(config) if ep not in config.endpoints
|
||||
ep for ep in config.llama_server_endpoints if ep not in config.endpoints
|
||||
]
|
||||
for ep in _warm_endpoints:
|
||||
try:
|
||||
|
|
@ -419,10 +403,29 @@ async def startup_event() -> None:
|
|||
flush_task = asyncio.create_task(flush_buffer())
|
||||
await init_llm_cache(config)
|
||||
|
||||
# Start the aissurance compliance plugin (non-critical; guarded by config).
|
||||
global compliance_plugin
|
||||
if config.compliance.enabled:
|
||||
try:
|
||||
compliance_plugin = AissurancePlugin(config.compliance)
|
||||
await compliance_plugin.start()
|
||||
except Exception as e:
|
||||
print(f"[startup] aissurance compliance plugin failed to start: {e}")
|
||||
compliance_plugin = None
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event() -> None:
|
||||
await close_all_sse_queues()
|
||||
|
||||
# Stop the compliance plugin first (best-effort buffer flush on its own).
|
||||
global compliance_plugin
|
||||
if compliance_plugin is not None:
|
||||
try:
|
||||
await compliance_plugin.stop()
|
||||
except Exception as e:
|
||||
print(f"[shutdown] Error stopping aissurance compliance plugin: {e}")
|
||||
compliance_plugin = None
|
||||
|
||||
# Stop background tasks first so they stop touching the DB before we close it.
|
||||
for t in (token_worker_task, flush_task):
|
||||
if t is not None:
|
||||
|
|
|
|||
121
routing.py
121
routing.py
|
|
@ -32,9 +32,6 @@ from backends.health import _is_fresh
|
|||
from backends.normalize import (
|
||||
is_ext_openai_endpoint,
|
||||
is_openai_compatible,
|
||||
is_llama_server,
|
||||
is_anthropic_endpoint,
|
||||
llama_endpoints,
|
||||
get_tracking_model,
|
||||
)
|
||||
from backends.probe import fetch
|
||||
|
|
@ -96,24 +93,13 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
"""
|
||||
config = get_config()
|
||||
# 1️⃣ Gather advertised‑model sets for all endpoints concurrently
|
||||
# Include config.endpoints plus any llama-server / llama-swap endpoints
|
||||
llama_eps_extra = [ep for ep in llama_endpoints(config) if ep not in config.endpoints]
|
||||
anthropic_eps_extra = [ep for ep in config.anthropic_endpoints if ep not in config.endpoints]
|
||||
all_endpoints = config.endpoints + llama_eps_extra + anthropic_eps_extra
|
||||
# Include both config.endpoints and config.llama_server_endpoints
|
||||
llama_eps_extra = [ep for ep in config.llama_server_endpoints if ep not in config.endpoints]
|
||||
all_endpoints = config.endpoints + llama_eps_extra
|
||||
|
||||
# Build the probe tasks in the SAME order as ``all_endpoints`` so the
|
||||
# gathered results stay aligned for the ``zip(all_endpoints, advertised_sets)``
|
||||
# below. (A previous partition into non-OpenAI / OpenAI groups reordered the
|
||||
# tasks relative to ``all_endpoints``, pairing each endpoint with another
|
||||
# endpoint's advertised models — so an OpenAI-compatible backend listed
|
||||
# before an Ollama one would inherit the Ollama model set and get a request
|
||||
# for a model it cannot serve, 404ing. See issue #128.)
|
||||
def _advertised_task(ep: str):
|
||||
if is_openai_compatible(ep) or is_anthropic_endpoint(ep):
|
||||
return fetch.available_models(ep, config.api_keys.get(ep))
|
||||
return fetch.available_models(ep)
|
||||
|
||||
tag_tasks = [_advertised_task(ep) for ep in all_endpoints]
|
||||
tag_tasks = [fetch.available_models(ep) for ep in config.endpoints if not is_openai_compatible(ep)]
|
||||
tag_tasks += [fetch.available_models(ep, config.api_keys.get(ep)) for ep in config.endpoints if is_openai_compatible(ep)]
|
||||
tag_tasks += [fetch.available_models(ep, config.api_keys.get(ep)) for ep in llama_eps_extra]
|
||||
advertised_sets = await asyncio.gather(*tag_tasks)
|
||||
|
||||
# 2️⃣ Filter endpoints that advertise the requested model
|
||||
|
|
@ -128,7 +114,7 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
model_without_latest = model.split(":latest")[0]
|
||||
candidate_endpoints = [
|
||||
ep for ep, models in zip(all_endpoints, advertised_sets)
|
||||
if model_without_latest in models and (is_ext_openai_endpoint(ep) or is_llama_server(ep))
|
||||
if model_without_latest in models and (is_ext_openai_endpoint(ep) or ep in config.llama_server_endpoints)
|
||||
]
|
||||
if not candidate_endpoints:
|
||||
# Only add :latest suffix if model doesn't already have a version suffix
|
||||
|
|
@ -216,43 +202,6 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
def utilization_ratio(ep: str) -> float:
|
||||
return tracking_usage(ep) / get_max_connections(ep)
|
||||
|
||||
def total_load(ep: str) -> int:
|
||||
"""Sum of in-flight requests across *all* models on the endpoint."""
|
||||
return sum(usage_counts.get(ep, {}).values())
|
||||
|
||||
# How many models each candidate currently has *resident* (from the
|
||||
# /api/ps probe). With infinite keep-alive a model stays loaded long
|
||||
# after its in-flight count drops to zero, so this is the signal that
|
||||
# spreads *distinct* models across backends.
|
||||
ep_loaded_counts = {
|
||||
ep: len(models) for ep, models in zip(candidate_endpoints, loaded_sets)
|
||||
}
|
||||
|
||||
def loaded_count(ep: str) -> int:
|
||||
return ep_loaded_counts.get(ep, 0)
|
||||
|
||||
def pick_least_loaded(eps: list[str]) -> str:
|
||||
"""Pick the least-committed endpoint, breaking ties at random.
|
||||
|
||||
Ordering key is ``(total_load, loaded_count)``:
|
||||
|
||||
* ``total_load`` (in-flight requests across *all* models) keeps a
|
||||
request off a backend already busy with a *different* model —
|
||||
otherwise the per-model count reads zero everywhere and the
|
||||
ranking is discarded (cold model B landing on the box serving A).
|
||||
* ``loaded_count`` (number of *resident* models) then spreads
|
||||
distinct models across backends. Two different cold models (27b,
|
||||
35b) requested back-to-back must not pile onto the same box: once
|
||||
27b is resident there, that box has loaded_count 1 while the idle
|
||||
backends have 0, so the next cold model prefers an empty backend
|
||||
even though every backend reports zero in-flight load.
|
||||
|
||||
``random.choice`` only breaks genuine ties on both keys, so a single
|
||||
idle cluster still distributes the very first cold model evenly."""
|
||||
best = min((total_load(ep), loaded_count(ep)) for ep in eps)
|
||||
tied = [ep for ep in eps if (total_load(ep), loaded_count(ep)) == best]
|
||||
return random.choice(tied)
|
||||
|
||||
# Priority map: position in all_endpoints list (lower = higher priority)
|
||||
ep_priority = {ep: i for i, ep in enumerate(all_endpoints)}
|
||||
|
||||
|
|
@ -286,11 +235,15 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
loaded_and_free.sort(key=utilization_ratio)
|
||||
selected = loaded_and_free[0]
|
||||
else:
|
||||
# All endpoints here already have the model loaded, so there
|
||||
# is no model-switching cost to optimise for. Pick the least
|
||||
# *total*-loaded one (tie broken at random) so we steer away
|
||||
# from a backend busy serving other models.
|
||||
selected = pick_least_loaded(loaded_and_free)
|
||||
# Sort ascending for load balancing — all endpoints here already have the
|
||||
# model loaded, so there is no model-switching cost to optimise for.
|
||||
loaded_and_free.sort(key=tracking_usage)
|
||||
# When all candidates are equally idle, randomise to avoid always picking
|
||||
# the first entry in a stable sort.
|
||||
if all(tracking_usage(ep) == 0 for ep in loaded_and_free):
|
||||
selected = random.choice(loaded_and_free)
|
||||
else:
|
||||
selected = loaded_and_free[0]
|
||||
else:
|
||||
# 4️⃣ Endpoints among the candidates that simply have a free slot
|
||||
endpoints_with_free_slot = [
|
||||
|
|
@ -304,10 +257,14 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
endpoints_with_free_slot.sort(key=utilization_ratio)
|
||||
selected = endpoints_with_free_slot[0]
|
||||
else:
|
||||
# Prefer the endpoint with the lowest *total* load so the
|
||||
# cold-start cost lands on genuinely idle hardware rather
|
||||
# than a backend already busy with a different model.
|
||||
selected = pick_least_loaded(endpoints_with_free_slot)
|
||||
# Sort by total endpoint load (ascending) to prefer idle endpoints.
|
||||
endpoints_with_free_slot.sort(
|
||||
key=lambda ep: sum(usage_counts.get(ep, {}).values())
|
||||
)
|
||||
if all(tracking_usage(ep) == 0 for ep in endpoints_with_free_slot):
|
||||
selected = random.choice(endpoints_with_free_slot)
|
||||
else:
|
||||
selected = endpoints_with_free_slot[0]
|
||||
else:
|
||||
# 5️⃣ All candidate endpoints are saturated – pick the least-busy one (will queue)
|
||||
if config.priority_routing:
|
||||
|
|
@ -323,23 +280,15 @@ async def choose_endpoint(model: str, reserve: bool = True,
|
|||
if reserve:
|
||||
usage_counts[selected][tracking_model] += 1
|
||||
snapshot = _capture_snapshot()
|
||||
# The slot is now reserved. Any failure (including CancelledError on client
|
||||
# disconnect) between here and `return` would otherwise leak it — the caller
|
||||
# never receives (endpoint, tracking_model) and so can never decrement it.
|
||||
try:
|
||||
if snapshot is not None:
|
||||
await _distribute_snapshot(snapshot)
|
||||
# Record / refresh affinity *after* releasing usage_lock.
|
||||
if reserve and config.conversation_affinity and affinity_key:
|
||||
expires_at = time.monotonic() + config.conversation_affinity_ttl
|
||||
async with _affinity_lock:
|
||||
_affinity_map[affinity_key] = (selected, model, expires_at)
|
||||
if len(_affinity_map) > _AFFINITY_MAX_ENTRIES:
|
||||
now = time.monotonic()
|
||||
for k in [k for k, v in _affinity_map.items() if v[2] < now]:
|
||||
_affinity_map.pop(k, None)
|
||||
except BaseException:
|
||||
if reserve:
|
||||
await decrement_usage(selected, tracking_model)
|
||||
raise
|
||||
if snapshot is not None:
|
||||
await _distribute_snapshot(snapshot)
|
||||
# Record / refresh affinity *after* releasing usage_lock.
|
||||
if reserve and config.conversation_affinity and affinity_key:
|
||||
expires_at = time.monotonic() + config.conversation_affinity_ttl
|
||||
async with _affinity_lock:
|
||||
_affinity_map[affinity_key] = (selected, model, expires_at)
|
||||
if len(_affinity_map) > _AFFINITY_MAX_ENTRIES:
|
||||
now = time.monotonic()
|
||||
for k in [k for k, v in _affinity_map.items() if v[2] < now]:
|
||||
_affinity_map.pop(k, None)
|
||||
return selected, tracking_model
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
230
test/test_compliance.py
Normal 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
|
||||
|
|
@ -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=[])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1,477 +0,0 @@
|
|||
"""Tests for the Anthropic Messages API support (api/messages.py + requests/anthropic.py).
|
||||
|
||||
Covers the pure translation layer, the translated (Ollama-style) and native
|
||||
(configured-anthropic-endpoint) backend paths, thinking→reasoning mapping,
|
||||
streaming event shape, token counting, and the nomyo-cache reflection.
|
||||
"""
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from types import SimpleNamespace as NS
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import orjson
|
||||
import pytest
|
||||
|
||||
import router
|
||||
from api import messages as api_messages
|
||||
from requests import anthropic as at
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Pure translation unit tests (no app / no I/O)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TestRequestTranslation:
|
||||
def test_system_string(self):
|
||||
chat = at.anthropic_messages_to_chat("be brief", [{"role": "user", "content": "hi"}])
|
||||
assert chat[0] == {"role": "system", "content": "be brief"}
|
||||
assert chat[1] == {"role": "user", "content": "hi"}
|
||||
|
||||
def test_system_blocks_join(self):
|
||||
chat = at.anthropic_messages_to_chat(
|
||||
[{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
||||
[{"role": "user", "content": "hi"}])
|
||||
assert chat[0] == {"role": "system", "content": "a\n\nb"}
|
||||
|
||||
def test_image_base64_to_data_url(self):
|
||||
chat = at.anthropic_messages_to_chat(None, [{
|
||||
"role": "user", "content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAA"}},
|
||||
{"type": "text", "text": "what?"},
|
||||
]}])
|
||||
parts = chat[0]["content"]
|
||||
assert parts[0] == {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}}
|
||||
assert parts[1] == {"type": "text", "text": "what?"}
|
||||
|
||||
def test_tool_use_and_result_roundtrip(self):
|
||||
chat = at.anthropic_messages_to_chat(None, [
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "text", "text": "calling"},
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "get", "input": {"x": 1}}]},
|
||||
{"role": "user", "content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_1", "content": "42"}]},
|
||||
])
|
||||
assert chat[0]["role"] == "assistant"
|
||||
assert chat[0]["content"] == "calling"
|
||||
assert chat[0]["tool_calls"][0]["id"] == "toolu_1"
|
||||
assert chat[0]["tool_calls"][0]["function"]["arguments"] == '{"x":1}'
|
||||
assert chat[1] == {"role": "tool", "tool_call_id": "toolu_1", "content": "42"}
|
||||
|
||||
def test_tool_result_blocks_flattened(self):
|
||||
chat = at.anthropic_messages_to_chat(None, [{
|
||||
"role": "user", "content": [
|
||||
{"type": "tool_result", "tool_use_id": "t1",
|
||||
"content": [{"type": "text", "text": "line"}]}]}])
|
||||
assert chat[0] == {"role": "tool", "tool_call_id": "t1", "content": "line"}
|
||||
|
||||
def test_tools_and_choice(self):
|
||||
sp = at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 10, "tools": [{"name": "get", "description": "d",
|
||||
"input_schema": {"type": "object"}}],
|
||||
"tool_choice": {"type": "tool", "name": "get"}},
|
||||
[], "m")
|
||||
assert sp["tools"] == [{"type": "function", "function": {
|
||||
"name": "get", "description": "d", "parameters": {"type": "object"}}}]
|
||||
assert sp["tool_choice"] == {"type": "function", "function": {"name": "get"}}
|
||||
|
||||
def test_tool_choice_variants(self):
|
||||
assert at.tool_choice_anthropic_to_chat({"type": "auto"}) == "auto"
|
||||
assert at.tool_choice_anthropic_to_chat({"type": "any"}) == "required"
|
||||
assert at.tool_choice_anthropic_to_chat({"type": "none"}) == "none"
|
||||
|
||||
def test_server_tools_dropped(self):
|
||||
# web_search has a type but no input_schema → not runnable locally
|
||||
assert at.tools_anthropic_to_chat([{"type": "web_search_20260209", "name": "web_search"}]) is None
|
||||
|
||||
def test_stop_sequences_and_sampling(self):
|
||||
sp = at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 5, "stop_sequences": ["X"], "temperature": 0.2, "top_p": 0.9}, [], "m")
|
||||
assert sp["stop"] == ["X"]
|
||||
assert sp["temperature"] == 0.2
|
||||
assert sp["top_p"] == 0.9
|
||||
|
||||
def test_thinking_maps_to_reasoning_effort(self):
|
||||
assert at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 1000}}, [], "m"
|
||||
)["reasoning_effort"] == "low"
|
||||
assert at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 4096}}, [], "m"
|
||||
)["reasoning_effort"] == "medium"
|
||||
assert at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 5, "thinking": {"type": "enabled", "budget_tokens": 20000}}, [], "m"
|
||||
)["reasoning_effort"] == "high"
|
||||
assert "reasoning_effort" not in at.anthropic_to_chat_send_params(
|
||||
{"max_tokens": 5, "thinking": {"type": "disabled"}}, [], "m")
|
||||
|
||||
|
||||
class TestResponseTranslation:
|
||||
def test_content_blocks_order(self):
|
||||
blocks = at.chat_message_to_content_blocks({
|
||||
"role": "assistant", "reasoning_content": "hmm", "content": "answer",
|
||||
"tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": '{"a":1}'}}]})
|
||||
assert [b["type"] for b in blocks] == ["thinking", "text", "tool_use"]
|
||||
assert blocks[0]["thinking"] == "hmm"
|
||||
assert blocks[1]["text"] == "answer"
|
||||
assert blocks[2] == {"type": "tool_use", "id": "c1", "name": "f", "input": {"a": 1}}
|
||||
|
||||
def test_stop_reason_mapping(self):
|
||||
assert at.finish_reason_to_stop_reason("stop") == "end_turn"
|
||||
assert at.finish_reason_to_stop_reason("length") == "max_tokens"
|
||||
assert at.finish_reason_to_stop_reason("tool_calls") == "tool_use"
|
||||
assert at.finish_reason_to_stop_reason("stop", has_tool_use=True) == "tool_use"
|
||||
|
||||
def test_usage_mapping_no_cache(self):
|
||||
u = at.usage_chat_to_anthropic({"prompt_tokens": 7, "completion_tokens": 3})
|
||||
assert u == {"input_tokens": 7, "output_tokens": 3,
|
||||
"cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}
|
||||
|
||||
def test_usage_mapping_cached_tokens_subtracted(self):
|
||||
# Backend reports 5 of 7 prompt tokens served from its prefix cache.
|
||||
u = at.usage_chat_to_anthropic({
|
||||
"prompt_tokens": 7, "completion_tokens": 3,
|
||||
"prompt_tokens_details": {"cached_tokens": 5}})
|
||||
assert u == {"input_tokens": 2, "output_tokens": 3,
|
||||
"cache_creation_input_tokens": 0, "cache_read_input_tokens": 5}
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Streaming translator (pure)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _chunk(content=None, tool_calls=None, reasoning=None, finish_reason=None):
|
||||
delta = NS(content=content, tool_calls=tool_calls)
|
||||
if reasoning is not None:
|
||||
delta.reasoning_content = reasoning
|
||||
return NS(choices=[NS(delta=delta, finish_reason=finish_reason)], usage=None)
|
||||
|
||||
|
||||
def _usage_chunk(p, c, cached=None):
|
||||
usage = NS(prompt_tokens=p, completion_tokens=c)
|
||||
if cached is not None:
|
||||
usage.prompt_tokens_details = NS(cached_tokens=cached)
|
||||
return NS(choices=[], usage=usage)
|
||||
|
||||
|
||||
async def _collect(translator, gen):
|
||||
frames = []
|
||||
async for sse in translator.events(gen):
|
||||
frames.append(sse.decode())
|
||||
return _parse_sse("".join(frames))
|
||||
|
||||
|
||||
def _parse_sse(text):
|
||||
out = []
|
||||
for frame in text.strip().split("\n\n"):
|
||||
if not frame.strip():
|
||||
continue
|
||||
etype = data = None
|
||||
for line in frame.splitlines():
|
||||
if line.startswith("event: "):
|
||||
etype = line[len("event: "):]
|
||||
elif line.startswith("data: "):
|
||||
data = orjson.loads(line[len("data: "):])
|
||||
out.append((etype, data))
|
||||
return out
|
||||
|
||||
|
||||
class TestStreamTranslator:
|
||||
async def test_text_stream(self):
|
||||
async def gen():
|
||||
yield _chunk(content="Hel")
|
||||
yield _chunk(content="lo", finish_reason="stop")
|
||||
yield _usage_chunk(3, 5)
|
||||
tr = at.ChatToMessagesStream("msg_1", "m")
|
||||
events = await _collect(tr, gen())
|
||||
types = [e[0] for e in events]
|
||||
assert types[0] == "message_start"
|
||||
assert types[-1] == "message_stop"
|
||||
assert "content_block_start" in types and "content_block_stop" in types
|
||||
text = "".join(d["delta"]["text"] for t, d in events
|
||||
if t == "content_block_delta" and d["delta"]["type"] == "text_delta")
|
||||
assert text == "Hello"
|
||||
md = [d for t, d in events if t == "message_delta"][0]
|
||||
assert md["delta"]["stop_reason"] == "end_turn"
|
||||
assert md["usage"]["output_tokens"] == 5
|
||||
assert tr.content_blocks == [{"type": "text", "text": "Hello"}]
|
||||
|
||||
async def test_cached_tokens_surface_as_cache_read(self):
|
||||
async def gen():
|
||||
yield _chunk(content="hi", finish_reason="stop")
|
||||
yield _usage_chunk(10, 2, cached=6)
|
||||
tr = at.ChatToMessagesStream("msg_1", "m")
|
||||
events = await _collect(tr, gen())
|
||||
md = [d for t, d in events if t == "message_delta"][0]
|
||||
assert md["usage"]["cache_read_input_tokens"] == 6
|
||||
assert md["usage"]["input_tokens"] == 4 # 10 prompt − 6 cached
|
||||
assert md["usage"]["cache_creation_input_tokens"] == 0
|
||||
|
||||
async def test_thinking_then_text(self):
|
||||
async def gen():
|
||||
yield _chunk(reasoning="think ")
|
||||
yield _chunk(reasoning="more")
|
||||
yield _chunk(content="answer", finish_reason="stop")
|
||||
tr = at.ChatToMessagesStream("msg_1", "m")
|
||||
events = await _collect(tr, gen())
|
||||
# thinking block (index 0) then text block (index 1)
|
||||
starts = [(d["index"], d["content_block"]["type"])
|
||||
for t, d in events if t == "content_block_start"]
|
||||
assert starts == [(0, "thinking"), (1, "text")]
|
||||
think = "".join(d["delta"]["thinking"] for t, d in events
|
||||
if t == "content_block_delta" and d["delta"]["type"] == "thinking_delta")
|
||||
assert think == "think more"
|
||||
assert tr.content_blocks[0] == {"type": "thinking", "thinking": "think more"}
|
||||
|
||||
async def test_tool_call_stream(self):
|
||||
tc0 = NS(index=0, id="call_1", function=NS(name="lookup", arguments='{"q":'))
|
||||
tc1 = NS(index=0, id=None, function=NS(name=None, arguments='"hi"}'))
|
||||
|
||||
async def gen():
|
||||
yield _chunk(tool_calls=[tc0])
|
||||
yield _chunk(tool_calls=[tc1], finish_reason="tool_calls")
|
||||
yield _usage_chunk(4, 2)
|
||||
tr = at.ChatToMessagesStream("msg_1", "m")
|
||||
events = await _collect(tr, gen())
|
||||
assert any(t == "content_block_start" and d["content_block"]["type"] == "tool_use"
|
||||
for t, d in events)
|
||||
partial = "".join(d["delta"]["partial_json"] for t, d in events
|
||||
if t == "content_block_delta" and d["delta"]["type"] == "input_json_delta")
|
||||
assert partial == '{"q":"hi"}'
|
||||
md = [d for t, d in events if t == "message_delta"][0]
|
||||
assert md["delta"]["stop_reason"] == "tool_use"
|
||||
assert tr.content_blocks[-1] == {"type": "tool_use", "id": "call_1",
|
||||
"name": "lookup", "input": {"q": "hi"}}
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Cache-hit replay
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
class TestCacheReplay:
|
||||
def test_message_object_to_sse_roundtrip(self):
|
||||
msg = at.build_message_object(
|
||||
message_id="msg_1", model="m",
|
||||
content_blocks=[{"type": "text", "text": "hi"}],
|
||||
stop_reason="end_turn",
|
||||
usage=at.usage_chat_to_anthropic({"prompt_tokens": 2, "completion_tokens": 1}))
|
||||
events = _parse_sse(at.message_object_to_sse(msg).decode())
|
||||
types = [e[0] for e in events]
|
||||
assert types[0] == "message_start"
|
||||
assert types[-1] == "message_stop"
|
||||
text = "".join(d["delta"]["text"] for t, d in events
|
||||
if t == "content_block_delta" and d["delta"]["type"] == "text_delta")
|
||||
assert text == "hi"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Route-level tests
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@contextmanager
|
||||
def _enter(*cms):
|
||||
with ExitStack() as stack:
|
||||
for cm in cms:
|
||||
stack.enter_context(cm)
|
||||
yield
|
||||
|
||||
|
||||
def _fake_completion(content="hello", usage=(3, 5), reasoning=None, tool_calls=None,
|
||||
finish_reason="stop", cached=None):
|
||||
md = {"role": "assistant", "content": content}
|
||||
if reasoning is not None:
|
||||
md["reasoning_content"] = reasoning
|
||||
if tool_calls is not None:
|
||||
md["tool_calls"] = tool_calls
|
||||
msg = MagicMock()
|
||||
msg.model_dump.return_value = md
|
||||
usage_dump = {"prompt_tokens": usage[0], "completion_tokens": usage[1],
|
||||
"total_tokens": sum(usage)}
|
||||
if cached is not None:
|
||||
usage_dump["prompt_tokens_details"] = {"cached_tokens": cached}
|
||||
usage_obj = MagicMock()
|
||||
usage_obj.model_dump.return_value = usage_dump
|
||||
return NS(choices=[NS(message=msg, finish_reason=finish_reason)], usage=usage_obj)
|
||||
|
||||
|
||||
def _patch_backend(native=False, endpoint="http://ollama:11434", cache=None):
|
||||
return (
|
||||
patch.object(api_messages, "choose_endpoint",
|
||||
AsyncMock(return_value=(endpoint, "test-model:latest"))),
|
||||
patch.object(api_messages, "decrement_usage", AsyncMock()),
|
||||
patch.object(api_messages, "is_anthropic_endpoint", return_value=native),
|
||||
patch.object(api_messages, "_make_openai_client", return_value=MagicMock()),
|
||||
patch.object(api_messages, "get_llm_cache", return_value=cache),
|
||||
)
|
||||
|
||||
|
||||
class TestTranslatedRoute:
|
||||
async def test_nonstream(self, client):
|
||||
with _enter(*_patch_backend(native=False),
|
||||
patch.object(api_messages, "create_chat_with_retries",
|
||||
AsyncMock(return_value=_fake_completion("hello world")))):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "test-model", "max_tokens": 100,
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["type"] == "message"
|
||||
assert body["role"] == "assistant"
|
||||
assert body["content"] == [{"type": "text", "text": "hello world"}]
|
||||
assert body["stop_reason"] == "end_turn"
|
||||
assert body["usage"]["input_tokens"] == 3 and body["usage"]["output_tokens"] == 5
|
||||
assert body["id"].startswith("msg_")
|
||||
|
||||
async def test_nonstream_cached_tokens(self, client):
|
||||
with _enter(*_patch_backend(native=False),
|
||||
patch.object(api_messages, "create_chat_with_retries",
|
||||
AsyncMock(return_value=_fake_completion(
|
||||
"hi", usage=(10, 4), cached=6)))):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "test-model", "max_tokens": 100,
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
u = resp.json()["usage"]
|
||||
assert u["cache_read_input_tokens"] == 6
|
||||
assert u["input_tokens"] == 4
|
||||
assert u["cache_creation_input_tokens"] == 0
|
||||
|
||||
async def test_missing_max_tokens_400(self, client):
|
||||
with _enter(*_patch_backend(native=False)):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "m", "messages": [{"role": "user", "content": "hi"}]})
|
||||
assert resp.status_code == 400
|
||||
|
||||
async def test_nonstream_tool_use(self, client):
|
||||
tc = [{"id": "c1", "function": {"name": "get", "arguments": '{"a":1}'}}]
|
||||
with _enter(*_patch_backend(native=False),
|
||||
patch.object(api_messages, "create_chat_with_retries",
|
||||
AsyncMock(return_value=_fake_completion(
|
||||
content=None, tool_calls=tc, finish_reason="tool_calls")))):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "m", "max_tokens": 50,
|
||||
"messages": [{"role": "user", "content": "call it"}]})
|
||||
body = resp.json()
|
||||
assert body["stop_reason"] == "tool_use"
|
||||
assert body["content"][0] == {"type": "tool_use", "id": "c1", "name": "get", "input": {"a": 1}}
|
||||
|
||||
async def test_stream_event_sequence(self, client):
|
||||
async def _text():
|
||||
yield _chunk(content="Hi", finish_reason="stop")
|
||||
yield _usage_chunk(3, 2)
|
||||
with _enter(*_patch_backend(native=False),
|
||||
patch.object(api_messages, "create_chat_with_retries",
|
||||
AsyncMock(return_value=_text()))):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "m", "max_tokens": 50, "stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert resp.headers["content-type"].startswith("text/event-stream")
|
||||
events = _parse_sse(resp.content.decode())
|
||||
types = [e[0] for e in events]
|
||||
assert types[0] == "message_start" and types[-1] == "message_stop"
|
||||
text = "".join(d["delta"]["text"] for t, d in events
|
||||
if t == "content_block_delta" and d["delta"]["type"] == "text_delta")
|
||||
assert text == "Hi"
|
||||
|
||||
async def test_thinking_passed_as_reasoning_effort(self, client):
|
||||
captured = {}
|
||||
|
||||
async def _spy(oclient, send_params, endpoint, model, tracking_model):
|
||||
captured.update(send_params)
|
||||
return _fake_completion("ok")
|
||||
|
||||
with _enter(*_patch_backend(native=False),
|
||||
patch.object(api_messages, "create_chat_with_retries", _spy)):
|
||||
await client.post("/v1/messages",
|
||||
json={"model": "m", "max_tokens": 50,
|
||||
"thinking": {"type": "enabled", "budget_tokens": 20000},
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert captured["reasoning_effort"] == "high"
|
||||
|
||||
|
||||
class TestCacheRoute:
|
||||
async def test_hit_reports_cache_read(self, client):
|
||||
stored = at.build_message_object(
|
||||
message_id="msg_old", model="m",
|
||||
content_blocks=[{"type": "text", "text": "cached"}], stop_reason="end_turn",
|
||||
usage=at.usage_chat_to_anthropic({"prompt_tokens": 9, "completion_tokens": 4}))
|
||||
fake_cache = MagicMock()
|
||||
fake_cache.get_chat = AsyncMock(return_value=orjson.dumps(stored))
|
||||
with _enter(*_patch_backend(native=False, cache=fake_cache)):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "m", "max_tokens": 50, "nomyo": {"cache": True},
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
body = resp.json()
|
||||
assert body["content"] == [{"type": "text", "text": "cached"}]
|
||||
assert body["usage"]["cache_read_input_tokens"] == 9
|
||||
assert body["usage"]["input_tokens"] == 0
|
||||
assert body["id"].startswith("msg_") and body["id"] != "msg_old"
|
||||
|
||||
|
||||
class TestNativeRoute:
|
||||
def _fake_client(self, *, post_return=None, stream_frames=None, status=200):
|
||||
client = MagicMock()
|
||||
if post_return is not None:
|
||||
resp = MagicMock()
|
||||
resp.status_code = status
|
||||
resp.json.return_value = post_return
|
||||
client.post = AsyncMock(return_value=resp)
|
||||
if stream_frames is not None:
|
||||
class _Stream:
|
||||
async def __aenter__(self_):
|
||||
return self_
|
||||
async def __aexit__(self_, *a):
|
||||
return False
|
||||
async def aiter_bytes(self_):
|
||||
for f in stream_frames:
|
||||
yield f
|
||||
client.stream = MagicMock(return_value=_Stream())
|
||||
return client
|
||||
|
||||
async def test_nonstream_passthrough(self, client):
|
||||
upstream = {"id": "msg_upstream", "type": "message", "role": "assistant",
|
||||
"content": [{"type": "text", "text": "native hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 2, "output_tokens": 3}}
|
||||
fake = self._fake_client(post_return=upstream)
|
||||
with _enter(*_patch_backend(native=True, endpoint="https://api.anthropic.com"),
|
||||
patch.object(api_messages, "_anthropic_http_client", return_value=fake)):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "claude-x", "max_tokens": 50,
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["content"][0]["text"] == "native hi"
|
||||
# request forwarded verbatim with stream disabled, nomyo stripped
|
||||
sent = fake.post.call_args.kwargs["json"]
|
||||
assert sent["stream"] is False and "nomyo" not in sent
|
||||
headers = fake.post.call_args.kwargs["headers"]
|
||||
assert headers["x-api-key"] and headers["anthropic-version"]
|
||||
|
||||
async def test_stream_passthrough(self, client):
|
||||
frames = [
|
||||
b'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":5}}}\n\n',
|
||||
b'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":7}}\n\n',
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
]
|
||||
fake = self._fake_client(stream_frames=frames)
|
||||
tracked = []
|
||||
with _enter(*_patch_backend(native=True, endpoint="https://api.anthropic.com"),
|
||||
patch.object(api_messages, "_anthropic_http_client", return_value=fake),
|
||||
patch.object(api_messages, "_track",
|
||||
AsyncMock(side_effect=lambda *a: tracked.append(a)))):
|
||||
resp = await client.post("/v1/messages",
|
||||
json={"model": "claude-x", "max_tokens": 50, "stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}]})
|
||||
body = resp.content.decode()
|
||||
assert "message_start" in body and "message_stop" in body
|
||||
# usage parsed out of the proxied stream for token tracking
|
||||
assert tracked and tracked[0][2] == 5 and tracked[0][3] == 7
|
||||
|
||||
|
||||
class TestCountTokens:
|
||||
async def test_local_estimate(self, client):
|
||||
with _enter(patch.object(api_messages, "choose_endpoint",
|
||||
AsyncMock(return_value=("http://ollama:11434", "m"))),
|
||||
patch.object(api_messages, "is_anthropic_endpoint", return_value=False)):
|
||||
resp = await client.post("/v1/messages/count_tokens",
|
||||
json={"model": "m",
|
||||
"messages": [{"role": "user", "content": "count me"}]})
|
||||
assert resp.status_code == 200
|
||||
assert isinstance(resp.json()["input_tokens"], int)
|
||||
assert resp.json()["input_tokens"] > 0
|
||||
|
|
@ -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"
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue