Coverage for src/lilbee/providers/fleet/client.py: 100%
684 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Thin httpx client for one llama-server OpenAI endpoint (local inference)."""
3from __future__ import annotations
5import base64
6import contextlib
7import json
8import logging
9import math
10import ssl
11import threading
12import time
13from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
14from concurrent.futures import ThreadPoolExecutor
15from typing import Any, Literal, TypedDict, TypeVar, overload
17import httpx
18import numpy as np
19import numpy.typing as npt
21from lilbee.core.config import cfg
22from lilbee.core.vectors import Vector
23from lilbee.providers.base import (
24 THINK_CLOSE_TAG,
25 THINK_OPEN_TAG,
26 ChatResult,
27 ChatToolResult,
28 ClosableIterator,
29 FinishReason,
30 ProviderError,
31 ProviderErrorKind,
32 StreamFinish,
33 TokenUsage,
34 ToolCall,
35 ToolCallDelta,
36)
37from lilbee.providers.fleet.adapters import LLM_RERANK_CONCURRENCY
38from lilbee.providers.fleet.normalize import ChatMessage, to_alternating
39from lilbee.providers.roles import RerankMode
41_PROVIDER_NAME = "llama-server"
43# Fleet clients only ever talk to a loopback llama-server over plain HTTP, so TLS is
44# never negotiated. httpx still builds a default SSL context per client (loads the
45# system CA bundle, ~13 ms each and slower on macOS via the keychain), which is pure
46# overhead paid on every fleet reload. Build one minimal context and share it so a
47# reload doesn't reload the CA bundle for each replica.
48_LOOPBACK_SSL_CONTEXT = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
49_LOOPBACK_SSL_CONTEXT.check_hostname = False
50_LOOPBACK_SSL_CONTEXT.verify_mode = ssl.CERT_NONE
51# Reranker pair format: query and candidate are joined with this separator into
52# one document so a cross-encoder GGUF scores the pair as a single sequence.
53_RERANK_PAIR_SEPARATOR = "</s></s>"
54# LLM reranker: score each candidate by the yes/no first-token logprob.
55_LLM_RERANK_PROMPT = (
56 "Judge whether the document is relevant to the query. "
57 "Answer with only 'yes' or 'no'.\n\nQuery: {query}\nDocument: {document}"
58)
59_LLM_RERANK_TOP_LOGPROBS = 20
60_YES_LABEL = "yes"
61_NO_LABEL = "no"
62_LLM_RERANK_NO_VERDICT_ERROR = (
63 "The reranker model never answered 'yes' or 'no', so its relevance scores are "
64 "unusable. Its chat template does not fit the relevance prompt. Choose a GGUF "
65 "built for reranking, set reranker_type to cross_encoder, or adapt "
66 "reranker_prompt to the model's expected format."
67)
68# Max sequences per /v1/embeddings request. Like the in-process backstop, a
69# batch is bounded by BOTH the token budget (the server's n_batch, == token_cap)
70# and this sequence count: a corpus of many tiny chunks would otherwise pack one
71# request past the server's batch/sequence limit and trip a 500.
72_EMBED_N_SEQ_MAX = 64
73# Estimate a chunk's token count from its character length so the bulk embed
74# path packs sub-batches without a /tokenize round-trip per input. The factor is
75# held below the corpus average (data.extract.chunk.CHARS_PER_TOKEN, 4 for Latin text)
76# so the estimate over-counts tokens and a sub-batch never packs past the
77# server's n_batch (== token_cap). Rerank pairs are token-dense (the separator
78# is several tokens in a few chars), so an under-count can slip an over-cap
79# pair through; the rerank path re-truncates exactly on the server's overflow
80# error instead of paying a /tokenize round trip per pair up front.
81_EMBED_EST_CHARS_PER_TOKEN = 3
82# llama-swap's error body when the spawned llama-server exited before serving.
83_UPSTREAM_DIED_MARKER = "exited prematurely"
84# llama-server's exit line when the port lilbee picked was taken by the time
85# the server bound it (the pick-then-bind gap spans the whole lazy-spawn wait,
86# so a passing ephemeral connection can occupy it). The occupation is
87# transient: llama-swap re-spawns on the next request against the same port
88# and normally binds.
89_BIND_FAILURE_MARKER = "couldn't bind HTTP server socket"
90# llama-server's 500 body when one input exceeds the physical batch (n_batch).
91_BATCH_OVERFLOW_MARKER = "too large to process"
94class _ChatToolSpecFunction(TypedDict, total=False):
95 """The ``function`` payload of an OpenAI tool definition (wire shape)."""
97 name: str
98 description: str
99 parameters: dict[str, Any]
102class ChatTool(TypedDict, total=False):
103 """One OpenAI tool definition sent in a chat request (wire shape)."""
105 type: str
106 function: _ChatToolSpecFunction
109# Some GGUF chat templates (Mistral-Nemo, Cohere command-r) reject a standard
110# OpenAI tool exchange: they require plain user/assistant turns to alternate and
111# raise a Jinja exception on the tool role or two same-role turns in a row. Rather
112# than fail a real request and parse the engine's error text, the client probes
113# the live template once per server with this representative tool exchange: if the
114# server rejects it as sent but renders the to_alternating() form, the model is
115# flagged so every later request is reshaped up front. Two assistant tool-call
116# turns separated by tool results is the minimal shape that trips strict
117# alternation; max_tokens=1 keeps the probe to template rendering, not generation.
118_ALTERNATION_PROBE_TOOLS: list[ChatTool] = [
119 {
120 "type": "function",
121 "function": {
122 "name": "probe",
123 "description": "Probe whether the chat template renders a tool exchange.",
124 # A single declared property (rather than an empty object) so a grammar
125 # that requires at least one parameter still renders the probe call.
126 "parameters": {
127 "type": "object",
128 "properties": {"query": {"type": "string"}},
129 "required": ["query"],
130 },
131 },
132 }
133]
134_ALTERNATION_PROBE_MESSAGES: list[ChatMessage] = [
135 {"role": "system", "content": "You are a helpful assistant."},
136 {"role": "user", "content": "Look something up."},
137 {
138 "role": "assistant",
139 "content": "",
140 "tool_calls": [
141 {
142 "id": "probe-1",
143 "type": "function",
144 "function": {"name": "probe", "arguments": '{"query": "x"}'},
145 }
146 ],
147 },
148 {"role": "tool", "tool_call_id": "probe-1", "content": "first result"},
149 {
150 "role": "assistant",
151 "content": "",
152 "tool_calls": [
153 {
154 "id": "probe-2",
155 "type": "function",
156 "function": {"name": "probe", "arguments": '{"query": "x"}'},
157 }
158 ],
159 },
160 {"role": "tool", "tool_call_id": "probe-2", "content": "second result"},
161 {"role": "user", "content": "Summarize."},
162]
163_ALTERNATION_PROBE_OPTIONS = {"max_tokens": 1}
164# The probe holds _alternation_lock across its request, so it uses a short, bounded
165# timeout rather than the chat default: a slow/wedged replica yields an inconclusive
166# (transient) result and a re-probe instead of blocking every first chat on the lock.
167_ALTERNATION_PROBE_TIMEOUT_S = 30.0
168_UPSTREAM_LOG_TAIL_CHARS = 2000
169_UPSTREAM_LOG_TIMEOUT_S = 2.0
170# Enough reads to cover one replay of llama-swap's 100KB per-model ring at
171# httpx's 64KB ceiling per chunk, with room to spare. The route keeps streaming
172# live lines afterwards, so without a bound this would cost the read timeout on
173# every death.
174_UPSTREAM_LOG_MAX_CHUNKS = 8
176log = logging.getLogger(__name__)
179def _estimate_tokens(text: str) -> int:
180 """Conservative (over-counting) token estimate from character length."""
181 return max(1, -(-len(text) // _EMBED_EST_CHARS_PER_TOKEN))
184def _raise_for_status(resp: httpx.Response) -> None:
185 """Raise including the server's error body, which ``raise_for_status`` drops.
187 A llama-server failure otherwise surfaces as a bare "Internal Server Error"
188 with no cause; the response body carries the actual reason (oversize prompt,
189 decode failure, ...), which both diagnosis and the user-facing error need.
190 """
191 if resp.is_success:
192 return
193 resp.read() # streaming responses aren't read yet; a no-op for buffered ones
194 body = resp.text.strip()
195 # llama-server reports an oversize prompt/conversation as a 400 whose body
196 # carries the "exceed_context_size_error" type. Tag it CONTEXT_OVERFLOW with a
197 # user-facing message so the chat route returns a clean context_length_exceeded
198 # (400) instead of a generic internal_error -- a long conversation that fills
199 # the window then reads as "too long", not "Internal server error".
200 if resp.status_code == _HTTP_BAD_REQUEST and (
201 "exceed_context_size" in body.lower() or "context size" in body.lower()
202 ):
203 raise ProviderError(
204 "The conversation exceeds this model's context window. "
205 "Start a new conversation or shorten the input.",
206 provider=_PROVIDER_NAME,
207 kind=ProviderErrorKind.CONTEXT_OVERFLOW,
208 )
209 # A 429 (slots full) is transient: a cold replica fleet rejects the first ingest
210 # fan-out until its slots load. Tag RATE_LIMIT so the caller backs off and retries
211 # instead of dropping the input.
212 if resp.status_code == _HTTP_TOO_MANY_REQUESTS:
213 raise ProviderError(
214 "llama-server is busy (HTTP 429); replicas may still be warming.",
215 provider=_PROVIDER_NAME,
216 kind=ProviderErrorKind.RATE_LIMIT,
217 )
218 detail = f": {body[:600]}" if body else ""
219 kind = _classify_error(resp.status_code, body)
220 # llama-swap masks a dead server as "exited prematurely"; surface the server's
221 # own captured output (a missing CUDA runtime, a model load failure, a bind
222 # error) so the real exit reason reaches the caller, not only the log.
223 if _UPSTREAM_DIED_MARKER in body:
224 tail = _upstream_failure_tail(resp)
225 if tail:
226 detail = f"{detail}\nupstream server output:\n{tail}"
227 classified = classify_upstream_death(tail)
228 if classified is not None:
229 kind = classified
230 raise ProviderError(
231 f"llama-server returned HTTP {resp.status_code}{detail}",
232 provider=_PROVIDER_NAME,
233 kind=kind,
234 )
237# What the engine prints when a device allocation fails during load. Every
238# backend words it differently and all of them mean the same thing: the plan
239# asked for more memory than the device had. Taken from the emit sites in
240# upstream rather than guessed, and matched lowercased.
241#
242# One entry covers CUDA, HIP and MUSA: the vendor headers #define cudaMalloc to
243# their own allocator, but the log string in ggml-cuda.cu is a literal, so an
244# AMD or Moore Threads build still prints "cudaMalloc failed". A separate
245# hipMalloc marker would match nothing.
246#
247# Vulkan is the one that needs its own wording. It is where every AMD and Intel
248# GPU lands, and it says neither "out of memory" nor "failed to allocate".
249_OOM_MARKERS: tuple[str, ...] = (
250 "out of memory",
251 "failed to allocate", # Metal's buffer failure, and most generic paths
252 "cudamalloc failed", # CUDA, HIP and MUSA alike
253 "device memory allocation of size", # ggml-vulkan's fatal allocation failure
254 "outofdevicememory", # a vk::OutOfDeviceMemoryError that reached the log
255 "unable to allocate",
256 "insufficient memory",
257 "out_of_device_memory", # SYCL, which exits through the runtime's own code
258 "out_of_resources",
259)
261# Lines that report a failure the engine then works around. ggml-vulkan warns
262# that pinned memory could not be allocated and falls back to unpinned, and that
263# text matches an allocation marker word for word, so a later unrelated death
264# would be read as a memory shortfall and answered with a context reduction.
265_SURVIVABLE_LINE_MARKERS: tuple[str, ...] = ("warning:", "warn:")
268def classify_upstream_death(tail: str) -> ProviderErrorKind | None:
269 """The kind of failure an engine's dying output describes, or ``None``.
271 ``CAPACITY`` for a load that ran out of device memory: retrying the identical
272 launch respawns it into a crash loop, while a smaller context might fit.
273 ``PORT_CONFLICT`` for losing the port-bind race, which is worth retrying
274 because the retry re-drives llama-swap's spawn, and worth naming because a
275 port held for good needs a different port rather than another attempt at the
276 same one. ``None`` leaves the existing classification alone rather than
277 guessing at an unfamiliar death.
278 """
279 fatal = [
280 line
281 for line in tail.lower().splitlines()
282 if not any(marker in line for marker in _SURVIVABLE_LINE_MARKERS)
283 ]
284 if any(marker in line for line in fatal for marker in _OOM_MARKERS):
285 return ProviderErrorKind.CAPACITY
286 if _BIND_FAILURE_MARKER in tail:
287 return ProviderErrorKind.PORT_CONFLICT
288 return None
291# Deaths a role's rebuild can fix, and a retry against the same launch cannot: a
292# memory shortfall needs a smaller plan, a held port needs a different port. Both
293# come from rebuilding the role, which re-plans and re-picks.
294_REBUILDABLE_KINDS = frozenset({ProviderErrorKind.CAPACITY, ProviderErrorKind.PORT_CONFLICT})
297def is_load_capacity_failure(exc: BaseException) -> bool:
298 """True when *exc* is an engine that died for lack of device memory on load."""
299 return isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.CAPACITY
302def is_rebuildable_failure(exc: BaseException) -> bool:
303 """True when rebuilding the role is what stands a chance, not another retry."""
304 return isinstance(exc, ProviderError) and exc.kind in _REBUILDABLE_KINDS
307def _classify_error(status_code: int, body: str) -> ProviderErrorKind:
308 """Error kind from a llama-server/llama-swap error status and body.
310 An input past the server's n_batch is a 500 whose body says "too large to
311 process" (CONTEXT_OVERFLOW, so the embed path re-truncates exactly); a dead
312 upstream is CONNECTION, so the router can mark the replica unhealthy. The
313 body markers win over the status: llama-swap reports a died upstream under
314 gateway statuses too, and that case needs the failover path, not a retry
315 against the same dead server (except a bind-race death, which
316 ``_raise_for_status`` upgrades to SERVER once the upstream tail proves it).
317 A bare gateway error (502/503/504) is a
318 momentarily-unreachable upstream -- restarting, OOM-killed, mid-swap -- so
319 it is SERVER, which the busy retry treats as transient.
320 """
321 if _BATCH_OVERFLOW_MARKER in body:
322 return ProviderErrorKind.CONTEXT_OVERFLOW
323 if _UPSTREAM_DIED_MARKER in body:
324 return ProviderErrorKind.CONNECTION
325 if status_code in _TRANSIENT_GATEWAY_STATUSES:
326 return ProviderErrorKind.SERVER
327 return ProviderErrorKind.UNKNOWN
330def is_connection_failure(exc: Exception) -> bool:
331 """Whether *exc* signals a dead/unreachable replica rather than a model error."""
332 if isinstance(exc, httpx.TransportError):
333 return True
334 # isinstance: only ProviderError carries a kind; other exceptions pass through.
335 return isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.CONNECTION
338def _is_transient_probe_failure(exc: Exception) -> bool:
339 """Whether a probe failure is transient (dead replica, busy 429, gateway
340 error), not a template verdict. A cold replica 429s or 502s its first
341 traffic, so neither response must be read as the template rejecting the
342 exchange."""
343 if is_connection_failure(exc):
344 return True
345 return isinstance(exc, ProviderError) and exc.kind in _TRANSIENT_KINDS
348def _upstream_failure_tail(resp: httpx.Response) -> str:
349 """Return (and log) the dead upstream's recent output, or empty when unreadable."""
350 with contextlib.suppress(httpx.HTTPError, json.JSONDecodeError, KeyError, TypeError):
351 base = str(resp.request.url).split("/v1/")[0]
352 model = json.loads(resp.request.content)["model"]
353 tail = _fetch_log_tail(f"{base}/logs/stream/{model}")
354 if tail:
355 log.warning("%s exited prematurely; recent server output:\n%s", model, tail)
356 return tail
357 return ""
360def _fetch_log_tail(url: str) -> str:
361 """The last ``_UPSTREAM_LOG_TAIL_CHARS`` of llama-swap's log stream for one model.
363 The stream replays the upstream's buffered output then stays open; the read
364 timeout is the cutoff once the replay is drained.
365 """
366 chunks: list[str] = []
367 with (
368 contextlib.suppress(httpx.HTTPError),
369 httpx.stream("GET", url, timeout=_UPSTREAM_LOG_TIMEOUT_S) as stream,
370 ):
371 for taken, chunk in enumerate(stream.iter_text(), start=1):
372 # Bounded by chunk count, not by the tail size. llama-swap replays a
373 # model's whole ring in one write and httpx hands over at most 64KB
374 # at a time, so stopping at the first chunk past the tail size
375 # returned the head of a warm model's log, where the fatal last line
376 # never is. Only the tail is retained as the replay goes by, and the
377 # route streams live lines afterwards, so the count is what keeps
378 # this from waiting out the timeout on every death.
379 chunks.append(chunk)
380 chunks = ["".join(chunks)[-_UPSTREAM_LOG_TAIL_CHARS:]]
381 if taken >= _UPSTREAM_LOG_MAX_CHUNKS:
382 break
383 return "".join(chunks)[-_UPSTREAM_LOG_TAIL_CHARS:]
386# llama-server L2-normalizes pooled embeddings by default (embd_normalize=2);
387# every embeddings request sends embd_normalize=-1 so the engine returns raw
388# vectors, and so a rank-pooling rerank score (a single value per pair) is not
389# collapsed to +-1 by normalization. The server only exposes this per request
390# body, not as a startup flag.
391_EMBD_NORMALIZE_NONE = -1
392# Vectors come back as a base64 float32 buffer: parsing thousands of JSON float
393# literals per batch is CPU-bound and holds the GIL, which caps embedding
394# throughput below what the GPUs can feed regardless of how many are dispatching.
395_EMBED_ENCODING_FORMAT = "base64"
396# The engine writes the raw float buffer in host byte order; supported targets are
397# all little-endian.
398_EMBED_VECTOR_DTYPE = "<f4"
399# Rank pooling puts the pair's relevance score in the vector's first slot.
400_RANK_SCORE_INDEX = 0
401_UNREADABLE_EMBEDDING_ERROR = (
402 "The embedding server returned vectors lilbee could not read. Update the "
403 "inference engine: base64 embedding responses need llama-server b4391 or newer."
404)
405_NO_RERANK_SCORE_ERROR = "The reranker returned no relevance score for a candidate."
406_HEALTH_PATH = "/health"
407_CHAT_PATH = "/v1/chat/completions"
408_EMBED_PATH = "/v1/embeddings"
409_TOKENIZE_PATH = "/tokenize"
410_DETOKENIZE_PATH = "/detokenize"
411# llama-swap proxies native (non-OpenAI) llama.cpp routes only under
412# /upstream/<model>/...; the bare /tokenize path 404s (it routes /v1/* by the
413# body's model field, but a native route carries no such field).
414_UPSTREAM_PREFIX = "/upstream"
415# Match the in-process tokenizer call (llm.tokenize(text, add_bos=True, special=False)):
416# the server adds BOS via add_special and leaves special-token strings unparsed.
417_TOKENIZE_ADD_SPECIAL = True
418_TOKENIZE_PARSE_SPECIAL = False
419_HTTP_OK = 200
420_HTTP_BAD_REQUEST = 400
421_HTTP_TOO_MANY_REQUESTS = 429
422# Gateway statuses llama-swap returns while an upstream is unreachable
423# (502 crashing/restarting, 503 unavailable, 504 gateway timeout). The request
424# succeeds once the upstream is back, so these must never terminalize a call.
425_TRANSIENT_GATEWAY_STATUSES = frozenset({502, 503, 504})
426# Error kinds the busy retry treats as transient: a 429 (slots still loading)
427# and a bare gateway error (upstream momentarily unreachable) both clear on
428# their own once the server is ready again.
429_TRANSIENT_KINDS = frozenset(
430 {ProviderErrorKind.RATE_LIMIT, ProviderErrorKind.SERVER, ProviderErrorKind.PORT_CONFLICT}
431)
432_DONE_SENTINEL = "[DONE]"
433_DATA_PREFIX = "data:"
434_DEFAULT_TIMEOUT_S = 300.0
435# Short, separate timeout for /health: a server can wedge under heavy prompt
436# processing, and readiness/monitor polls must not block on the request timeout.
437_HEALTH_TIMEOUT_S = 5.0
438# Retry a server-busy (HTTP 429) response with exponential backoff (capped): a
439# cold replica fleet 429s the first fan-out until its slots load. Interactive
440# callers fail fast after this short budget (~15s).
441_BUSY_RETRIES = 6
442_BUSY_BACKOFF_BASE_S = 0.5
443_BUSY_BACKOFF_MAX_S = 8.0
444# Bulk embed ingest is background work, so it waits out a full cold start rather
445# than dropping files: an 8B embedder warming while a large chat model loads on
446# neighboring cards can take well past the interactive budget. Capped backoff
447# keeps the total near ~80s (0.5+1+2+4 then 8 each), which covers a real warmup.
448_EMBED_BUSY_RETRIES = 14
449# Half-open recovery: a replica marked unhealthy becomes routable again after
450# this cool-down. Recovery is probe-by-traffic and unmetered: every concurrent
451# caller sees it routable once cooled down (a success restores it, another
452# connection failure re-stamps the cool-down).
453_UNHEALTHY_RETRY_S = 30.0
454_T = TypeVar("_T")
457class ChatDeadlineError(ProviderError):
458 """A bounded chat exceeded its caller-supplied total wall-clock deadline.
460 Distinct from a transport/server error so a deadline-bounded caller (vision
461 OCR) can word its own timeout message and skip failover without matching
462 error strings. Its ``UNKNOWN`` kind keeps it out of ``is_connection_failure``.
463 """
466def retry_on_busy(
467 call: Callable[[], _T], *, retries: int = _BUSY_RETRIES, deadline: float | None = None
468) -> _T:
469 """Run *call*, retrying transient failures (429, gateway errors) with capped backoff.
471 A cold replica fleet 429s the first fan-out until its slots load, and a
472 replica restarting mid-run answers 502 until it is back; backing off and
473 retrying turns both drops into successes. With a *deadline*
474 (``time.monotonic`` epoch) the retry waits out the server until that
475 deadline -- a page on a deep OCR queue keeps waiting for a genuinely free
476 slot instead of dropping after a fixed budget. Without one, *retries* bounds
477 the attempts. Non-transient errors (and the final still-failing response)
478 propagate.
479 """
480 delay = _BUSY_BACKOFF_BASE_S
481 attempt = 0
482 while True:
483 try:
484 return call()
485 except ProviderError as exc:
486 if exc.kind not in _TRANSIENT_KINDS:
487 raise
488 attempt += 1
489 exhausted = (
490 time.monotonic() + delay >= deadline if deadline is not None else attempt >= retries
491 )
492 if exhausted:
493 raise
494 time.sleep(delay)
495 delay = min(delay * 2, _BUSY_BACKOFF_MAX_S)
498class LlamaServerClient:
499 """Calls one llama-server's OpenAI surface. Tracks in-flight requests so the
500 fleet router can pick the least-busy replica."""
502 def __init__(
503 self,
504 base_url: str,
505 model: str,
506 *,
507 http: httpx.Client | None = None,
508 token_cap: int | None = None,
509 timeout: float = _DEFAULT_TIMEOUT_S,
510 rerank_mode: RerankMode | None = None,
511 inline_reasoning: bool = False,
512 embed_busy_deadline_s: float | None = None,
513 ) -> None:
514 self._base = base_url.rstrip("/")
515 self._model = model
516 # Cold-load budget (seconds) the embed path waits out a still-warming replica
517 # before dropping the input, in place of the short attempt cap. Set on the
518 # EMBED-role client to the same ceiling llama-swap keeps the server alive for,
519 # so a bulk ingest never gives up while the replica is legitimately loading.
520 # None (rerank, chat, vision, self-check) keeps the fixed interactive budget.
521 self._embed_busy_deadline_s = embed_busy_deadline_s
522 self._http = http or httpx.Client(
523 base_url=self._base, timeout=timeout, verify=_LOOPBACK_SSL_CONTEXT
524 )
525 self._owns_http = http is None
526 # Chat-role clients re-inline server-extracted reasoning as <think> text;
527 # the other roles (vision OCR) keep dropping it, as their servers already did.
528 self._inline_reasoning = inline_reasoning
529 # Per-slot context for embed/rerank servers: inputs longer than this are
530 # token-truncated (via the server's tokenizer) before embedding, mirroring
531 # the in-process backstop. None for chat/vision, which don't truncate inputs.
532 self._token_cap = token_cap
533 # LLM => score candidates by yes/no logprob; None/cross-encoder => rank pooling.
534 self._rerank_mode = rerank_mode
535 self.in_flight = 0
536 self._in_flight_lock = threading.Lock()
537 # Live SSE responses, so a cancel can sever the transport from another
538 # thread: a reader blocked in iter_lines cannot see a cooperative
539 # cancel flag, but closing its response unblocks it with an error.
540 self._active_streams: set[httpx.Response] = set()
541 # Whether this server's chat template needs OpenAI tool exchanges reshaped
542 # into strict user/assistant alternation. Determined lazily by a one-time
543 # probe of the live template (see _prepare_chat_messages); None until then.
544 # A client is bound to one model for its lifetime, so the template (hence
545 # the verdict) is fixed once determined.
546 self._needs_alternation: bool | None = None
547 self._alternation_lock = threading.Lock()
548 # Routing health: cleared on a connection-level failure (see _UNHEALTHY_RETRY_S).
549 self._healthy = True
550 # Monotonic stamp of the last mark_unhealthy; consulted only while unhealthy.
551 self._unhealthy_since = 0.0
553 @property
554 def healthy(self) -> bool:
555 """Routable: healthy, or unhealthy past the ``_UNHEALTHY_RETRY_S`` cool-down."""
556 with self._in_flight_lock:
557 if self._healthy:
558 return True
559 return time.monotonic() - self._unhealthy_since >= _UNHEALTHY_RETRY_S
561 def mark_unhealthy(self) -> None:
562 """Record a connection-level failure so the router skips this replica."""
563 with self._in_flight_lock:
564 self._healthy = False
565 self._unhealthy_since = time.monotonic()
567 def mark_healthy(self) -> None:
568 """Restore the replica to the routing pool after a successful call."""
569 with self._in_flight_lock:
570 self._healthy = True
572 def reserve(self) -> None:
573 """Mark a routed request assigned to this replica, at selection time.
575 The router balances on ``in_flight`` but the per-request tracking only
576 bumps it once the HTTP call starts. Under a bulk ingest many threads pick
577 a replica at the same instant, all see the momentarily-idlest one at the
578 same low count, and pile onto it (a thundering herd that leaves the other
579 cards idle). Reserving at selection makes the assignment visible to the
580 next picker so requests spread across replicas. Paired with :meth:`release`.
581 """
582 with self._in_flight_lock:
583 self.in_flight += 1
585 def release(self) -> None:
586 """Release a reservation taken by :meth:`reserve`."""
587 with self._in_flight_lock:
588 self.in_flight -= 1
590 def health(self) -> bool:
591 """True iff ``GET /health`` returns 200 (liveness, not readiness)."""
592 try:
593 resp = self._http.get(_HEALTH_PATH, timeout=_HEALTH_TIMEOUT_S)
594 except httpx.HTTPError:
595 return False
596 return resp.status_code == _HTTP_OK
598 @overload
599 def chat(
600 self,
601 messages: Sequence[Mapping[str, Any]],
602 *,
603 options: dict[str, Any] | None = None,
604 stream: Literal[False] = False,
605 timeout: float | None = None,
606 ) -> str: ...
608 @overload
609 def chat(
610 self,
611 messages: Sequence[Mapping[str, Any]],
612 *,
613 options: dict[str, Any] | None = None,
614 stream: Literal[True],
615 timeout: float | None = None,
616 ) -> Iterator[str]: ...
618 @overload
619 def chat(
620 self,
621 messages: Sequence[Mapping[str, Any]],
622 *,
623 options: dict[str, Any] | None = None,
624 stream: bool,
625 timeout: float | None = None,
626 ) -> str | Iterator[str]: ...
628 def chat(
629 self,
630 messages: Sequence[Mapping[str, Any]],
631 *,
632 options: dict[str, Any] | None = None,
633 stream: bool = False,
634 timeout: float | None = None,
635 ) -> str | Iterator[str]:
636 """Chat completion. Returns the full text, or a token iterator if streaming.
638 ``messages`` accepts both plain ``{role, content: str}`` and multipart
639 ``content`` lists (vision image parts), so the vision path reuses this.
640 ``timeout`` overrides the client default for either path, so a
641 caller-enforced deadline (vision OCR) ends the request itself.
642 """
643 payload: dict[str, Any] = {"model": self._model, "messages": messages, **(options or {})}
644 request_timeout = timeout if timeout is not None else httpx.USE_CLIENT_DEFAULT
645 if stream:
646 return self._chat_stream(payload, request_timeout)
647 with self._track():
648 resp = self._http.post(
649 _CHAT_PATH, json={**payload, "stream": False}, timeout=request_timeout
650 )
651 _raise_for_status(resp)
652 # content is null for a refusal / content-filter stop / empty completion;
653 # coerce to "" (like chat_result/chat_tools) so callers never see "None".
654 return _inline_message_reasoning(
655 resp.json()["choices"][0]["message"], enabled=self._inline_reasoning
656 )
658 def _chat_stream(
659 self, payload: dict[str, Any], timeout: Any = httpx.USE_CLIENT_DEFAULT
660 ) -> Iterator[str]:
661 inliner = _ThinkInliner(enabled=self._inline_reasoning)
662 with (
663 self._track(),
664 self._http.stream(
665 "POST", _CHAT_PATH, json={**payload, "stream": True}, timeout=timeout
666 ) as resp,
667 ):
668 _raise_for_status(resp)
669 with self._abortable(resp):
670 for line in resp.iter_lines():
671 delta = inliner.feed(*_parse_sse_deltas(line))
672 if delta:
673 yield delta
674 tail = inliner.finish()
675 if tail:
676 yield tail
678 def chat_bounded(
679 self,
680 messages: Sequence[Mapping[str, Any]],
681 *,
682 options: dict[str, Any] | None = None,
683 deadline_s: float,
684 ) -> str:
685 """Stream a chat completion and return its text, bounded by a total deadline.
687 httpx float timeouts are per-phase (connect/read/...), never a total
688 budget, so a steadily trickling upstream can pin a worker past its
689 deadline. Streaming in the caller's own thread and checking a monotonic
690 deadline per frame bounds total time: on expiry the ``with`` block closes
691 the stream (releasing the in-flight slot) and raises
692 :class:`ChatDeadlineError`.
693 """
694 payload: dict[str, Any] = {"model": self._model, "messages": messages, **(options or {})}
695 deadline = time.monotonic() + deadline_s
696 inliner = _ThinkInliner(enabled=self._inline_reasoning)
697 parts: list[str] = []
698 with (
699 self._track(),
700 self._http.stream("POST", _CHAT_PATH, json={**payload, "stream": True}) as resp,
701 ):
702 _raise_for_status(resp)
703 for line in resp.iter_lines():
704 if time.monotonic() >= deadline:
705 raise ChatDeadlineError(
706 f"llama-server chat exceeded its {deadline_s:.0f}s deadline.",
707 provider=_PROVIDER_NAME,
708 )
709 parts.append(inliner.feed(*_parse_sse_deltas(line)))
710 parts.append(inliner.finish())
711 return "".join(parts)
713 def chat_tools(
714 self,
715 messages: Sequence[Mapping[str, Any]],
716 *,
717 tools: list[dict[str, Any]],
718 tool_choice: str | dict[str, Any] | None = None,
719 options: dict[str, Any] | None = None,
720 ) -> ChatToolResult:
721 """Non-streaming chat with function tools; returns content + any tool calls.
723 The server is launched with ``--jinja`` so it parses the model's native
724 tool-call syntax into structured ``message.tool_calls``. When a model
725 instead emits a bare-JSON call as content (a native miss), recover it.
726 """
727 payload: dict[str, Any] = {
728 "model": self._model,
729 "messages": self._prepare_chat_messages(messages),
730 "tools": tools,
731 "stream": False,
732 **(options or {}),
733 }
734 if tool_choice is not None:
735 payload["tool_choice"] = tool_choice
736 with self._track():
737 resp = self._http.post(_CHAT_PATH, json=payload)
738 _raise_for_status(resp)
739 message = resp.json()["choices"][0]["message"]
740 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
741 native = _parse_native_tool_calls(message.get("tool_calls"))
742 if native:
743 return ChatToolResult(content=content, tool_calls=native)
744 return _recover_bare_json_tool_calls(content)
746 def chat_result(
747 self,
748 messages: Sequence[Mapping[str, Any]],
749 *,
750 tools: list[dict[str, Any]] | None = None,
751 tool_choice: str | dict[str, Any] | None = None,
752 options: dict[str, Any] | None = None,
753 ) -> ChatResult:
754 """Non-streaming chat returning text, tool calls, and a finish reason.
756 The server is launched with ``--jinja`` so it parses the model's native
757 tool-call syntax into structured ``message.tool_calls``. When a model
758 instead emits a bare-JSON call as content (a native miss), recover it
759 and report ``tool_calls`` as the finish reason. Messages are reshaped to
760 strict alternation up front when this server's template needs it (see
761 :meth:`_prepare_chat_messages`).
762 """
763 payload = self._chat_payload(
764 self._prepare_chat_messages(messages), tools, tool_choice, options, stream=False
765 )
766 with self._track():
767 resp = self._http.post(_CHAT_PATH, json=payload)
768 _raise_for_status(resp)
769 body = dict(resp.json())
770 choice = body["choices"][0]
771 usage = _usage_from_body(body) or TokenUsage()
772 message = choice["message"]
773 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
774 finish_reason = _coerce_finish_reason(choice.get("finish_reason"))
775 native = _parse_native_tool_calls(message.get("tool_calls"))
776 if native:
777 return ChatResult(
778 text=content,
779 tool_calls=tuple(native),
780 finish_reason=finish_reason,
781 usage=usage,
782 )
783 recovered = _recover_bare_json_tool_calls(content)
784 if recovered.tool_calls:
785 return ChatResult(
786 text=recovered.content,
787 tool_calls=tuple(recovered.tool_calls),
788 finish_reason=FinishReason.TOOL_CALLS,
789 usage=usage,
790 )
791 return ChatResult(text=content, tool_calls=(), finish_reason=finish_reason, usage=usage)
793 def chat_stream_items(
794 self,
795 messages: Sequence[Mapping[str, Any]],
796 *,
797 tools: list[dict[str, Any]] | None = None,
798 tool_choice: str | dict[str, Any] | None = None,
799 options: dict[str, Any] | None = None,
800 ) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
801 """Stream text tokens and tool-call deltas from the server's OpenAI SSE.
803 Each SSE chunk's ``choices[0].delta`` carries a ``content`` token and/or
804 a ``tool_calls`` array; both are surfaced as :data:`ChatStreamItem`
805 frames (text strings and :class:`ToolCallDelta`). The dispatch's stream
806 translator accumulates the deltas by ``index``. Messages are reshaped to
807 strict alternation up front when this server's template needs it (see
808 :meth:`_prepare_chat_messages`), so the open never fails on a template
809 that rejects the raw tool exchange.
811 Not a generator: the up-front probe runs when this is called, not deferred
812 to the first iteration, matching the eager non-stream paths.
814 A model that emits a tool call as bare-JSON text instead of native
815 ``tool_calls`` (a native miss, as on the non-stream paths) is recovered by
816 wrapping the raw frames; see :func:`_recover_bare_json_stream`.
817 """
818 prepared = self._prepare_chat_messages(messages)
819 return _recover_bare_json_stream(
820 self._open_chat_stream(prepared, tools, tool_choice, options)
821 )
823 def _open_chat_stream(
824 self,
825 messages: Sequence[Mapping[str, Any]],
826 tools: list[dict[str, Any]] | None,
827 tool_choice: str | dict[str, Any] | None,
828 options: dict[str, Any] | None,
829 ) -> Iterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
830 """Open one SSE chat stream and yield its frames; raises before the first frame."""
831 payload = self._chat_payload(messages, tools, tool_choice, options, stream=True)
832 inliner = _ThinkInliner(enabled=self._inline_reasoning)
833 with (
834 self._track(),
835 self._http.stream("POST", _CHAT_PATH, json=payload) as resp,
836 ):
837 _raise_for_status(resp)
838 with self._abortable(resp):
839 for line in resp.iter_lines():
840 yield from _parse_sse_stream_items(line, inliner)
841 tail = inliner.finish()
842 if tail:
843 yield tail
845 def _chat_payload(
846 self,
847 messages: Sequence[Mapping[str, Any]],
848 tools: Sequence[Mapping[str, Any]] | None,
849 tool_choice: str | dict[str, Any] | None,
850 options: dict[str, Any] | None,
851 *,
852 stream: bool,
853 ) -> dict[str, Any]:
854 """Build the chat-completions request body shared by the stream and non-stream paths."""
855 payload: dict[str, Any] = {"model": self._model, "messages": messages, "stream": stream}
856 if stream:
857 # include_usage makes llama-server emit a final SSE chunk carrying the
858 # token usage (with an empty choices list) just before [DONE].
859 payload["stream_options"] = {"include_usage": True}
860 if tools is not None:
861 payload["tools"] = tools
862 if tool_choice is not None:
863 payload["tool_choice"] = tool_choice
864 payload.update(options or {})
865 return payload
867 def _prepare_chat_messages(
868 self, messages: Sequence[Mapping[str, Any]]
869 ) -> Sequence[Mapping[str, Any]]:
870 """Reshape *messages* to strict alternation when this server's template needs it.
872 The need is detected once per server by :meth:`_ensure_alternation_probed`
873 (a binary accept/reject of a representative tool exchange against the live
874 template), then cached, so real requests are normalized up front rather
875 than failing and retrying.
876 """
877 self._ensure_alternation_probed()
878 if self._needs_alternation:
879 return to_alternating([dict(m) for m in messages])
880 return messages
882 def _ensure_alternation_probed(self) -> None:
883 """Probe the live template once to learn whether it needs alternation.
885 Caches only a conclusive verdict: a transient unreachable server leaves
886 the flag unset so the next request re-probes rather than locking in a
887 wrong answer.
888 """
889 if self._needs_alternation is not None:
890 return
891 with self._alternation_lock:
892 if self._needs_alternation is not None:
893 return
894 verdict = self._probe_alternation()
895 if verdict is not None:
896 self._needs_alternation = verdict
898 def _probe_alternation(self) -> bool | None:
899 """Whether the template needs normalization: ``None`` when undetermined.
901 Renders the probe exchange as sent; if the template accepts it, no
902 normalization is needed. If it rejects it, normalization is needed only
903 when the reshaped exchange is accepted. A transient failure on either
904 render is inconclusive (``None``) so no verdict is cached; a genuine
905 rejection of both forms is a conclusive ``False`` (the template fault is
906 unrelated to alternation, so reshaping would not help).
907 """
908 raw = self._chat_probe(_ALTERNATION_PROBE_MESSAGES)
909 if raw is None:
910 return None # transient; stay undetermined so the next request re-probes
911 if raw:
912 return False # the template renders the raw OpenAI exchange as sent
913 reshaped = self._chat_probe(to_alternating([dict(m) for m in _ALTERNATION_PROBE_MESSAGES]))
914 if reshaped is None:
915 return None # transient on the reshape probe; stay undetermined
916 return reshaped
918 def _chat_probe(self, messages: Sequence[Mapping[str, Any]]) -> bool | None:
919 """Post the probe exchange: ``True`` rendered, ``False`` rejected, ``None`` undetermined.
921 A connection failure or a server-busy (HTTP 429) response is transient and
922 unrelated to the template, so it is undetermined: only a clean render or a
923 genuine rejection is a verdict the caller may cache.
924 """
925 payload = self._chat_payload(
926 messages, _ALTERNATION_PROBE_TOOLS, None, _ALTERNATION_PROBE_OPTIONS, stream=False
927 )
928 try:
929 with self._track():
930 resp = self._http.post(
931 _CHAT_PATH, json=payload, timeout=_ALTERNATION_PROBE_TIMEOUT_S
932 )
933 _raise_for_status(resp)
934 except (ProviderError, httpx.TransportError) as exc:
935 return None if _is_transient_probe_failure(exc) else False
936 return True
938 def embed(self, texts: list[str]) -> list[Vector]:
939 """Embed a batch via ``/v1/embeddings``."""
940 if not texts:
941 # Match the in-process embedder; the server rejects an empty input.
942 return []
943 vectors: list[Vector] = []
944 for sub_batch in self._truncate_and_subbatch(texts, estimate=True):
945 data = self._embed_subbatch(sub_batch)
946 vectors.extend(_embedding_vector(item) for item in data)
947 return vectors
949 def _embed_subbatch(self, sub_batch: list[str]) -> list[dict[str, Any]]:
950 """Embed one estimate-budgeted sub-batch, re-truncating exactly on overflow.
952 ``_estimate_tokens`` is char-based and can under-count token-dense inputs
953 (XML, code), so an estimate-trusted input may still exceed the server's
954 context. On that error -- and only that -- redo the batch with exact
955 server-side tokenization, which truncates the oversize input to the cap.
956 """
957 try:
958 return self._embeddings_call(sub_batch)
959 except ProviderError as exc:
960 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW:
961 raise
962 data: list[dict[str, Any]] = []
963 for exact in self._truncate_and_subbatch(sub_batch, estimate=False):
964 data.extend(self._embeddings_call(exact))
965 return data
967 def rerank(self, query: str, candidates: list[str]) -> list[float]:
968 """Relevance scores via rank-pooling embeddings.
970 The server runs with ``--pooling rank``; we send ``query</s></s>candidate``
971 pairs to ``/v1/embeddings`` and read each item's first embedding value as the
972 score, so the ``/v1/rerank`` template-dependency (and its zero-output failure
973 modes) is moot.
975 All pairs go out in one request per ``_EMBED_N_SEQ_MAX`` sequences: the
976 server queues one task per input, so splitting a query-time pool into
977 per-pair requests only adds HTTP round trips.
978 """
979 if not candidates:
980 return []
981 if self._rerank_mode is RerankMode.LLM:
982 return self._rerank_llm(query, candidates)
983 pairs = [
984 self._fit_estimated(f"{query}{_RERANK_PAIR_SEPARATOR}{candidate}")
985 for candidate in candidates
986 ]
987 scores: list[float] = []
988 for start in range(0, len(pairs), _EMBED_N_SEQ_MAX):
989 data = self._rerank_batch(pairs[start : start + _EMBED_N_SEQ_MAX])
990 scores.extend(_rerank_score(item) for item in data)
991 return scores
993 def _fit_estimated(self, text: str) -> str:
994 """Truncate *text* to the token cap only when its char estimate exceeds it."""
995 if self._token_cap is None:
996 return text
997 return self._fit_input(text, self._token_cap, estimate=True)[0]
999 def _rerank_batch(self, batch: list[str]) -> list[dict[str, Any]]:
1000 """Score one rerank batch, redoing it with exact truncation on overflow.
1002 The char estimate under-counts token-dense pairs, so an over-cap pair
1003 can slip through untruncated; the server rejects it as CONTEXT_OVERFLOW
1004 and the batch is redone with exact server-side tokenization.
1005 """
1006 try:
1007 return self._embeddings_call(batch)
1008 except ProviderError as exc:
1009 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW or self._token_cap is None:
1010 raise
1011 cap = self._token_cap
1012 exact = [self._fit_input(text, cap, estimate=False)[0] for text in batch]
1013 return self._embeddings_call(exact)
1015 def _rerank_llm(self, query: str, candidates: list[str]) -> list[float]:
1016 """Score each candidate by an LLM's yes/no first-token logprob.
1018 Raises ``ProviderError`` when no candidate yields a verdict.
1019 """
1020 template = cfg.reranker_prompt or _LLM_RERANK_PROMPT
1021 workers = min(LLM_RERANK_CONCURRENCY, len(candidates))
1022 with ThreadPoolExecutor(max_workers=workers) as pool:
1023 scores = list(pool.map(lambda c: self._llm_rerank_one(template, query, c), candidates))
1024 if all(score is None for score in scores):
1025 raise ProviderError(_LLM_RERANK_NO_VERDICT_ERROR, provider=_PROVIDER_NAME)
1026 return [0.0 if score is None else score for score in scores]
1028 def _llm_rerank_one(self, template: str, query: str, candidate: str) -> float | None:
1029 """One chat request scoring a single candidate's relevance to the query."""
1030 content = template.format(query=query, document=candidate)
1031 payload = {
1032 "model": self._model,
1033 "messages": [{"role": "user", "content": content}],
1034 "max_tokens": 1,
1035 "temperature": 0,
1036 "logprobs": True,
1037 "top_logprobs": _LLM_RERANK_TOP_LOGPROBS,
1038 "stream": False,
1039 # Scoring reads the first generated token; a thinking template would
1040 # spend it opening a <think> block instead of answering.
1041 "chat_template_kwargs": {"enable_thinking": False},
1042 }
1044 def _call() -> dict[str, Any]:
1045 with self._track():
1046 resp = self._http.post(_CHAT_PATH, json=payload)
1047 _raise_for_status(resp)
1048 return dict(resp.json())
1050 return _llm_rerank_score(_first_token_top_logprobs(retry_on_busy(_call)))
1052 def _embeddings_call(self, inputs: list[str]) -> list[dict[str, Any]]:
1053 """POST one already-budgeted sub-batch to ``/v1/embeddings``; return its data."""
1055 def _call() -> list[dict[str, Any]]:
1056 with self._track():
1057 resp = self._http.post(
1058 _EMBED_PATH,
1059 json={
1060 "model": self._model,
1061 "input": inputs,
1062 "embd_normalize": _EMBD_NORMALIZE_NONE,
1063 "encoding_format": _EMBED_ENCODING_FORMAT,
1064 },
1065 )
1066 _raise_for_status(resp)
1067 data = resp.json()["data"]
1068 if len(data) != len(inputs):
1069 raise ProviderError(
1070 f"Embedder returned {len(data)} vectors for {len(inputs)} inputs",
1071 provider=_PROVIDER_NAME,
1072 )
1073 return list(data)
1075 # Bulk ingest can afford to wait out a cold-start warmup rather than drop
1076 # files. With a cold-load deadline (the EMBED-role client) the retry waits
1077 # out a still-loading replica for the full budget llama-swap keeps it alive,
1078 # instead of dropping the file after the fixed attempt cap; without one the
1079 # fixed count bounds an interactive caller.
1080 if self._embed_busy_deadline_s is not None:
1081 return retry_on_busy(_call, deadline=time.monotonic() + self._embed_busy_deadline_s)
1082 return retry_on_busy(_call, retries=_EMBED_BUSY_RETRIES)
1084 def _truncate_and_subbatch(self, texts: list[str], *, estimate: bool) -> list[list[str]]:
1085 """Token-truncate over-cap inputs, then pack into server-sized sub-batches.
1087 An input longer than ``token_cap`` (the server's per-slot context /
1088 n_batch) is truncated to it via the server's tokenizer, since the server
1089 cannot split a pooled embedding sequence. Inputs are then grouped so each
1090 request stays within both the token budget and ``_EMBED_N_SEQ_MAX``
1091 sequences, bounding per-request size and the busy-retry window. No cap
1092 (chat/vision) sends a single batch untouched.
1094 When ``estimate`` is set the per-input token count comes from
1095 :func:`_estimate_tokens`, and ``/tokenize`` is consulted only for the
1096 rare input whose estimate exceeds the cap -- eliminating a round-trip per
1097 chunk during bulk ingest. ``estimate=False`` (the overflow redo) tokenizes
1098 every input exactly.
1099 """
1100 if self._token_cap is None:
1101 return [texts]
1102 cap = self._token_cap
1103 batches: list[list[str]] = []
1104 current: list[str] = []
1105 current_tokens = 0
1106 for text in texts:
1107 item, item_tokens = self._fit_input(text, cap, estimate=estimate)
1108 if current and (current_tokens + item_tokens > cap or len(current) >= _EMBED_N_SEQ_MAX):
1109 batches.append(current)
1110 current = []
1111 current_tokens = 0
1112 current.append(item)
1113 current_tokens += item_tokens
1114 if current:
1115 batches.append(current)
1116 return batches
1118 def _fit_input(self, text: str, cap: int, *, estimate: bool) -> tuple[str, int]:
1119 """Return ``(input, token_count)`` for one sequence, truncating if over cap.
1121 Estimation short-circuits the common case: an estimate within the cap is
1122 trusted (no ``/tokenize``); only an over-cap estimate is confirmed against
1123 the server tokenizer and truncated if it really exceeds the cap.
1124 """
1125 if estimate:
1126 est = _estimate_tokens(text)
1127 if est <= cap:
1128 return text, est
1129 tokens = self._tokenize(text)
1130 if len(tokens) > cap:
1131 log.warning("Truncating oversize embed input: %d tokens > cap %d", len(tokens), cap)
1132 return self._detokenize(tokens[:cap]), cap
1133 return text, max(1, len(tokens))
1135 def _native_route(self, suffix: str) -> str:
1136 """Path for a native (non-OpenAI) llama-server route through llama-swap.
1138 llama-swap proxies these only under ``/upstream/<model>/...``; the model
1139 is carried in the path, not the body (unlike the ``/v1`` OpenAI routes).
1140 """
1141 return f"{_UPSTREAM_PREFIX}/{self._model}{suffix}"
1143 def _tokenize(self, text: str) -> list[int]:
1144 resp = self._http.post(
1145 self._native_route(_TOKENIZE_PATH),
1146 json={
1147 "content": text,
1148 "add_special": _TOKENIZE_ADD_SPECIAL,
1149 "parse_special": _TOKENIZE_PARSE_SPECIAL,
1150 },
1151 )
1152 _raise_for_status(resp)
1153 return list(resp.json()["tokens"])
1155 def count_tokens(self, text: str) -> int:
1156 """Number of tokens *text* encodes to under the server's tokenizer."""
1157 return len(self._tokenize(text))
1159 def _detokenize(self, tokens: list[int]) -> str:
1160 resp = self._http.post(self._native_route(_DETOKENIZE_PATH), json={"tokens": tokens})
1161 _raise_for_status(resp)
1162 return str(resp.json()["content"])
1164 @contextlib.contextmanager
1165 def _abortable(self, resp: httpx.Response) -> Generator[None]:
1166 """Expose *resp* to ``abort_streams`` for the duration of its read loop."""
1167 with self._in_flight_lock:
1168 self._active_streams.add(resp)
1169 try:
1170 yield
1171 finally:
1172 with self._in_flight_lock:
1173 self._active_streams.discard(resp)
1175 def abort_streams(self) -> None:
1176 """Sever every in-flight SSE response on this replica.
1178 Closing the response from another thread unblocks a reader stuck in
1179 ``iter_lines`` with a stream error, which unwinds its worker;
1180 llama-server stops generating when the connection drops.
1181 """
1182 with self._in_flight_lock:
1183 streams = list(self._active_streams)
1184 for resp in streams:
1185 with contextlib.suppress(Exception):
1186 resp.close()
1188 def close(self) -> None:
1189 """Close the underlying client if this instance created it."""
1190 if self._owns_http:
1191 self._http.close()
1193 def _track(self) -> _InFlight:
1194 return _InFlight(self)
1197class _InFlight:
1198 """Context manager that atomically bumps the owner's in-flight counter.
1200 ``+= 1`` is a read-modify-write, so concurrent chat/embed calls would corrupt
1201 the counter the router balances on; the client's lock makes it atomic.
1202 """
1204 def __init__(self, client: LlamaServerClient) -> None:
1205 self._client = client
1207 def __enter__(self) -> None:
1208 with self._client._in_flight_lock:
1209 self._client.in_flight += 1
1211 def __exit__(self, *_exc: object) -> None:
1212 with self._client._in_flight_lock:
1213 self._client.in_flight -= 1
1216def _embedding_vector(item: dict[str, Any]) -> npt.NDArray[np.float32]:
1217 """Decode one ``/v1/embeddings`` item's vector from its base64 float buffer."""
1218 embedding = item.get("embedding")
1219 # Untyped server JSON: a non-string means the encoding format was not honored.
1220 if not isinstance(embedding, str):
1221 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME)
1222 try:
1223 return np.frombuffer(base64.b64decode(embedding), dtype=_EMBED_VECTOR_DTYPE)
1224 except ValueError as exc:
1225 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME) from exc
1228def _rerank_score(item: dict[str, Any]) -> float:
1229 """Pull one relevance score from a rank-pooling ``/v1/embeddings`` item."""
1230 vector = _embedding_vector(item)
1231 if not vector.size:
1232 raise ProviderError(_NO_RERANK_SCORE_ERROR, provider=_PROVIDER_NAME)
1233 return float(vector[_RANK_SCORE_INDEX])
1236def _first_token_top_logprobs(response: dict[str, Any]) -> list[dict[str, Any]]:
1237 """The first generated token's top_logprobs list from a chat completion, or []."""
1238 choices = response.get("choices") or []
1239 if not choices:
1240 return []
1241 content = (choices[0].get("logprobs") or {}).get("content") or []
1242 if not content:
1243 return []
1244 return list(content[0].get("top_logprobs") or [])
1247def _llm_rerank_score(top_logprobs: list[dict[str, Any]]) -> float | None:
1248 """Softmax of the yes vs no logprobs in a token's top_logprobs (case/space-insensitive).
1250 ``None`` when neither verdict appears, distinct from the 0.0 of a confident "no".
1251 """
1252 yes_lp: float | None = None
1253 no_lp: float | None = None
1254 for entry in top_logprobs:
1255 token = str(entry.get("token", "")).strip().lower()
1256 logprob = float(entry.get("logprob", 0.0))
1257 if token == _YES_LABEL and (yes_lp is None or logprob > yes_lp):
1258 yes_lp = logprob
1259 elif token == _NO_LABEL and (no_lp is None or logprob > no_lp):
1260 no_lp = logprob
1261 if yes_lp is None:
1262 return None if no_lp is None else 0.0
1263 if no_lp is None:
1264 return math.exp(yes_lp)
1265 yes_e, no_e = math.exp(yes_lp), math.exp(no_lp)
1266 return yes_e / (yes_e + no_e)
1269def _parse_sse_deltas(line: str) -> tuple[str, str]:
1270 """Extract the (reasoning, content) deltas from one OpenAI SSE line."""
1271 if not line.startswith(_DATA_PREFIX):
1272 return "", ""
1273 body = line[len(_DATA_PREFIX) :].strip()
1274 if not body or body == _DONE_SENTINEL:
1275 return "", ""
1276 try:
1277 obj = json.loads(body)
1278 except json.JSONDecodeError:
1279 return "", ""
1280 choices = obj.get("choices") or []
1281 if not choices:
1282 return "", ""
1283 delta = choices[0].get("delta") or {}
1284 return str(delta.get("reasoning_content") or ""), str(delta.get("content") or "")
1287class _ThinkInliner:
1288 """Re-inlines server-extracted reasoning deltas as inline ``<think>`` text.
1290 The server parses each model's reasoning format natively (``--reasoning-format``)
1291 and streams it as ``reasoning_content``; lilbee's pipeline speaks inline
1292 ``<think>`` text, so the chat boundary opens the tag on the first reasoning
1293 delta and closes it when the answer starts (or at end of stream). Disabled
1294 (the non-chat roles), reasoning is dropped and content passes through, matching
1295 the server-extracted default those roles already ran with.
1296 """
1298 def __init__(self, *, enabled: bool) -> None:
1299 self._enabled = enabled
1300 self._in_think = False
1302 def feed(self, reasoning: str, content: str) -> str:
1303 if not self._enabled:
1304 return content
1305 parts: list[str] = []
1306 if reasoning:
1307 if not self._in_think:
1308 self._in_think = True
1309 parts.append(THINK_OPEN_TAG)
1310 parts.append(reasoning)
1311 if content:
1312 if self._in_think:
1313 self._in_think = False
1314 parts.append(THINK_CLOSE_TAG)
1315 parts.append(content)
1316 return "".join(parts)
1318 def finish(self) -> str:
1319 """Close an unterminated think block at end of stream."""
1320 if self._in_think:
1321 self._in_think = False
1322 return THINK_CLOSE_TAG
1323 return ""
1326def _inline_message_reasoning(message: Mapping[str, Any], *, enabled: bool) -> str:
1327 """A non-streaming message's text with any extracted reasoning re-inlined."""
1328 content = str(message.get("content") or "")
1329 reasoning = str(message.get("reasoning_content") or "") if enabled else ""
1330 if reasoning:
1331 return f"{THINK_OPEN_TAG}{reasoning}{THINK_CLOSE_TAG}{content}"
1332 return content
1335def _usage_from_body(body: Mapping[str, Any]) -> TokenUsage | None:
1336 """Read the ``usage`` block of an OpenAI response, or ``None`` if absent.
1338 llama-server reports ``prompt_tokens`` / ``completion_tokens``; a missing or
1339 malformed block yields ``None`` so callers can decide between a zero default
1340 (non-streaming) and skipping the frame (streaming terminator).
1341 """
1342 usage = body.get("usage")
1343 if not isinstance(usage, Mapping):
1344 return None
1345 prompt = usage.get("prompt_tokens")
1346 completion = usage.get("completion_tokens")
1347 return TokenUsage(
1348 prompt_tokens=prompt if isinstance(prompt, int) else 0,
1349 completion_tokens=completion if isinstance(completion, int) else 0,
1350 )
1353def _coerce_finish_reason(raw: Any) -> FinishReason:
1354 """Map a server-supplied finish_reason string to :class:`FinishReason`."""
1355 return FinishReason.coerce(raw)
1358def _tool_call_delta_from_chunk(call: Mapping[str, Any], *, fallback_index: int) -> ToolCallDelta:
1359 """Map one streaming ``delta.tool_calls`` entry to a :class:`ToolCallDelta`.
1361 Mirrors the SDK path: ``id`` / ``name`` arrive on the opener and accumulate
1362 by ``index``; empty strings normalise to ``None`` so the dispatch's stream
1363 translator (which gates on ``is not None``) does not emit spurious openers.
1364 """
1365 raw_index = call.get("index")
1366 index = raw_index if isinstance(raw_index, int) else fallback_index
1367 call_id = call.get("id")
1368 fn = call.get("function")
1369 raw_name = fn.get("name") if isinstance(fn, Mapping) else None
1370 raw_args = fn.get("arguments") if isinstance(fn, Mapping) else None
1371 return ToolCallDelta(
1372 index=index,
1373 id=str(call_id) if call_id else None,
1374 name=str(raw_name) if raw_name else None,
1375 arguments_delta=str(raw_args) if raw_args else None,
1376 )
1379def _parse_sse_stream_items(
1380 line: str, inliner: _ThinkInliner
1381) -> Iterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1382 """Yield text tokens, tool-call deltas, and the finish frame from one SSE line.
1384 A chunk can carry a ``content`` token, a ``reasoning_content`` token (routed
1385 through *inliner*), a ``tool_calls`` delta array, or a mix; each is yielded as
1386 its own :data:`ChatStreamItem` frame. The chunk that closes the turn carries
1387 ``choices[0].finish_reason``, surfaced as a :class:`StreamFinish` so the
1388 dispatch reports ``length`` (and friends), not just the default end-of-turn.
1389 """
1390 if not line.startswith(_DATA_PREFIX):
1391 return
1392 body = line[len(_DATA_PREFIX) :].strip()
1393 if not body or body == _DONE_SENTINEL:
1394 return
1395 try:
1396 obj = json.loads(body)
1397 except json.JSONDecodeError:
1398 return
1399 choices = obj.get("choices") or []
1400 if not choices:
1401 # The include_usage terminator chunk has an empty choices list and the
1402 # token totals on a top-level ``usage`` block; surface it as the final
1403 # frame so the dispatch can attach real counts to the stream.
1404 usage = _usage_from_body(obj)
1405 if usage is not None:
1406 yield usage
1407 return
1408 delta = choices[0].get("delta") or {}
1409 text = inliner.feed(str(delta.get("reasoning_content") or ""), str(delta.get("content") or ""))
1410 if text:
1411 yield text
1412 raw_calls = delta.get("tool_calls") or []
1413 for i, call in enumerate(raw_calls):
1414 if isinstance(call, Mapping):
1415 yield _tool_call_delta_from_chunk(call, fallback_index=i)
1416 raw_finish = choices[0].get("finish_reason")
1417 if raw_finish is not None:
1418 yield StreamFinish(reason=_coerce_finish_reason(raw_finish))
1421def _arguments_to_str(arguments: Any) -> str:
1422 """Normalize a tool-call ``arguments`` value to a JSON string (OpenAI's shape)."""
1423 if isinstance(arguments, str):
1424 return arguments
1425 if arguments is None:
1426 return "{}"
1427 return json.dumps(arguments)
1430def _parse_native_tool_calls(raw: Any) -> list[ToolCall]:
1431 """Map a response's ``message.tool_calls`` array to :class:`ToolCall` objects.
1433 Reads the OpenAI shape (``{"id", "function": {"name", "arguments"}}``) that
1434 ``--jinja`` produces. Malformed or nameless entries are skipped.
1435 """
1436 if not isinstance(raw, list):
1437 return []
1438 calls: list[ToolCall] = []
1439 for idx, entry in enumerate(raw):
1440 if not isinstance(entry, Mapping):
1441 continue
1442 fn = entry.get("function")
1443 if not isinstance(fn, Mapping):
1444 continue
1445 name = fn.get("name")
1446 if not isinstance(name, str) or not name:
1447 continue
1448 call_id = entry.get("id")
1449 calls.append(
1450 ToolCall(
1451 id=call_id if isinstance(call_id, str) and call_id else f"call_{idx}",
1452 name=name,
1453 arguments=_arguments_to_str(fn.get("arguments")),
1454 )
1455 )
1456 return calls
1459def _bare_call_from_mapping(obj: Mapping[str, Any], *, index: int) -> ToolCall | None:
1460 """Build a ToolCall from a bare ``{"name", "arguments"|"parameters"}`` object."""
1461 name = obj.get("name")
1462 if not isinstance(name, str) or not name:
1463 return None
1464 arguments = obj.get("arguments")
1465 if arguments is None:
1466 arguments = obj.get("parameters")
1467 return ToolCall(id=f"call_{index}", name=name, arguments=_arguments_to_str(arguments))
1470def _recover_bare_json_tool_calls(content: str) -> ChatToolResult:
1471 """Recover a tool call a model emitted as bare-JSON content (a native miss).
1473 Some models ignore the tool-call protocol and print ``{"name": ...,
1474 "arguments": {...}}`` (or a list of them) as the message body. When the whole
1475 content parses as such, treat it as the call(s) and clear the text; otherwise
1476 return the content unchanged with no calls.
1477 """
1478 stripped = content.strip()
1479 if not stripped or stripped[0] not in "{[":
1480 return ChatToolResult(content=content, tool_calls=[])
1481 try:
1482 parsed = json.loads(stripped)
1483 except json.JSONDecodeError:
1484 return ChatToolResult(content=content, tool_calls=[])
1485 entries = parsed if isinstance(parsed, list) else [parsed]
1486 calls = [
1487 call
1488 for idx, entry in enumerate(entries)
1489 if isinstance(entry, Mapping) and (call := _bare_call_from_mapping(entry, index=idx))
1490 ]
1491 if not calls:
1492 return ChatToolResult(content=content, tool_calls=[])
1493 return ChatToolResult(content="", tool_calls=calls)
1496# Leading non-whitespace characters that mark streamed text as a potential bare
1497# JSON tool call (an object or an array of them); any other first char is plain
1498# text and streams through untouched.
1499_BARE_CALL_OPENERS = "{["
1502def _tool_call_delta_from_recovered(call: ToolCall, index: int) -> ToolCallDelta:
1503 """Shape a recovered bare-JSON :class:`ToolCall` as a single streaming delta.
1505 Mirrors :func:`_tool_call_delta_from_chunk`: id and name ride the opener (the
1506 only frame for a recovered call), the arguments JSON is the lone
1507 ``arguments_delta``, and the position is the index.
1508 """
1509 return ToolCallDelta(
1510 index=index,
1511 id=call.id or None,
1512 name=call.name or None,
1513 arguments_delta=call.arguments or None,
1514 )
1517def _recover_bare_json_stream(
1518 items: Iterator[str | ToolCallDelta | TokenUsage | StreamFinish],
1519) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1520 """Wrap a raw chat stream to recover a tool call emitted as bare-JSON text.
1522 Some small models print ``{"name": ..., "arguments": {...}}`` as content
1523 instead of native ``tool_calls``; the non-stream paths recover this via
1524 :func:`_recover_bare_json_tool_calls`. This applies the same recovery to the
1525 stream, but only when the model emitted no native :class:`ToolCallDelta` and
1526 the streamed text looks like a bare call from its first character. Normal text
1527 still streams token by token: once the buffered head proves not to be a bare
1528 call it is flushed and all later text passes straight through.
1529 """
1530 buffer = "" # leading text held back as a potential bare call until resolved
1531 # True once leading text has streamed as plain (or a native call was seen):
1532 # past that, a later '{'/'[' is content, not a bare call -- never buffer again.
1533 committed = False
1534 try:
1535 for item in items:
1536 if isinstance(item, ToolCallDelta):
1537 yield from _flush_plain(buffer)
1538 buffer, committed = "", True
1539 yield item
1540 elif isinstance(item, TokenUsage | StreamFinish):
1541 yield from _recover_buffer(buffer)
1542 buffer = ""
1543 yield item
1544 elif committed or _passthrough_text(buffer, item):
1545 yield from _flush_plain(buffer)
1546 buffer = ""
1547 committed = True
1548 yield item
1549 else:
1550 buffer += item
1551 yield from _recover_buffer(buffer)
1552 finally:
1553 # Forward close to the source generator: if a consumer closes this
1554 # wrapper mid-stream, a plain for-loop would not propagate GeneratorExit
1555 # to *items*, leaking the underlying HTTP stream and its in_flight slot.
1556 # Suppress teardown errors (httpx stream close can raise) so they don't
1557 # mask the exception that triggered this finally.
1558 if isinstance(items, Generator):
1559 with contextlib.suppress(Exception):
1560 items.close()
1563def _passthrough_text(buffer: str, text: str) -> bool:
1564 """Whether *text* should stream through directly rather than buffer.
1566 True once the accumulated head's first non-whitespace char is known and is not
1567 a bare-call opener (plain text): the buffer is empty in that case, so the
1568 caller yields *text* as is. While the head is all whitespace, or once it opens
1569 with ``{``/``[``, the text is buffered (False) pending recovery.
1570 """
1571 head = (buffer + text).lstrip()
1572 return bool(head) and head[0] not in _BARE_CALL_OPENERS
1575def _flush_plain(buffer: str) -> Iterator[str]:
1576 """Yield buffered leading text verbatim (it was not a bare call after all)."""
1577 if buffer:
1578 yield buffer
1581def _recover_buffer(buffer: str) -> Iterator[str | ToolCallDelta]:
1582 """Resolve the buffered leading text at a terminator or end of stream.
1584 The buffer reaching here was held as a potential bare call (text starting with
1585 ``{``/``[`` and no native call seen). Run :func:`_recover_bare_json_tool_calls`:
1586 emit one delta per recovered call, or yield the text unchanged when it only
1587 happened to start with ``{``/``[`` but is not a call.
1588 """
1589 if not buffer:
1590 return
1591 recovered = _recover_bare_json_tool_calls(buffer)
1592 if not recovered.tool_calls:
1593 yield buffer
1594 return
1595 for index, call in enumerate(recovered.tool_calls):
1596 yield _tool_call_delta_from_recovered(call, index)