Coverage for src/lilbee/providers/fleet/client.py: 100%
716 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +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 on_prefill: Callable[[tuple[int, int] | None], None] | None = None,
514 ) -> None:
515 self._base = base_url.rstrip("/")
516 self._model = model
517 # Prefill observer: called with (processed, total) per engine progress
518 # frame during a streamed chat's prompt processing, then None once the
519 # first generated frame proves the prefill is over.
520 self._on_prefill = on_prefill
521 # Cold-load budget (seconds) the embed path waits out a still-warming replica
522 # before dropping the input, in place of the short attempt cap. Set on the
523 # EMBED-role client to the same ceiling llama-swap keeps the server alive for,
524 # so a bulk ingest never gives up while the replica is legitimately loading.
525 # None (rerank, chat, vision, self-check) keeps the fixed interactive budget.
526 self._embed_busy_deadline_s = embed_busy_deadline_s
527 self._http = http or httpx.Client(
528 base_url=self._base, timeout=timeout, verify=_LOOPBACK_SSL_CONTEXT
529 )
530 self._owns_http = http is None
531 # Chat-role clients re-inline server-extracted reasoning as <think> text;
532 # the other roles (vision OCR) keep dropping it, as their servers already did.
533 self._inline_reasoning = inline_reasoning
534 # Per-slot context for embed/rerank servers: inputs longer than this are
535 # token-truncated (via the server's tokenizer) before embedding, mirroring
536 # the in-process backstop. None for chat/vision, which don't truncate inputs.
537 self._token_cap = token_cap
538 # LLM => score candidates by yes/no logprob; None/cross-encoder => rank pooling.
539 self._rerank_mode = rerank_mode
540 self.in_flight = 0
541 self._in_flight_lock = threading.Lock()
542 # Live SSE responses, so a cancel can sever the transport from another
543 # thread: a reader blocked in iter_lines cannot see a cooperative
544 # cancel flag, but closing its response unblocks it with an error.
545 self._active_streams: set[httpx.Response] = set()
546 # Whether this server's chat template needs OpenAI tool exchanges reshaped
547 # into strict user/assistant alternation. Determined lazily by a one-time
548 # probe of the live template (see _prepare_chat_messages); None until then.
549 # A client is bound to one model for its lifetime, so the template (hence
550 # the verdict) is fixed once determined.
551 self._needs_alternation: bool | None = None
552 self._alternation_lock = threading.Lock()
553 # Routing health: cleared on a connection-level failure (see _UNHEALTHY_RETRY_S).
554 self._healthy = True
555 # Monotonic stamp of the last mark_unhealthy; consulted only while unhealthy.
556 self._unhealthy_since = 0.0
558 @property
559 def healthy(self) -> bool:
560 """Routable: healthy, or unhealthy past the ``_UNHEALTHY_RETRY_S`` cool-down."""
561 with self._in_flight_lock:
562 if self._healthy:
563 return True
564 return time.monotonic() - self._unhealthy_since >= _UNHEALTHY_RETRY_S
566 def mark_unhealthy(self) -> None:
567 """Record a connection-level failure so the router skips this replica."""
568 with self._in_flight_lock:
569 self._healthy = False
570 self._unhealthy_since = time.monotonic()
572 def mark_healthy(self) -> None:
573 """Restore the replica to the routing pool after a successful call."""
574 with self._in_flight_lock:
575 self._healthy = True
577 def reserve(self) -> None:
578 """Mark a routed request assigned to this replica, at selection time.
580 The router balances on ``in_flight`` but the per-request tracking only
581 bumps it once the HTTP call starts. Under a bulk ingest many threads pick
582 a replica at the same instant, all see the momentarily-idlest one at the
583 same low count, and pile onto it (a thundering herd that leaves the other
584 cards idle). Reserving at selection makes the assignment visible to the
585 next picker so requests spread across replicas. Paired with :meth:`release`.
586 """
587 with self._in_flight_lock:
588 self.in_flight += 1
590 def release(self) -> None:
591 """Release a reservation taken by :meth:`reserve`."""
592 with self._in_flight_lock:
593 self.in_flight -= 1
595 def health(self) -> bool:
596 """True iff ``GET /health`` returns 200 (liveness, not readiness)."""
597 try:
598 resp = self._http.get(_HEALTH_PATH, timeout=_HEALTH_TIMEOUT_S)
599 except httpx.HTTPError:
600 return False
601 return resp.status_code == _HTTP_OK
603 @overload
604 def chat(
605 self,
606 messages: Sequence[Mapping[str, Any]],
607 *,
608 options: dict[str, Any] | None = None,
609 stream: Literal[False] = False,
610 timeout: float | None = None,
611 ) -> str: ...
613 @overload
614 def chat(
615 self,
616 messages: Sequence[Mapping[str, Any]],
617 *,
618 options: dict[str, Any] | None = None,
619 stream: Literal[True],
620 timeout: float | None = None,
621 ) -> Iterator[str]: ...
623 @overload
624 def chat(
625 self,
626 messages: Sequence[Mapping[str, Any]],
627 *,
628 options: dict[str, Any] | None = None,
629 stream: bool,
630 timeout: float | None = None,
631 ) -> str | Iterator[str]: ...
633 def chat(
634 self,
635 messages: Sequence[Mapping[str, Any]],
636 *,
637 options: dict[str, Any] | None = None,
638 stream: bool = False,
639 timeout: float | None = None,
640 ) -> str | Iterator[str]:
641 """Chat completion. Returns the full text, or a token iterator if streaming.
643 ``messages`` accepts both plain ``{role, content: str}`` and multipart
644 ``content`` lists (vision image parts), so the vision path reuses this.
645 ``timeout`` overrides the client default for either path, so a
646 caller-enforced deadline (vision OCR) ends the request itself.
647 """
648 payload: dict[str, Any] = {"model": self._model, "messages": messages, **(options or {})}
649 request_timeout = timeout if timeout is not None else httpx.USE_CLIENT_DEFAULT
650 if stream:
651 return self._chat_stream(payload, request_timeout)
652 with self._track():
653 resp = self._http.post(
654 _CHAT_PATH, json={**payload, "stream": False}, timeout=request_timeout
655 )
656 _raise_for_status(resp)
657 # content is null for a refusal / content-filter stop / empty completion;
658 # coerce to "" (like chat_result/chat_tools) so callers never see "None".
659 return _inline_message_reasoning(
660 resp.json()["choices"][0]["message"], enabled=self._inline_reasoning
661 )
663 def _chat_stream(
664 self, payload: dict[str, Any], timeout: Any = httpx.USE_CLIENT_DEFAULT
665 ) -> Iterator[str]:
666 inliner = _ThinkInliner(enabled=self._inline_reasoning)
667 with (
668 self._track(),
669 self._http.stream(
670 "POST", _CHAT_PATH, json={**payload, "stream": True}, timeout=timeout
671 ) as resp,
672 ):
673 _raise_for_status(resp)
674 with self._abortable(resp):
675 for line in resp.iter_lines():
676 delta = inliner.feed(*_parse_sse_deltas(line))
677 if delta:
678 yield delta
679 tail = inliner.finish()
680 if tail:
681 yield tail
683 def chat_bounded(
684 self,
685 messages: Sequence[Mapping[str, Any]],
686 *,
687 options: dict[str, Any] | None = None,
688 deadline_s: float,
689 ) -> str:
690 """Stream a chat completion and return its text, bounded by a total deadline.
692 httpx float timeouts are per-phase (connect/read/...), never a total
693 budget, so a steadily trickling upstream can pin a worker past its
694 deadline. Streaming in the caller's own thread and checking a monotonic
695 deadline per frame bounds total time: on expiry the ``with`` block closes
696 the stream (releasing the in-flight slot) and raises
697 :class:`ChatDeadlineError`.
698 """
699 payload: dict[str, Any] = {"model": self._model, "messages": messages, **(options or {})}
700 deadline = time.monotonic() + deadline_s
701 inliner = _ThinkInliner(enabled=self._inline_reasoning)
702 parts: list[str] = []
703 with (
704 self._track(),
705 self._http.stream("POST", _CHAT_PATH, json={**payload, "stream": True}) as resp,
706 ):
707 _raise_for_status(resp)
708 for line in resp.iter_lines():
709 if time.monotonic() >= deadline:
710 raise ChatDeadlineError(
711 f"llama-server chat exceeded its {deadline_s:.0f}s deadline.",
712 provider=_PROVIDER_NAME,
713 )
714 parts.append(inliner.feed(*_parse_sse_deltas(line)))
715 parts.append(inliner.finish())
716 return "".join(parts)
718 def chat_tools(
719 self,
720 messages: Sequence[Mapping[str, Any]],
721 *,
722 tools: list[dict[str, Any]],
723 tool_choice: str | dict[str, Any] | None = None,
724 options: dict[str, Any] | None = None,
725 ) -> ChatToolResult:
726 """Non-streaming chat with function tools; returns content + any tool calls.
728 The server is launched with ``--jinja`` so it parses the model's native
729 tool-call syntax into structured ``message.tool_calls``. When a model
730 instead emits a bare-JSON call as content (a native miss), recover it.
731 """
732 payload: dict[str, Any] = {
733 "model": self._model,
734 "messages": self._prepare_chat_messages(messages),
735 "tools": tools,
736 "stream": False,
737 **(options or {}),
738 }
739 if tool_choice is not None:
740 payload["tool_choice"] = tool_choice
741 with self._track():
742 resp = self._http.post(_CHAT_PATH, json=payload)
743 _raise_for_status(resp)
744 message = resp.json()["choices"][0]["message"]
745 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
746 native = _parse_native_tool_calls(message.get("tool_calls"))
747 if native:
748 return ChatToolResult(content=content, tool_calls=native)
749 return _recover_bare_json_tool_calls(content)
751 def chat_result(
752 self,
753 messages: Sequence[Mapping[str, Any]],
754 *,
755 tools: list[dict[str, Any]] | None = None,
756 tool_choice: str | dict[str, Any] | None = None,
757 options: dict[str, Any] | None = None,
758 ) -> ChatResult:
759 """Non-streaming chat returning text, tool calls, and a finish reason.
761 The server is launched with ``--jinja`` so it parses the model's native
762 tool-call syntax into structured ``message.tool_calls``. When a model
763 instead emits a bare-JSON call as content (a native miss), recover it
764 and report ``tool_calls`` as the finish reason. Messages are reshaped to
765 strict alternation up front when this server's template needs it (see
766 :meth:`_prepare_chat_messages`).
767 """
768 payload = self._chat_payload(
769 self._prepare_chat_messages(messages), tools, tool_choice, options, stream=False
770 )
771 with self._track():
772 resp = self._http.post(_CHAT_PATH, json=payload)
773 _raise_for_status(resp)
774 body = dict(resp.json())
775 choice = body["choices"][0]
776 usage = _usage_from_body(body) or TokenUsage()
777 message = choice["message"]
778 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
779 finish_reason = _coerce_finish_reason(choice.get("finish_reason"))
780 native = _parse_native_tool_calls(message.get("tool_calls"))
781 if native:
782 return ChatResult(
783 text=content,
784 tool_calls=tuple(native),
785 finish_reason=finish_reason,
786 usage=usage,
787 )
788 recovered = _recover_bare_json_tool_calls(content)
789 if recovered.tool_calls:
790 return ChatResult(
791 text=recovered.content,
792 tool_calls=tuple(recovered.tool_calls),
793 finish_reason=FinishReason.TOOL_CALLS,
794 usage=usage,
795 )
796 return ChatResult(text=content, tool_calls=(), finish_reason=finish_reason, usage=usage)
798 def chat_stream_items(
799 self,
800 messages: Sequence[Mapping[str, Any]],
801 *,
802 tools: list[dict[str, Any]] | None = None,
803 tool_choice: str | dict[str, Any] | None = None,
804 options: dict[str, Any] | None = None,
805 ) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
806 """Stream text tokens and tool-call deltas from the server's OpenAI SSE.
808 Each SSE chunk's ``choices[0].delta`` carries a ``content`` token and/or
809 a ``tool_calls`` array; both are surfaced as :data:`ChatStreamItem`
810 frames (text strings and :class:`ToolCallDelta`). The dispatch's stream
811 translator accumulates the deltas by ``index``. Messages are reshaped to
812 strict alternation up front when this server's template needs it (see
813 :meth:`_prepare_chat_messages`), so the open never fails on a template
814 that rejects the raw tool exchange.
816 Not a generator: the up-front probe runs when this is called, not deferred
817 to the first iteration, matching the eager non-stream paths.
819 A model that emits a tool call as bare-JSON text instead of native
820 ``tool_calls`` (a native miss, as on the non-stream paths) is recovered by
821 wrapping the raw frames; see :func:`_recover_bare_json_stream`.
822 """
823 prepared = self._prepare_chat_messages(messages)
824 return _recover_bare_json_stream(
825 self._open_chat_stream(prepared, tools, tool_choice, options)
826 )
828 def _open_chat_stream(
829 self,
830 messages: Sequence[Mapping[str, Any]],
831 tools: list[dict[str, Any]] | None,
832 tool_choice: str | dict[str, Any] | None,
833 options: dict[str, Any] | None,
834 ) -> Iterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
835 """Open one SSE chat stream and yield its frames; raises before the first frame."""
836 payload = self._chat_payload(messages, tools, tool_choice, options, stream=True)
837 inliner = _ThinkInliner(enabled=self._inline_reasoning)
838 on_prefill = self._on_prefill
839 prefilling = False
840 try:
841 with (
842 self._track(),
843 self._http.stream("POST", _CHAT_PATH, json=payload) as resp,
844 ):
845 _raise_for_status(resp)
846 with self._abortable(resp):
847 for line in resp.iter_lines():
848 if on_prefill is not None:
849 progress = _prefill_progress(line)
850 if progress is not None:
851 prefilling = True
852 on_prefill(progress)
853 for item in _parse_sse_stream_items(line, inliner):
854 if prefilling and on_prefill is not None:
855 # The first real frame proves prefill is over.
856 prefilling = False
857 on_prefill(None)
858 yield item
859 tail = inliner.finish()
860 if tail:
861 yield tail
862 finally:
863 # A stream that dies or is closed mid-prefill must not leave a
864 # stale in-progress reading on the status surface.
865 if prefilling and on_prefill is not None:
866 on_prefill(None)
868 def _chat_payload(
869 self,
870 messages: Sequence[Mapping[str, Any]],
871 tools: Sequence[Mapping[str, Any]] | None,
872 tool_choice: str | dict[str, Any] | None,
873 options: dict[str, Any] | None,
874 *,
875 stream: bool,
876 ) -> dict[str, Any]:
877 """Build the chat-completions request body shared by the stream and non-stream paths."""
878 payload: dict[str, Any] = {"model": self._model, "messages": messages, "stream": stream}
879 if stream:
880 # include_usage makes llama-server emit a final SSE chunk carrying the
881 # token usage (with an empty choices list) just before [DONE].
882 payload["stream_options"] = {"include_usage": True}
883 # return_progress makes llama-server stream prompt_progress frames
884 # during prefill, so a long first turn is observable server-side.
885 payload["return_progress"] = True
886 if tools is not None:
887 payload["tools"] = tools
888 if tool_choice is not None:
889 payload["tool_choice"] = tool_choice
890 payload.update(options or {})
891 return payload
893 def _prepare_chat_messages(
894 self, messages: Sequence[Mapping[str, Any]]
895 ) -> Sequence[Mapping[str, Any]]:
896 """Reshape *messages* to strict alternation when this server's template needs it.
898 The need is detected once per server by :meth:`_ensure_alternation_probed`
899 (a binary accept/reject of a representative tool exchange against the live
900 template), then cached, so real requests are normalized up front rather
901 than failing and retrying.
902 """
903 self._ensure_alternation_probed()
904 if self._needs_alternation:
905 return to_alternating([dict(m) for m in messages])
906 return messages
908 def _ensure_alternation_probed(self) -> None:
909 """Probe the live template once to learn whether it needs alternation.
911 Caches only a conclusive verdict: a transient unreachable server leaves
912 the flag unset so the next request re-probes rather than locking in a
913 wrong answer.
914 """
915 if self._needs_alternation is not None:
916 return
917 with self._alternation_lock:
918 if self._needs_alternation is not None:
919 return
920 verdict = self._probe_alternation()
921 if verdict is not None:
922 self._needs_alternation = verdict
924 def _probe_alternation(self) -> bool | None:
925 """Whether the template needs normalization: ``None`` when undetermined.
927 Renders the probe exchange as sent; if the template accepts it, no
928 normalization is needed. If it rejects it, normalization is needed only
929 when the reshaped exchange is accepted. A transient failure on either
930 render is inconclusive (``None``) so no verdict is cached; a genuine
931 rejection of both forms is a conclusive ``False`` (the template fault is
932 unrelated to alternation, so reshaping would not help).
933 """
934 raw = self._chat_probe(_ALTERNATION_PROBE_MESSAGES)
935 if raw is None:
936 return None # transient; stay undetermined so the next request re-probes
937 if raw:
938 return False # the template renders the raw OpenAI exchange as sent
939 reshaped = self._chat_probe(to_alternating([dict(m) for m in _ALTERNATION_PROBE_MESSAGES]))
940 if reshaped is None:
941 return None # transient on the reshape probe; stay undetermined
942 return reshaped
944 def _chat_probe(self, messages: Sequence[Mapping[str, Any]]) -> bool | None:
945 """Post the probe exchange: ``True`` rendered, ``False`` rejected, ``None`` undetermined.
947 A connection failure or a server-busy (HTTP 429) response is transient and
948 unrelated to the template, so it is undetermined: only a clean render or a
949 genuine rejection is a verdict the caller may cache.
950 """
951 payload = self._chat_payload(
952 messages, _ALTERNATION_PROBE_TOOLS, None, _ALTERNATION_PROBE_OPTIONS, stream=False
953 )
954 try:
955 with self._track():
956 resp = self._http.post(
957 _CHAT_PATH, json=payload, timeout=_ALTERNATION_PROBE_TIMEOUT_S
958 )
959 _raise_for_status(resp)
960 except (ProviderError, httpx.TransportError) as exc:
961 return None if _is_transient_probe_failure(exc) else False
962 return True
964 def embed(self, texts: list[str]) -> list[Vector]:
965 """Embed a batch via ``/v1/embeddings``."""
966 if not texts:
967 # Match the in-process embedder; the server rejects an empty input.
968 return []
969 vectors: list[Vector] = []
970 for sub_batch in self._truncate_and_subbatch(texts, estimate=True):
971 data = self._embed_subbatch(sub_batch)
972 vectors.extend(_embedding_vector(item) for item in data)
973 return vectors
975 def _embed_subbatch(self, sub_batch: list[str]) -> list[dict[str, Any]]:
976 """Embed one estimate-budgeted sub-batch, re-truncating exactly on overflow.
978 ``_estimate_tokens`` is char-based and can under-count token-dense inputs
979 (XML, code), so an estimate-trusted input may still exceed the server's
980 context. On that error -- and only that -- redo the batch with exact
981 server-side tokenization, which truncates the oversize input to the cap.
982 """
983 try:
984 return self._embeddings_call(sub_batch)
985 except ProviderError as exc:
986 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW:
987 raise
988 data: list[dict[str, Any]] = []
989 for exact in self._truncate_and_subbatch(sub_batch, estimate=False):
990 data.extend(self._embeddings_call(exact))
991 return data
993 def rerank(self, query: str, candidates: list[str]) -> list[float]:
994 """Relevance scores via rank-pooling embeddings.
996 The server runs with ``--pooling rank``; we send ``query</s></s>candidate``
997 pairs to ``/v1/embeddings`` and read each item's first embedding value as the
998 score, so the ``/v1/rerank`` template-dependency (and its zero-output failure
999 modes) is moot.
1001 All pairs go out in one request per ``_EMBED_N_SEQ_MAX`` sequences: the
1002 server queues one task per input, so splitting a query-time pool into
1003 per-pair requests only adds HTTP round trips.
1004 """
1005 if not candidates:
1006 return []
1007 if self._rerank_mode is RerankMode.LLM:
1008 return self._rerank_llm(query, candidates)
1009 pairs = [
1010 self._fit_estimated(f"{query}{_RERANK_PAIR_SEPARATOR}{candidate}")
1011 for candidate in candidates
1012 ]
1013 scores: list[float] = []
1014 for start in range(0, len(pairs), _EMBED_N_SEQ_MAX):
1015 data = self._rerank_batch(pairs[start : start + _EMBED_N_SEQ_MAX])
1016 scores.extend(_rerank_score(item) for item in data)
1017 return scores
1019 def _fit_estimated(self, text: str) -> str:
1020 """Truncate *text* to the token cap only when its char estimate exceeds it."""
1021 if self._token_cap is None:
1022 return text
1023 return self._fit_input(text, self._token_cap, estimate=True)[0]
1025 def _rerank_batch(self, batch: list[str]) -> list[dict[str, Any]]:
1026 """Score one rerank batch, redoing it with exact truncation on overflow.
1028 The char estimate under-counts token-dense pairs, so an over-cap pair
1029 can slip through untruncated; the server rejects it as CONTEXT_OVERFLOW
1030 and the batch is redone with exact server-side tokenization.
1031 """
1032 try:
1033 return self._embeddings_call(batch)
1034 except ProviderError as exc:
1035 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW or self._token_cap is None:
1036 raise
1037 cap = self._token_cap
1038 exact = [self._fit_input(text, cap, estimate=False)[0] for text in batch]
1039 return self._embeddings_call(exact)
1041 def _rerank_llm(self, query: str, candidates: list[str]) -> list[float]:
1042 """Score each candidate by an LLM's yes/no first-token logprob.
1044 Raises ``ProviderError`` when no candidate yields a verdict.
1045 """
1046 template = cfg.reranker_prompt or _LLM_RERANK_PROMPT
1047 workers = min(LLM_RERANK_CONCURRENCY, len(candidates))
1048 with ThreadPoolExecutor(max_workers=workers) as pool:
1049 scores = list(pool.map(lambda c: self._llm_rerank_one(template, query, c), candidates))
1050 if all(score is None for score in scores):
1051 raise ProviderError(_LLM_RERANK_NO_VERDICT_ERROR, provider=_PROVIDER_NAME)
1052 return [0.0 if score is None else score for score in scores]
1054 def _llm_rerank_one(self, template: str, query: str, candidate: str) -> float | None:
1055 """One chat request scoring a single candidate's relevance to the query."""
1056 content = template.format(query=query, document=candidate)
1057 payload = {
1058 "model": self._model,
1059 "messages": [{"role": "user", "content": content}],
1060 "max_tokens": 1,
1061 "temperature": 0,
1062 "logprobs": True,
1063 "top_logprobs": _LLM_RERANK_TOP_LOGPROBS,
1064 "stream": False,
1065 # Scoring reads the first generated token; a thinking template would
1066 # spend it opening a <think> block instead of answering.
1067 "chat_template_kwargs": {"enable_thinking": False},
1068 }
1070 def _call() -> dict[str, Any]:
1071 with self._track():
1072 resp = self._http.post(_CHAT_PATH, json=payload)
1073 _raise_for_status(resp)
1074 return dict(resp.json())
1076 return _llm_rerank_score(_first_token_top_logprobs(retry_on_busy(_call)))
1078 def _embeddings_call(self, inputs: list[str]) -> list[dict[str, Any]]:
1079 """POST one already-budgeted sub-batch to ``/v1/embeddings``; return its data."""
1081 def _call() -> list[dict[str, Any]]:
1082 with self._track():
1083 resp = self._http.post(
1084 _EMBED_PATH,
1085 json={
1086 "model": self._model,
1087 "input": inputs,
1088 "embd_normalize": _EMBD_NORMALIZE_NONE,
1089 "encoding_format": _EMBED_ENCODING_FORMAT,
1090 },
1091 )
1092 _raise_for_status(resp)
1093 data = resp.json()["data"]
1094 if len(data) != len(inputs):
1095 raise ProviderError(
1096 f"Embedder returned {len(data)} vectors for {len(inputs)} inputs",
1097 provider=_PROVIDER_NAME,
1098 )
1099 return list(data)
1101 # Bulk ingest can afford to wait out a cold-start warmup rather than drop
1102 # files. With a cold-load deadline (the EMBED-role client) the retry waits
1103 # out a still-loading replica for the full budget llama-swap keeps it alive,
1104 # instead of dropping the file after the fixed attempt cap; without one the
1105 # fixed count bounds an interactive caller.
1106 if self._embed_busy_deadline_s is not None:
1107 return retry_on_busy(_call, deadline=time.monotonic() + self._embed_busy_deadline_s)
1108 return retry_on_busy(_call, retries=_EMBED_BUSY_RETRIES)
1110 def _truncate_and_subbatch(self, texts: list[str], *, estimate: bool) -> list[list[str]]:
1111 """Token-truncate over-cap inputs, then pack into server-sized sub-batches.
1113 An input longer than ``token_cap`` (the server's per-slot context /
1114 n_batch) is truncated to it via the server's tokenizer, since the server
1115 cannot split a pooled embedding sequence. Inputs are then grouped so each
1116 request stays within both the token budget and ``_EMBED_N_SEQ_MAX``
1117 sequences, bounding per-request size and the busy-retry window. No cap
1118 (chat/vision) sends a single batch untouched.
1120 When ``estimate`` is set the per-input token count comes from
1121 :func:`_estimate_tokens`, and ``/tokenize`` is consulted only for the
1122 rare input whose estimate exceeds the cap -- eliminating a round-trip per
1123 chunk during bulk ingest. ``estimate=False`` (the overflow redo) tokenizes
1124 every input exactly.
1125 """
1126 if self._token_cap is None:
1127 return [texts]
1128 cap = self._token_cap
1129 batches: list[list[str]] = []
1130 current: list[str] = []
1131 current_tokens = 0
1132 for text in texts:
1133 item, item_tokens = self._fit_input(text, cap, estimate=estimate)
1134 if current and (current_tokens + item_tokens > cap or len(current) >= _EMBED_N_SEQ_MAX):
1135 batches.append(current)
1136 current = []
1137 current_tokens = 0
1138 current.append(item)
1139 current_tokens += item_tokens
1140 if current:
1141 batches.append(current)
1142 return batches
1144 def _fit_input(self, text: str, cap: int, *, estimate: bool) -> tuple[str, int]:
1145 """Return ``(input, token_count)`` for one sequence, truncating if over cap.
1147 Estimation short-circuits the common case: an estimate within the cap is
1148 trusted (no ``/tokenize``); only an over-cap estimate is confirmed against
1149 the server tokenizer and truncated if it really exceeds the cap.
1150 """
1151 if estimate:
1152 est = _estimate_tokens(text)
1153 if est <= cap:
1154 return text, est
1155 tokens = self._tokenize(text)
1156 if len(tokens) > cap:
1157 log.warning("Truncating oversize embed input: %d tokens > cap %d", len(tokens), cap)
1158 return self._detokenize(tokens[:cap]), cap
1159 return text, max(1, len(tokens))
1161 def _native_route(self, suffix: str) -> str:
1162 """Path for a native (non-OpenAI) llama-server route through llama-swap.
1164 llama-swap proxies these only under ``/upstream/<model>/...``; the model
1165 is carried in the path, not the body (unlike the ``/v1`` OpenAI routes).
1166 """
1167 return f"{_UPSTREAM_PREFIX}/{self._model}{suffix}"
1169 def _tokenize(self, text: str) -> list[int]:
1170 resp = self._http.post(
1171 self._native_route(_TOKENIZE_PATH),
1172 json={
1173 "content": text,
1174 "add_special": _TOKENIZE_ADD_SPECIAL,
1175 "parse_special": _TOKENIZE_PARSE_SPECIAL,
1176 },
1177 )
1178 _raise_for_status(resp)
1179 return list(resp.json()["tokens"])
1181 def count_tokens(self, text: str) -> int:
1182 """Number of tokens *text* encodes to under the server's tokenizer."""
1183 return len(self._tokenize(text))
1185 def _detokenize(self, tokens: list[int]) -> str:
1186 resp = self._http.post(self._native_route(_DETOKENIZE_PATH), json={"tokens": tokens})
1187 _raise_for_status(resp)
1188 return str(resp.json()["content"])
1190 @contextlib.contextmanager
1191 def _abortable(self, resp: httpx.Response) -> Generator[None]:
1192 """Expose *resp* to ``abort_streams`` for the duration of its read loop."""
1193 with self._in_flight_lock:
1194 self._active_streams.add(resp)
1195 try:
1196 yield
1197 finally:
1198 with self._in_flight_lock:
1199 self._active_streams.discard(resp)
1201 def abort_streams(self) -> None:
1202 """Sever every in-flight SSE response on this replica.
1204 Closing the response from another thread unblocks a reader stuck in
1205 ``iter_lines`` with a stream error, which unwinds its worker;
1206 llama-server stops generating when the connection drops.
1207 """
1208 with self._in_flight_lock:
1209 streams = list(self._active_streams)
1210 for resp in streams:
1211 with contextlib.suppress(Exception):
1212 resp.close()
1214 def close(self) -> None:
1215 """Close the underlying client if this instance created it."""
1216 if self._owns_http:
1217 self._http.close()
1219 def _track(self) -> _InFlight:
1220 return _InFlight(self)
1223class _InFlight:
1224 """Context manager that atomically bumps the owner's in-flight counter.
1226 ``+= 1`` is a read-modify-write, so concurrent chat/embed calls would corrupt
1227 the counter the router balances on; the client's lock makes it atomic.
1228 """
1230 def __init__(self, client: LlamaServerClient) -> None:
1231 self._client = client
1233 def __enter__(self) -> None:
1234 with self._client._in_flight_lock:
1235 self._client.in_flight += 1
1237 def __exit__(self, *_exc: object) -> None:
1238 with self._client._in_flight_lock:
1239 self._client.in_flight -= 1
1242def _embedding_vector(item: dict[str, Any]) -> npt.NDArray[np.float32]:
1243 """Decode one ``/v1/embeddings`` item's vector from its base64 float buffer."""
1244 embedding = item.get("embedding")
1245 # Untyped server JSON: a non-string means the encoding format was not honored.
1246 if not isinstance(embedding, str):
1247 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME)
1248 try:
1249 return np.frombuffer(base64.b64decode(embedding), dtype=_EMBED_VECTOR_DTYPE)
1250 except ValueError as exc:
1251 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME) from exc
1254def _rerank_score(item: dict[str, Any]) -> float:
1255 """Pull one relevance score from a rank-pooling ``/v1/embeddings`` item."""
1256 vector = _embedding_vector(item)
1257 if not vector.size:
1258 raise ProviderError(_NO_RERANK_SCORE_ERROR, provider=_PROVIDER_NAME)
1259 return float(vector[_RANK_SCORE_INDEX])
1262def _first_token_top_logprobs(response: dict[str, Any]) -> list[dict[str, Any]]:
1263 """The first generated token's top_logprobs list from a chat completion, or []."""
1264 choices = response.get("choices") or []
1265 if not choices:
1266 return []
1267 content = (choices[0].get("logprobs") or {}).get("content") or []
1268 if not content:
1269 return []
1270 return list(content[0].get("top_logprobs") or [])
1273def _llm_rerank_score(top_logprobs: list[dict[str, Any]]) -> float | None:
1274 """Softmax of the yes vs no logprobs in a token's top_logprobs (case/space-insensitive).
1276 ``None`` when neither verdict appears, distinct from the 0.0 of a confident "no".
1277 """
1278 yes_lp: float | None = None
1279 no_lp: float | None = None
1280 for entry in top_logprobs:
1281 token = str(entry.get("token", "")).strip().lower()
1282 logprob = float(entry.get("logprob", 0.0))
1283 if token == _YES_LABEL and (yes_lp is None or logprob > yes_lp):
1284 yes_lp = logprob
1285 elif token == _NO_LABEL and (no_lp is None or logprob > no_lp):
1286 no_lp = logprob
1287 if yes_lp is None:
1288 return None if no_lp is None else 0.0
1289 if no_lp is None:
1290 return math.exp(yes_lp)
1291 yes_e, no_e = math.exp(yes_lp), math.exp(no_lp)
1292 return yes_e / (yes_e + no_e)
1295def _parse_sse_deltas(line: str) -> tuple[str, str]:
1296 """Extract the (reasoning, content) deltas from one OpenAI SSE line."""
1297 if not line.startswith(_DATA_PREFIX):
1298 return "", ""
1299 body = line[len(_DATA_PREFIX) :].strip()
1300 if not body or body == _DONE_SENTINEL:
1301 return "", ""
1302 try:
1303 obj = json.loads(body)
1304 except json.JSONDecodeError:
1305 return "", ""
1306 choices = obj.get("choices") or []
1307 if not choices:
1308 return "", ""
1309 delta = choices[0].get("delta") or {}
1310 return str(delta.get("reasoning_content") or ""), str(delta.get("content") or "")
1313class _ThinkInliner:
1314 """Re-inlines server-extracted reasoning deltas as inline ``<think>`` text.
1316 The server parses each model's reasoning format natively (``--reasoning-format``)
1317 and streams it as ``reasoning_content``; lilbee's pipeline speaks inline
1318 ``<think>`` text, so the chat boundary opens the tag on the first reasoning
1319 delta and closes it when the answer starts (or at end of stream). Disabled
1320 (the non-chat roles), reasoning is dropped and content passes through, matching
1321 the server-extracted default those roles already ran with.
1322 """
1324 def __init__(self, *, enabled: bool) -> None:
1325 self._enabled = enabled
1326 self._in_think = False
1328 def feed(self, reasoning: str, content: str) -> str:
1329 if not self._enabled:
1330 return content
1331 parts: list[str] = []
1332 if reasoning:
1333 if not self._in_think:
1334 self._in_think = True
1335 parts.append(THINK_OPEN_TAG)
1336 parts.append(reasoning)
1337 if content:
1338 if self._in_think:
1339 self._in_think = False
1340 parts.append(THINK_CLOSE_TAG)
1341 parts.append(content)
1342 return "".join(parts)
1344 def finish(self) -> str:
1345 """Close an unterminated think block at end of stream."""
1346 if self._in_think:
1347 self._in_think = False
1348 return THINK_CLOSE_TAG
1349 return ""
1352def _inline_message_reasoning(message: Mapping[str, Any], *, enabled: bool) -> str:
1353 """A non-streaming message's text with any extracted reasoning re-inlined."""
1354 content = str(message.get("content") or "")
1355 reasoning = str(message.get("reasoning_content") or "") if enabled else ""
1356 if reasoning:
1357 return f"{THINK_OPEN_TAG}{reasoning}{THINK_CLOSE_TAG}{content}"
1358 return content
1361def _usage_from_body(body: Mapping[str, Any]) -> TokenUsage | None:
1362 """Read the ``usage`` block of an OpenAI response, or ``None`` if absent.
1364 llama-server reports ``prompt_tokens`` / ``completion_tokens``; a missing or
1365 malformed block yields ``None`` so callers can decide between a zero default
1366 (non-streaming) and skipping the frame (streaming terminator).
1367 """
1368 usage = body.get("usage")
1369 if not isinstance(usage, Mapping):
1370 return None
1371 prompt = usage.get("prompt_tokens")
1372 completion = usage.get("completion_tokens")
1373 return TokenUsage(
1374 prompt_tokens=prompt if isinstance(prompt, int) else 0,
1375 completion_tokens=completion if isinstance(completion, int) else 0,
1376 )
1379def _coerce_finish_reason(raw: Any) -> FinishReason:
1380 """Map a server-supplied finish_reason string to :class:`FinishReason`."""
1381 return FinishReason.coerce(raw)
1384def _tool_call_delta_from_chunk(call: Mapping[str, Any], *, fallback_index: int) -> ToolCallDelta:
1385 """Map one streaming ``delta.tool_calls`` entry to a :class:`ToolCallDelta`.
1387 Mirrors the SDK path: ``id`` / ``name`` arrive on the opener and accumulate
1388 by ``index``; empty strings normalise to ``None`` so the dispatch's stream
1389 translator (which gates on ``is not None``) does not emit spurious openers.
1390 """
1391 raw_index = call.get("index")
1392 index = raw_index if isinstance(raw_index, int) else fallback_index
1393 call_id = call.get("id")
1394 fn = call.get("function")
1395 raw_name = fn.get("name") if isinstance(fn, Mapping) else None
1396 raw_args = fn.get("arguments") if isinstance(fn, Mapping) else None
1397 return ToolCallDelta(
1398 index=index,
1399 id=str(call_id) if call_id else None,
1400 name=str(raw_name) if raw_name else None,
1401 arguments_delta=str(raw_args) if raw_args else None,
1402 )
1405_PROMPT_PROGRESS_KEY = "prompt_progress"
1408def _prefill_progress(line: str) -> tuple[int, int] | None:
1409 """The ``(processed, total)`` of one SSE line's ``prompt_progress``, or None.
1411 llama-server emits the block on streamed chats that opt in via
1412 ``return_progress``; the substring pre-check keeps the extra JSON parse off
1413 every ordinary token line.
1414 """
1415 if _PROMPT_PROGRESS_KEY not in line or not line.startswith(_DATA_PREFIX):
1416 return None
1417 body = line[len(_DATA_PREFIX) :].strip()
1418 try:
1419 obj = json.loads(body)
1420 except json.JSONDecodeError:
1421 return None
1422 progress = obj.get(_PROMPT_PROGRESS_KEY)
1423 if not isinstance(progress, Mapping):
1424 return None
1425 try:
1426 return int(progress["processed"]), int(progress["total"])
1427 except (KeyError, TypeError, ValueError):
1428 return None
1431def _parse_sse_stream_items(
1432 line: str, inliner: _ThinkInliner
1433) -> Iterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1434 """Yield text tokens, tool-call deltas, and the finish frame from one SSE line.
1436 A chunk can carry a ``content`` token, a ``reasoning_content`` token (routed
1437 through *inliner*), a ``tool_calls`` delta array, or a mix; each is yielded as
1438 its own :data:`ChatStreamItem` frame. The chunk that closes the turn carries
1439 ``choices[0].finish_reason``, surfaced as a :class:`StreamFinish` so the
1440 dispatch reports ``length`` (and friends), not just the default end-of-turn.
1441 """
1442 if not line.startswith(_DATA_PREFIX):
1443 return
1444 body = line[len(_DATA_PREFIX) :].strip()
1445 if not body or body == _DONE_SENTINEL:
1446 return
1447 try:
1448 obj = json.loads(body)
1449 except json.JSONDecodeError:
1450 return
1451 choices = obj.get("choices") or []
1452 if not choices:
1453 # The include_usage terminator chunk has an empty choices list and the
1454 # token totals on a top-level ``usage`` block; surface it as the final
1455 # frame so the dispatch can attach real counts to the stream.
1456 usage = _usage_from_body(obj)
1457 if usage is not None:
1458 yield usage
1459 return
1460 delta = choices[0].get("delta") or {}
1461 text = inliner.feed(str(delta.get("reasoning_content") or ""), str(delta.get("content") or ""))
1462 if text:
1463 yield text
1464 raw_calls = delta.get("tool_calls") or []
1465 for i, call in enumerate(raw_calls):
1466 if isinstance(call, Mapping):
1467 yield _tool_call_delta_from_chunk(call, fallback_index=i)
1468 raw_finish = choices[0].get("finish_reason")
1469 if raw_finish is not None:
1470 yield StreamFinish(reason=_coerce_finish_reason(raw_finish))
1473def _arguments_to_str(arguments: Any) -> str:
1474 """Normalize a tool-call ``arguments`` value to a JSON string (OpenAI's shape)."""
1475 if isinstance(arguments, str):
1476 return arguments
1477 if arguments is None:
1478 return "{}"
1479 return json.dumps(arguments)
1482def _parse_native_tool_calls(raw: Any) -> list[ToolCall]:
1483 """Map a response's ``message.tool_calls`` array to :class:`ToolCall` objects.
1485 Reads the OpenAI shape (``{"id", "function": {"name", "arguments"}}``) that
1486 ``--jinja`` produces. Malformed or nameless entries are skipped.
1487 """
1488 if not isinstance(raw, list):
1489 return []
1490 calls: list[ToolCall] = []
1491 for idx, entry in enumerate(raw):
1492 if not isinstance(entry, Mapping):
1493 continue
1494 fn = entry.get("function")
1495 if not isinstance(fn, Mapping):
1496 continue
1497 name = fn.get("name")
1498 if not isinstance(name, str) or not name:
1499 continue
1500 call_id = entry.get("id")
1501 calls.append(
1502 ToolCall(
1503 id=call_id if isinstance(call_id, str) and call_id else f"call_{idx}",
1504 name=name,
1505 arguments=_arguments_to_str(fn.get("arguments")),
1506 )
1507 )
1508 return calls
1511def _bare_call_from_mapping(obj: Mapping[str, Any], *, index: int) -> ToolCall | None:
1512 """Build a ToolCall from a bare ``{"name", "arguments"|"parameters"}`` object."""
1513 name = obj.get("name")
1514 if not isinstance(name, str) or not name:
1515 return None
1516 arguments = obj.get("arguments")
1517 if arguments is None:
1518 arguments = obj.get("parameters")
1519 return ToolCall(id=f"call_{index}", name=name, arguments=_arguments_to_str(arguments))
1522def _recover_bare_json_tool_calls(content: str) -> ChatToolResult:
1523 """Recover a tool call a model emitted as bare-JSON content (a native miss).
1525 Some models ignore the tool-call protocol and print ``{"name": ...,
1526 "arguments": {...}}`` (or a list of them) as the message body. When the whole
1527 content parses as such, treat it as the call(s) and clear the text; otherwise
1528 return the content unchanged with no calls.
1529 """
1530 stripped = content.strip()
1531 if not stripped or stripped[0] not in "{[":
1532 return ChatToolResult(content=content, tool_calls=[])
1533 try:
1534 parsed = json.loads(stripped)
1535 except json.JSONDecodeError:
1536 return ChatToolResult(content=content, tool_calls=[])
1537 entries = parsed if isinstance(parsed, list) else [parsed]
1538 calls = [
1539 call
1540 for idx, entry in enumerate(entries)
1541 if isinstance(entry, Mapping) and (call := _bare_call_from_mapping(entry, index=idx))
1542 ]
1543 if not calls:
1544 return ChatToolResult(content=content, tool_calls=[])
1545 return ChatToolResult(content="", tool_calls=calls)
1548# Leading non-whitespace characters that mark streamed text as a potential bare
1549# JSON tool call (an object or an array of them); any other first char is plain
1550# text and streams through untouched.
1551_BARE_CALL_OPENERS = "{["
1554def _tool_call_delta_from_recovered(call: ToolCall, index: int) -> ToolCallDelta:
1555 """Shape a recovered bare-JSON :class:`ToolCall` as a single streaming delta.
1557 Mirrors :func:`_tool_call_delta_from_chunk`: id and name ride the opener (the
1558 only frame for a recovered call), the arguments JSON is the lone
1559 ``arguments_delta``, and the position is the index.
1560 """
1561 return ToolCallDelta(
1562 index=index,
1563 id=call.id or None,
1564 name=call.name or None,
1565 arguments_delta=call.arguments or None,
1566 )
1569def _recover_bare_json_stream(
1570 items: Iterator[str | ToolCallDelta | TokenUsage | StreamFinish],
1571) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1572 """Wrap a raw chat stream to recover a tool call emitted as bare-JSON text.
1574 Some small models print ``{"name": ..., "arguments": {...}}`` as content
1575 instead of native ``tool_calls``; the non-stream paths recover this via
1576 :func:`_recover_bare_json_tool_calls`. This applies the same recovery to the
1577 stream, but only when the model emitted no native :class:`ToolCallDelta` and
1578 the streamed text looks like a bare call from its first character. Normal text
1579 still streams token by token: once the buffered head proves not to be a bare
1580 call it is flushed and all later text passes straight through.
1581 """
1582 buffer = "" # leading text held back as a potential bare call until resolved
1583 # True once leading text has streamed as plain (or a native call was seen):
1584 # past that, a later '{'/'[' is content, not a bare call -- never buffer again.
1585 committed = False
1586 try:
1587 for item in items:
1588 if isinstance(item, ToolCallDelta):
1589 yield from _flush_plain(buffer)
1590 buffer, committed = "", True
1591 yield item
1592 elif isinstance(item, TokenUsage | StreamFinish):
1593 yield from _recover_buffer(buffer)
1594 buffer = ""
1595 yield item
1596 elif committed or _passthrough_text(buffer, item):
1597 yield from _flush_plain(buffer)
1598 buffer = ""
1599 committed = True
1600 yield item
1601 else:
1602 buffer += item
1603 yield from _recover_buffer(buffer)
1604 finally:
1605 # Forward close to the source generator: if a consumer closes this
1606 # wrapper mid-stream, a plain for-loop would not propagate GeneratorExit
1607 # to *items*, leaking the underlying HTTP stream and its in_flight slot.
1608 # Suppress teardown errors (httpx stream close can raise) so they don't
1609 # mask the exception that triggered this finally.
1610 if isinstance(items, Generator):
1611 with contextlib.suppress(Exception):
1612 items.close()
1615def _passthrough_text(buffer: str, text: str) -> bool:
1616 """Whether *text* should stream through directly rather than buffer.
1618 True once the accumulated head's first non-whitespace char is known and is not
1619 a bare-call opener (plain text): the buffer is empty in that case, so the
1620 caller yields *text* as is. While the head is all whitespace, or once it opens
1621 with ``{``/``[``, the text is buffered (False) pending recovery.
1622 """
1623 head = (buffer + text).lstrip()
1624 return bool(head) and head[0] not in _BARE_CALL_OPENERS
1627def _flush_plain(buffer: str) -> Iterator[str]:
1628 """Yield buffered leading text verbatim (it was not a bare call after all)."""
1629 if buffer:
1630 yield buffer
1633def _recover_buffer(buffer: str) -> Iterator[str | ToolCallDelta]:
1634 """Resolve the buffered leading text at a terminator or end of stream.
1636 The buffer reaching here was held as a potential bare call (text starting with
1637 ``{``/``[`` and no native call seen). Run :func:`_recover_bare_json_tool_calls`:
1638 emit one delta per recovered call, or yield the text unchanged when it only
1639 happened to start with ``{``/``[`` but is not a call.
1640 """
1641 if not buffer:
1642 return
1643 recovered = _recover_bare_json_tool_calls(buffer)
1644 if not recovered.tool_calls:
1645 yield buffer
1646 return
1647 for index, call in enumerate(recovered.tool_calls):
1648 yield _tool_call_delta_from_recovered(call, index)