Jaroslav Beneš
9e2caeac48
A streamed completion carries no token counts unless you ask for them, and
`stream_options: {include_usage: true}` is how. Not every server implements
it, and an unknown key is a 400 from some -- the same hazard as sending a
tools array to an endpoint without support. So it is asked for once per base
URL per process, and an endpoint that refuses is remembered and retried
without it. The retry is safe because the status is checked before a single
line is read: nothing has been yielded, so there is nothing to duplicate.
chunk_usage() reads the resulting chunk. It needed no change to the loop
above it: a usage chunk carries `choices: []`, which is exactly the shape
delta_text, delta_reasoning, delta_tool_calls and finish_reason have always
returned early on. All-zero counts are treated as absent, because some
servers attach zeros to every chunk and the real numbers only at the end.
services/tokens.py is the fallback for endpoints that never report: four
characters to a token, counting the tools array because thirteen schemas is
a meaningful slice of a short window, and counting nothing for an image
because its cost depends on tiling and an invented number would be worse
than the omission. Crude on purpose -- a real tokeniser means one per model
family, for a figure that is displayed beside a tilde.
Nothing uses any of this yet.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
417 lines
16 KiB
Python
417 lines
16 KiB
Python
"""Client for OpenAI-compatible chat endpoints.
|
|
|
|
Deliberately plain httpx rather than the official SDK. The target is not just
|
|
api.openai.com but LM Studio, vLLM, llama.cpp, Ollama's compatibility layer,
|
|
OpenRouter and anything else exposing /v1 -- and they differ in small ways. A
|
|
thin client passes request parameters through untouched and is tolerant about
|
|
what comes back, which is exactly what talking to all of them requires.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from lembas.config import settings
|
|
from lembas.db.models import Connection
|
|
from lembas.services.crypto import decrypt
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMError(Exception):
|
|
"""An upstream failure with a message fit to show a user.
|
|
|
|
Every failure path in this module raises this rather than letting an httpx
|
|
or JSON exception escape, so callers have exactly one thing to catch and
|
|
the chat UI always has something intelligible to display.
|
|
"""
|
|
|
|
def __init__(self, message: str, *, status_code: int | None = None) -> None:
|
|
super().__init__(message)
|
|
self.message = message
|
|
self.status_code = status_code
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Endpoint:
|
|
"""Everything needed to call a connection, with the key already decrypted.
|
|
|
|
A frozen snapshot rather than the ORM object because streaming outlives the
|
|
request that started it, and a detached SQLAlchemy instance is a trap.
|
|
"""
|
|
|
|
base_url: str
|
|
api_key: str
|
|
extra_headers: dict[str, str]
|
|
name: str = ""
|
|
|
|
@classmethod
|
|
def from_connection(cls, connection: Connection) -> Endpoint:
|
|
return cls(
|
|
base_url=connection.base_url.rstrip("/"),
|
|
api_key=decrypt(connection.api_key_encrypted),
|
|
extra_headers=dict(connection.extra_headers_json or {}),
|
|
name=connection.name,
|
|
)
|
|
|
|
def url(self, path: str) -> str:
|
|
# Accept both "http://host:1234" and "http://host:1234/v1" so users do
|
|
# not have to guess which form this expects.
|
|
base = self.base_url
|
|
if not base.endswith("/v1") and "/v1/" not in base:
|
|
base = f"{base}/v1"
|
|
return f"{base}/{path.lstrip('/')}"
|
|
|
|
def headers(self) -> dict[str, str]:
|
|
headers = {"Content-Type": "application/json", **self.extra_headers}
|
|
# Local endpoints frequently need no key at all; sending an empty
|
|
# bearer token makes some of them reject the request outright.
|
|
if self.api_key:
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
return headers
|
|
|
|
|
|
def describe_http_error(exc: httpx.HTTPStatusError) -> str:
|
|
"""Turn an upstream error response into something worth reading.
|
|
|
|
Public because the audio and search clients talk to the same class of
|
|
server and want the same translation; LLMError stays the one thing a
|
|
caller has to catch.
|
|
|
|
Providers put the useful part in wildly different places, so try the common
|
|
shapes before falling back to the raw body.
|
|
"""
|
|
status = exc.response.status_code
|
|
detail = ""
|
|
try:
|
|
payload = exc.response.json()
|
|
if isinstance(payload, dict):
|
|
error = payload.get("error")
|
|
if isinstance(error, dict):
|
|
detail = error.get("message", "")
|
|
elif isinstance(error, str):
|
|
detail = error
|
|
detail = detail or payload.get("message", "") or payload.get("detail", "")
|
|
except (ValueError, json.JSONDecodeError):
|
|
detail = exc.response.text[:300]
|
|
|
|
friendly = {
|
|
401: "The API key was rejected.",
|
|
403: "The API key is not permitted to use this model.",
|
|
404: "The endpoint or model was not found.",
|
|
429: "Rate limited by the provider.",
|
|
}.get(status)
|
|
|
|
if friendly and detail:
|
|
return f"{friendly} {detail}"
|
|
return friendly or detail or f"The endpoint returned HTTP {status}."
|
|
|
|
|
|
def wrap_transport_error(exc: httpx.RequestError, endpoint: Endpoint) -> LLMError:
|
|
if isinstance(exc, httpx.ConnectError):
|
|
return LLMError(
|
|
f"Could not reach {endpoint.base_url}. Is the endpoint running and "
|
|
f"the URL correct?"
|
|
)
|
|
if isinstance(exc, httpx.TimeoutException):
|
|
return LLMError(
|
|
f"{endpoint.base_url} did not respond within "
|
|
f"{settings.request_timeout:.0f}s."
|
|
)
|
|
return LLMError(f"Could not reach {endpoint.base_url}: {exc}")
|
|
|
|
|
|
async def list_models(endpoint: Endpoint) -> list[dict[str, Any]]:
|
|
"""Fetch the models a connection advertises via GET /v1/models."""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.get(endpoint.url("models"), headers=endpoint.headers())
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
except httpx.HTTPStatusError as exc:
|
|
raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
|
|
except httpx.RequestError as exc:
|
|
raise wrap_transport_error(exc, endpoint) from exc
|
|
except ValueError as exc:
|
|
raise LLMError("The endpoint returned a response that was not JSON.") from exc
|
|
|
|
# The spec says {"data": [...]}, but some servers return a bare list.
|
|
entries = payload.get("data", payload) if isinstance(payload, dict) else payload
|
|
if not isinstance(entries, list):
|
|
raise LLMError("The endpoint's model list was not in the expected format.")
|
|
|
|
models = []
|
|
for entry in entries:
|
|
if isinstance(entry, dict) and entry.get("id"):
|
|
models.append(entry)
|
|
elif isinstance(entry, str):
|
|
models.append({"id": entry})
|
|
return models
|
|
|
|
|
|
# Where the runners that bother to advertise a context length put it. There is
|
|
# no standard field, so this is a list of what the common ones actually emit.
|
|
_CONTEXT_KEYS = ("context_length", "max_model_len", "context_window", "max_context_length")
|
|
|
|
# Below the first, the number is not a context length; above the second it is a
|
|
# typo or a different unit. Either way, better to record nothing than a wrong
|
|
# figure a percentage would then be computed from.
|
|
MIN_CONTEXT = 256
|
|
MAX_CONTEXT = 10_000_000
|
|
|
|
|
|
def context_from(entry: dict[str, Any]) -> int:
|
|
"""A model's context length as advertised by /v1/models, or 0 if it is not.
|
|
|
|
Strings are accepted because some servers quote the number, but only when
|
|
they are digits alone -- "8192 tokens" is a label, not a measurement.
|
|
"""
|
|
candidates = [entry.get(key) for key in _CONTEXT_KEYS]
|
|
meta = entry.get("meta")
|
|
if isinstance(meta, dict):
|
|
candidates += [meta.get("n_ctx"), *(meta.get(key) for key in _CONTEXT_KEYS)]
|
|
|
|
for value in candidates:
|
|
if isinstance(value, bool):
|
|
continue
|
|
if isinstance(value, str):
|
|
value = value.strip()
|
|
if not value.isdigit():
|
|
continue
|
|
value = int(value)
|
|
if isinstance(value, int) and MIN_CONTEXT <= value <= MAX_CONTEXT:
|
|
return value
|
|
return 0
|
|
|
|
|
|
# Endpoints that rejected `stream_options`, so it is asked for once per base URL
|
|
# per process and then never again. Not persisted: it is a property of whatever
|
|
# is running there now, and a restart is the right time to find out afresh.
|
|
_NO_STREAM_OPTIONS: set[str] = set()
|
|
|
|
|
|
async def stream_chat(
|
|
endpoint: Endpoint,
|
|
payload: dict[str, Any],
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
"""Stream a chat completion, yielding each parsed SSE data object.
|
|
|
|
Yields the raw upstream chunks; interpreting them is the caller's job. The
|
|
terminating "[DONE]" sentinel is consumed here and not yielded.
|
|
|
|
`stream_options` asks for the final usage chunk, which is the only way to
|
|
learn what a streamed reply actually cost. Not every server implements it,
|
|
and an unknown key is a 400 from some of them -- the same hazard as sending
|
|
a `tools` array to an endpoint without support. So it is asked for once,
|
|
and an endpoint that refuses is remembered and never asked again. Retrying
|
|
is safe because the status is checked before a single line is read: nothing
|
|
has been yielded, so there is nothing to duplicate.
|
|
"""
|
|
wants_usage = endpoint.base_url not in _NO_STREAM_OPTIONS
|
|
|
|
try:
|
|
async for chunk in _stream_once(endpoint, payload, usage=wants_usage):
|
|
yield chunk
|
|
except LLMError as exc:
|
|
if not wants_usage or exc.status_code not in (400, 422):
|
|
raise
|
|
_NO_STREAM_OPTIONS.add(endpoint.base_url)
|
|
log.info(
|
|
"%s rejected stream_options; token counts will be estimated there",
|
|
endpoint.base_url,
|
|
)
|
|
async for chunk in _stream_once(endpoint, payload, usage=False):
|
|
yield chunk
|
|
|
|
|
|
async def _stream_once(
|
|
endpoint: Endpoint,
|
|
payload: dict[str, Any],
|
|
*,
|
|
usage: bool,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
body = {**payload, "stream": True}
|
|
if usage:
|
|
body["stream_options"] = {"include_usage": True}
|
|
|
|
try:
|
|
async with (
|
|
httpx.AsyncClient(timeout=settings.request_timeout) as client,
|
|
client.stream(
|
|
"POST",
|
|
endpoint.url("chat/completions"),
|
|
headers=endpoint.headers(),
|
|
json=body,
|
|
) as response,
|
|
):
|
|
if response.status_code >= 400:
|
|
# The body has not been read yet on a streaming response, and
|
|
# the error detail is in it.
|
|
await response.aread()
|
|
response.raise_for_status()
|
|
|
|
async for line in response.aiter_lines():
|
|
line = line.strip()
|
|
if not line or line.startswith(":"):
|
|
continue # keep-alive comment
|
|
if not line.startswith("data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
return
|
|
try:
|
|
yield json.loads(data)
|
|
except json.JSONDecodeError:
|
|
# A malformed frame is not worth killing a reply over.
|
|
log.warning("skipping unparseable SSE frame: %.120s", data)
|
|
continue
|
|
|
|
except httpx.HTTPStatusError as exc:
|
|
raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
|
|
except httpx.RequestError as exc:
|
|
raise wrap_transport_error(exc, endpoint) from exc
|
|
|
|
|
|
async def complete(endpoint: Endpoint, payload: dict[str, Any]) -> str:
|
|
"""Non-streaming completion. Used for short internal calls like auto-titling."""
|
|
body = {**payload, "stream": False}
|
|
try:
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
response = await client.post(
|
|
endpoint.url("chat/completions"), headers=endpoint.headers(), json=body
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
except httpx.HTTPStatusError as exc:
|
|
raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
|
|
except httpx.RequestError as exc:
|
|
raise wrap_transport_error(exc, endpoint) from exc
|
|
except ValueError as exc:
|
|
raise LLMError("The endpoint returned a response that was not JSON.") from exc
|
|
|
|
try:
|
|
return data["choices"][0]["message"]["content"] or ""
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise LLMError("The endpoint returned no completion.") from exc
|
|
|
|
|
|
def delta_reasoning(chunk: dict[str, Any]) -> str:
|
|
"""Pull a reasoning delta out of one streamed chunk.
|
|
|
|
Providers disagree on the field name -- llama.cpp, llama-swap and vLLM use
|
|
``reasoning_content``, some others just ``reasoning`` -- so both are read.
|
|
Models that emit ``<think>`` tags inline in ``content`` instead are handled
|
|
by lembas.services.reasoning.
|
|
"""
|
|
try:
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
return ""
|
|
delta = choices[0].get("delta") or {}
|
|
for field in ("reasoning_content", "reasoning"):
|
|
value = delta.get(field)
|
|
if isinstance(value, str) and value:
|
|
return value
|
|
return ""
|
|
except (AttributeError, TypeError):
|
|
return ""
|
|
|
|
|
|
def delta_tool_calls(chunk: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""Pull tool-call fragments out of one streamed chunk.
|
|
|
|
Each entry carries an ``index`` and, across chunks, a name that arrives
|
|
once and an ``arguments`` string that arrives in pieces. Reassembling them
|
|
is lembas.services.tools.ToolCallAccumulator's job; this only extracts.
|
|
"""
|
|
try:
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
return []
|
|
calls = (choices[0].get("delta") or {}).get("tool_calls")
|
|
return calls if isinstance(calls, list) else []
|
|
except (AttributeError, TypeError):
|
|
return []
|
|
|
|
|
|
def finish_reason(chunk: dict[str, Any]) -> str:
|
|
"""Why the model stopped, when the chunk says so.
|
|
|
|
``tool_calls`` here is the signal that the reply is not an answer but a
|
|
request to run something and come back. Some servers send ``stop`` even
|
|
when they emitted tool calls, so the accumulator's contents are the real
|
|
authority and this is only a hint.
|
|
"""
|
|
try:
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
return ""
|
|
return choices[0].get("finish_reason") or ""
|
|
except (AttributeError, TypeError):
|
|
return ""
|
|
|
|
|
|
def chunk_usage(chunk: dict[str, Any]) -> dict[str, int] | None:
|
|
"""Token counts from a usage chunk, or None if this is not one.
|
|
|
|
A usage chunk carries `choices: []`, which is exactly the shape delta_text,
|
|
delta_reasoning, delta_tool_calls and finish_reason all return early on --
|
|
they have always tolerated it, so nothing else needs to change to let one
|
|
through.
|
|
|
|
Fields are read defensively because "the endpoint returned something odd"
|
|
must never be the reason a reply fails; a bad shape simply means no counts.
|
|
"""
|
|
try:
|
|
raw = chunk.get("usage")
|
|
except AttributeError:
|
|
return None
|
|
if not isinstance(raw, dict):
|
|
return None
|
|
|
|
counts: dict[str, int] = {}
|
|
for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
|
|
value = raw.get(key)
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
continue
|
|
if value >= 0:
|
|
counts[key] = int(value)
|
|
|
|
# Some servers send a usage object of zeros on every chunk and the real
|
|
# numbers only at the end. All-zero is indistinguishable from that, and
|
|
# treating it as an answer would freeze the count at nothing.
|
|
if not counts or not any(counts.values()):
|
|
return None
|
|
counts.setdefault(
|
|
"total_tokens", counts.get("prompt_tokens", 0) + counts.get("completion_tokens", 0)
|
|
)
|
|
return counts
|
|
|
|
|
|
def delta_text(chunk: dict[str, Any]) -> str:
|
|
"""Pull the text out of one streamed chunk, tolerating provider variation."""
|
|
try:
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
return ""
|
|
delta = choices[0].get("delta") or {}
|
|
content = delta.get("content")
|
|
if isinstance(content, str):
|
|
return content
|
|
# Some providers send content as a list of typed parts even in deltas.
|
|
if isinstance(content, list):
|
|
return "".join(
|
|
part.get("text", "")
|
|
for part in content
|
|
if isinstance(part, dict) and part.get("type") == "text"
|
|
)
|
|
return ""
|
|
except (AttributeError, TypeError):
|
|
return ""
|