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