Coverage for src/local_deep_research/llm/llm_registry.py: 100%
71 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""Registry for custom LangChain LLMs.
3This module provides a registry for registering and managing custom LangChain
4LLMs that can be used with Local Deep Research.
6Entries are keyed per user. Built-in providers (auto-discovered at import via
7``discover_providers()``) and SDK/MCP/benchmark registrations that carry no
8username live in a shared namespace and resolve for everyone. A registration
9made with a ``username`` lives in that user's namespace and is resolved for
10that user first, so one user's custom LLM can neither shadow a built-in
11provider for anyone else nor leak into another user's registry listing.
12"""
14import threading
15from typing import Callable, Dict, Optional, Union
17from langchain.chat_models.base import BaseChatModel
18from loguru import logger
20# Sentinel key for the shared/global namespace (see module docstring).
21_SHARED_NAMESPACE: Optional[str] = None
23_LLMType = Union[BaseChatModel, Callable[..., BaseChatModel]]
26class LLMRegistry:
27 """Thread-safe, per-user registry for custom LangChain LLMs."""
29 def __init__(self):
30 # {namespace: {normalized_name: llm}}. ``namespace`` is a username
31 # or ``_SHARED_NAMESPACE`` for shared/global entries (built-ins +
32 # username-less registrations).
33 self._llms: Dict[Optional[str], Dict[str, _LLMType]] = {}
34 self._lock = threading.Lock()
36 @staticmethod
37 def _ns(username: Optional[str]) -> Optional[str]:
38 """Normalize a username to a namespace key.
40 A falsy username maps to the shared namespace, so legacy/global
41 callers keep a single shared namespace and never crash for want of
42 a username.
43 """
44 return username or _SHARED_NAMESPACE
46 def register(
47 self,
48 name: str,
49 llm: _LLMType,
50 username: Optional[str] = None,
51 ) -> None:
52 """Register a custom LLM.
54 Args:
55 name: Unique name for the LLM (case-insensitive)
56 llm: Either a BaseChatModel instance or a factory function that returns one
57 username: Owner of the registration. When omitted the entry is
58 stored in the shared namespace (visible to all callers).
60 Raises:
61 ValueError: ``llm`` is None. A stored ``None`` would silently
62 miss in ``get()``'s ``found is not None`` check and fall
63 through to the shared namespace, letting a misregistered
64 entry resolve to an unrelated (possibly built-in) provider.
65 """
66 if llm is None:
67 raise ValueError(f"Cannot register LLM '{name}': value is None")
68 ns = self._ns(username)
69 with self._lock:
70 # Normalize name to lowercase for case-insensitive storage
71 normalized_name = name.lower()
72 bucket = self._llms.setdefault(ns, {})
73 if normalized_name in bucket:
74 logger.warning(f"Overwriting existing LLM: {name}")
75 bucket[normalized_name] = llm
76 logger.info(
77 f"Registered custom LLM: {name} (normalized: {normalized_name})"
78 )
80 # Completes the CRUD API surface for the registry.
81 # Used in tests to verify cleanup behavior.
82 def unregister(self, name: str, username: Optional[str] = None) -> None:
83 """Unregister a custom LLM.
85 Args:
86 name: Name of the LLM to unregister (case-insensitive)
87 username: Owner namespace to remove from (None = shared)
88 """
89 ns = self._ns(username)
90 with self._lock:
91 normalized_name = name.lower()
92 bucket = self._llms.get(ns)
93 if bucket and normalized_name in bucket:
94 del bucket[normalized_name]
95 logger.info(f"Unregistered custom LLM: {name}")
97 def get(
98 self, name: str, username: Optional[str] = None
99 ) -> Optional[_LLMType]:
100 """Get a registered LLM.
102 Resolution order: the caller's own namespace, then the shared
103 namespace (which holds built-in providers). A user's registration
104 never resolves for another user.
106 Args:
107 name: Name of the LLM to retrieve (case-insensitive)
108 username: Requesting user (None resolves the shared namespace only)
110 Returns:
111 The LLM instance/factory or None if not found
112 """
113 ns = self._ns(username)
114 with self._lock:
115 normalized_name = name.lower()
116 if ns is not _SHARED_NAMESPACE:
117 found = self._llms.get(ns, {}).get(normalized_name)
118 if found is not None:
119 return found
120 return self._llms.get(_SHARED_NAMESPACE, {}).get(normalized_name)
122 def is_registered(self, name: str, username: Optional[str] = None) -> bool:
123 """Check if an LLM is registered and resolvable for the caller.
125 Args:
126 name: Name to check (case-insensitive)
127 username: Requesting user (None resolves the shared namespace only)
129 Returns:
130 True if registered, False otherwise
131 """
132 return self.get(name, username=username) is not None
134 # Used in test assertions to verify registry state;
135 # part of public API for plugin authors.
136 def list_registered(self, username: Optional[str] = None) -> list[str]:
137 """Get list of LLM names visible to the caller.
139 Returns the caller's own registrations plus shared ones (built-in
140 providers); never another user's names.
142 Order is DETERMINISTIC: the caller's own entries first (insertion
143 order), then shared entries not shadowed by an own-namespace name.
144 This mirrors ``get()``'s own-namespace-first resolution (a same-named
145 own entry shadows the shared built-in, so it is listed once, under
146 own) and replaces the previous ``set``-union, whose ``list(set(...))``
147 order was nondeterministic run to run.
149 Args:
150 username: Requesting user (None lists the shared namespace only)
152 Returns:
153 List of registered LLM names
154 """
155 ns = self._ns(username)
156 with self._lock:
157 shared = self._llms.get(_SHARED_NAMESPACE, {})
158 if ns is _SHARED_NAMESPACE:
159 return list(shared.keys())
160 own = self._llms.get(ns, {})
161 return list(own.keys()) + [
162 name for name in shared if name not in own
163 ]
165 # Used in 7+ test files' autouse fixtures for test isolation
166 # (64+ tests depend on this to reset global state between runs).
167 def clear(self, username: Optional[str] = None) -> None:
168 """Clear registered LLMs.
170 With no ``username`` every namespace is cleared — including the
171 shared built-in providers — matching the historic reset behavior
172 the test-isolation fixtures depend on. With a ``username`` only
173 that user's namespace is cleared.
174 """
175 # TODO: hook clear(username=<deleted user>) into a user-deletion
176 # flow so a removed user's registered LLMs are evicted from this
177 # in-process registry. No user-deletion hook exists yet, so
178 # eviction-on-deletion is out of scope here; this is the seam.
179 with self._lock:
180 if username is None:
181 self._llms.clear()
182 logger.info("Cleared all registered custom LLMs")
183 else:
184 self._llms.pop(self._ns(username), None)
185 logger.info("Cleared registered custom LLMs for one user")
188# Global registry instance
189_llm_registry = LLMRegistry()
192# Public API functions
193def register_llm(
194 name: str,
195 llm: _LLMType,
196 username: Optional[str] = None,
197) -> None:
198 """Register a custom LLM in the registry.
200 Args:
201 name: Unique name for the LLM
202 llm: Either a BaseChatModel instance or a factory function
203 username: Owner of the registration (None = shared namespace)
204 """
205 _llm_registry.register(name, llm, username=username)
208def unregister_llm(name: str, username: Optional[str] = None) -> None:
209 """Unregister a custom LLM from the registry.
211 Args:
212 name: Name of the LLM to unregister
213 username: Owner namespace to remove from (None = shared)
214 """
215 _llm_registry.unregister(name, username=username)
218def get_llm_from_registry(
219 name: str,
220 username: Optional[str] = None,
221) -> Optional[_LLMType]:
222 """Get a registered LLM from the registry.
224 Args:
225 name: Name of the LLM to retrieve
226 username: Requesting user (None resolves the shared namespace only)
228 Returns:
229 The LLM instance/factory or None if not found
230 """
231 return _llm_registry.get(name, username=username)
234def is_llm_registered(name: str, username: Optional[str] = None) -> bool:
235 """Check if an LLM is registered in the registry.
237 Args:
238 name: Name to check
239 username: Requesting user (None resolves the shared namespace only)
241 Returns:
242 True if registered, False otherwise
243 """
244 return _llm_registry.is_registered(name, username=username)
247def list_registered_llms(username: Optional[str] = None) -> list[str]:
248 """Get list of registered LLM names visible to the caller.
250 Args:
251 username: Requesting user (None lists the shared namespace only)
253 Returns:
254 List of registered LLM names
255 """
256 return _llm_registry.list_registered(username=username)
259def clear_llm_registry(username: Optional[str] = None) -> None:
260 """Clear registered LLMs (all namespaces when no username is given)."""
261 _llm_registry.clear(username=username)