Coverage for src/lilbee/providers/fleet/client.py: 100%
851 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Thin httpx client for one llama-server OpenAI endpoint (local inference)."""
3from __future__ import annotations
5import asyncio
6import base64
7import contextlib
8import json
9import logging
10import math
11import socket
12import ssl
13import threading
14import time
15from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
16from concurrent.futures import ThreadPoolExecutor
17from pathlib import Path
18from typing import Any, Literal, TypedDict, TypeVar, overload
20import httpx
21import numpy as np
22import numpy.typing as npt
24from lilbee.core.config import cfg
25from lilbee.core.vectors import Vector
26from lilbee.providers.base import (
27 THINK_CLOSE_TAG,
28 THINK_OPEN_TAG,
29 ChatResult,
30 ChatToolResult,
31 ClosableIterator,
32 FinishReason,
33 ProviderError,
34 ProviderErrorKind,
35 StreamFinish,
36 TokenUsage,
37 ToolCall,
38 ToolCallDelta,
39)
40from lilbee.providers.fleet.adapters import LLM_RERANK_CONCURRENCY
41from lilbee.providers.fleet.normalize import ChatMessage, to_alternating
42from lilbee.providers.roles import RerankMode
43from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
45_PROVIDER_NAME = "llama-server"
46# Shown when a stream breaks after its first frame. The client cannot know why.
47_STREAM_INTERRUPTED_MESSAGE = (
48 "The model server stopped answering and the answer is cut. Send your message again."
49)
51# Fleet clients only ever talk to a loopback llama-server over plain HTTP, so TLS is
52# never negotiated. httpx still builds a default SSL context per client (loads the
53# system CA bundle, ~13 ms each and slower on macOS via the keychain), which is pure
54# overhead paid on every fleet reload. Build one minimal context and share it so a
55# reload doesn't reload the CA bundle for each replica.
56_LOOPBACK_SSL_CONTEXT = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
57_LOOPBACK_SSL_CONTEXT.check_hostname = False
58_LOOPBACK_SSL_CONTEXT.verify_mode = ssl.CERT_NONE
59# Reranker pair format: query and candidate are joined with this separator into
60# one document so a cross-encoder GGUF scores the pair as a single sequence.
61_RERANK_PAIR_SEPARATOR = "</s></s>"
62# LLM reranker: score each candidate by the yes/no first-token logprob.
63_LLM_RERANK_PROMPT = (
64 "Judge whether the document is relevant to the query. "
65 "Answer with only 'yes' or 'no'.\n\nQuery: {query}\nDocument: {document}"
66)
67_LLM_RERANK_TOP_LOGPROBS = 20
68_YES_LABEL = "yes"
69_NO_LABEL = "no"
70_LLM_RERANK_NO_VERDICT_ERROR = (
71 "The reranker model never answered 'yes' or 'no', so its relevance scores are "
72 "unusable. Its chat template does not fit the relevance prompt. Choose a GGUF "
73 "built for reranking, set reranker_type to cross_encoder, or adapt "
74 "reranker_prompt to the model's expected format."
75)
76# Max sequences per /v1/embeddings request. Like the in-process backstop, a
77# batch is bounded by BOTH the token budget (the server's n_batch, == token_cap)
78# and this sequence count: a corpus of many tiny chunks would otherwise pack one
79# request past the server's batch/sequence limit and trip a 500.
80_EMBED_N_SEQ_MAX = 64
81# Estimate a chunk's token count from its character length so the bulk embed
82# path packs sub-batches without a /tokenize round-trip per input. The factor is
83# held below the corpus average (data.extract.chunk.CHARS_PER_TOKEN, 4 for Latin text)
84# so the estimate over-counts tokens and a sub-batch never packs past the
85# server's n_batch (== token_cap). Rerank pairs are token-dense (the separator
86# is several tokens in a few chars), so an under-count can slip an over-cap
87# pair through; the rerank path re-truncates exactly on the server's overflow
88# error instead of paying a /tokenize round trip per pair up front.
89_EMBED_EST_CHARS_PER_TOKEN = 3
90# llama-swap's error body when the spawned llama-server exited before serving.
91_UPSTREAM_DIED_MARKER = "exited prematurely"
92# llama-server's exit line when the port lilbee picked was taken by the time
93# the server bound it (the pick-then-bind gap spans the whole lazy-spawn wait,
94# so a passing ephemeral connection can occupy it). The occupation is
95# transient: llama-swap re-spawns on the next request against the same port
96# and normally binds.
97_BIND_FAILURE_MARKER = "couldn't bind HTTP server socket"
98# llama-server's 500 body when one input exceeds the physical batch (n_batch).
99_BATCH_OVERFLOW_MARKER = "too large to process"
102class _ChatToolSpecFunction(TypedDict, total=False):
103 """The ``function`` payload of an OpenAI tool definition (wire shape)."""
105 name: str
106 description: str
107 parameters: dict[str, Any]
110class ChatTool(TypedDict, total=False):
111 """One OpenAI tool definition sent in a chat request (wire shape)."""
113 type: str
114 function: _ChatToolSpecFunction
117# Some GGUF chat templates (Mistral-Nemo, Cohere command-r) reject a standard
118# OpenAI tool exchange: they require plain user/assistant turns to alternate and
119# raise a Jinja exception on the tool role or two same-role turns in a row. Rather
120# than fail a real request and parse the engine's error text, the client probes
121# the live template once per server with this representative tool exchange: if the
122# server rejects it as sent but renders the to_alternating() form, the model is
123# flagged so every later request is reshaped up front. Two assistant tool-call
124# turns separated by tool results is the minimal shape that trips strict
125# alternation; max_tokens=1 keeps the probe to template rendering, not generation.
126_ALTERNATION_PROBE_TOOLS: list[ChatTool] = [
127 {
128 "type": "function",
129 "function": {
130 "name": "probe",
131 "description": "Probe whether the chat template renders a tool exchange.",
132 # A single declared property (rather than an empty object) so a grammar
133 # that requires at least one parameter still renders the probe call.
134 "parameters": {
135 "type": "object",
136 "properties": {"query": {"type": "string"}},
137 "required": ["query"],
138 },
139 },
140 }
141]
142_ALTERNATION_PROBE_MESSAGES: list[ChatMessage] = [
143 {"role": "system", "content": "You are a helpful assistant."},
144 {"role": "user", "content": "Look something up."},
145 {
146 "role": "assistant",
147 "content": "",
148 "tool_calls": [
149 {
150 "id": "probe-1",
151 "type": "function",
152 "function": {"name": "probe", "arguments": '{"query": "x"}'},
153 }
154 ],
155 },
156 {"role": "tool", "tool_call_id": "probe-1", "content": "first result"},
157 {
158 "role": "assistant",
159 "content": "",
160 "tool_calls": [
161 {
162 "id": "probe-2",
163 "type": "function",
164 "function": {"name": "probe", "arguments": '{"query": "x"}'},
165 }
166 ],
167 },
168 {"role": "tool", "tool_call_id": "probe-2", "content": "second result"},
169 {"role": "user", "content": "Summarize."},
170]
171_ALTERNATION_PROBE_OPTIONS = {"max_tokens": 1}
172# The probe holds _alternation_lock across its request, so it uses a short, bounded
173# timeout rather than the chat default: a slow/wedged replica yields an inconclusive
174# (transient) result and a re-probe instead of blocking every first chat on the lock.
175_ALTERNATION_PROBE_TIMEOUT_S = 30.0
176_UPSTREAM_LOG_TAIL_CHARS = 2000
177_UPSTREAM_LOG_TIMEOUT_S = 2.0
178# Enough reads to cover one replay of llama-swap's 100KB per-model ring at
179# httpx's 64KB ceiling per chunk, with room to spare. The route keeps streaming
180# live lines afterwards, so without a bound this would cost the read timeout on
181# every death.
182_UPSTREAM_LOG_MAX_CHUNKS = 8
184log = logging.getLogger(__name__)
187def _estimate_tokens(text: str) -> int:
188 """Conservative (over-counting) token estimate from character length."""
189 return max(1, -(-len(text) // _EMBED_EST_CHARS_PER_TOKEN))
192_UPSTREAM_DEATH_MESSAGES: dict[ProviderErrorKind, str] = {
193 ProviderErrorKind.CAPACITY: "The engine ran out of memory while loading {model}",
194 ProviderErrorKind.PORT_CONFLICT: "The engine could not start {model}: its port was taken",
195}
196_UPSTREAM_DEATH_DEFAULT = "The engine for {model} stopped before it could answer"
197_BROKEN_MODEL_FILE_MESSAGE = (
198 "The model file for {model} is incomplete or damaged. "
199 "Run 'lilbee model rm {model}', then 'lilbee model pull {model}'."
200)
201_BROKEN_UNNAMED_MODEL_FILE_MESSAGE = (
202 "The model file is incomplete or damaged. "
203 "Run 'lilbee model list', remove that model with 'lilbee model rm', then pull it again."
204)
205# What llama.cpp prints when a GGUF ends before its tensor data does.
206_BROKEN_FILE_MARKERS: tuple[str, ...] = (
207 "model is corrupted or incomplete",
208 "is not within the file bounds",
209)
212def _upstream_death_message(
213 kind: ProviderErrorKind, status_code: int, model: str | None, engine_log: Path | None
214) -> str:
215 """One plain line for a model server that exited, built from what lilbee knows."""
216 subject = model if model is not None else "the model"
217 if kind is ProviderErrorKind.MODEL_FILE_BROKEN:
218 if model is None:
219 return _BROKEN_UNNAMED_MODEL_FILE_MESSAGE
220 return _BROKEN_MODEL_FILE_MESSAGE.format(model=model)
221 cause = _UPSTREAM_DEATH_MESSAGES.get(kind, _UPSTREAM_DEATH_DEFAULT).format(model=subject)
222 message = f"{cause} (HTTP {status_code})."
223 if engine_log is None:
224 return message
225 return f"{message} The full engine output is in {engine_log}."
228# What the engine prints when a device allocation fails during load. Every
229# backend words it differently and all of them mean the same thing: the plan
230# asked for more memory than the device had. Taken from the emit sites in
231# upstream rather than guessed, and matched lowercased.
232#
233# One entry covers CUDA, HIP and MUSA: the vendor headers #define cudaMalloc to
234# their own allocator, but the log string in ggml-cuda.cu is a literal, so an
235# AMD or Moore Threads build still prints "cudaMalloc failed". A separate
236# hipMalloc marker would match nothing.
237#
238# Vulkan is the one that needs its own wording. It is where every AMD and Intel
239# GPU lands, and it says neither "out of memory" nor "failed to allocate".
240_OOM_MARKERS: tuple[str, ...] = (
241 "out of memory",
242 "failed to allocate", # Metal's buffer failure, and most generic paths
243 "cudamalloc failed", # CUDA, HIP and MUSA alike
244 "device memory allocation of size", # ggml-vulkan's fatal allocation failure
245 "outofdevicememory", # a vk::OutOfDeviceMemoryError that reached the log
246 "unable to allocate",
247 "insufficient memory",
248 "out_of_device_memory", # SYCL, which exits through the runtime's own code
249 "out_of_resources",
250)
252# Lines that report a failure the engine then works around. ggml-vulkan warns
253# that pinned memory could not be allocated and falls back to unpinned, and that
254# text matches an allocation marker word for word, so a later unrelated death
255# would be read as a memory shortfall and answered with a context reduction.
256_SURVIVABLE_LINE_MARKERS: tuple[str, ...] = ("warning:", "warn:")
259def classify_upstream_death(tail: str) -> ProviderErrorKind | None:
260 """The kind of failure an engine's dying output describes, or ``None``.
262 ``CAPACITY`` for a load that ran out of device memory: retrying the identical
263 launch respawns it into a crash loop, while a smaller context might fit.
264 ``PORT_CONFLICT`` for losing the port-bind race, which is worth retrying
265 because the retry re-drives llama-swap's spawn, and worth naming because a
266 port held for good needs a different port rather than another attempt at the
267 same one. ``None`` leaves the existing classification alone rather than
268 guessing at an unfamiliar death.
269 """
270 fatal = [
271 line
272 for line in tail.lower().splitlines()
273 if not any(marker in line for marker in _SURVIVABLE_LINE_MARKERS)
274 ]
275 # A damaged file is checked first: no memory or port change makes it load.
276 if any(marker in line for line in fatal for marker in _BROKEN_FILE_MARKERS):
277 return ProviderErrorKind.MODEL_FILE_BROKEN
278 if any(marker in line for line in fatal for marker in _OOM_MARKERS):
279 return ProviderErrorKind.CAPACITY
280 if _BIND_FAILURE_MARKER in tail:
281 return ProviderErrorKind.PORT_CONFLICT
282 return None
285# Deaths a role's rebuild can fix, and a retry against the same launch cannot: a
286# memory shortfall needs a smaller plan, a held port needs a different port. Both
287# come from rebuilding the role, which re-plans and re-picks.
288_REBUILDABLE_KINDS = frozenset({ProviderErrorKind.CAPACITY, ProviderErrorKind.PORT_CONFLICT})
291def is_load_capacity_failure(exc: BaseException) -> bool:
292 """True when *exc* is an engine that died for lack of device memory on load."""
293 return isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.CAPACITY
296def is_rebuildable_failure(exc: BaseException) -> bool:
297 """True when rebuilding the role is what stands a chance, not another retry."""
298 return isinstance(exc, ProviderError) and exc.kind in _REBUILDABLE_KINDS
301def _classify_error(status_code: int, body: str) -> ProviderErrorKind:
302 """Error kind from a llama-server/llama-swap error status and body.
304 An input past the server's n_batch is a 500 whose body says "too large to
305 process" (CONTEXT_OVERFLOW, so the embed path re-truncates exactly); a dead
306 upstream is CONNECTION, so the router can mark the replica unhealthy. The
307 body markers win over the status: llama-swap reports a died upstream under
308 gateway statuses too, and that case needs the failover path, not a retry
309 against the same dead server (except a bind-race death, which
310 ``_raise_for_status`` upgrades to SERVER once the upstream tail proves it).
311 A bare gateway error (502/503/504) is a
312 momentarily-unreachable upstream -- restarting, OOM-killed, mid-swap -- so
313 it is SERVER, which the gateway retry treats as transient.
314 """
315 if _BATCH_OVERFLOW_MARKER in body:
316 return ProviderErrorKind.CONTEXT_OVERFLOW
317 if _UPSTREAM_DIED_MARKER in body:
318 return ProviderErrorKind.CONNECTION
319 if status_code in _TRANSIENT_GATEWAY_STATUSES:
320 return ProviderErrorKind.SERVER
321 return ProviderErrorKind.UNKNOWN
324def is_connection_failure(exc: Exception) -> bool:
325 """Whether *exc* signals a dead/unreachable replica rather than a model error."""
326 if isinstance(exc, httpx.TransportError):
327 return True
328 # isinstance: only ProviderError carries a kind; other exceptions pass through.
329 return isinstance(exc, ProviderError) and exc.kind is ProviderErrorKind.CONNECTION
332def _is_transient_probe_failure(exc: Exception) -> bool:
333 """Whether a probe failure is transient (dead replica, busy 429, gateway
334 error), not a template verdict. llama-swap 429s past its concurrency limit
335 and a restarting replica 502s, so neither response must be read as the
336 template rejecting the exchange."""
337 if is_connection_failure(exc):
338 return True
339 return isinstance(exc, ProviderError) and exc.kind in _TRANSIENT_KINDS
342def _upstream_failure_tail(resp: httpx.Response) -> str:
343 """Return (and log) the dead upstream's recent output, or empty when unreadable."""
344 with contextlib.suppress(httpx.HTTPError, json.JSONDecodeError, KeyError, TypeError):
345 base = str(resp.request.url).split("/v1/")[0]
346 model = json.loads(resp.request.content)["model"]
347 tail = _fetch_log_tail(f"{base}/logs/stream/{model}")
348 if tail:
349 log.warning("%s exited prematurely; recent server output:\n%s", model, tail)
350 return tail
351 return ""
354def _fetch_log_tail(url: str) -> str:
355 """The last ``_UPSTREAM_LOG_TAIL_CHARS`` of llama-swap's log stream for one model.
357 The stream replays the upstream's buffered output then stays open; the read
358 timeout is the cutoff once the replay is drained.
359 """
360 chunks: list[str] = []
361 with (
362 contextlib.suppress(httpx.HTTPError),
363 httpx.stream("GET", url, timeout=_UPSTREAM_LOG_TIMEOUT_S) as stream,
364 ):
365 for taken, chunk in enumerate(stream.iter_text(), start=1):
366 # Bounded by chunk count, not by the tail size. llama-swap replays a
367 # model's whole ring in one write and httpx hands over at most 64KB
368 # at a time, so stopping at the first chunk past the tail size
369 # returned the head of a warm model's log, where the fatal last line
370 # never is. Only the tail is retained as the replay goes by, and the
371 # route streams live lines afterwards, so the count is what keeps
372 # this from waiting out the timeout on every death.
373 chunks.append(chunk)
374 chunks = ["".join(chunks)[-_UPSTREAM_LOG_TAIL_CHARS:]]
375 if taken >= _UPSTREAM_LOG_MAX_CHUNKS:
376 break
377 return "".join(chunks)[-_UPSTREAM_LOG_TAIL_CHARS:]
380# llama-server L2-normalizes pooled embeddings by default (embd_normalize=2);
381# every embeddings request sends embd_normalize=-1 so the engine returns raw
382# vectors, and so a rank-pooling rerank score (a single value per pair) is not
383# collapsed to +-1 by normalization. The server only exposes this per request
384# body, not as a startup flag.
385_EMBD_NORMALIZE_NONE = -1
386# Vectors come back as a base64 float32 buffer: parsing thousands of JSON float
387# literals per batch is CPU-bound and holds the GIL, which caps embedding
388# throughput below what the GPUs can feed regardless of how many are dispatching.
389_EMBED_ENCODING_FORMAT = "base64"
390# The engine writes the raw float buffer in host byte order; supported targets are
391# all little-endian.
392_EMBED_VECTOR_DTYPE = "<f4"
393# Rank pooling puts the pair's relevance score in the vector's first slot.
394_RANK_SCORE_INDEX = 0
395_UNREADABLE_EMBEDDING_ERROR = (
396 "The embedding server returned vectors lilbee could not read. Update the "
397 "inference engine: base64 embedding responses need llama-server b4391 or newer."
398)
399_NO_RERANK_SCORE_ERROR = "The reranker returned no relevance score for a candidate."
400_HEALTH_PATH = "/health"
401_CHAT_PATH = "/v1/chat/completions"
402_EMBED_PATH = "/v1/embeddings"
403_TOKENIZE_PATH = "/tokenize"
404_DETOKENIZE_PATH = "/detokenize"
405_APPLY_TEMPLATE_PATH = "/apply-template"
406# llama-swap proxies native (non-OpenAI) llama.cpp routes only under
407# /upstream/<model>/...; the bare /tokenize path 404s (it routes /v1/* by the
408# body's model field, but a native route carries no such field).
409_UPSTREAM_PREFIX = "/upstream"
410# Match the in-process tokenizer call (llm.tokenize(text, add_bos=True, special=False)):
411# the server adds BOS via add_special, and the chunk-sizing call leaves
412# special-token strings unparsed; the chat-prompt count overrides that, because
413# the render carries them.
414_TOKENIZE_ADD_SPECIAL = True
415_TOKENIZE_PARSE_SPECIAL = False
416_HTTP_OK = 200
417_HTTP_BAD_REQUEST = 400
418_HTTP_NOT_FOUND = 404
419_HTTP_METHOD_NOT_ALLOWED = 405
420_HTTP_TOO_MANY_REQUESTS = 429
421# Gateway statuses llama-swap returns while an upstream is unreachable
422# (502 crashing/restarting, 503 unavailable, 504 gateway timeout). The request
423# succeeds once the upstream is back, so these must never terminalize a call.
424_TRANSIENT_GATEWAY_STATUSES = frozenset({502, 503, 504})
425# Error kinds the gateway retry re-sends: a bare gateway error (upstream
426# momentarily unreachable) and a lost port-bind race clear on their own once
427# the server is back. A 429 is llama-swap's concurrency limit, which the swap
428# config sizes for several lilbee processes at full load, so it is reported,
429# not re-sent.
430_RETRY_KINDS = frozenset({ProviderErrorKind.SERVER, ProviderErrorKind.PORT_CONFLICT})
431# Error kinds that leave a template probe undetermined: the retry kinds plus a 429.
432_TRANSIENT_KINDS = _RETRY_KINDS | {ProviderErrorKind.RATE_LIMIT}
433_DONE_SENTINEL = "[DONE]"
434_DATA_PREFIX = "data:"
435_DEFAULT_TIMEOUT_S = 300.0
436# How often an abortable chat checks its cancel signal.
437_CANCEL_POLL_S = 0.1
438# Short, separate timeout for /health: a server can wedge under heavy prompt
439# processing, and readiness/monitor polls must not block on the request timeout.
440_HEALTH_TIMEOUT_S = 5.0
441# Retry a gateway error with exponential backoff (capped). llama-swap holds a
442# request through a cold load, so a gateway error is a replica restarting, and
443# callers fail after this short budget (~15s).
444_GATEWAY_RETRIES = 6
445_GATEWAY_BACKOFF_BASE_S = 0.5
446_GATEWAY_BACKOFF_MAX_S = 8.0
447# Half-open recovery: a replica marked unhealthy becomes routable again after
448# this cool-down. Recovery is probe-by-traffic and unmetered: every concurrent
449# caller sees it routable once cooled down (a success restores it, another
450# connection failure re-stamps the cool-down).
451_UNHEALTHY_RETRY_S = 30.0
452_T = TypeVar("_T")
455def _route_is_absent(resp: httpx.Response) -> bool:
456 """Whether *resp* means the route is missing rather than the request bad.
458 A proxy that knows the path but not the method answers 405. An engine too old
459 to carry the route answers 404 with a plain body, while llama-swap answers
460 404 with a JSON error envelope for a model it cannot route -- a
461 misconfiguration the caller must see, not one to answer with an estimate.
462 """
463 if resp.status_code == _HTTP_METHOD_NOT_ALLOWED:
464 return True
465 if resp.status_code != _HTTP_NOT_FOUND:
466 return False
467 resp.read() # streaming responses aren't read yet; a no-op for buffered ones
468 try:
469 body = resp.json()
470 except ValueError:
471 return True
472 return not (isinstance(body, dict) and "error" in body)
475class ChatDeadlineError(ProviderError):
476 """A bounded chat exceeded its caller-supplied total wall-clock deadline.
478 Distinct from a transport/server error so a deadline-bounded caller (vision
479 OCR) can word its own timeout message and skip failover without matching
480 error strings. Its ``UNKNOWN`` kind keeps it out of ``is_connection_failure``.
481 """
484class StreamInterruptedError(ProviderError):
485 """A chat stream broke after it had sent frames, so its answer is cut."""
487 def __init__(self) -> None:
488 super().__init__(_STREAM_INTERRUPTED_MESSAGE, provider=_PROVIDER_NAME)
491def retry_on_gateway_error(
492 call: Callable[[], _T], *, retries: int = _GATEWAY_RETRIES, deadline: float | None = None
493) -> _T:
494 """Run *call*, retrying gateway errors and a lost port-bind race with capped backoff.
496 A replica restarting mid-run answers 502 until it is back; backing off and
497 retrying turns that drop into a success. With a *deadline*
498 (``time.monotonic`` epoch) the retry waits out the server until that
499 deadline; without one, *retries* bounds the attempts. Every other error,
500 a 429 included, and the final still-failing response propagate.
501 """
502 delay = _GATEWAY_BACKOFF_BASE_S
503 attempt = 0
504 while True:
505 try:
506 return call()
507 except ProviderError as exc:
508 if exc.kind not in _RETRY_KINDS:
509 raise
510 attempt += 1
511 exhausted = (
512 time.monotonic() + delay >= deadline if deadline is not None else attempt >= retries
513 )
514 if exhausted:
515 raise
516 time.sleep(delay)
517 delay = min(delay * 2, _GATEWAY_BACKOFF_MAX_S)
520class LlamaServerClient:
521 """Calls one llama-server's OpenAI surface. Tracks in-flight requests so the
522 fleet router can pick the least-busy replica."""
524 def __init__(
525 self,
526 base_url: str,
527 model: str,
528 *,
529 http: httpx.Client | None = None,
530 async_transport: httpx.AsyncBaseTransport | None = None,
531 token_cap: int | None = None,
532 timeout: float = _DEFAULT_TIMEOUT_S,
533 rerank_mode: RerankMode | None = None,
534 inline_reasoning: bool = False,
535 on_prefill: Callable[[tuple[int, int] | None], None] | None = None,
536 engine_log: Path | None = None,
537 model_ref: str | None = None,
538 ) -> None:
539 self._base = base_url.rstrip("/")
540 # The file this server's own engine writes its output to; None when the
541 # launch writes none or no swap manager fronts the client (the self-check).
542 self.engine_log = engine_log
543 self._model = model
544 # The model as the user knows it (the registry ref); the replica id is not
545 # something a person recognises. None names it "the model".
546 self.model_ref = model_ref
547 # Prefill observer: called with (processed, total) per engine progress
548 # frame during a streamed chat's prompt processing, then None once the
549 # first generated frame proves the prefill is over.
550 self._on_prefill = on_prefill
551 self._http = http or httpx.Client(
552 base_url=self._base, timeout=timeout, verify=_LOOPBACK_SSL_CONTEXT
553 )
554 self._timeout = timeout
555 self._async_transport = async_transport
556 self._owns_http = http is None
557 # Chat-role clients re-inline server-extracted reasoning as <think> text;
558 # the other roles (vision OCR) keep dropping it, as their servers already did.
559 self._inline_reasoning = inline_reasoning
560 # Per-slot context for embed/rerank servers: inputs longer than this are
561 # token-truncated (via the server's tokenizer) before embedding, mirroring
562 # the in-process backstop. None for chat/vision, which don't truncate inputs.
563 self._token_cap = token_cap
564 # LLM => score candidates by yes/no logprob; None/cross-encoder => rank pooling.
565 self._rerank_mode = rerank_mode
566 self.in_flight = 0
567 self._in_flight_lock = threading.Lock()
568 # Called each time a request takes in_flight to zero; see notify_when_idle.
569 self._on_idle: Callable[[], None] | None = None
570 # Live SSE responses, so a cancel can sever the transport from another
571 # thread: a reader blocked in iter_lines cannot see a cooperative
572 # cancel flag, but a shutdown of its socket unblocks it with an error.
573 self._active_streams: set[httpx.Response] = set()
574 # Whether this server's chat template needs OpenAI tool exchanges reshaped
575 # into strict user/assistant alternation. Determined lazily by a one-time
576 # probe of the live template (see _prepare_chat_messages); None until then.
577 # A client is bound to one model for its lifetime, so the template (hence
578 # the verdict) is fixed once determined.
579 self._needs_alternation: bool | None = None
580 self._alternation_lock = threading.Lock()
581 # Routing health: cleared on a connection-level failure (see _UNHEALTHY_RETRY_S).
582 self._healthy = True
583 # Monotonic stamp of the last mark_unhealthy; consulted only while unhealthy.
584 self._unhealthy_since = 0.0
586 @property
587 def healthy(self) -> bool:
588 """Routable: healthy, or unhealthy past the ``_UNHEALTHY_RETRY_S`` cool-down."""
589 with self._in_flight_lock:
590 if self._healthy:
591 return True
592 return time.monotonic() - self._unhealthy_since >= _UNHEALTHY_RETRY_S
594 def _raise_for_status(self, resp: httpx.Response) -> None:
595 """Raise the plain error for a failed *resp*, naming this client's model and log.
597 The one place a client reports a failed response, so no call site can skip the
598 model name and the engine log. ``raise_for_status`` drops the server's error
599 body; the body carries the real reason (oversize prompt, decode failure, a
600 dead server), which both diagnosis and the user-facing error need.
601 """
602 if resp.is_success:
603 return
604 resp.read() # streaming responses aren't read yet; a no-op for buffered ones
605 body = resp.text.strip()
606 # llama-server reports an oversize prompt/conversation as a 400 whose body
607 # carries the "exceed_context_size_error" type. Tag it CONTEXT_OVERFLOW with a
608 # user-facing message so the chat route returns a clean context_length_exceeded
609 # (400) instead of a generic internal_error -- a long conversation that fills
610 # the window then reads as "too long", not "Internal server error".
611 if resp.status_code == _HTTP_BAD_REQUEST and (
612 "exceed_context_size" in body.lower() or "context size" in body.lower()
613 ):
614 raise ProviderError(
615 "The conversation exceeds this model's context window. "
616 "Start a new conversation or shorten the input.",
617 provider=_PROVIDER_NAME,
618 kind=ProviderErrorKind.CONTEXT_OVERFLOW,
619 )
620 # llama-swap answers 429 once its concurrency limit for the model is reached.
621 # The swap config sizes that limit for several lilbee processes at full load,
622 # so a request past it is reported rather than re-sent into the same load.
623 if resp.status_code == _HTTP_TOO_MANY_REQUESTS:
624 raise ProviderError(
625 "The model is busy with too many requests. Try again in a moment.",
626 provider=_PROVIDER_NAME,
627 kind=ProviderErrorKind.RATE_LIMIT,
628 )
629 kind = _classify_error(resp.status_code, body)
630 # llama-swap masks a dead server as "exited prematurely". The server's own
631 # captured output goes to the log (in _upstream_failure_tail) and decides the
632 # kind; the message the user sees stays one plain line.
633 if _UPSTREAM_DIED_MARKER in body:
634 classified = classify_upstream_death(_upstream_failure_tail(resp))
635 if classified is not None:
636 kind = classified
637 raise ProviderError(
638 _upstream_death_message(kind, resp.status_code, self.model_ref, self.engine_log),
639 provider=_PROVIDER_NAME,
640 kind=kind,
641 )
642 detail = f": {body[:600]}" if body else ""
643 raise ProviderError(
644 f"llama-server returned HTTP {resp.status_code}{detail}",
645 provider=_PROVIDER_NAME,
646 kind=kind,
647 )
649 def mark_unhealthy(self) -> None:
650 """Record a connection-level failure so the router skips this replica."""
651 with self._in_flight_lock:
652 self._healthy = False
653 self._unhealthy_since = time.monotonic()
655 def mark_healthy(self) -> None:
656 """Restore the replica to the routing pool after a successful call."""
657 with self._in_flight_lock:
658 self._healthy = True
660 def reserve(self) -> None:
661 """Mark a routed request assigned to this replica, at selection time.
663 The router balances on ``in_flight`` but the per-request tracking only
664 bumps it once the HTTP call starts. Under a bulk ingest many threads pick
665 a replica at the same instant, all see the momentarily-idlest one at the
666 same low count, and pile onto it (a thundering herd that leaves the other
667 cards idle). Reserving at selection makes the assignment visible to the
668 next picker so requests spread across replicas. Paired with :meth:`release`.
669 """
670 with self._in_flight_lock:
671 self.in_flight += 1
673 def release(self) -> None:
674 """Release a reservation taken by :meth:`reserve`."""
675 self._leave()
677 def notify_when_idle(self, callback: Callable[[], None]) -> None:
678 """Call *callback* each time a request ends and leaves ``in_flight`` at zero.
680 The callback runs on the thread that ended the request, after the count
681 lock is released. That thread can hold any lock of its caller, so the
682 callback must not block on one.
683 """
684 with self._in_flight_lock:
685 self._on_idle = callback
687 def _leave(self) -> None:
688 """End one request; call the idle callback when it was the last one."""
689 with self._in_flight_lock:
690 self.in_flight -= 1
691 callback = self._on_idle if self.in_flight == 0 else None
692 if callback is not None:
693 callback()
695 def health(self) -> bool:
696 """True iff ``GET /health`` returns 200 (liveness, not readiness)."""
697 try:
698 resp = self._http.get(_HEALTH_PATH, timeout=_HEALTH_TIMEOUT_S)
699 except httpx.HTTPError:
700 return False
701 return resp.status_code == _HTTP_OK
703 @overload
704 def chat(
705 self,
706 messages: Sequence[Mapping[str, Any]],
707 *,
708 options: dict[str, Any] | None = None,
709 stream: Literal[False] = False,
710 timeout: float | None = None,
711 ) -> str: ...
713 @overload
714 def chat(
715 self,
716 messages: Sequence[Mapping[str, Any]],
717 *,
718 options: dict[str, Any] | None = None,
719 stream: Literal[True],
720 timeout: float | None = None,
721 ) -> Iterator[str]: ...
723 @overload
724 def chat(
725 self,
726 messages: Sequence[Mapping[str, Any]],
727 *,
728 options: dict[str, Any] | None = None,
729 stream: bool,
730 timeout: float | None = None,
731 ) -> str | Iterator[str]: ...
733 def chat(
734 self,
735 messages: Sequence[Mapping[str, Any]],
736 *,
737 options: dict[str, Any] | None = None,
738 stream: bool = False,
739 timeout: float | None = None,
740 ) -> str | Iterator[str]:
741 """Chat completion. Returns the full text, or a token iterator if streaming.
743 ``messages`` accepts both plain ``{role, content: str}`` and multipart
744 ``content`` lists (vision image parts), so the vision path reuses this.
745 ``timeout`` overrides the client default for either path, so a
746 caller-enforced deadline (vision OCR) ends the request itself.
747 """
748 payload: dict[str, Any] = {"model": self._model, "messages": messages, **(options or {})}
749 request_timeout = timeout if timeout is not None else httpx.USE_CLIENT_DEFAULT
750 if stream:
751 return self._chat_stream(payload, request_timeout)
752 with self._track():
753 resp = self._http.post(
754 _CHAT_PATH, json={**payload, "stream": False}, timeout=request_timeout
755 )
756 self._raise_for_status(resp)
757 # content is null for a refusal / content-filter stop / empty completion;
758 # coerce to "" (like chat_result/chat_tools) so callers never see "None".
759 return _inline_message_reasoning(
760 resp.json()["choices"][0]["message"], enabled=self._inline_reasoning
761 )
763 def _chat_stream(
764 self, payload: dict[str, Any], timeout: Any = httpx.USE_CLIENT_DEFAULT
765 ) -> Iterator[str]:
766 inliner = _ThinkInliner(enabled=self._inline_reasoning)
767 with (
768 self._track(),
769 self._http.stream(
770 "POST", _CHAT_PATH, json={**payload, "stream": True}, timeout=timeout
771 ) as resp,
772 ):
773 self._raise_for_status(resp)
774 self._abortable(resp)
775 for line in resp.iter_lines():
776 delta = inliner.feed(*_parse_sse_deltas(line))
777 if delta:
778 yield delta
779 tail = inliner.finish()
780 if tail:
781 yield tail
783 def chat_abortable(
784 self,
785 messages: Sequence[Mapping[str, Any]],
786 *,
787 options: dict[str, Any] | None = None,
788 deadline_s: float | None = None,
789 cancel: CancelSignal | None = None,
790 ) -> str:
791 """Stream a chat completion and return its text, aborted by a deadline or *cancel*.
793 The request runs as an asyncio task on a private loop, so cancelling it
794 closes the connection whether or not response headers have arrived, and
795 llama-server stops generating for a closed stream. A passed *deadline_s*
796 (total seconds, ``None`` for no limit) raises :class:`ChatDeadlineError`;
797 a set *cancel* raises ``TaskCancelledError``.
798 """
799 payload = {"model": self._model, "messages": messages, **(options or {}), "stream": True}
800 with self._track():
801 return asyncio.run(self._chat_until_stopped(payload, deadline_s, cancel))
803 async def _chat_until_stopped(
804 self, payload: dict[str, Any], deadline_s: float | None, cancel: CancelSignal | None
805 ) -> str:
806 async with httpx.AsyncClient(
807 base_url=self._base,
808 timeout=self._timeout,
809 verify=_LOOPBACK_SSL_CONTEXT,
810 transport=self._async_transport,
811 ) as http:
812 chat = asyncio.ensure_future(self._stream_text(http, payload))
813 try:
814 async with asyncio.timeout(deadline_s):
815 while not chat.done():
816 if cancel is not None and cancel.is_set():
817 raise TaskCancelledError
818 await asyncio.wait({chat}, timeout=_CANCEL_POLL_S)
819 except TimeoutError:
820 raise ChatDeadlineError(
821 f"llama-server chat exceeded its {deadline_s:g}s deadline.",
822 provider=_PROVIDER_NAME,
823 ) from None
824 finally:
825 chat.cancel()
826 await asyncio.gather(chat, return_exceptions=True)
827 return chat.result()
829 async def _stream_text(self, http: httpx.AsyncClient, payload: dict[str, Any]) -> str:
830 inliner = _ThinkInliner(enabled=self._inline_reasoning)
831 parts: list[str] = []
832 async with http.stream("POST", _CHAT_PATH, json=payload) as resp:
833 if not resp.is_success:
834 await resp.aread()
835 self._raise_for_status(resp)
836 async for line in resp.aiter_lines():
837 parts.append(inliner.feed(*_parse_sse_deltas(line)))
838 parts.append(inliner.finish())
839 return "".join(parts)
841 def chat_tools(
842 self,
843 messages: Sequence[Mapping[str, Any]],
844 *,
845 tools: list[dict[str, Any]],
846 tool_choice: str | dict[str, Any] | None = None,
847 options: dict[str, Any] | None = None,
848 ) -> ChatToolResult:
849 """Non-streaming chat with function tools; returns content + any tool calls.
851 The server is launched with ``--jinja`` so it parses the model's native
852 tool-call syntax into structured ``message.tool_calls``. When a model
853 instead emits a bare-JSON call as content (a native miss), recover it.
854 """
855 payload: dict[str, Any] = {
856 "model": self._model,
857 "messages": self._prepare_chat_messages(messages),
858 "tools": tools,
859 "stream": False,
860 **(options or {}),
861 }
862 if tool_choice is not None:
863 payload["tool_choice"] = tool_choice
864 with self._track():
865 resp = self._http.post(_CHAT_PATH, json=payload)
866 self._raise_for_status(resp)
867 message = resp.json()["choices"][0]["message"]
868 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
869 native = _parse_native_tool_calls(message.get("tool_calls"))
870 if native:
871 return ChatToolResult(content=content, tool_calls=native)
872 return _recover_bare_json_tool_calls(content)
874 def chat_result(
875 self,
876 messages: Sequence[Mapping[str, Any]],
877 *,
878 tools: list[dict[str, Any]] | None = None,
879 tool_choice: str | dict[str, Any] | None = None,
880 options: dict[str, Any] | None = None,
881 ) -> ChatResult:
882 """Non-streaming chat returning text, tool calls, and a finish reason.
884 The server is launched with ``--jinja`` so it parses the model's native
885 tool-call syntax into structured ``message.tool_calls``. When a request
886 that offered tools gets a bare-JSON call as content (a native miss),
887 recover it and report ``tool_calls`` as the finish reason; without tools
888 a JSON answer is the answer. Messages are reshaped to
889 strict alternation up front when this server's template needs it (see
890 :meth:`_prepare_chat_messages`).
891 """
892 payload = self._chat_payload(
893 self._prepare_chat_messages(messages), tools, tool_choice, options, stream=False
894 )
895 with self._track():
896 resp = self._http.post(_CHAT_PATH, json=payload)
897 self._raise_for_status(resp)
898 body = dict(resp.json())
899 choice = body["choices"][0]
900 usage = _usage_from_body(body) or TokenUsage()
901 message = choice["message"]
902 content = _inline_message_reasoning(message, enabled=self._inline_reasoning)
903 finish_reason = _coerce_finish_reason(choice.get("finish_reason"))
904 native = _parse_native_tool_calls(message.get("tool_calls"))
905 if native:
906 return ChatResult(
907 text=content,
908 tool_calls=tuple(native),
909 finish_reason=finish_reason,
910 usage=usage,
911 )
912 recovered = _recover_bare_json_tool_calls(content) if tools else None
913 if recovered is not None and recovered.tool_calls:
914 return ChatResult(
915 text=recovered.content,
916 tool_calls=tuple(recovered.tool_calls),
917 finish_reason=FinishReason.TOOL_CALLS,
918 usage=usage,
919 )
920 return ChatResult(text=content, tool_calls=(), finish_reason=finish_reason, usage=usage)
922 def chat_stream_items(
923 self,
924 messages: Sequence[Mapping[str, Any]],
925 *,
926 tools: list[dict[str, Any]] | None = None,
927 tool_choice: str | dict[str, Any] | None = None,
928 options: dict[str, Any] | None = None,
929 ) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
930 """Stream text tokens and tool-call deltas from the server's OpenAI SSE.
932 Each SSE chunk's ``choices[0].delta`` carries a ``content`` token and/or
933 a ``tool_calls`` array; both are surfaced as :data:`ChatStreamItem`
934 frames (text strings and :class:`ToolCallDelta`). The dispatch's stream
935 translator accumulates the deltas by ``index``. Messages are reshaped to
936 strict alternation up front when this server's template needs it (see
937 :meth:`_prepare_chat_messages`), so the open never fails on a template
938 that rejects the raw tool exchange.
940 Not a generator: the up-front probe runs when this is called, not deferred
941 to the first iteration, matching the eager non-stream paths.
943 When the request offered tools, a model that emits a tool call as
944 bare-JSON text instead of native ``tool_calls`` (a native miss, as on the
945 non-stream paths) is recovered by wrapping the raw frames; see
946 :func:`_recover_bare_json_stream`. Without tools the frames pass through
947 as they arrive.
948 """
949 prepared = self._prepare_chat_messages(messages)
950 stream = self._open_chat_stream(prepared, tools, tool_choice, options)
951 return _recover_bare_json_stream(stream) if tools else stream
953 def _open_chat_stream(
954 self,
955 messages: Sequence[Mapping[str, Any]],
956 tools: list[dict[str, Any]] | None,
957 tool_choice: str | dict[str, Any] | None,
958 options: dict[str, Any] | None,
959 ) -> Generator[str | ToolCallDelta | TokenUsage | StreamFinish]:
960 """Open one SSE chat stream and yield its frames; raises before the first frame."""
961 payload = self._chat_payload(messages, tools, tool_choice, options, stream=True)
962 inliner = _ThinkInliner(enabled=self._inline_reasoning)
963 on_prefill = self._on_prefill
964 prefilling = False
965 started = False
966 ended = False
967 try:
968 with (
969 self._track(),
970 self._http.stream("POST", _CHAT_PATH, json=payload) as resp,
971 ):
972 self._raise_for_status(resp)
973 self._abortable(resp)
974 for line in resp.iter_lines():
975 if on_prefill is not None:
976 progress = _prefill_progress(line)
977 if progress is not None:
978 prefilling = True
979 on_prefill(progress)
980 for item in _parse_sse_stream_items(line, inliner):
981 if prefilling and on_prefill is not None:
982 # The first real frame proves prefill is over.
983 prefilling = False
984 on_prefill(None)
985 started = True
986 ended |= isinstance(item, StreamFinish)
987 yield item
988 ended |= _is_done_line(line)
989 tail = inliner.finish()
990 if tail:
991 yield tail
992 _raise_if_cut(started=started, ended=ended)
993 except httpx.TransportError as exc:
994 if not started:
995 raise
996 raise StreamInterruptedError from exc
997 finally:
998 # A stream that dies or is closed mid-prefill must not leave a
999 # stale in-progress reading on the status surface.
1000 if prefilling and on_prefill is not None:
1001 on_prefill(None)
1003 def _chat_payload(
1004 self,
1005 messages: Sequence[Mapping[str, Any]],
1006 tools: Sequence[Mapping[str, Any]] | None,
1007 tool_choice: str | dict[str, Any] | None,
1008 options: dict[str, Any] | None,
1009 *,
1010 stream: bool,
1011 ) -> dict[str, Any]:
1012 """Build the chat-completions request body shared by the stream and non-stream paths."""
1013 payload: dict[str, Any] = {"model": self._model, "messages": messages, "stream": stream}
1014 if stream:
1015 # include_usage makes llama-server emit a final SSE chunk carrying the
1016 # token usage (with an empty choices list) just before [DONE].
1017 payload["stream_options"] = {"include_usage": True}
1018 # return_progress makes llama-server stream prompt_progress frames
1019 # during prefill, so a long first turn is observable server-side.
1020 payload["return_progress"] = True
1021 if tools is not None:
1022 payload["tools"] = tools
1023 if tool_choice is not None:
1024 payload["tool_choice"] = tool_choice
1025 payload.update(options or {})
1026 return payload
1028 def _prepare_chat_messages(
1029 self, messages: Sequence[Mapping[str, Any]]
1030 ) -> Sequence[Mapping[str, Any]]:
1031 """Reshape *messages* to strict alternation when this server's template needs it.
1033 The need is detected once per server by :meth:`_ensure_alternation_probed`
1034 (a binary accept/reject of a representative tool exchange against the live
1035 template), then cached, so real requests are normalized up front rather
1036 than failing and retrying.
1037 """
1038 self._ensure_alternation_probed()
1039 if self._needs_alternation:
1040 return to_alternating([dict(m) for m in messages])
1041 return messages
1043 def _ensure_alternation_probed(self) -> None:
1044 """Probe the live template once to learn whether it needs alternation.
1046 Caches only a conclusive verdict: a transient unreachable server leaves
1047 the flag unset so the next request re-probes rather than locking in a
1048 wrong answer.
1049 """
1050 if self._needs_alternation is not None:
1051 return
1052 with self._alternation_lock:
1053 if self._needs_alternation is not None:
1054 return
1055 verdict = self._probe_alternation()
1056 if verdict is not None:
1057 self._needs_alternation = verdict
1059 def _probe_alternation(self) -> bool | None:
1060 """Whether the template needs normalization: ``None`` when undetermined.
1062 Renders the probe exchange as sent; if the template accepts it, no
1063 normalization is needed. If it rejects it, normalization is needed only
1064 when the reshaped exchange is accepted. A transient failure on either
1065 render is inconclusive (``None``) so no verdict is cached; a genuine
1066 rejection of both forms is a conclusive ``False`` (the template fault is
1067 unrelated to alternation, so reshaping would not help).
1068 """
1069 raw = self._chat_probe(_ALTERNATION_PROBE_MESSAGES)
1070 if raw is None:
1071 return None # transient; stay undetermined so the next request re-probes
1072 if raw:
1073 return False # the template renders the raw OpenAI exchange as sent
1074 reshaped = self._chat_probe(to_alternating([dict(m) for m in _ALTERNATION_PROBE_MESSAGES]))
1075 if reshaped is None:
1076 return None # transient on the reshape probe; stay undetermined
1077 return reshaped
1079 def _chat_probe(self, messages: Sequence[Mapping[str, Any]]) -> bool | None:
1080 """Post the probe exchange: ``True`` rendered, ``False`` rejected, ``None`` undetermined.
1082 A connection failure or a server-busy (HTTP 429) response is transient and
1083 unrelated to the template, so it is undetermined: only a clean render or a
1084 genuine rejection is a verdict the caller may cache.
1085 """
1086 payload = self._chat_payload(
1087 messages, _ALTERNATION_PROBE_TOOLS, None, _ALTERNATION_PROBE_OPTIONS, stream=False
1088 )
1089 try:
1090 with self._track():
1091 resp = self._http.post(
1092 _CHAT_PATH, json=payload, timeout=_ALTERNATION_PROBE_TIMEOUT_S
1093 )
1094 self._raise_for_status(resp)
1095 except (ProviderError, httpx.TransportError) as exc:
1096 return None if _is_transient_probe_failure(exc) else False
1097 return True
1099 def embed(self, texts: list[str]) -> list[Vector]:
1100 """Embed a batch via ``/v1/embeddings``."""
1101 if not texts:
1102 # Match the in-process embedder; the server rejects an empty input.
1103 return []
1104 vectors: list[Vector] = []
1105 for sub_batch in self._truncate_and_subbatch(texts, estimate=True):
1106 data = self._embed_subbatch(sub_batch)
1107 vectors.extend(_embedding_vector(item) for item in data)
1108 return vectors
1110 def _embed_subbatch(self, sub_batch: list[str]) -> list[dict[str, Any]]:
1111 """Embed one estimate-budgeted sub-batch, re-truncating exactly on overflow.
1113 ``_estimate_tokens`` is char-based and can under-count token-dense inputs
1114 (XML, code), so an estimate-trusted input may still exceed the server's
1115 context. On that error -- and only that -- redo the batch with exact
1116 server-side tokenization, which truncates the oversize input to the cap.
1117 """
1118 try:
1119 return self._embeddings_call(sub_batch)
1120 except ProviderError as exc:
1121 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW:
1122 raise
1123 data: list[dict[str, Any]] = []
1124 for exact in self._truncate_and_subbatch(sub_batch, estimate=False):
1125 data.extend(self._embeddings_call(exact))
1126 return data
1128 def rerank(self, query: str, candidates: list[str]) -> list[float]:
1129 """Relevance scores via rank-pooling embeddings.
1131 The server runs with ``--pooling rank``; we send ``query</s></s>candidate``
1132 pairs to ``/v1/embeddings`` and read each item's first embedding value as the
1133 score, so the ``/v1/rerank`` template-dependency (and its zero-output failure
1134 modes) is moot.
1136 All pairs go out in one request per ``_EMBED_N_SEQ_MAX`` sequences: the
1137 server queues one task per input, so splitting a query-time pool into
1138 per-pair requests only adds HTTP round trips.
1139 """
1140 if not candidates:
1141 return []
1142 if self._rerank_mode is RerankMode.LLM:
1143 return self._rerank_llm(query, candidates)
1144 pairs = [
1145 self._fit_estimated(f"{query}{_RERANK_PAIR_SEPARATOR}{candidate}")
1146 for candidate in candidates
1147 ]
1148 scores: list[float] = []
1149 for start in range(0, len(pairs), _EMBED_N_SEQ_MAX):
1150 data = self._rerank_batch(pairs[start : start + _EMBED_N_SEQ_MAX])
1151 scores.extend(_rerank_score(item) for item in data)
1152 return scores
1154 def _fit_estimated(self, text: str) -> str:
1155 """Truncate *text* to the token cap only when its char estimate exceeds it."""
1156 if self._token_cap is None:
1157 return text
1158 return self._fit_input(text, self._token_cap, estimate=True)[0]
1160 def _rerank_batch(self, batch: list[str]) -> list[dict[str, Any]]:
1161 """Score one rerank batch, redoing it with exact truncation on overflow.
1163 The char estimate under-counts token-dense pairs, so an over-cap pair
1164 can slip through untruncated; the server rejects it as CONTEXT_OVERFLOW
1165 and the batch is redone with exact server-side tokenization.
1166 """
1167 try:
1168 return self._embeddings_call(batch)
1169 except ProviderError as exc:
1170 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW or self._token_cap is None:
1171 raise
1172 cap = self._token_cap
1173 exact = [self._fit_input(text, cap, estimate=False)[0] for text in batch]
1174 return self._embeddings_call(exact)
1176 def _rerank_llm(self, query: str, candidates: list[str]) -> list[float]:
1177 """Score each candidate by an LLM's yes/no first-token logprob.
1179 Raises ``ProviderError`` when no candidate yields a verdict.
1180 """
1181 template = cfg.reranker_prompt or _LLM_RERANK_PROMPT
1182 workers = min(LLM_RERANK_CONCURRENCY, len(candidates))
1183 with ThreadPoolExecutor(max_workers=workers) as pool:
1184 scores = list(pool.map(lambda c: self._llm_rerank_one(template, query, c), candidates))
1185 if all(score is None for score in scores):
1186 raise ProviderError(_LLM_RERANK_NO_VERDICT_ERROR, provider=_PROVIDER_NAME)
1187 return [0.0 if score is None else score for score in scores]
1189 def _llm_rerank_one(self, template: str, query: str, candidate: str) -> float | None:
1190 """One chat request scoring a single candidate's relevance to the query."""
1191 content = template.format(query=query, document=candidate)
1192 payload = {
1193 "model": self._model,
1194 "messages": [{"role": "user", "content": content}],
1195 "max_tokens": 1,
1196 "temperature": 0,
1197 "logprobs": True,
1198 "top_logprobs": _LLM_RERANK_TOP_LOGPROBS,
1199 "stream": False,
1200 # Scoring reads the first generated token; a thinking template would
1201 # spend it opening a <think> block instead of answering.
1202 "chat_template_kwargs": {"enable_thinking": False},
1203 }
1205 def _call() -> dict[str, Any]:
1206 with self._track():
1207 resp = self._http.post(_CHAT_PATH, json=payload)
1208 self._raise_for_status(resp)
1209 return dict(resp.json())
1211 return _llm_rerank_score(_first_token_top_logprobs(retry_on_gateway_error(_call)))
1213 def _embeddings_call(self, inputs: list[str]) -> list[dict[str, Any]]:
1214 """POST one already-budgeted sub-batch to ``/v1/embeddings``; return its data."""
1216 def _call() -> list[dict[str, Any]]:
1217 with self._track():
1218 resp = self._http.post(
1219 _EMBED_PATH,
1220 json={
1221 "model": self._model,
1222 "input": inputs,
1223 "embd_normalize": _EMBD_NORMALIZE_NONE,
1224 "encoding_format": _EMBED_ENCODING_FORMAT,
1225 },
1226 )
1227 self._raise_for_status(resp)
1228 data = resp.json()["data"]
1229 if len(data) != len(inputs):
1230 raise ProviderError(
1231 f"Embedder returned {len(data)} vectors for {len(inputs)} inputs",
1232 provider=_PROVIDER_NAME,
1233 )
1234 return list(data)
1236 return retry_on_gateway_error(_call)
1238 def _truncate_and_subbatch(self, texts: list[str], *, estimate: bool) -> list[list[str]]:
1239 """Token-truncate over-cap inputs, then pack into server-sized sub-batches.
1241 An input longer than ``token_cap`` (the server's per-slot context /
1242 n_batch) is truncated to it via the server's tokenizer, since the server
1243 cannot split a pooled embedding sequence. Inputs are then grouped so each
1244 request stays within both the token budget and ``_EMBED_N_SEQ_MAX``
1245 sequences, bounding per-request size and the gateway-retry window. No cap
1246 (chat/vision) sends a single batch untouched.
1248 When ``estimate`` is set the per-input token count comes from
1249 :func:`_estimate_tokens`, and ``/tokenize`` is consulted only for the
1250 rare input whose estimate exceeds the cap -- eliminating a round-trip per
1251 chunk during bulk ingest. ``estimate=False`` (the overflow redo) tokenizes
1252 every input exactly.
1253 """
1254 if self._token_cap is None:
1255 return [texts]
1256 cap = self._token_cap
1257 batches: list[list[str]] = []
1258 current: list[str] = []
1259 current_tokens = 0
1260 for text in texts:
1261 item, item_tokens = self._fit_input(text, cap, estimate=estimate)
1262 if current and (current_tokens + item_tokens > cap or len(current) >= _EMBED_N_SEQ_MAX):
1263 batches.append(current)
1264 current = []
1265 current_tokens = 0
1266 current.append(item)
1267 current_tokens += item_tokens
1268 if current:
1269 batches.append(current)
1270 return batches
1272 def _fit_input(self, text: str, cap: int, *, estimate: bool) -> tuple[str, int]:
1273 """Return ``(input, token_count)`` for one sequence, truncating if over cap.
1275 Estimation short-circuits the common case: an estimate within the cap is
1276 trusted (no ``/tokenize``); only an over-cap estimate is confirmed against
1277 the server tokenizer and truncated if it really exceeds the cap.
1278 """
1279 if estimate:
1280 est = _estimate_tokens(text)
1281 if est <= cap:
1282 return text, est
1283 tokens = self._tokenize(text)
1284 if len(tokens) > cap:
1285 log.warning("Truncating oversize embed input: %d tokens > cap %d", len(tokens), cap)
1286 return self._detokenize(tokens[:cap]), cap
1287 return text, max(1, len(tokens))
1289 def _native_route(self, suffix: str) -> str:
1290 """Path for a native (non-OpenAI) llama-server route through llama-swap.
1292 llama-swap proxies these only under ``/upstream/<model>/...``; the model
1293 is carried in the path, not the body (unlike the ``/v1`` OpenAI routes).
1294 """
1295 return f"{_UPSTREAM_PREFIX}/{self._model}{suffix}"
1297 def _tokenize(self, text: str, *, parse_special: bool = _TOKENIZE_PARSE_SPECIAL) -> list[int]:
1298 """Token ids for *text*; a gateway error is retried like an embedding."""
1300 def _call() -> list[int]:
1301 resp = self._http.post(
1302 self._native_route(_TOKENIZE_PATH),
1303 json={
1304 "content": text,
1305 "add_special": _TOKENIZE_ADD_SPECIAL,
1306 "parse_special": parse_special,
1307 },
1308 )
1309 self._raise_for_status(resp)
1310 return list(resp.json()["tokens"])
1312 return retry_on_gateway_error(_call)
1314 def count_tokens(self, text: str) -> int:
1315 """Number of tokens *text* encodes to under the server's tokenizer."""
1316 return len(self._tokenize(text))
1318 def count_chat_prompt_tokens(
1319 self,
1320 messages: Sequence[Mapping[str, Any]],
1321 *,
1322 tools: Sequence[Mapping[str, Any]] | None = None,
1323 tool_choice: str | dict[str, Any] | None = None,
1324 options: dict[str, Any] | None = None,
1325 ) -> int:
1326 """Tokens the server prefills for this prompt, template applied.
1328 The template's role markers and its tool-call preamble are part of the
1329 prompt the model reads, so the rendered text is what gets tokenized.
1330 Special-token strings in the rendered text are parsed, matching how the
1331 server tokenizes its own prompt.
1332 """
1333 rendered = self._render_chat_prompt(messages, tools, tool_choice, options)
1334 return len(self._tokenize(rendered, parse_special=True))
1336 def _render_chat_prompt(
1337 self,
1338 messages: Sequence[Mapping[str, Any]],
1339 tools: Sequence[Mapping[str, Any]] | None,
1340 tool_choice: str | dict[str, Any] | None,
1341 options: dict[str, Any] | None,
1342 ) -> str:
1343 """The prompt text this server's template renders for a chat body.
1345 Sends the body the chat paths build, messages reshaped to strict
1346 alternation when this server's template needs it (see
1347 :meth:`_prepare_chat_messages`), so the render cannot diverge from the
1348 prompt a chat call prefills. Raises ``NotImplementedError`` when the
1349 server has no such route, which is the one case a caller can answer with
1350 an estimate instead.
1351 """
1352 payload = self._chat_payload(
1353 self._prepare_chat_messages(messages), tools, tool_choice, options, stream=False
1354 )
1356 def _call() -> str:
1357 resp = self._http.post(self._native_route(_APPLY_TEMPLATE_PATH), json=payload)
1358 if _route_is_absent(resp):
1359 raise NotImplementedError(
1360 f"This inference engine has no {_APPLY_TEMPLATE_PATH} route."
1361 )
1362 self._raise_for_status(resp)
1363 return str(resp.json()["prompt"])
1365 return retry_on_gateway_error(_call)
1367 def _detokenize(self, tokens: list[int]) -> str:
1368 resp = self._http.post(self._native_route(_DETOKENIZE_PATH), json={"tokens": tokens})
1369 self._raise_for_status(resp)
1370 return str(resp.json()["content"])
1372 def _abortable(self, resp: httpx.Response) -> None:
1373 """Expose *resp* to ``abort_streams`` until the response closes."""
1374 # A response read whole on arrival has no connection left and no close
1375 # to come; httpx types the body sync-or-async, and a sync client's is sync.
1376 if resp.is_closed or not isinstance(resp.stream, httpx.SyncByteStream):
1377 return
1378 resp.stream = _AbortableStream(resp.stream, resp, self)
1379 with self._in_flight_lock:
1380 self._active_streams.add(resp)
1382 def abort_streams(self) -> None:
1383 """Sever every in-flight SSE response on this replica.
1385 A socket shutdown sends the server a FIN and makes a reader blocked in
1386 ``iter_lines`` read end-of-stream at once, so it raises a stream error
1387 and closes its own response; llama-server stops generating when the
1388 connection drops. Closing the socket from another thread does neither.
1389 A response with no socket (a mock transport), or whose shutdown fails
1390 because the peer is gone, is closed instead.
1392 The shutdowns run under the registry lock, and a response leaves the
1393 registry under that lock before its close returns the connection to
1394 the pool, so a shutdown never reaches a pooled socket. The fallback
1395 closes run after the lock is released, because a close takes it.
1396 """
1397 unsevered: list[httpx.Response] = []
1398 with self._in_flight_lock:
1399 for resp in self._active_streams:
1400 with contextlib.suppress(Exception):
1401 if not _shut_down_socket(resp):
1402 unsevered.append(resp)
1403 for resp in unsevered:
1404 with contextlib.suppress(Exception):
1405 resp.close()
1407 def close(self) -> None:
1408 """Close the underlying client if this instance created it."""
1409 if self._owns_http:
1410 self._http.close()
1412 def _track(self) -> _InFlight:
1413 return _InFlight(self)
1416class _InFlight:
1417 """Context manager that atomically bumps the owner's in-flight counter.
1419 ``+= 1`` is a read-modify-write, so concurrent chat/embed calls would corrupt
1420 the counter the router balances on; the client's lock makes it atomic.
1421 """
1423 def __init__(self, client: LlamaServerClient) -> None:
1424 self._client = client
1426 def __enter__(self) -> None:
1427 with self._client._in_flight_lock:
1428 self._client.in_flight += 1
1430 def __exit__(self, *_exc: object) -> None:
1431 self._client._leave()
1434class _AbortableStream(httpx.SyncByteStream):
1435 """A response body that leaves the abort registry before it releases its connection.
1437 httpx closes a response from inside ``iter_raw`` at the last chunk, and the
1438 close of httpcore's pool stream hands the connection to the next request.
1439 Deregistering in this close, ahead of the inner close, keeps a late
1440 ``abort_streams`` off a connection another request may now use.
1441 """
1443 def __init__(
1444 self, inner: httpx.SyncByteStream, resp: httpx.Response, client: LlamaServerClient
1445 ) -> None:
1446 self._inner = inner
1447 self._resp = resp
1448 self._client = client
1449 self._closed = False
1451 def __iter__(self) -> Iterator[bytes]:
1452 # httpcore's iterators release the connection when they are finalized.
1453 # The local keeps the inner iterator alive until this close has run:
1454 # an exception at the yield pops a ``for`` loop's iterator before the
1455 # except body, and ``yield from`` closes it first on GeneratorExit.
1456 chunks = iter(self._inner)
1457 try:
1458 while (chunk := next(chunks, None)) is not None:
1459 yield chunk
1460 except BaseException:
1461 self.close()
1462 raise
1464 def close(self) -> None:
1465 with self._client._in_flight_lock:
1466 if self._closed:
1467 return
1468 self._closed = True
1469 self._client._active_streams.discard(self._resp)
1470 self._inner.close()
1473def _shut_down_socket(resp: httpx.Response) -> bool:
1474 """Shut down *resp*'s socket; False when it has no socket or the shutdown fails."""
1475 network_stream = resp.extensions.get("network_stream")
1476 sock = network_stream.get_extra_info("socket") if network_stream else None
1477 if sock is None:
1478 return False
1479 try:
1480 sock.shutdown(socket.SHUT_RDWR)
1481 except OSError:
1482 return False
1483 return True
1486def _embedding_vector(item: dict[str, Any]) -> npt.NDArray[np.float32]:
1487 """Decode one ``/v1/embeddings`` item's vector from its base64 float buffer."""
1488 embedding = item.get("embedding")
1489 # Untyped server JSON: a non-string means the encoding format was not honored.
1490 if not isinstance(embedding, str):
1491 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME)
1492 try:
1493 return np.frombuffer(base64.b64decode(embedding), dtype=_EMBED_VECTOR_DTYPE)
1494 except ValueError as exc:
1495 raise ProviderError(_UNREADABLE_EMBEDDING_ERROR, provider=_PROVIDER_NAME) from exc
1498def _rerank_score(item: dict[str, Any]) -> float:
1499 """Pull one relevance score from a rank-pooling ``/v1/embeddings`` item."""
1500 vector = _embedding_vector(item)
1501 if not vector.size:
1502 raise ProviderError(_NO_RERANK_SCORE_ERROR, provider=_PROVIDER_NAME)
1503 return float(vector[_RANK_SCORE_INDEX])
1506def _first_token_top_logprobs(response: dict[str, Any]) -> list[dict[str, Any]]:
1507 """The first generated token's top_logprobs list from a chat completion, or []."""
1508 choices = response.get("choices") or []
1509 if not choices:
1510 return []
1511 content = (choices[0].get("logprobs") or {}).get("content") or []
1512 if not content:
1513 return []
1514 return list(content[0].get("top_logprobs") or [])
1517def _llm_rerank_score(top_logprobs: list[dict[str, Any]]) -> float | None:
1518 """Softmax of the yes vs no logprobs in a token's top_logprobs (case/space-insensitive).
1520 ``None`` when neither verdict appears, distinct from the 0.0 of a confident "no".
1521 """
1522 yes_lp: float | None = None
1523 no_lp: float | None = None
1524 for entry in top_logprobs:
1525 token = str(entry.get("token", "")).strip().lower()
1526 logprob = float(entry.get("logprob", 0.0))
1527 if token == _YES_LABEL and (yes_lp is None or logprob > yes_lp):
1528 yes_lp = logprob
1529 elif token == _NO_LABEL and (no_lp is None or logprob > no_lp):
1530 no_lp = logprob
1531 if yes_lp is None:
1532 return None if no_lp is None else 0.0
1533 if no_lp is None:
1534 return math.exp(yes_lp)
1535 yes_e, no_e = math.exp(yes_lp), math.exp(no_lp)
1536 return yes_e / (yes_e + no_e)
1539def _parse_sse_deltas(line: str) -> tuple[str, str]:
1540 """Extract the (reasoning, content) deltas from one OpenAI SSE line."""
1541 if not line.startswith(_DATA_PREFIX):
1542 return "", ""
1543 body = line[len(_DATA_PREFIX) :].strip()
1544 if not body or body == _DONE_SENTINEL:
1545 return "", ""
1546 try:
1547 obj = json.loads(body)
1548 except json.JSONDecodeError:
1549 return "", ""
1550 choices = obj.get("choices") or []
1551 if not choices:
1552 return "", ""
1553 delta = choices[0].get("delta") or {}
1554 return str(delta.get("reasoning_content") or ""), str(delta.get("content") or "")
1557class _ThinkInliner:
1558 """Re-inlines server-extracted reasoning deltas as inline ``<think>`` text.
1560 The server parses each model's reasoning format natively (``--reasoning-format``)
1561 and streams it as ``reasoning_content``; lilbee's pipeline speaks inline
1562 ``<think>`` text, so the chat boundary opens the tag on the first reasoning
1563 delta and closes it when the answer starts (or at end of stream). Disabled
1564 (the non-chat roles), reasoning is dropped and content passes through, matching
1565 the server-extracted default those roles already ran with.
1566 """
1568 def __init__(self, *, enabled: bool) -> None:
1569 self._enabled = enabled
1570 self._in_think = False
1572 def feed(self, reasoning: str, content: str) -> str:
1573 if not self._enabled:
1574 return content
1575 parts: list[str] = []
1576 if reasoning:
1577 if not self._in_think:
1578 self._in_think = True
1579 parts.append(THINK_OPEN_TAG)
1580 parts.append(reasoning)
1581 if content:
1582 if self._in_think:
1583 self._in_think = False
1584 parts.append(THINK_CLOSE_TAG)
1585 parts.append(content)
1586 return "".join(parts)
1588 def finish(self) -> str:
1589 """Close an unterminated think block at end of stream."""
1590 if self._in_think:
1591 self._in_think = False
1592 return THINK_CLOSE_TAG
1593 return ""
1596def _inline_message_reasoning(message: Mapping[str, Any], *, enabled: bool) -> str:
1597 """A non-streaming message's text with any extracted reasoning re-inlined."""
1598 content = str(message.get("content") or "")
1599 reasoning = str(message.get("reasoning_content") or "") if enabled else ""
1600 if reasoning:
1601 return f"{THINK_OPEN_TAG}{reasoning}{THINK_CLOSE_TAG}{content}"
1602 return content
1605def _cached_prompt_tokens(usage: Mapping[str, Any]) -> int:
1606 """Read ``prompt_tokens_details.cached_tokens``, or 0 when the engine omits it."""
1607 details = usage.get("prompt_tokens_details")
1608 if not isinstance(details, Mapping):
1609 return 0
1610 cached = details.get("cached_tokens")
1611 return cached if isinstance(cached, int) else 0
1614def _usage_from_body(body: Mapping[str, Any]) -> TokenUsage | None:
1615 """Read the ``usage`` block of an OpenAI response, or ``None`` if absent.
1617 llama-server reports ``prompt_tokens`` / ``completion_tokens``; a missing or
1618 malformed block yields ``None`` so callers can decide between a zero default
1619 (non-streaming) and skipping the frame (streaming terminator).
1620 """
1621 usage = body.get("usage")
1622 if not isinstance(usage, Mapping):
1623 return None
1624 prompt = usage.get("prompt_tokens")
1625 completion = usage.get("completion_tokens")
1626 return TokenUsage(
1627 prompt_tokens=prompt if isinstance(prompt, int) else 0,
1628 completion_tokens=completion if isinstance(completion, int) else 0,
1629 cached_prompt_tokens=_cached_prompt_tokens(usage),
1630 )
1633def _coerce_finish_reason(raw: Any) -> FinishReason:
1634 """Map a server-supplied finish_reason string to :class:`FinishReason`."""
1635 return FinishReason.coerce(raw)
1638def _tool_call_delta_from_chunk(call: Mapping[str, Any], *, fallback_index: int) -> ToolCallDelta:
1639 """Map one streaming ``delta.tool_calls`` entry to a :class:`ToolCallDelta`.
1641 Mirrors the SDK path: ``id`` / ``name`` arrive on the opener and accumulate
1642 by ``index``; empty strings normalise to ``None`` so the dispatch's stream
1643 translator (which gates on ``is not None``) does not emit spurious openers.
1644 """
1645 raw_index = call.get("index")
1646 index = raw_index if isinstance(raw_index, int) else fallback_index
1647 call_id = call.get("id")
1648 fn = call.get("function")
1649 raw_name = fn.get("name") if isinstance(fn, Mapping) else None
1650 raw_args = fn.get("arguments") if isinstance(fn, Mapping) else None
1651 return ToolCallDelta(
1652 index=index,
1653 id=str(call_id) if call_id else None,
1654 name=str(raw_name) if raw_name else None,
1655 arguments_delta=str(raw_args) if raw_args else None,
1656 )
1659_PROMPT_PROGRESS_KEY = "prompt_progress"
1662def _prefill_progress(line: str) -> tuple[int, int] | None:
1663 """The ``(processed, total)`` of one SSE line's ``prompt_progress``, or None.
1665 llama-server emits the block on streamed chats that opt in via
1666 ``return_progress``; the substring pre-check keeps the extra JSON parse off
1667 every ordinary token line.
1668 """
1669 if _PROMPT_PROGRESS_KEY not in line or not line.startswith(_DATA_PREFIX):
1670 return None
1671 body = line[len(_DATA_PREFIX) :].strip()
1672 try:
1673 obj = json.loads(body)
1674 except json.JSONDecodeError:
1675 return None
1676 progress = obj.get(_PROMPT_PROGRESS_KEY)
1677 if not isinstance(progress, Mapping):
1678 return None
1679 try:
1680 return int(progress["processed"]), int(progress["total"])
1681 except (KeyError, TypeError, ValueError):
1682 return None
1685def _raise_if_cut(*, started: bool, ended: bool) -> None:
1686 """Raise when a started stream's body ended with neither a finish frame nor ``[DONE]``.
1688 The engine died but its proxy closed the body cleanly, so no transport error
1689 was raised. A stream with no frame at all is left to the caller's retry path.
1690 """
1691 if started and not ended:
1692 raise StreamInterruptedError
1695def _is_done_line(line: str) -> bool:
1696 """Whether *line* is the SSE ``data: [DONE]`` sentinel that closes a stream."""
1697 return _sse_data_body(line) == _DONE_SENTINEL
1700def _sse_data_body(line: str) -> str | None:
1701 """The stripped payload of an SSE ``data:`` line, or ``None`` for any other line."""
1702 if not line.startswith(_DATA_PREFIX):
1703 return None
1704 return line[len(_DATA_PREFIX) :].strip()
1707def _parse_sse_stream_items(
1708 line: str, inliner: _ThinkInliner
1709) -> Iterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1710 """Yield text tokens, tool-call deltas, and the finish frame from one SSE line.
1712 A chunk can carry a ``content`` token, a ``reasoning_content`` token (routed
1713 through *inliner*), a ``tool_calls`` delta array, or a mix; each is yielded as
1714 its own :data:`ChatStreamItem` frame. The chunk that closes the turn carries
1715 ``choices[0].finish_reason``, surfaced as a :class:`StreamFinish` so the
1716 dispatch reports ``length`` (and friends), not just the default end-of-turn.
1717 """
1718 body = _sse_data_body(line)
1719 if not body or body == _DONE_SENTINEL:
1720 return
1721 try:
1722 obj = json.loads(body)
1723 except json.JSONDecodeError:
1724 return
1725 choices = obj.get("choices") or []
1726 if not choices:
1727 # The include_usage terminator chunk has an empty choices list and the
1728 # token totals on a top-level ``usage`` block; surface it as the final
1729 # frame so the dispatch can attach real counts to the stream.
1730 usage = _usage_from_body(obj)
1731 if usage is not None:
1732 yield usage
1733 return
1734 delta = choices[0].get("delta") or {}
1735 text = inliner.feed(str(delta.get("reasoning_content") or ""), str(delta.get("content") or ""))
1736 if text:
1737 yield text
1738 raw_calls = delta.get("tool_calls") or []
1739 for i, call in enumerate(raw_calls):
1740 if isinstance(call, Mapping):
1741 yield _tool_call_delta_from_chunk(call, fallback_index=i)
1742 raw_finish = choices[0].get("finish_reason")
1743 if raw_finish is not None:
1744 yield StreamFinish(reason=_coerce_finish_reason(raw_finish))
1747def _arguments_to_str(arguments: Any) -> str:
1748 """Normalize a tool-call ``arguments`` value to a JSON string (OpenAI's shape)."""
1749 if isinstance(arguments, str):
1750 return arguments
1751 if arguments is None:
1752 return "{}"
1753 return json.dumps(arguments)
1756def _parse_native_tool_calls(raw: Any) -> list[ToolCall]:
1757 """Map a response's ``message.tool_calls`` array to :class:`ToolCall` objects.
1759 Reads the OpenAI shape (``{"id", "function": {"name", "arguments"}}``) that
1760 ``--jinja`` produces. Malformed or nameless entries are skipped.
1761 """
1762 if not isinstance(raw, list):
1763 return []
1764 calls: list[ToolCall] = []
1765 for idx, entry in enumerate(raw):
1766 if not isinstance(entry, Mapping):
1767 continue
1768 fn = entry.get("function")
1769 if not isinstance(fn, Mapping):
1770 continue
1771 name = fn.get("name")
1772 if not isinstance(name, str) or not name:
1773 continue
1774 call_id = entry.get("id")
1775 calls.append(
1776 ToolCall(
1777 id=call_id if isinstance(call_id, str) and call_id else f"call_{idx}",
1778 name=name,
1779 arguments=_arguments_to_str(fn.get("arguments")),
1780 )
1781 )
1782 return calls
1785def _bare_call_from_mapping(obj: Mapping[str, Any], *, index: int) -> ToolCall | None:
1786 """Build a ToolCall from a bare ``{"name", "arguments"|"parameters"}`` object."""
1787 name = obj.get("name")
1788 if not isinstance(name, str) or not name:
1789 return None
1790 arguments = obj.get("arguments")
1791 if arguments is None:
1792 arguments = obj.get("parameters")
1793 return ToolCall(id=f"call_{index}", name=name, arguments=_arguments_to_str(arguments))
1796def _recover_bare_json_tool_calls(content: str) -> ChatToolResult:
1797 """Recover a tool call a model emitted as bare-JSON content (a native miss).
1799 Some models ignore the tool-call protocol and print ``{"name": ...,
1800 "arguments": {...}}`` (or a list of them) as the message body. When the whole
1801 content parses as such, treat it as the call(s) and clear the text; otherwise
1802 return the content unchanged with no calls.
1803 """
1804 stripped = content.strip()
1805 if not stripped or stripped[0] not in "{[":
1806 return ChatToolResult(content=content, tool_calls=[])
1807 try:
1808 parsed = json.loads(stripped)
1809 except json.JSONDecodeError:
1810 return ChatToolResult(content=content, tool_calls=[])
1811 entries = parsed if isinstance(parsed, list) else [parsed]
1812 calls = [
1813 call
1814 for idx, entry in enumerate(entries)
1815 if isinstance(entry, Mapping) and (call := _bare_call_from_mapping(entry, index=idx))
1816 ]
1817 if not calls:
1818 return ChatToolResult(content=content, tool_calls=[])
1819 return ChatToolResult(content="", tool_calls=calls)
1822# Leading non-whitespace characters that mark streamed text as a potential bare
1823# JSON tool call (an object or an array of them); any other first char is plain
1824# text and streams through untouched.
1825_BARE_CALL_OPENERS = "{["
1828def _tool_call_delta_from_recovered(call: ToolCall, index: int) -> ToolCallDelta:
1829 """Shape a recovered bare-JSON :class:`ToolCall` as a single streaming delta.
1831 Mirrors :func:`_tool_call_delta_from_chunk`: id and name ride the opener (the
1832 only frame for a recovered call), the arguments JSON is the lone
1833 ``arguments_delta``, and the position is the index.
1834 """
1835 return ToolCallDelta(
1836 index=index,
1837 id=call.id or None,
1838 name=call.name or None,
1839 arguments_delta=call.arguments or None,
1840 )
1843def _recover_bare_json_stream(
1844 items: Iterator[str | ToolCallDelta | TokenUsage | StreamFinish],
1845) -> ClosableIterator[str | ToolCallDelta | TokenUsage | StreamFinish]:
1846 """Wrap a raw chat stream to recover a tool call emitted as bare-JSON text.
1848 Some small models print ``{"name": ..., "arguments": {...}}`` as content
1849 instead of native ``tool_calls``; the non-stream paths recover this via
1850 :func:`_recover_bare_json_tool_calls`. This applies the same recovery to the
1851 stream, but only when the model emitted no native :class:`ToolCallDelta` and
1852 the streamed text looks like a bare call from its first character. Normal text
1853 still streams token by token: once the buffered head proves not to be a bare
1854 call it is flushed and all later text passes straight through.
1855 """
1856 buffer = "" # leading text held back as a potential bare call until resolved
1857 # True once leading text has streamed as plain (or a native call was seen):
1858 # past that, a later '{'/'[' is content, not a bare call -- never buffer again.
1859 committed = False
1860 try:
1861 for item in items:
1862 if isinstance(item, ToolCallDelta):
1863 yield from _flush_plain(buffer)
1864 buffer, committed = "", True
1865 yield item
1866 elif isinstance(item, TokenUsage | StreamFinish):
1867 yield from _recover_buffer(buffer)
1868 buffer = ""
1869 yield item
1870 elif committed or _passthrough_text(buffer, item):
1871 yield from _flush_plain(buffer)
1872 buffer = ""
1873 committed = True
1874 yield item
1875 else:
1876 buffer += item
1877 yield from _recover_buffer(buffer)
1878 finally:
1879 # Forward close to the source generator: if a consumer closes this
1880 # wrapper mid-stream, a plain for-loop would not propagate GeneratorExit
1881 # to *items*, leaking the underlying HTTP stream and its in_flight slot.
1882 # Suppress teardown errors (httpx stream close can raise) so they don't
1883 # mask the exception that triggered this finally.
1884 if isinstance(items, Generator):
1885 with contextlib.suppress(Exception):
1886 items.close()
1889def _passthrough_text(buffer: str, text: str) -> bool:
1890 """Whether *text* should stream through directly rather than buffer.
1892 True once the accumulated head's first non-whitespace char is known and is not
1893 a bare-call opener (plain text): the buffer is empty in that case, so the
1894 caller yields *text* as is. While the head is all whitespace, or once it opens
1895 with ``{``/``[``, the text is buffered (False) pending recovery.
1896 """
1897 head = (buffer + text).lstrip()
1898 return bool(head) and head[0] not in _BARE_CALL_OPENERS
1901def _flush_plain(buffer: str) -> Iterator[str]:
1902 """Yield buffered leading text verbatim (it was not a bare call after all)."""
1903 if buffer:
1904 yield buffer
1907def _recover_buffer(buffer: str) -> Iterator[str | ToolCallDelta]:
1908 """Resolve the buffered leading text at a terminator or end of stream.
1910 The buffer reaching here was held as a potential bare call (text starting with
1911 ``{``/``[`` and no native call seen). Run :func:`_recover_bare_json_tool_calls`:
1912 emit one delta per recovered call, or yield the text unchanged when it only
1913 happened to start with ``{``/``[`` but is not a call.
1914 """
1915 if not buffer:
1916 return
1917 recovered = _recover_bare_json_tool_calls(buffer)
1918 if not recovered.tool_calls:
1919 yield buffer
1920 return
1921 for index, call in enumerate(recovered.tool_calls):
1922 yield _tool_call_delta_from_recovered(call, index)