Coverage for src/local_deep_research/web_search_engines/retriever_registry.py: 100%
86 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Registry for dynamically registering LangChain retrievers as search engines.
3"""
5from threading import Lock
6from typing import Dict, Optional
8from langchain_core.retrievers import BaseRetriever
9from ..security.secure_logging import logger
11# Sentinel key for the shared/global namespace. Retrievers registered
12# without a username (SDK, MCP, benchmarks, legacy internal callers) live
13# here and remain visible to everyone, preserving the historic single-
14# namespace behavior. Per-user registrations live under their username so
15# one user can neither shadow nor see another user's retrievers.
16_SHARED_NAMESPACE: Optional[str] = None
19class RetrieverRegistry:
20 """
21 Thread-safe registry for LangChain retrievers.
23 This allows users to register retrievers programmatically and use them
24 as search engines within LDR.
26 Entries are keyed per user: a registration made with a ``username`` is
27 stored in that user's namespace, and a registration made without one
28 goes to a shared namespace. Reads resolve the caller's own namespace
29 first and fall back to the shared namespace, so a user's registration
30 can never shadow or leak into another user's engine list.
31 """
33 def __init__(self):
34 # {namespace: {name: retriever}}. ``namespace`` is a username or
35 # ``_SHARED_NAMESPACE`` for shared/global entries.
36 self._retrievers: Dict[Optional[str], Dict[str, BaseRetriever]] = {}
37 # Parallel {namespace: {name: metadata}} map, e.g.
38 # {"is_local": True}. Kept separate from _retrievers so get()
39 # keeps returning the bare retriever object (callers + tests
40 # depend on `get() is retriever`).
41 self._metadata: Dict[Optional[str], Dict[str, dict]] = {}
42 self._lock = Lock()
44 @staticmethod
45 def _ns(username: Optional[str]) -> Optional[str]:
46 """Normalize a username to a namespace key.
48 A falsy username (``None`` / empty string) maps to the shared
49 namespace so legacy/global callers keep a single shared namespace
50 and never crash for want of a username.
51 """
52 return username or _SHARED_NAMESPACE
54 def register(
55 self,
56 name: str,
57 retriever: BaseRetriever,
58 is_local: bool = True,
59 username: Optional[str] = None,
60 ) -> None:
61 """
62 Register a retriever with a given name.
64 Args:
65 name: Name to register the retriever under
66 retriever: LangChain BaseRetriever instance
67 is_local: Whether this retriever reads local/private data
68 (a user's vector store / KB) rather than the public
69 internet. Defaults to True because programmatically
70 registered retrievers are almost always private corpora;
71 the egress policy uses this to decide whether the
72 retriever may run under PUBLIC_ONLY / PRIVATE_ONLY scopes.
73 username: Owner of the registration. When omitted the entry is
74 stored in the shared namespace (visible to all callers).
76 Raises:
77 ValueError: ``retriever`` is None. A stored ``None`` would
78 silently miss in ``get()``'s ``found is not None`` check and
79 fall through to the shared namespace, letting a
80 misregistered entry resolve to an unrelated retriever.
81 """
82 if retriever is None:
83 raise ValueError(
84 f"Cannot register retriever '{name}': value is None"
85 )
86 ns = self._ns(username)
87 with self._lock:
88 self._retrievers.setdefault(ns, {})[name] = retriever
89 self._metadata.setdefault(ns, {})[name] = {"is_local": is_local}
90 logger.info(
91 f"Registered retriever '{name}' of type {type(retriever).__name__} "
92 f"(is_local={is_local})"
93 )
95 def register_multiple(
96 self,
97 retrievers: Dict[str, "BaseRetriever | dict"],
98 is_local: bool = True,
99 username: Optional[str] = None,
100 ) -> None:
101 """
102 Register multiple retrievers at once.
104 Args:
105 retrievers: Either ``{name: retriever}`` (uses ``is_local``
106 arg as a single default) OR a richer
107 ``{name: {"retriever": …, "is_local": bool}}`` mapping
108 so different retrievers can carry different
109 classifications in one call. Mixing both shapes inside
110 the same dict is allowed.
111 is_local: Default classification when an entry uses the
112 bare-retriever shape. The public programmatic API
113 (api/research_functions.py) calls this; we default to
114 True here ONLY because the old API contract was True.
115 CLI / programmatic callers passing a cloud-backed
116 retriever MUST use the dict shape with
117 ``is_local=False`` or they will egress under
118 PRIVATE_ONLY (silent classification lie).
119 username: Owner of the registrations. When omitted the
120 entries are stored in the shared namespace.
122 Raises:
123 ValueError: an entry resolves to a None retriever (bare shape
124 ``{name: None}`` or dict shape with a missing/None
125 ``"retriever"`` key). A stored ``None`` would silently miss
126 in ``get()``'s ``found is not None`` check and fall through
127 to the shared namespace — see ``register``'s docstring.
128 """
129 ns = self._ns(username)
130 # Two-pass so the batch is atomic: validate every entry (resolving
131 # its retriever + is_local) BEFORE mutating any state, so a None
132 # value raises without leaving a partially-registered batch behind.
133 validated = []
134 for name, entry in retrievers.items():
135 if isinstance(entry, dict):
136 retriever = entry.get("retriever")
137 entry_is_local = bool(entry.get("is_local", is_local))
138 else:
139 retriever = entry
140 entry_is_local = is_local
141 if retriever is None:
142 raise ValueError(
143 f"Cannot register retriever '{name}': value is None"
144 )
145 validated.append((name, retriever, entry_is_local))
146 with self._lock:
147 bucket = self._retrievers.setdefault(ns, {})
148 meta_bucket = self._metadata.setdefault(ns, {})
149 for name, retriever, entry_is_local in validated:
150 bucket[name] = retriever
151 meta_bucket[name] = {"is_local": entry_is_local}
152 logger.info(
153 f"Registered retriever '{name}' of type {type(retriever).__name__} "
154 f"(is_local={entry_is_local})"
155 )
157 def get(
158 self, name: str, username: Optional[str] = None
159 ) -> Optional[BaseRetriever]:
160 """
161 Get a registered retriever by name.
163 Resolution order: the caller's own namespace, then the shared
164 namespace. A user's registration never resolves for another user.
166 Args:
167 name: Name of the retriever
168 username: Requesting user (None resolves the shared namespace only)
170 Returns:
171 The retriever if found, None otherwise
172 """
173 ns = self._ns(username)
174 with self._lock:
175 if ns is not _SHARED_NAMESPACE:
176 found = self._retrievers.get(ns, {}).get(name)
177 if found is not None:
178 return found
179 return self._retrievers.get(_SHARED_NAMESPACE, {}).get(name)
181 def unregister(self, name: str, username: Optional[str] = None) -> None:
182 """
183 Remove a registered retriever from the given user's namespace.
185 Args:
186 name: Name of the retriever to remove
187 username: Owner namespace to remove from (None = shared)
188 """
189 ns = self._ns(username)
190 with self._lock:
191 bucket = self._retrievers.get(ns)
192 if bucket and name in bucket:
193 del bucket[name]
194 self._metadata.get(ns, {}).pop(name, None)
195 logger.info(f"Unregistered retriever '{name}'")
197 def clear(self, username: Optional[str] = None) -> None:
198 """Clear registered retrievers.
200 With no ``username`` every namespace is cleared (used by test
201 isolation fixtures to reset all global state). With a ``username``
202 only that user's namespace is cleared.
203 """
204 # TODO: hook clear(username=<deleted user>) into a user-deletion
205 # flow so a removed user's registered retrievers are evicted from
206 # this in-process registry. No user-deletion hook exists yet, so
207 # eviction-on-deletion is out of scope here; this is the seam.
208 with self._lock:
209 if username is None:
210 count = sum(len(bucket) for bucket in self._retrievers.values())
211 self._retrievers.clear()
212 self._metadata.clear()
213 else:
214 ns = self._ns(username)
215 count = len(self._retrievers.get(ns, {}))
216 self._retrievers.pop(ns, None)
217 self._metadata.pop(ns, None)
218 logger.info(f"Cleared {count} registered retrievers")
220 def is_registered(self, name: str, username: Optional[str] = None) -> bool:
221 """
222 Check if a retriever is registered and resolvable for the caller.
224 Args:
225 name: Name of the retriever
226 username: Requesting user (None resolves the shared namespace only)
228 Returns:
229 True if registered, False otherwise
230 """
231 return self.get(name, username=username) is not None
233 def list_registered(self, username: Optional[str] = None) -> list[str]:
234 """
235 Get list of retriever names visible to the caller.
237 Returns the caller's own registrations plus shared ones; never
238 another user's names.
240 Order is DETERMINISTIC: the caller's own entries first (insertion
241 order), then shared entries not shadowed by an own-namespace name.
242 This mirrors ``get()``'s own-namespace-first resolution (a same-named
243 own entry shadows the shared one, so it is listed once, under own) and
244 replaces the previous ``set``-union, whose ``list(set(...))`` order was
245 nondeterministic run to run.
247 Args:
248 username: Requesting user (None lists the shared namespace only)
250 Returns:
251 List of retriever names
252 """
253 ns = self._ns(username)
254 with self._lock:
255 shared = self._retrievers.get(_SHARED_NAMESPACE, {})
256 if ns is _SHARED_NAMESPACE:
257 return list(shared.keys())
258 own = self._retrievers.get(ns, {})
259 return list(own.keys()) + [
260 name for name in shared if name not in own
261 ]
263 def get_metadata(
264 self, name: str, username: Optional[str] = None
265 ) -> Optional[Dict]:
266 """Return policy-relevant metadata for a registered retriever.
268 Returns a dict like ``{"is_local": True}`` for a registered
269 retriever, or ``None`` when the retriever is unknown. A
270 registered retriever with no recorded classification yields an
271 empty dict, which ``evaluate_retriever`` treats as "unclassified"
272 and fails closed under any non-BOTH scope.
274 The egress policy's ``evaluate_retriever`` consults this hook to
275 decide whether a retriever may run under the active scope.
277 ``username`` scopes the lookup the same way ``get`` does: the
278 caller's own namespace first, then the shared namespace, so a
279 user's classification is never read for another user.
280 """
281 ns = self._ns(username)
282 with self._lock:
283 if ns is not _SHARED_NAMESPACE and name in self._retrievers.get(
284 ns, {}
285 ):
286 # Copy so callers can't mutate the stored classification.
287 return dict(self._metadata.get(ns, {}).get(name, {}))
288 if name in self._retrievers.get(_SHARED_NAMESPACE, {}):
289 return dict(
290 self._metadata.get(_SHARED_NAMESPACE, {}).get(name, {})
291 )
292 return None
295# Global registry instance
296retriever_registry = RetrieverRegistry()