Coverage for src/lilbee/providers/sdk_llm_provider.py: 100%
181 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""SDK-agnostic LLM provider implementing the public ``LLMProvider`` Protocol.
3``SdkLLMProvider`` owns the semantic layer: the API key each request
4carries, option translation, model-ref parsing, error wrapping, and lazy
5one-shot backend initialization (``configure_logging`` on first use).
6It 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
17from collections.abc import Callable
18from typing import Any, Literal, overload
20import numpy as np
22from lilbee.core.config import cfg
23from lilbee.core.vectors import Vector
24from lilbee.providers.base import (
25 ChatMessage,
26 ChatResult,
27 ChatStreamItem,
28 ChatToolResult,
29 ClosableIterator,
30 FinishReason,
31 LLMProvider,
32 ProviderError,
33 StreamFinish,
34 ToolCall,
35 ToolCallDelta,
36 require_role_ref,
37)
38from lilbee.providers.key_check import provider_api_key_in_use
39from lilbee.providers.local_servers import LOCAL_SERVER_KEYS, OLLAMA
40from lilbee.providers.local_servers.config_urls import base_url_for, configured_local_servers
41from lilbee.providers.model_ref import ProviderModelRef, parse_model_ref, translate_options
42from lilbee.providers.roles import WorkerRole
43from lilbee.providers.sdk_backend import (
44 CompletionRequest,
45 EmbeddingRequest,
46 LlmSdkBackend,
47 RerankRequest,
48)
49from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
50from lilbee.runtime.daemon_call import DaemonCall
52log = logging.getLogger(__name__)
54_VISION_OCR_TIMEOUT_ERROR = "Vision OCR did not finish within {timeout:g} seconds."
57def _api_base_for(ref: ProviderModelRef) -> str | None:
58 """Endpoint for a local-server ref; ``None`` for hosted APIs (no base needed)."""
59 if ref.provider in LOCAL_SERVER_KEYS:
60 return base_url_for(ref.provider)
61 return None
64def _api_key_for(ref: ProviderModelRef) -> str | None:
65 """A hosted ref sends only the key the key check tests; Ollama sends none.
67 Other local-server refs send ``llm_api_key``.
68 """
69 if ref.is_api:
70 return provider_api_key_in_use(ref.provider)
71 if ref.provider == OLLAMA.key:
72 return None
73 return cfg.llm_api_key or None
76class SdkLLMProvider(LLMProvider):
77 """Provider that delegates SDK calls to an ``LlmSdkBackend``."""
79 def __init__(self, backend: LlmSdkBackend) -> None:
80 self._backend = backend
81 self._initialized = False
83 def _ensure_initialized(self) -> None:
84 """Run ``configure_logging(suppress_debug=cfg.json_mode)`` once, before any SDK call."""
85 if self._initialized:
86 return
87 try:
88 self._backend.configure_logging(suppress_debug=cfg.json_mode)
89 except (ImportError, AttributeError):
90 log.debug("backend.configure_logging failed", exc_info=True)
91 self._initialized = True
93 def embed(self, texts: list[str]) -> list[Vector]:
94 """Embed texts via the configured backend, converting its JSON floats to float32."""
95 self._ensure_initialized()
96 ref = parse_model_ref(
97 require_role_ref(cfg.embedding_model, WorkerRole.EMBED, provider="sdk")
98 )
99 request = EmbeddingRequest(
100 ref=ref,
101 inputs=texts,
102 api_base=_api_base_for(ref),
103 api_key=_api_key_for(ref),
104 )
105 try:
106 result = self._backend.embed(request)
107 except ProviderError:
108 raise
109 except Exception as exc:
110 raise ProviderError(
111 f"Embedding failed: {exc}", provider=self._backend.provider_name
112 ) from exc
113 return [np.asarray(vec, dtype=np.float32) for vec in result.vectors]
115 def count_tokens(self, text: str) -> int:
116 """Cloud SDK backends expose no local tokenizer, so chunk sizing falls back
117 to a character estimate (see LilbeeTokenizerBackend)."""
118 raise NotImplementedError("SDK backends have no local tokenizer for chunk sizing")
120 def count_chat_prompt_tokens(
121 self,
122 messages: list[ChatMessage],
123 *,
124 options: dict[str, Any] | None = None,
125 model: str | None = None,
126 tools: list[dict[str, Any]] | None = None,
127 tool_choice: str | dict[str, Any] | None = None,
128 ) -> int:
129 """Cloud SDK backends render no prompt locally, so an exact count is unavailable."""
130 raise NotImplementedError("SDK backends have no tokenizer for chat prompt counts")
132 @overload
133 def chat(
134 self,
135 messages: list[dict[str, str]],
136 *,
137 stream: Literal[False] = False,
138 options: dict[str, Any] | None = None,
139 model: str | None = None,
140 tools: list[dict[str, Any]] | None = None,
141 tool_choice: str | dict[str, Any] | None = None,
142 ) -> ChatResult: ...
144 @overload
145 def chat(
146 self,
147 messages: list[dict[str, str]],
148 *,
149 stream: Literal[True],
150 options: dict[str, Any] | None = None,
151 model: str | None = None,
152 tools: list[dict[str, Any]] | None = None,
153 tool_choice: str | dict[str, Any] | None = None,
154 ) -> ClosableIterator[ChatStreamItem]: ...
156 def chat(
157 self,
158 messages: list[dict[str, str]],
159 *,
160 stream: bool = False,
161 options: dict[str, Any] | None = None,
162 model: str | None = None,
163 tools: list[dict[str, Any]] | None = None,
164 tool_choice: str | dict[str, Any] | None = None,
165 ) -> ChatResult | ClosableIterator[ChatStreamItem]:
166 """Chat completion via the configured backend.
168 Non-streaming returns a :class:`ChatResult` carrying the assistant
169 text, any tool-call frames the model emitted, and a finish reason.
170 Streaming yields :data:`ChatStreamItem` frames (text tokens and
171 tool-call deltas).
172 """
173 self._ensure_initialized()
174 ref = parse_model_ref(
175 require_role_ref(model or cfg.chat_model, WorkerRole.CHAT, provider="sdk")
176 )
177 if tools and not self.supports_tools(model or cfg.chat_model):
178 chosen = model or cfg.chat_model
179 raise ProviderError(
180 f"Model {chosen!r} does not support tool calls. Pick a different "
181 f"chat model that advertises tool support, or remove tools from "
182 f"the request.",
183 provider=self._backend.provider_name,
184 )
185 translated = translate_options(options, ref) if options else {}
186 if tools is not None:
187 translated["tools"] = tools
188 if tool_choice is not None:
189 translated["tool_choice"] = tool_choice
190 request = CompletionRequest(
191 ref=ref,
192 messages=list(messages),
193 options=translated,
194 api_base=_api_base_for(ref),
195 api_key=_api_key_for(ref),
196 )
197 if stream:
198 return self._chat_stream(request)
199 try:
200 result = self._backend.complete(request)
201 except ProviderError:
202 raise
203 except Exception as exc:
204 raise ProviderError(
205 f"Chat failed: {exc}", provider=self._backend.provider_name
206 ) from exc
207 return ChatResult(
208 text=result.content,
209 tool_calls=tuple(
210 ToolCall(id=tc.id, name=tc.name, arguments=tc.arguments) for tc in result.tool_calls
211 ),
212 finish_reason=FinishReason.coerce(result.finish_reason),
213 usage=result.usage,
214 )
216 def supports_tools(self, model_ref: str) -> bool:
217 """Delegate to the backend's ``supports_tools`` probe."""
218 return self._backend.supports_tools(model_ref)
220 def chat_with_tools(
221 self,
222 messages: list[ChatMessage],
223 *,
224 tools: list[dict[str, Any]],
225 tool_choice: str | dict[str, Any] | None = None,
226 options: dict[str, Any] | None = None,
227 model: str | None = None,
228 ) -> ChatToolResult:
229 """Tool-calling chat for remote/SDK models.
231 The base stub raises, but this backend advertises tool support and ``chat``
232 already forwards tools/tool_choice, so route through it instead of refusing.
233 """
234 result = self.chat(
235 # Pass each message through whole: a tool conversation carries
236 # ``tool_calls`` / ``tool_call_id`` / ``name`` that link an assistant
237 # call to its result, and stripping to role+content breaks that chain.
238 [dict(m) for m in messages],
239 stream=False,
240 options=options,
241 model=model,
242 tools=tools,
243 tool_choice=tool_choice,
244 )
245 return ChatToolResult(content=result.text, tool_calls=list(result.tool_calls))
247 def _chat_stream(self, request: CompletionRequest) -> ClosableIterator[ChatStreamItem]:
248 """Yield content tokens and tool-call deltas from a streaming completion.
250 Exceptions surfaced by the backend at either call time or during
251 iteration are re-raised as ``ProviderError`` so callers always
252 see a consistent error type.
253 """
254 try:
255 stream = self._backend.complete_stream(request)
256 for chunk in stream:
257 if chunk.content:
258 yield chunk.content
259 for delta in chunk.tool_call_deltas:
260 yield ToolCallDelta(
261 index=delta.index,
262 id=delta.id,
263 name=delta.name,
264 arguments_delta=delta.arguments_delta,
265 )
266 if chunk.finish_reason is not None:
267 # The closing chunk's finish_reason lets the dispatch report
268 # length/stop, matching the non-streaming path.
269 yield StreamFinish(reason=FinishReason.coerce(chunk.finish_reason))
270 if chunk.usage is not None:
271 yield chunk.usage
272 except ProviderError:
273 raise
274 except Exception as exc:
275 raise ProviderError(
276 f"Chat failed: {exc}", provider=self._backend.provider_name
277 ) from exc
279 def vision_ocr(
280 self,
281 png_bytes: bytes,
282 model: str,
283 prompt: str = "",
284 *,
285 timeout: float | None = None,
286 cancel: CancelSignal | None = None,
287 ) -> str:
288 """OCR via a multipart chat completion; ``timeout`` enforced on a daemon thread."""
289 from lilbee.vision import build_vision_messages, resolve_ocr_prompt
291 if cancel is not None and cancel.is_set():
292 raise TaskCancelledError
294 messages = build_vision_messages(prompt or resolve_ocr_prompt(model), png_bytes)
295 if timeout and timeout > 0:
296 # The caller is freed at the deadline; a wedged call's thread lives until
297 # the backend httpx timeout, and exit does not wait for it.
298 call = DaemonCall(
299 lambda: self.chat(messages, stream=False, model=model), name="vision-ocr"
300 )
301 if not call.wait(timeout):
302 raise TimeoutError(_VISION_OCR_TIMEOUT_ERROR.format(timeout=timeout))
303 result = call.result()
304 else:
305 result = self.chat(messages, stream=False, model=model)
306 if not isinstance(result, ChatResult):
307 raise ProviderError(
308 f"Vision OCR returned non-text response ({type(result).__name__}).",
309 provider=self._backend.provider_name,
310 )
311 return result.text
313 def vision_slot_capacity(self) -> int | None:
314 """Hosted backends have no local OCR slots; the caller estimates."""
315 return None
317 def list_models(self) -> list[str]:
318 """List models across every configured local server.
320 A single unreachable server is logged and skipped so its outage does not
321 drop the models served by the other reachable servers.
322 """
323 names: list[str] = []
324 for spec, base_url in configured_local_servers():
325 try:
326 names.extend(
327 self._backend.list_models(spec=spec, base_url=base_url, api_key=cfg.llm_api_key)
328 )
329 except NotImplementedError:
330 continue
331 except Exception as exc:
332 log.debug("Skipping unreachable local server %s: %s", spec.key, exc)
333 return names
335 def list_chat_models(self, provider: str) -> list[str]:
336 """List frontier chat models known to the backend for *provider*.
338 Initializes the backend first so ``cfg.json_mode`` suppression is
339 applied before the SDK import inside the backend runs.
340 """
341 self._ensure_initialized()
342 try:
343 return self._backend.list_chat_models(provider)
344 except NotImplementedError:
345 return []
346 except ProviderError:
347 raise
348 except Exception as exc:
349 raise ProviderError(
350 f"Listing chat models failed: {exc}", provider=self._backend.provider_name
351 ) from exc
353 def pull_model(self, model: str, *, on_progress: Callable[..., Any] | None = None) -> None:
354 """Pull a model via the backend."""
355 try:
356 base_url = _api_base_for(parse_model_ref(model)) or ""
357 self._backend.pull_model(model, base_url=base_url, on_progress=on_progress)
358 except NotImplementedError as exc:
359 raise ProviderError(
360 f"Cannot pull model {model!r}: backend does not support pulling",
361 provider=self._backend.provider_name,
362 ) from exc
363 except ProviderError:
364 raise
365 except Exception as exc:
366 raise ProviderError(
367 f"Cannot pull model {model!r}: {exc}", provider=self._backend.provider_name
368 ) from exc
370 def show_model(self, model: str) -> dict[str, Any] | None:
371 """Return model metadata, or None when unsupported or not found."""
372 try:
373 base_url = _api_base_for(parse_model_ref(model)) or ""
374 return self._backend.show_model(model, base_url=base_url)
375 except NotImplementedError:
376 return None
377 except ProviderError:
378 raise
379 except Exception as exc:
380 raise ProviderError(
381 f"Showing model {model!r} failed: {exc}", provider=self._backend.provider_name
382 ) from exc
384 def get_capabilities(self, model: str) -> list[str]:
385 """Return capability tags from ``show_model`` output, or ``[]``."""
386 info = self.show_model(model)
387 if info is None:
388 return []
389 caps = info.get("capabilities", [])
390 return caps if isinstance(caps, list) else []
392 def rerank(self, query: str, candidates: list[str]) -> list[float]:
393 """Rerank candidates via the SDK backend using ``cfg.reranker_model``."""
394 if not candidates:
395 return []
396 self._ensure_initialized()
397 ref = parse_model_ref(
398 require_role_ref(cfg.reranker_model, WorkerRole.RERANK, provider="sdk")
399 )
400 request = RerankRequest(
401 ref=ref,
402 query=query,
403 candidates=candidates,
404 api_base=_api_base_for(ref),
405 api_key=_api_key_for(ref),
406 )
407 try:
408 result = self._backend.rerank(request)
409 except ProviderError:
410 raise
411 except Exception as exc:
412 raise ProviderError(
413 f"Rerank failed: {exc}", provider=self._backend.provider_name
414 ) from exc
415 return result.scores
417 def supports_rerank(self) -> bool:
418 """SDK-backed rerank is available when the underlying SDK is importable."""
419 return self._backend.available()
421 def available(self) -> bool:
422 """Return True when the configured SDK backend can service catalog calls."""
423 return self._backend.available()
425 def shutdown(self) -> None:
426 """SDK-backed providers hold no lilbee-side resources."""