Coverage for src/lilbee/providers/base.py: 100%

1"""Base protocol and exceptions for LLM providers."""

3from __future__ import annotations

5from collections.abc import Callable, Iterator

6from dataclasses import dataclass

7from enum import StrEnum

8from pathlib import Path

9from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, overload, runtime_checkable

11from pydantic import BaseModel

13from lilbee.core.vectors import Vector

14from lilbee.providers.roles import WorkerRole

16if TYPE_CHECKING:

17 from lilbee.providers.warm_progress import WarmProgress

19T_co = TypeVar("T_co", covariant=True)

21# The inline reasoning markers lilbee's pipeline speaks. A provider whose server

22# extracts reasoning into a separate field re-inlines it with these tags at the

23# client boundary, so every downstream consumer parses one format.

24# What every provider holds back from the context window for the answer it is

25# about to generate, plus a small margin for the chat template's own framing.

26# One owner for both numbers: the fleet ENFORCES this default budget (rejecting a

27# prompt that exceeds it), while retrieval FITS its context to it. An explicit

28# over-large output reservation (num_predict past the default) is clamped back to

29# this default rather than rejected, so an agent that over-reserves still fits.

30# When the two disagreed, retrieval assembled prompts up to the margin larger than

31# the engine would accept, and a grounded turn failed with a 400 the caller could

32# do nothing about.

33GENERATION_RESERVE_TOKENS = 1024

34CONTEXT_WINDOW_MARGIN_TOKENS = 128

36# Chars-per-token assumed when BUDGETING context, deliberately harsher than the

37# display estimator's 4: dense OCR/legal text tokenizes at ~2.5-3 chars per token.

38BUDGET_CHARS_PER_TOKEN = 3

41def prompt_token_budget(ctx: int, num_predict: int | None = None) -> int:

42 """Tokens a prompt may occupy in a *ctx*-token window, reserve and margin removed."""

43 return ctx - (num_predict or GENERATION_RESERVE_TOKENS) - CONTEXT_WINDOW_MARGIN_TOKENS

46def estimate_budget_tokens(text: str) -> int:

47 """Conservative token cost of *text* for budgeting (see BUDGET_CHARS_PER_TOKEN)."""

