khive_runtime/runtime/engine_binding.rs
1//! Ordered configured provider binding and programmatic registration.
2
3use super::*;
4
5impl KhiveRuntime {
6 /// Return the name of the default embedding model (empty string if none configured).
7 pub fn default_embedder_name(&self) -> &str {
8 self.default_embedder_name.as_ref()
9 }
10
11 /// Resolve a model name (or `None` for the default) to an `EmbeddingModel`.
12 ///
13 /// This compatibility API only represents built-in lattice models.
14 /// Use named embedding APIs for custom providers.
15 /// Returns `UnknownModel` if the name is not in the registry, or
16 /// `Unconfigured` if `None` is passed and no default model is set.
17 pub fn resolve_embedding_model(&self, name: Option<&str>) -> RuntimeResult<EmbeddingModel> {
18 let model = match name {
19 Some(raw) => parse_embedding_model_alias(raw)
20 .ok_or_else(|| crate::RuntimeError::UnknownModel(raw.to_string()))?,
21 None => self
22 .config
23 .embedding_model
24 .ok_or_else(|| crate::RuntimeError::Unconfigured("embedding_model".into()))?,
25 };
26 let key = model.to_string();
27 if request_excludes_embedder(&key) {
28 return Err(crate::RuntimeError::UnknownModel(
29 name.unwrap_or_else(|| self.default_embedder_name())
30 .to_string(),
31 ));
32 }
33 let contains = self
34 .embedder_registry
35 .read()
36 .map(|reg| reg.contains(&key))
37 .unwrap_or(false);
38 if contains {
39 Ok(model)
40 } else {
41 Err(crate::RuntimeError::UnknownModel(
42 name.unwrap_or_else(|| self.default_embedder_name())
43 .to_string(),
44 ))
45 }
46 }
47
48 /// Names of all registered embedding models in this runtime.
49 ///
50 /// Includes both built-in lattice models and any custom embedders
51 /// registered by packs via [`register_embedder`](Self::register_embedder).
52 /// Useful for operations that must touch every model's storage (e.g.,
53 /// scoped vector deletion on note delete). The default model is included.
54 pub fn registered_embedding_model_names(&self) -> Vec<String> {
55 self.embedder_registry
56 .read()
57 .map(|reg| {
58 reg.names()
59 .into_iter()
60 .filter(|name| !request_excludes_embedder(name))
61 .collect()
62 })
63 .unwrap_or_default()
64 }
65
66 /// Get the lazily-initialized embedding service for the named model.
67 ///
68 /// Accepts both built-in lattice model names (e.g. `"all-minilm-l6-v2"`,
69 /// `"paraphrase"`) and custom provider names registered via
70 /// [`register_embedder`](Self::register_embedder).
71 ///
72 /// For lattice model names, aliases (e.g. `"paraphrase"`) are resolved to
73 /// their canonical key before looking up the registry. For custom providers
74 /// the name must match exactly as supplied during registration.
75 ///
76 /// First call for any name loads the underlying service (cold start cost);
77 /// subsequent calls are cheap (registry caches the `Arc`).
78 pub async fn embedder(&self, name: &str) -> RuntimeResult<Arc<dyn EmbeddingService>> {
79 Ok(self.embedder_inner(name, None).await?.0)
80 }
81
82 pub(crate) async fn embedder_with_token(
83 &self,
84 token: &NamespaceToken,
85 name: &str,
86 ) -> RuntimeResult<Arc<dyn EmbeddingService>> {
87 Ok(self.embedder_inner(name, Some(token)).await?.0)
88 }
89
90 /// Resolve the service and its document-preparation attestation from the
91 /// same registry entry. Cloning it also avoids holding a registry lock
92 /// across the asynchronous service resolution.
93 pub(crate) async fn embedder_with_input_attestation(
94 &self,
95 name: &str,
96 token: Option<&NamespaceToken>,
97 ) -> RuntimeResult<(Arc<dyn EmbeddingService>, bool)> {
98 self.embedder_inner(name, token).await
99 }
100
101 /// Register a custom embedding provider with this runtime.
102 ///
103 /// The provider is added to the shared [`EmbedderRegistry`] so all clones
104 /// of this runtime see the new provider immediately. Setup may replace a
105 /// provider before it is selected for resolution. Later duplicates are
106 /// refused and logged; use [`try_register_embedder`](Self::try_register_embedder)
107 /// when the caller needs to handle registration errors.
108 ///
109 /// Packs should call this from [`crate::PackRuntime::register_embedders`] (the
110 /// hook is invoked by the transport during pack initialisation, before the
111 /// first verb dispatch).
112 ///
113 /// [`EmbedderRegistry`]: crate::embedder_registry::EmbedderRegistry
114 pub fn register_embedder(
115 &self,
116 provider: impl crate::embedder_registry::EmbedderProvider + 'static,
117 ) {
118 if let Err(error) = self.try_register_embedder(provider) {
119 tracing::warn!(target: "khive_runtime::runtime", %error, "embedder registration refused");
120 }
121 }
122
123 /// Register a custom embedding provider and return serving-duplicate or lock errors.
124 ///
125 /// Unlike [`register_embedder`](Self::register_embedder), this method lets
126 /// callers fail initialization when registration cannot be completed.
127 pub fn try_register_embedder(
128 &self,
129 provider: impl crate::embedder_registry::EmbedderProvider + 'static,
130 ) -> RuntimeResult<()> {
131 if let Some(model) = parse_embedding_model_alias(provider.name()) {
132 if provider.dimensions() != model.dimensions() {
133 return Err(RuntimeError::InvalidInput(format!(
134 "embedding provider `{}` has dimensions {}, but its built-in storage binding requires {}",
135 provider.name(), provider.dimensions(), model.dimensions()
136 )));
137 }
138 }
139 if let Some(engine) = self
140 .config
141 .engines
142 .as_ref()
143 .and_then(|engines| engines.iter().find(|engine| engine.name == provider.name()))
144 {
145 engine
146 .check_dimensions(provider.dimensions())
147 .map_err(|error| RuntimeError::InvalidInput(error.to_string()))?;
148 }
149 let mut registry = self
150 .embedder_registry
151 .write()
152 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?;
153 registry.register(provider)
154 }
155
156 /// Install a deterministic backend for exact-input provenance tests.
157 /// The test adapter, not the supplied backend, owns lattice passage
158 /// prefixing; this API is absent unless `test-internals` is enabled.
159 #[cfg(feature = "test-internals")]
160 pub fn register_test_audited_embedder(
161 &self,
162 model: EmbeddingModel,
163 provider: impl crate::embedder_registry::EmbedderProvider + 'static,
164 ) {
165 self.embedder_registry
166 .write()
167 .expect("test embedder registry lock")
168 .register_test_audited(model, provider);
169 }
170
171 /// Replace the provider under an already-bound engine name in a test host.
172 /// See [`EmbedderRegistry::replace_bound_for_test`] for the refusals; this
173 /// API is absent unless `test-internals` is enabled.
174 ///
175 /// [`EmbedderRegistry::replace_bound_for_test`]: crate::embedder_registry::EmbedderRegistry::replace_bound_for_test
176 #[cfg(feature = "test-internals")]
177 pub fn replace_bound_embedder_for_test(
178 &self,
179 provider: impl crate::embedder_registry::EmbedderProvider + 'static,
180 ) -> RuntimeResult<()> {
181 self.embedder_registry
182 .write()
183 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?
184 .replace_bound_for_test(provider)
185 }
186
187 /// Bind every configured engine after host and pack provider registration.
188 ///
189 /// Low-level hosts call this before serving. Validation is atomic within
190 /// this shared registry and reads provider metadata only: it does not build
191 /// services, access storage, or change historical embedding spaces. A
192 /// successful binding freezes those provider names. Repeating the same
193 /// binding is a no-op; a failed attempt can be retried after registration.
194 /// Explicitly named registered providers remain independently accessible.
195 pub fn finalize_embedding_engines(&self) -> RuntimeResult<()> {
196 self.initialize_embedding_engines(|| {})
197 }
198
199 /// Canonical configured participation order, distinct from all registered names.
200 ///
201 /// Returns `Unconfigured` until the host finalizes binding. An explicitly
202 /// empty configuration binds an empty list and has no default provider.
203 pub fn bound_embedding_engine_names(&self) -> RuntimeResult<Vec<String>> {
204 self.embedder_registry
205 .read()
206 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?
207 .bound_engines()
208 .map(|engines| engines.iter().map(|engine| engine.name.clone()).collect())
209 .ok_or_else(|| RuntimeError::Unconfigured("embedding engine binding".into()))
210 }
211
212 pub(crate) fn embedding_registry_identity(&self) -> usize {
213 Arc::as_ptr(&self.embedder_registry) as usize
214 }
215
216 pub(crate) fn initialize_embedding_engines(
217 &self,
218 register: impl FnOnce(),
219 ) -> RuntimeResult<()> {
220 let initialization = self
221 .embedder_registry
222 .read()
223 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?
224 .initialization_lock();
225 // Registration itself acquires the registry write lock. Serialize the
226 // lifecycle outside that lock so hooks can use the ordinary public API.
227 let _initialization = initialization
228 .lock()
229 .map_err(|_| RuntimeError::Internal("embedder initialization lock poisoned".into()))?;
230 let engines = self.config.configured_engines();
231 let already_bound = self
232 .embedder_registry
233 .read()
234 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?
235 .bound_engines()
236 .is_some();
237 if !already_bound && !engines.is_empty() {
238 register();
239 }
240 self.embedder_registry
241 .write()
242 .map_err(|_| RuntimeError::Internal("embedder registry lock poisoned".into()))?
243 .bind_engines(&engines)
244 }
245}