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

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

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable, Iterator 

6from pathlib import Path 

7from typing import Any, Protocol, TypeVar, runtime_checkable 

8 

9from pydantic import BaseModel 

10 

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

12 

13 

14@runtime_checkable 

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

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

17 

18 Streaming chat responses use this to guarantee the upstream model lock 

19 is released even when callers truncate the stream before exhaustion. 

20 Generators satisfy this implicitly; explicit wrappers (e.g. the llama-cpp 

21 chat-lock iterator) implement it directly. 

22 """ 

23 

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

25 

26 

27class LLMOptions(BaseModel): 

28 """Validated options passed to LLM providers. 

29 Only these fields are forwarded — everything else is rejected 

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

31 """ 

32 

33 temperature: float | None = None 

34 top_p: float | None = None 

35 top_k: int | None = None 

36 seed: int | None = None 

37 num_predict: int | None = None 

38 repeat_penalty: float | None = None 

39 num_ctx: int | None = None 

40 

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

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

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

44 

45 

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

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

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

49 

50 

51class ProviderError(Exception): 

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

53 

54 def __init__(self, message: str, *, provider: str = "") -> None: 

55 self.provider = provider 

56 super().__init__(message) 

57 

58 

59ChatMessage = dict[str, str] 

60 

61 

62class LLMProvider(Protocol): 

63 """Protocol for pluggable LLM backends.""" 

64 

65 def embed(self, texts: list[str]) -> list[list[float]]: 

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

67 ... 

68 

69 def chat( 

70 self, 

71 messages: list[ChatMessage], 

72 *, 

73 stream: bool = False, 

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

75 model: str | None = None, 

76 ) -> str | ClosableIterator[str]: 

77 """Chat completion. Returns str for non-stream, ClosableIterator[str] for stream.""" 

78 ... 

79 

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

81 """List available model identifiers.""" 

82 ... 

83 

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

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

86 

87 Returns ``[]`` when the backend has no static catalog (for 

88 example, native llama-cpp has no notion of external API catalogs). 

89 """ 

90 ... 

91 

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

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

94 ... 

95 

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

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

98 ... 

99 

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

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

102 

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

104 reporting or the model is not found. 

105 """ 

106 ... 

107 

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

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

110 

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

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

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

114 

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

116 Empty ``candidates`` returns ``[]``. 

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

118 reranking or ``cfg.reranker_model`` is empty. 

119 """ 

120 ... 

121 

122 def supports_rerank(self) -> bool: 

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

124 

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

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

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

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

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

130 non-empty value. 

131 """ 

132 return False 

133 

134 def shutdown(self) -> None: 

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

136 ... 

137 

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

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

140 return