48 return max(1, len(text) // BUDGET_CHARS_PER_TOKEN)

51THINK_OPEN_TAG = "<think>"

52THINK_CLOSE_TAG = "</think>"

55@runtime_checkable

56class ClosableIterator(Iterator[T_co], Protocol[T_co]):

57 """An iterator that releases resources when ``close()`` is called.

59 Streaming chat responses use this to guarantee upstream resources (the

60 fleet's in-flight request slot) are released even when callers truncate

61 the stream before exhaustion. Generators satisfy this implicitly.

62 """

64 def close(self) -> None: ...

67class LLMOptions(BaseModel):

68 """Validated options passed to LLM providers.

69 Only these fields are forwarded: everything else is rejected

70 to prevent injection of sensitive parameters like api_base or api_key.

71 """

73 temperature: float | None = None

74 top_p: float | None = None

75 top_k: int | None = None

76 seed: int | None = None

77 num_predict: int | None = None

78 repeat_penalty: float | None = None

79 frequency_penalty: float | None = None

80 presence_penalty: float | None = None

81 num_ctx: int | None = None

82 stop: list[str] | None = None

83 # Thinking-template control for structured internal calls (schema

84 # induction and similar): a small thinking model can burn its whole

85 # token budget inside <think> and emit nothing. llama-server maps this

86 # to chat_template_kwargs; hosted-API translators drop it.

87 think: bool | None = None

88 # Structured-output request for the internal calls that want a bare JSON

89 # value back. Must be listed here or the allowlist above silently drops it.

90 response_format: dict[str, Any] | None = None

92 def to_dict(self) -> dict[str, Any]:

93 """Return only non-None values as a dict."""

94 return {k: v for k, v in self.model_dump().items() if v is not None}

97def filter_options(options: dict[str, Any]) -> dict[str, Any]:

98 """Validate and filter generation options through LLMOptions model."""

99 return LLMOptions(**options).to_dict()

100

101

102def normalize_generation_options(options: dict[str, Any] | None) -> dict[str, Any]:

103 """Validate options and map them to the per-call set an OpenAI/llama-server body takes.

104

105 ``filter_options`` validates against :class:`LLMOptions`; ``num_predict`` then

106 becomes ``max_tokens`` and ``num_ctx`` is dropped (a model-load param, not a

107 per-call one). Shared by the fleet and SDK option translators so the mapping

108 lives in one place.

109 """

110 if not options:

111 return {}

112 filtered = filter_options(options)

113 if "num_predict" in filtered:

114 filtered["max_tokens"] = filtered.pop("num_predict")

115 filtered.pop("num_ctx", None)

116 return filtered

117

118

119class ProviderErrorKind(StrEnum):

120 """Provider-agnostic category of a failed provider call.

121

122 Classified by exception type at each backend boundary so callers can

123 branch on the kind instead of matching message strings (which are

124 provider-specific and drift between SDK versions).

125 """

126

127 AUTH = "auth"

128 RATE_LIMIT = "rate_limit"

129 CONTEXT_OVERFLOW = "context_overflow"

130 NOT_FOUND = "not_found"

131 BAD_REQUEST = "bad_request"

132 CONNECTION = "connection"

133 SERVER = "server"

134 CAPACITY = "capacity"

135 PORT_CONFLICT = "port_conflict"

136 UNKNOWN = "unknown"

137

138

139class ProviderError(Exception):

140 """Raised when an LLM provider operation fails.

141

142 ``kind`` is the provider-agnostic category; backends that can't classify a

143 failure leave it ``UNKNOWN``.

144 """

145

146 def __init__(

147 self,

148 message: str,

149 *,

150 provider: str = "",

151 kind: ProviderErrorKind = ProviderErrorKind.UNKNOWN,

152 ) -> None:

153 self.provider = provider

154 self.kind = kind

155 super().__init__(message)

156

157

158# Human word per role for the not-configured error ("embedding model", not "embed model").

159_ROLE_WORDS: dict[WorkerRole, str] = {

160 WorkerRole.CHAT: "chat",

161 WorkerRole.EMBED: "embedding",

162 WorkerRole.RERANK: "reranker",

163 WorkerRole.VISION: "vision",

164}

165

166

167def require_role_ref(ref: str, role: WorkerRole, *, provider: str = "") -> str:

168 """Reject an unconfigured role with a clean error instead of parsing ''."""

169 if not ref:

170 raise ProviderError(

171 f"No {_ROLE_WORDS[role]} model is configured. Pick one from the catalog "

172 f"or run 'lilbee model pull <model>'.",

173 provider=provider,

174 kind=ProviderErrorKind.NOT_FOUND,

175 )

176 return ref

177

178

179ChatMessage = dict[str, str]

180

181

182@dataclass(frozen=True)

183class ToolCall:

184 """One tool/function call the model requested.

185

186 ``arguments`` is the raw JSON-encoded argument object (OpenAI's shape), left

187 as a string so the caller decides how to parse and validate it. ``id`` is the

188 server-assigned call id, echoed back in the tool result message.

189 """

190

191 id: str

192 name: str

193 arguments: str

194

195

196@dataclass(frozen=True)

197class ChatToolResult:

198 """A chat turn that may carry tool calls alongside (or instead of) text.

199

200 ``tool_calls`` is empty for an ordinary text answer; ``content`` is empty when

201 the model returned only tool calls. Both can be populated when a model emits

202 commentary plus a call.

203 """

204

205 content: str

206 tool_calls: list[ToolCall]

207

208

209class FinishReason(StrEnum):

210 """Why a chat completion stopped, mirroring OpenAI's vocabulary."""

211

212 STOP = "stop"

213 LENGTH = "length"

214 TOOL_CALLS = "tool_calls"

215 CONTENT_FILTER = "content_filter"

216

217 @classmethod

218 def coerce(cls, raw: object) -> FinishReason:

219 """Map a backend-supplied finish_reason to a member, defaulting to STOP.

220

221 Both the streaming and non-streaming paths read finish_reason from the

222 backend; an unknown or non-string value (a model that omits it) falls

223 back to STOP so the dispatch reports an ordinary end-of-turn.

224 """

225 if isinstance(raw, str):

226 try:

227 return cls(raw)

228 except ValueError:

229 return cls.STOP

230 return cls.STOP

231

232

233@dataclass(frozen=True)

234class TokenUsage:

235 """Prompt / completion token counts for one chat call.

236

237 Defaults to zero so a backend that reports no usage block still yields a

238 well-formed result; the fleet populates these from llama-server's ``usage``.

239 """

240

241 prompt_tokens: int = 0

242 completion_tokens: int = 0

243

244

245@dataclass(frozen=True)

246class ChatResult:

247 """Structured result from a non-streaming chat call.

248

249 ``tool_calls`` is empty for an ordinary text answer; ``text`` is empty when

250 the model returned only tool calls. ``usage`` carries the backend's token

251 counts (zero when unreported). The canonical chat dispatch reads these to

252 build its OpenAI/Anthropic-shaped response.

253 """

254

255 text: str

256 tool_calls: tuple[ToolCall, ...]

257 finish_reason: FinishReason

258 usage: TokenUsage = TokenUsage()

259

260

261@dataclass(frozen=True)

262class ToolCallDelta:

263 """Partial tool-call delta in a streaming response, accumulated by ``index``.

264

265 ``id`` and ``name`` arrive on the opener frame for a call; ``arguments_delta``

266 accumulates across subsequent frames at the same ``index``.

267 """

268

269 index: int

270 id: str | None

271 name: str | None

272 arguments_delta: str | None

273

274

275@dataclass(frozen=True)

276class StreamFinish:

277 """Terminal frame carrying why a streaming chat call stopped.

278

279 Emitted once, near the end of the stream, so the dispatch can report the

280 same finish_reason the non-streaming path already surfaces, notably

281 ``length`` on a max_tokens truncation. Tool-call streams already infer

282 TOOL_USE from their deltas, so a finish frame never downgrades that.

283 """

284

285 reason: FinishReason

286

287

288ChatStreamItem = str | ToolCallDelta | TokenUsage | StreamFinish

289"""One frame yielded by a streaming chat call: text token, tool-call delta, the

290final token-usage summary, or the finish-reason terminator (each emitted once,

291last, when the backend reports them)."""

292

293

294class LLMProvider(Protocol):

295 """Protocol for pluggable LLM backends."""

296

297 def embed(self, texts: list[str]) -> list[Vector]:

298 """Embed a batch of texts, return list of vectors."""

299 ...

300

301 def count_tokens(self, text: str) -> int:

302 """Exact token count of *text* under the embedding model's tokenizer.

303

304 Raise ``NotImplementedError`` when the backend has no local tokenizer (cloud

305 SDK backends); token-budgeted chunk sizing then falls back to a character

306 estimate.

307 """

308 ...

309

310 @overload

311 def chat(

312 self,

313 messages: list[ChatMessage],

314 *,

315 stream: Literal[False] = False,

316 options: dict[str, Any] | None = None,

317 model: str | None = None,

318 tools: list[dict[str, Any]] | None = None,

319 tool_choice: str | dict[str, Any] | None = None,

320 ) -> ChatResult: ...

321

322 @overload

323 def chat(

324 self,

325 messages: list[ChatMessage],

326 *,

327 stream: Literal[True],

328 options: dict[str, Any] | None = None,

329 model: str | None = None,

330 tools: list[dict[str, Any]] | None = None,

331 tool_choice: str | dict[str, Any] | None = None,

332 ) -> ClosableIterator[ChatStreamItem]: ...

333

334 def chat(

335 self,

336 messages: list[ChatMessage],

337 *,

338 stream: bool = False,

339 options: dict[str, Any] | None = None,

340 model: str | None = None,

341 tools: list[dict[str, Any]] | None = None,

342 tool_choice: str | dict[str, Any] | None = None,

343 ) -> ChatResult | ClosableIterator[ChatStreamItem]:

344 """Chat completion.

345

346 Non-streaming returns a :class:`ChatResult` (assistant text, any

347 tool-call frames, and a finish reason). Streaming returns a

348 :class:`ClosableIterator` of :data:`ChatStreamItem` (text tokens

349 interleaved with :class:`ToolCallDelta` frames). ``tools`` is the

350 OpenAI function-tool list; ``tool_choice`` is ``"auto"`` / ``"none"`` /

351 ``"required"`` or a ``{"type": "function", ...}`` selector. A model

352 that lacks tool support returns an empty ``tool_calls`` / yields no

353 tool deltas rather than erroring.

354 """

355 ...

356

357 def supports_tools(self, model_ref: str) -> bool:

358 """Return True iff the backend can route tool calls for *model_ref*.

359

360 Default False so backends without a tool path are never offered tools;

361 tool-capable backends override this with a real probe.

362 """

363 return False

364

365 def chat_with_tools(

366 self,

367 messages: list[ChatMessage],

368 *,

369 tools: list[dict[str, Any]],

370 tool_choice: str | dict[str, Any] | None = None,

371 options: dict[str, Any] | None = None,

372 model: str | None = None,

373 ) -> ChatToolResult:

374 """Non-streaming chat that may return tool calls.

375

376 ``tools`` is the OpenAI function-tool list; ``tool_choice`` is ``"auto"``

377 / ``"none"`` / ``"required"`` or a specific ``{"type": "function", ...}``

378 selector. Backends without tool support raise :class:`ProviderError`.

379 """

380 raise ProviderError("This backend does not support tool calling.")

381

382 def vision_ocr(

383 self,

384 png_bytes: bytes,

385 model: str,

386 prompt: str = "",

387 *,

388 timeout: float | None = None,

389 ) -> str:

390 """OCR one page image; ``timeout`` seconds, ``None``/``0`` = no cap."""

391 ...

392

393 def vision_slot_capacity(self) -> int | None:

394 """Fitted concurrent-OCR slots if the vision fleet is running, else None.

395

396 The ingest fan-out uses this to size itself to the servers' real

397 continuous-batching capacity rather than the requested concurrency,

398 which a memory-constrained card cannot always fit. ``None`` means the

399 capacity isn't known yet (no local vision backend, or the fleet hasn't

400 started); the caller falls back to its own estimate.

401 """

402 ...

403

404 def list_models(self) -> list[str]:

405 """List available model identifiers."""

406 ...

407

408 def list_chat_models(self, provider: str) -> list[str]:

409 """List frontier chat models the provider is aware of for *provider*.

410

411 Returns the unfiltered upstream catalog (whatever litellm

412 exposes for API providers; an empty list for local backends

413 like the llama-server fleet that have no notion of external

414 catalogs).

415 """

416 ...

417

418 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:

419 """Download a model. Raises NotImplementedError if not supported."""

420 ...

421

422 def show_model(self, model: str) -> dict[str, Any] | None:

423 """Return model metadata, or None if backend doesn't expose it."""

424 ...

425

426 def get_capabilities(self, model: str) -> list[str]:

427 """Return capability tags (e.g. ``["completion", "vision"]``) for *model*.

428

429 Returns an empty list when the backend does not support capability

430 reporting or the model is not found.

431 """

432 ...

433

434 def rerank(self, query: str, candidates: list[str]) -> list[float]:

435 """Score *candidates* for their relevance to *query*, one float per candidate.

436

437 The backend resolves the reranker model from ``cfg.reranker_model``.

438 Callers MUST check ``cfg.reranker_model`` is non-empty before

439 calling; use :meth:`supports_rerank` for UI-render decisions.

440

441 Returns: list of floats in input order, higher = more relevant.

442 Empty ``candidates`` returns ``[]``.

443 Raises :class:`ProviderError` when the backend does not support

444 reranking, ``cfg.reranker_model`` is empty, or the model scored no

445 candidate. A backend must raise rather than return uniform scores,

446 which would silently preserve the caller's input order.

447 """

448 ...

449

450 def supports_rerank(self) -> bool:

451 """Capability probe: can this backend rerank *if* a model is configured?

452

453 Pure capability check, NOT "a reranker is currently active". An

454 empty ``cfg.reranker_model`` returns ``True`` so the settings UI

455 keeps the picker visible; callers that need to know whether

456 reranking is actually configured must check ``bool(cfg.reranker_model)``

457 separately. ``rerank()`` is the gated path that requires a

458 non-empty value.

459 """

460 return False

461

462 def shutdown(self) -> None:

463 """Release resources (e.g. background threads). No-op if nothing to clean up."""

464 ...

465

466 def invalidate_load_cache(self, model_path: Path | None = None) -> None:

467 """Drop loaded-model state; ``None`` evicts all, else only that path. No-op default."""

468 return

469

470 def drop_loaded_models_async(self) -> None:

471 """Drop all loaded-model state off the caller's thread. No-op default.

472

473 Like :meth:`invalidate_load_cache` with no path, but the teardown (which

474 stops every server and waits on each process) runs on a background thread

475 so a settings change that touches a role-agnostic load key never blocks

476 the UI / request thread. The next call rebuilds with current cfg.

477 """

478 self.invalidate_load_cache()

479

480 def warm_up_pool(self) -> None:

481 """Eagerly start the configured role servers so the first call lands warm.

482

483 Default no-op so providers without managed servers (SDK / routing

484 wrappers) can be passed to ``Services`` unchanged. Implemented by

485 :class:`FleetProvider` to spawn the chat / embed / rerank / vision

486 servers whose model is configured.

487 """

488 return

489

490 def cancel_inference(self) -> None:

491 """Interrupt any in-flight generation. No-op default.

492

493 The fleet engine severs its live chat streams (llama-server stops

494 generating when the connection drops); the SDK wrapper has nothing to

495 interrupt here.

496 """

497 return

498

499 def reload_role(self, role: WorkerRole, *, wait: bool = False) -> None:

500 """Drop and respawn just *role*'s model so it picks up changed cfg.

501

502 Default no-op for providers without per-role model servers. The fleet

503 respawns only that role's server; other roles and their in-flight work

504 are left untouched. ``wait=True`` blocks until the respawn finishes (for a

505 caller already off the event loop); the default returns immediately.

506 """

507 return

508

509 def reload_placement(self, *, wait: bool = False) -> None:

510 """Re-plan GPU placement with current cfg, restarting only moved roles.

511

512 Default no-op for providers without GPU-placed servers. The fleet diffs

513 the fresh plan against the running fleet and respawns only the roles

514 whose placement changed, so an untouched role's loaded model stays

515 resident. ``wait=True`` blocks until the restarted proxies are healthy.

516 """

517 return

518

519 def role_ready(self, role: WorkerRole) -> bool:

520 """Whether *role* has a healthy server now, without starting one.

521

522 Default ``True``: providers without managed servers (SDK / routing

523 wrappers) are always reachable. The fleet returns ``False`` while a role

524 is still cold-starting so surfaces can show a warming state.

525 """

526 del role

527 return True

528

529 def max_concurrent_chats(self) -> int:

530 """Upper bound on simultaneous chat generations this provider can serve.

531

532 Default ``1``: a single in-process model cannot take concurrent generate

533 calls, so chat is serialized. A server-backed provider that batches (the

534 fleet) overrides this with its slot capacity, so the chat admission gate

535 lets that many run at once instead of one at a time.

536 """

537 return 1

538

539 def served_chat_ctx(self) -> int | None:

540 """Per-slot context the active chat server runs with, or None if unknown.

541

542 A client trims its conversation to this so a long agentic session fits

543 the model's actual window instead of overflowing. Default ``None``:

544 providers without a managed context (SDK wrappers) advertise nothing.

545 """

546 return None

547

548 def warm_pending(self) -> bool:

549 """Whether a warm has been requested and has not finished.

550

551 True from the moment ``warm_up_pool`` accepts a warm until its background

552 work ends, so a surface can hold before the first phase is stamped. Default

553 ``False``: providers without managed servers never warm.

554 """

555 return False

556

557 def warm_progress(self) -> WarmProgress | None:

558 """Snapshot of the chat model's cold-load progress, or None when idle.

559

560 A launcher streams this to render a real progress bar while a large chat

561 model loads. Default ``None``: providers without a managed load (SDK /

562 routing wrappers) expose nothing, so a launcher falls back to a plain

563 spinner. The fleet returns live read / engine-load state.

564 """

565 return None

566

567 def add_spawn_listener(

568 self,

569 *,

570 on_spawning: Callable[[WorkerRole], None] | None = None,

571 on_spawned: Callable[[WorkerRole], None] | None = None,

572 ) -> None:

573 """Subscribe to server (re)spawn lifecycle events. No-op default.

574

575 The fleet calls ``on_spawning`` before a role's server starts and

576 ``on_spawned`` once it is healthy, so the TUI can surface cold-start and

577 reload progress. Providers without managed servers ignore it.

578 """

579 return