Coverage for src/lilbee/providers/sdk_llm_provider.py: 100%
177 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""SDK-agnostic LLM provider implementing the public ``LLMProvider`` Protocol.
3``SdkLLMProvider`` owns the semantic layer: auth key injection, option
4translation, model-ref parsing, error wrapping, and lazy one-shot
5backend initialization (``configure_logging`` + ``inject_provider_keys``
6on first use). It speaks to the underlying SDK exclusively through an
7``LlmSdkBackend``, so swapping SDKs is a one-file adapter change.
9Zero direct SDK imports live here. The adapter owns SDK-specific
10concerns like wire-format prefixes (``ollama/``) and OpenAI content-parts
11schema for image inputs.
12"""
14from __future__ import annotations
16import logging
17import os
18from collections.abc import Callable
19from pathlib import Path
20from typing import Any, Literal, overload
22import numpy as np
24from lilbee.core.config import cfg
25from lilbee.core.vectors import Vector
26from lilbee.providers.base import (
27 ChatMessage,
28 ChatResult,
29 ChatStreamItem,
30 ChatToolResult,
31 ClosableIterator,
32 FinishReason,
33 LLMProvider,
34 ProviderError,
35 StreamFinish,
36 ToolCall,
37 ToolCallDelta,
38 require_role_ref,
39)
40from lilbee.providers.local_servers import LOCAL_SERVER_KEYS
41from lilbee.providers.local_servers.config_urls import base_url_for, configured_local_servers
42from lilbee.providers.model_ref import ProviderModelRef, parse_model_ref, translate_options
43from lilbee.providers.roles import WorkerRole
44from lilbee.providers.sdk_backend import (
45 PROVIDER_KEYS,
46 CompletionRequest,
47 EmbeddingRequest,
48 LlmSdkBackend,
49 RerankRequest,
50)
52log = logging.getLogger(__name__)
55def _api_base_for(ref: ProviderModelRef) -> str | None:
56 """Endpoint for a local-server ref; ``None`` for hosted APIs (no base needed)."""
57 if ref.provider in LOCAL_SERVER_KEYS:
58 return base_url_for(ref.provider)
59 return None
62def inject_provider_keys() -> None:
63 """Copy per-provider API keys from config into ``os.environ``.
65 OpenAI-compatible SDKs read provider-specific env vars
66 (``OPENAI_API_KEY``, ``ANTHROPIC_API_KEY``, ...) at call time. This
67 bridges lilbee's config system to that convention. Explicit env
68 vars are never overwritten so users can still override via their
69 shell.
70 """
71 for _, cfg_field, env_var, _ in PROVIDER_KEYS:
72 # No default: every PROVIDER_KEYS field is a declared config attribute, so
73 # a typo in the table should surface as AttributeError, not silently read "".
74 value = getattr(cfg, cfg_field)
75 if value and not os.environ.get(env_var):
76 os.environ[env_var] = value
79class SdkLLMProvider(LLMProvider):
80 """Provider that delegates SDK calls to an ``LlmSdkBackend``."""
82 def __init__(
83 self,
84 backend: LlmSdkBackend,
85 *,
86 api_key: str = "",
87 ) -> None:
88 self._backend = backend
89 self._api_key = api_key
90 self._initialized = False
92 def _ensure_initialized(self) -> None:
93 """Apply one-shot backend setup before the first call.
95 Runs ``configure_logging(suppress_debug=cfg.json_mode)`` and
96 ``inject_provider_keys()`` exactly once, regardless of whether
97 the first operation is ``chat``, ``embed``, or a catalog query.
98 Both steps happen together because the backend's first SDK
99 import must see (a) the debug flag applied, and (b) per-provider
100 API keys in ``os.environ``.
101 """
102 if self._initialized:
103 return
104 try:
105 self._backend.configure_logging(suppress_debug=cfg.json_mode)
106 except (ImportError, AttributeError):
107 log.debug("backend.configure_logging failed", exc_info=True)
108 inject_provider_keys()
109 self._initialized = True
111 def embed(self, texts: list[str]) -> list[Vector]:
112 """Embed texts via the configured backend, converting its JSON floats to float32."""
113 self._ensure_initialized()
114 ref = parse_model_ref(
115 require_role_ref(cfg.embedding_model, WorkerRole.EMBED, provider="sdk")
116 )
117 request = EmbeddingRequest(
118 ref=ref,
119 inputs=texts,
120 api_base=_api_base_for(ref),
121 api_key=self._api_key or None,
122 )
123 try:
124 result = self._backend.embed(request)
125 except ProviderError:
126 raise
127 except Exception as exc:
128 raise ProviderError(
129 f"Embedding failed: {exc}", provider=self._backend.provider_name
130 ) from exc
131 return [np.asarray(vec, dtype=np.float32) for vec in result.vectors]
133 def count_tokens(self, text: str) -> int:
134 """Cloud SDK backends expose no local tokenizer, so chunk sizing falls back
135 to a character estimate (see LilbeeTokenizerBackend)."""
136 raise NotImplementedError("SDK backends have no local tokenizer for chunk sizing")
138 @overload
139 def chat(
140 self,
141 messages: list[dict[str, str]],
142 *,
143 stream: Literal[False] = False,
144 options: dict[str, Any] | None = None,
145 model: str | None = None,
146 tools: list[dict[str, Any]] | None = None,
147 tool_choice: str | dict[str, Any] | None = None,
148 ) -> ChatResult: ...
150 @overload
151 def chat(
152 self,
153 messages: list[dict[str, str]],
154 *,
155 stream: Literal[True],
156 options: dict[str, Any] | None = None,
157 model: str | None = None,
158 tools: list[dict[str, Any]] | None = None,
159 tool_choice: str | dict[str, Any] | None = None,
160 ) -> ClosableIterator[ChatStreamItem]: ...
162 def chat(
163 self,
164 messages: list[dict[str, str]],
165 *,
166 stream: bool = False,
167 options: dict[str, Any] | None = None,
168 model: str | None = None,
169 tools: list[dict[str, Any]] | None = None,
170 tool_choice: str | dict[str, Any] | None = None,
171 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
172 """Chat completion via the configured backend.
174 Non-streaming returns a :class:`ChatResult` carrying the assistant
175 text, any tool-call frames the model emitted, and a finish reason.
176 Streaming yields :data:`ChatStreamItem` frames (text tokens and
177 tool-call deltas).
178 """
179 self._ensure_initialized()
180 ref = parse_model_ref(
181 require_role_ref(model or cfg.chat_model, WorkerRole.CHAT, provider="sdk")
182 )
183 if tools and not self.supports_tools(model or cfg.chat_model):
184 chosen = model or cfg.chat_model
185 raise ProviderError(
186 f"Model {chosen!r} does not support tool calls. Pick a different "
187 f"chat model that advertises tool support, or remove tools from "
188 f"the request.",
189 provider=self._backend.provider_name,
190 )
191 translated = translate_options(options, ref) if options else {}
192 if tools is not None:
193 translated["tools"] = tools
194 if tool_choice is not None:
195 translated["tool_choice"] = tool_choice
196 request = CompletionRequest(
197 ref=ref,
198 messages=list(messages),
199 options=translated,
200 api_base=_api_base_for(ref),
201 api_key=self._api_key or None,
202 )
203 if stream:
204 return self._chat_stream(request)
205 try:
206 result = self._backend.complete(request)
207 except ProviderError:
208 raise
209 except Exception as exc:
210 raise ProviderError(
211 f"Chat failed: {exc}", provider=self._backend.provider_name
212 ) from exc
213 return ChatResult(
214 text=result.content,
215 tool_calls=tuple(
216 ToolCall(id=tc.id, name=tc.name, arguments=tc.arguments) for tc in result.tool_calls
217 ),
218 finish_reason=FinishReason.coerce(result.finish_reason),
219 )
221 def supports_tools(self, model_ref: str) -> bool:
222 """Delegate to the backend's ``supports_tools`` probe."""
223 return self._backend.supports_tools(model_ref)
225 def chat_with_tools(
226 self,
227 messages: list[ChatMessage],
228 *,
229 tools: list[dict[str, Any]],
230 tool_choice: str | dict[str, Any] | None = None,
231 options: dict[str, Any] | None = None,
232 model: str | None = None,
233 ) -> ChatToolResult:
234 """Tool-calling chat for remote/SDK models.
236 The base stub raises, but this backend advertises tool support and ``chat``
237 already forwards tools/tool_choice, so route through it instead of refusing.
238 """
239 result = self.chat(
240 # Pass each message through whole: a tool conversation carries
241 # ``tool_calls`` / ``tool_call_id`` / ``name`` that link an assistant
242 # call to its result, and stripping to role+content breaks that chain.
243 [dict(m) for m in messages],
244 stream=False,
245 options=options,
246 model=model,
247 tools=tools,
248 tool_choice=tool_choice,
249 )
250 return ChatToolResult(content=result.text, tool_calls=list(result.tool_calls))
252 def _chat_stream(self, request: CompletionRequest) -> ClosableIterator[ChatStreamItem]:
253 """Yield content tokens and tool-call deltas from a streaming completion.
255 Exceptions surfaced by the backend at either call time or during
256 iteration are re-raised as ``ProviderError`` so callers always
257 see a consistent error type.
258 """
259 try:
260 stream = self._backend.complete_stream(request)
261 for chunk in stream:
262 if chunk.content:
263 yield chunk.content
264 for delta in chunk.tool_call_deltas:
265 yield ToolCallDelta(
266 index=delta.index,
267 id=delta.id,
268 name=delta.name,
269 arguments_delta=delta.arguments_delta,
270 )
271 if chunk.finish_reason is not None:
272 # The closing chunk's finish_reason lets the dispatch report
273 # length/stop, matching the non-streaming path.
274 yield StreamFinish(reason=FinishReason.coerce(chunk.finish_reason))
275 except ProviderError:
276 raise
277 except Exception as exc:
278 raise ProviderError(
279 f"Chat failed: {exc}", provider=self._backend.provider_name
280 ) from exc
282 def vision_ocr(
283 self,
284 png_bytes: bytes,
285 model: str,
286 prompt: str = "",
287 *,
288 timeout: float | None = None,
289 ) -> str:
290 """OCR via a multipart chat completion; ``timeout`` enforced via thread pool."""
291 from lilbee.vision import build_vision_messages, resolve_ocr_prompt
293 messages = build_vision_messages(prompt or resolve_ocr_prompt(model), png_bytes)
294 if timeout and timeout > 0:
295 from concurrent.futures import ThreadPoolExecutor
297 # Don't use the context manager: its __exit__ shutdown(wait=True) would
298 # block until a hung call returns, so the caller would not be freed at
299 # the deadline. Shut down without waiting (matching the fleet OCR path);
300 # a wedged call's worker thread lives until the backend httpx timeout.
301 pool = ThreadPoolExecutor(max_workers=1)
302 try:
303 future = pool.submit(self.chat, messages, stream=False, model=model)
304 result = future.result(timeout=timeout)
305 finally:
306 pool.shutdown(wait=False, cancel_futures=True)
307 else:
308 result = self.chat(messages, stream=False, model=model)
309 if not isinstance(result, ChatResult):
310 raise ProviderError(
311 f"Vision OCR returned non-text response ({type(result).__name__}).",
312 provider=self._backend.provider_name,
313 )
314 return result.text
316 def vision_slot_capacity(self) -> int | None:
317 """Hosted backends have no local OCR slots; the caller estimates."""
318 return None
320 def list_models(self) -> list[str]:
321 """List models across every configured local server.
323 A single unreachable server is logged and skipped so its outage does not
324 drop the models served by the other reachable servers.
325 """
326 names: list[str] = []
327 for spec, base_url in configured_local_servers():
328 try:
329 names.extend(self._backend.list_models(base_url=base_url, api_key=self._api_key))
330 except NotImplementedError:
331 continue
332 except Exception as exc:
333 log.debug("Skipping unreachable local server %s: %s", spec.key, exc)
334 return names
336 def list_chat_models(self, provider: str) -> list[str]:
337 """List frontier chat models known to the backend for *provider*.
339 Initializes the backend first so ``cfg.json_mode`` suppression is
340 applied before the SDK import inside the backend runs.
341 """
342 self._ensure_initialized()
343 try:
344 return self._backend.list_chat_models(provider)
345 except NotImplementedError:
346 return []
347 except ProviderError:
348 raise
349 except Exception as exc:
350 raise ProviderError(
351 f"Listing chat models failed: {exc}", provider=self._backend.provider_name
352 ) from exc
354 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
355 """Pull a model via the backend."""
356 try:
357 base_url = _api_base_for(parse_model_ref(model)) or ""
358 self._backend.pull_model(model, base_url=base_url, on_progress=on_progress)
359 except NotImplementedError as exc:
360 raise ProviderError(
361 f"Cannot pull model {model!r}: backend does not support pulling",
362 provider=self._backend.provider_name,
363 ) from exc
364 except ProviderError:
365 raise
366 except Exception as exc:
367 raise ProviderError(
368 f"Cannot pull model {model!r}: {exc}", provider=self._backend.provider_name
369 ) from exc
371 def show_model(self, model: str) -> dict[str, Any] | None:
372 """Return model metadata, or None when unsupported or not found."""
373 try:
374 base_url = _api_base_for(parse_model_ref(model)) or ""
375 return self._backend.show_model(model, base_url=base_url)
376 except NotImplementedError:
377 return None
378 except ProviderError:
379 raise
380 except Exception as exc:
381 raise ProviderError(
382 f"Showing model {model!r} failed: {exc}", provider=self._backend.provider_name
383 ) from exc
385 def get_capabilities(self, model: str) -> list[str]:
386 """Return capability tags from ``show_model`` output, or ``[]``."""
387 info = self.show_model(model)
388 if info is None:
389 return []
390 caps = info.get("capabilities", [])
391 return caps if isinstance(caps, list) else []
393 def rerank(self, query: str, candidates: list[str]) -> list[float]:
394 """Rerank candidates via the SDK backend using ``cfg.reranker_model``."""
395 if not candidates:
396 return []
397 self._ensure_initialized()
398 ref = parse_model_ref(
399 require_role_ref(cfg.reranker_model, WorkerRole.RERANK, provider="sdk")
400 )
401 request = RerankRequest(
402 ref=ref,
403 query=query,
404 candidates=candidates,
405 api_base=_api_base_for(ref),
406 api_key=self._api_key or None,
407 )
408 try:
409 result = self._backend.rerank(request)
410 except ProviderError:
411 raise
412 except Exception as exc:
413 raise ProviderError(
414 f"Rerank failed: {exc}", provider=self._backend.provider_name
415 ) from exc
416 return result.scores
418 def supports_rerank(self) -> bool:
419 """SDK-backed rerank is available when the underlying SDK is importable."""
420 return self._backend.available()
422 def available(self) -> bool:
423 """Return True when the configured SDK backend can service catalog calls."""
424 return self._backend.available()
426 def shutdown(self) -> None:
427 """SDK-backed providers hold no lilbee-side resources."""
429 def invalidate_load_cache(self, model_path: Path | None = None) -> None:
430 """No-op: cloud backends have no local model cache to evict."""