1use std::cmp::Ordering;
2use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
3
4use serde::{Deserialize, Serialize};
5
6use crate::error::{DataError, Result};
7use crate::ids::{RepresentationId, TypeId};
8use crate::plan::{FitScope, PlanIssue};
9
10#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
11pub struct InputPortSpec {
12 pub name: String,
13 pub accepted_representations: Vec<RepresentationId>,
14 pub accepted_types: Vec<TypeId>,
15 pub rank: Option<usize>,
16 #[serde(default)]
17 pub multi_source: bool,
18 #[serde(default)]
19 pub optional: bool,
20}
21
22impl InputPortSpec {
23 pub fn validate(&self) -> Result<()> {
24 validate_name("input port", &self.name)?;
25 if self.accepted_representations.is_empty() {
26 return Err(DataError::Validation(format!(
27 "input port `{}` accepts no representations",
28 self.name
29 )));
30 }
31 if self.accepted_types.is_empty() {
32 return Err(DataError::Validation(format!(
33 "input port `{}` accepts no types",
34 self.name
35 )));
36 }
37 Ok(())
38 }
39}
40
41#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
42pub struct ModelInputSpec {
43 pub ports: Vec<InputPortSpec>,
44 #[serde(default)]
45 pub default_fusion: Option<serde_json::Value>,
46}
47
48impl ModelInputSpec {
49 pub fn validate(&self) -> Result<()> {
50 if self.ports.is_empty() {
51 return Err(DataError::Validation(
52 "model input spec contains no ports".to_string(),
53 ));
54 }
55 let mut names = BTreeSet::new();
56 for port in &self.ports {
57 port.validate()?;
58 if !names.insert(port.name.as_str()) {
59 return Err(DataError::Validation(format!(
60 "duplicate model input port `{}`",
61 port.name
62 )));
63 }
64 }
65 Ok(())
66 }
67}
68
69#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
70pub struct AdapterSpec {
71 pub id: String,
72 pub version: String,
73 pub input_type: TypeId,
74 pub input_representation: RepresentationId,
75 pub output_type: TypeId,
76 pub output_representation: RepresentationId,
77 pub cost: u64,
78 #[serde(default)]
79 pub lossy: bool,
80 #[serde(default)]
81 pub supervised: bool,
82 #[serde(default)]
83 pub stateful: bool,
84 #[serde(default = "default_true")]
85 pub deterministic: bool,
86 pub fit_scope: FitScope,
87 #[serde(default)]
88 pub params: BTreeMap<String, serde_json::Value>,
89}
90
91fn default_true() -> bool {
92 true
93}
94
95impl AdapterSpec {
96 pub fn validate(&self) -> Result<()> {
97 validate_name("adapter", &self.id)?;
98 validate_name("adapter version", &self.version)?;
99 if !self.deterministic {
100 return Err(DataError::Validation(format!(
101 "adapter `{}` is not deterministic",
102 self.id
103 )));
104 }
105 if self.stateful && self.fit_scope == FitScope::Stateless {
106 return Err(DataError::Validation(format!(
107 "stateful adapter `{}` cannot use stateless fit scope",
108 self.id
109 )));
110 }
111 Ok(())
112 }
113
114 fn source(&self) -> RepresentationNode {
115 RepresentationNode {
116 type_id: self.input_type.clone(),
117 representation_id: self.input_representation.clone(),
118 }
119 }
120
121 fn target(&self) -> RepresentationNode {
122 RepresentationNode {
123 type_id: self.output_type.clone(),
124 representation_id: self.output_representation.clone(),
125 }
126 }
127}
128
129#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
130pub struct PlanningPolicy {
131 #[serde(default)]
132 pub allow_lossy: bool,
133 #[serde(default)]
134 pub allow_stateful: bool,
135 #[serde(default)]
136 pub allow_supervised: bool,
137 #[serde(default)]
138 pub forbidden_adapters: BTreeSet<String>,
139 #[serde(default)]
140 pub preferred_adapters: BTreeSet<String>,
141 #[serde(default = "default_true")]
142 pub require_user_choice_on_ambiguity: bool,
143 pub max_hops: Option<usize>,
144}
145
146impl Default for PlanningPolicy {
147 fn default() -> Self {
148 Self {
149 allow_lossy: false,
150 allow_stateful: false,
151 allow_supervised: false,
152 forbidden_adapters: BTreeSet::new(),
153 preferred_adapters: BTreeSet::new(),
154 require_user_choice_on_ambiguity: true,
155 max_hops: None,
156 }
157 }
158}
159
160#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
161pub struct RepresentationNode {
162 pub type_id: TypeId,
163 pub representation_id: RepresentationId,
164}
165
166#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
167pub struct AdapterPath {
168 pub adapters: Vec<AdapterSpec>,
169 pub total_cost: u64,
170 pub effective_score: u64,
171}
172
173#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
174pub struct PathResolution {
175 pub path: Option<AdapterPath>,
176 #[serde(default)]
177 pub requires_user_choice: bool,
178 #[serde(default)]
179 pub issues: Vec<PlanIssue>,
180}
181
182impl PathResolution {
183 pub fn resolved(path: AdapterPath) -> Self {
184 Self {
185 path: Some(path),
186 requires_user_choice: false,
187 issues: Vec::new(),
188 }
189 }
190
191 pub fn unresolved(code: &str, message: String, choices: Vec<String>) -> Self {
192 Self {
193 path: None,
194 requires_user_choice: !choices.is_empty(),
195 issues: vec![PlanIssue {
196 code: code.to_string(),
197 message,
198 choices,
199 }],
200 }
201 }
202}
203
204#[derive(Clone, Debug, Default)]
205pub struct AdapterRegistry {
206 adapters: BTreeMap<String, AdapterSpec>,
207}
208
209#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
210pub struct AdapterRegistrySpec {
211 #[serde(default)]
212 pub adapters: Vec<AdapterSpec>,
213}
214
215impl AdapterRegistry {
216 pub fn new() -> Self {
217 Self::default()
218 }
219
220 pub fn from_spec(spec: AdapterRegistrySpec) -> Result<Self> {
221 let mut registry = Self::new();
222 for adapter in spec.adapters {
223 registry.register_adapter(adapter)?;
224 }
225 Ok(registry)
226 }
227
228 pub fn register_adapter(&mut self, adapter: AdapterSpec) -> Result<()> {
229 adapter.validate()?;
230 if self.adapters.contains_key(&adapter.id) {
231 return Err(DataError::Validation(format!(
232 "duplicate adapter id `{}`",
233 adapter.id
234 )));
235 }
236 self.adapters.insert(adapter.id.clone(), adapter);
237 Ok(())
238 }
239
240 pub fn adapters(&self) -> impl Iterator<Item = &AdapterSpec> {
241 self.adapters.values()
242 }
243
244 pub fn find_path(
245 &self,
246 source_type: &TypeId,
247 source_representation: &RepresentationId,
248 target_type: &TypeId,
249 target_representation: &RepresentationId,
250 policy: &PlanningPolicy,
251 ) -> PathResolution {
252 let start = RepresentationNode {
253 type_id: source_type.clone(),
254 representation_id: source_representation.clone(),
255 };
256 let goal = RepresentationNode {
257 type_id: target_type.clone(),
258 representation_id: target_representation.clone(),
259 };
260 if start == goal {
261 return PathResolution::resolved(AdapterPath {
262 adapters: Vec::new(),
263 total_cost: 0,
264 effective_score: 0,
265 });
266 }
267
268 let mut edges: BTreeMap<RepresentationNode, Vec<&AdapterSpec>> = BTreeMap::new();
269 for adapter in self.adapters.values() {
270 if policy.forbidden_adapters.contains(&adapter.id) {
271 continue;
272 }
273 if adapter.lossy && !policy.allow_lossy {
274 continue;
275 }
276 if adapter.stateful && !policy.allow_stateful {
277 continue;
278 }
279 if adapter.supervised && !policy.allow_supervised {
280 continue;
281 }
282 edges.entry(adapter.source()).or_default().push(adapter);
283 }
284
285 let mut heap = BinaryHeap::new();
286 heap.push(SearchState {
287 score: 0,
288 raw_cost: 0,
289 hops: 0,
290 node: start.clone(),
291 adapter_ids: Vec::new(),
292 });
293
294 let mut best_seen: BTreeMap<RepresentationNode, (u64, usize)> = BTreeMap::new();
295 best_seen.insert(start.clone(), (0, 0));
296 let mut best_goal: Option<(u64, usize, u64)> = None;
297 let mut goal_paths = Vec::new();
298
299 while let Some(state) = heap.pop() {
300 if let Some((best_score, best_hops, _)) = best_goal {
301 if (state.score, state.hops) > (best_score, best_hops) {
302 break;
303 }
304 }
305 if state.node == goal {
306 best_goal.get_or_insert((state.score, state.hops, state.raw_cost));
307 goal_paths.push(state.adapter_ids);
308 continue;
309 }
310 if policy
311 .max_hops
312 .is_some_and(|max_hops| state.hops >= max_hops)
313 {
314 continue;
315 }
316 let Some(next_edges) = edges.get(&state.node) else {
317 continue;
318 };
319 for adapter in next_edges {
320 if state.adapter_ids.iter().any(|id| id == &adapter.id) {
321 continue;
322 }
323 let next = adapter.target();
324 let score = state.score + adapter_score(adapter, policy);
325 let raw_cost = state.raw_cost + adapter.cost;
326 let hops = state.hops + 1;
327 if best_seen.get(&next).is_some_and(|(best_score, best_hops)| {
328 (score, hops) > (*best_score, *best_hops)
329 }) {
330 continue;
331 }
332 best_seen.insert(next.clone(), (score, hops));
333 let mut adapter_ids = state.adapter_ids.clone();
334 adapter_ids.push(adapter.id.clone());
335 heap.push(SearchState {
336 score,
337 raw_cost,
338 hops,
339 node: next,
340 adapter_ids,
341 });
342 }
343 }
344
345 if goal_paths.is_empty() {
346 return PathResolution::unresolved(
347 "no_path",
348 format!(
349 "no adapter path from `{}/{}` to `{}/{}`",
350 source_type, source_representation, target_type, target_representation
351 ),
352 Vec::new(),
353 );
354 }
355
356 goal_paths.sort();
357 goal_paths.dedup();
358 if goal_paths.len() > 1 && policy.require_user_choice_on_ambiguity {
359 let choices = goal_paths
360 .iter()
361 .map(|path| path.join(" -> "))
362 .collect::<Vec<_>>();
363 return PathResolution::unresolved(
364 "ambiguous_path",
365 "multiple equivalent adapter paths require user choice".to_string(),
366 choices,
367 );
368 }
369
370 let adapter_ids = goal_paths.remove(0);
371 let adapters = adapter_ids
372 .iter()
373 .map(|id| self.adapters.get(id).expect("path adapter exists").clone())
374 .collect::<Vec<_>>();
375 let total_cost = adapters.iter().map(|adapter| adapter.cost).sum();
376 let effective_score = adapters
377 .iter()
378 .map(|adapter| adapter_score(adapter, policy))
379 .sum();
380 PathResolution::resolved(AdapterPath {
381 adapters,
382 total_cost,
383 effective_score,
384 })
385 }
386}
387
388#[derive(Clone, Debug, Eq, PartialEq)]
389struct SearchState {
390 score: u64,
391 raw_cost: u64,
392 hops: usize,
393 node: RepresentationNode,
394 adapter_ids: Vec<String>,
395}
396
397impl Ord for SearchState {
398 fn cmp(&self, other: &Self) -> Ordering {
399 other
400 .score
401 .cmp(&self.score)
402 .then_with(|| other.hops.cmp(&self.hops))
403 .then_with(|| other.raw_cost.cmp(&self.raw_cost))
404 .then_with(|| other.node.cmp(&self.node))
405 .then_with(|| other.adapter_ids.cmp(&self.adapter_ids))
406 }
407}
408
409impl PartialOrd for SearchState {
410 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
411 Some(self.cmp(other))
412 }
413}
414
415fn adapter_score(adapter: &AdapterSpec, policy: &PlanningPolicy) -> u64 {
416 let mut score = adapter.cost.max(1);
417 if adapter.lossy {
418 score += 1_000_000;
419 }
420 if adapter.stateful {
421 score += 100_000;
422 }
423 if adapter.supervised {
424 score += 100_000;
425 }
426 if policy.preferred_adapters.contains(&adapter.id) {
427 score = score.saturating_sub(1);
428 }
429 score
430}
431
432fn validate_name(kind: &str, value: &str) -> Result<()> {
433 if value.trim().is_empty() {
434 return Err(DataError::Validation(format!("{kind} name is empty")));
435 }
436 if !value
437 .bytes()
438 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b':' | b'/'))
439 {
440 return Err(DataError::Validation(format!(
441 "{kind} `{value}` contains unsupported characters"
442 )));
443 }
444 Ok(())
445}
446
447#[cfg(test)]
448mod tests {
449 use super::*;
450
451 fn tid(value: &str) -> TypeId {
452 TypeId::new(value).unwrap()
453 }
454
455 fn rid(value: &str) -> RepresentationId {
456 RepresentationId::new(value).unwrap()
457 }
458
459 fn adapter(id: &str, input: &str, output: &str, cost: u64) -> AdapterSpec {
460 AdapterSpec {
461 id: id.to_string(),
462 version: "0.1.0".to_string(),
463 input_type: tid("dense_signal"),
464 input_representation: rid(input),
465 output_type: if output == "tabular_numeric" {
466 tid("table")
467 } else {
468 tid("dense_signal")
469 },
470 output_representation: rid(output),
471 cost,
472 lossy: false,
473 supervised: false,
474 stateful: false,
475 deterministic: true,
476 fit_scope: FitScope::Stateless,
477 params: BTreeMap::new(),
478 }
479 }
480
481 #[test]
482 fn validates_model_input_ports() {
483 let spec = ModelInputSpec {
484 ports: vec![InputPortSpec {
485 name: "X".to_string(),
486 accepted_representations: vec![rid("tabular_numeric")],
487 accepted_types: vec![tid("table")],
488 rank: Some(2),
489 multi_source: true,
490 optional: false,
491 }],
492 default_fusion: None,
493 };
494
495 assert!(spec.validate().is_ok());
496 }
497
498 #[test]
499 fn rejects_duplicate_adapter_ids() {
500 let mut registry = AdapterRegistry::new();
501 registry
502 .register_adapter(adapter(
503 "spectra.flatten",
504 "signal_1d",
505 "tabular_numeric",
506 1,
507 ))
508 .unwrap();
509
510 assert!(registry
511 .register_adapter(adapter(
512 "spectra.flatten",
513 "signal_1d",
514 "tabular_numeric",
515 1
516 ))
517 .is_err());
518 }
519
520 #[test]
521 fn path_selection_is_registration_order_independent() {
522 let mut left = AdapterRegistry::new();
523 left.register_adapter(adapter("a.to_mid", "signal_1d", "signal_mid", 1))
524 .unwrap();
525 left.register_adapter(adapter("b.to_tabular", "signal_mid", "tabular_numeric", 1))
526 .unwrap();
527 left.register_adapter(adapter("c.direct", "signal_1d", "tabular_numeric", 10))
528 .unwrap();
529
530 let mut right = AdapterRegistry::new();
531 right
532 .register_adapter(adapter("c.direct", "signal_1d", "tabular_numeric", 10))
533 .unwrap();
534 right
535 .register_adapter(adapter("b.to_tabular", "signal_mid", "tabular_numeric", 1))
536 .unwrap();
537 right
538 .register_adapter(adapter("a.to_mid", "signal_1d", "signal_mid", 1))
539 .unwrap();
540
541 let policy = PlanningPolicy::default();
542 let left_path = left
543 .find_path(
544 &tid("dense_signal"),
545 &rid("signal_1d"),
546 &tid("table"),
547 &rid("tabular_numeric"),
548 &policy,
549 )
550 .path
551 .unwrap();
552 let right_path = right
553 .find_path(
554 &tid("dense_signal"),
555 &rid("signal_1d"),
556 &tid("table"),
557 &rid("tabular_numeric"),
558 &policy,
559 )
560 .path
561 .unwrap();
562
563 assert_eq!(
564 left_path
565 .adapters
566 .iter()
567 .map(|adapter| adapter.id.as_str())
568 .collect::<Vec<_>>(),
569 vec!["a.to_mid", "b.to_tabular"]
570 );
571 assert_eq!(left_path, right_path);
572 }
573
574 #[test]
575 fn lossy_paths_are_refused_unless_allowed() {
576 let mut lossy = adapter("image.embedding", "signal_1d", "tabular_numeric", 1);
577 lossy.lossy = true;
578
579 let mut registry = AdapterRegistry::new();
580 registry.register_adapter(lossy).unwrap();
581
582 let refused = registry.find_path(
583 &tid("dense_signal"),
584 &rid("signal_1d"),
585 &tid("table"),
586 &rid("tabular_numeric"),
587 &PlanningPolicy::default(),
588 );
589 assert!(refused.path.is_none());
590
591 let allowed = registry.find_path(
592 &tid("dense_signal"),
593 &rid("signal_1d"),
594 &tid("table"),
595 &rid("tabular_numeric"),
596 &PlanningPolicy {
597 allow_lossy: true,
598 ..PlanningPolicy::default()
599 },
600 );
601 assert_eq!(allowed.path.unwrap().adapters[0].id, "image.embedding");
602 }
603
604 #[test]
605 fn equivalent_best_paths_require_user_choice() {
606 let mut registry = AdapterRegistry::new();
607 registry
608 .register_adapter(adapter("a.flatten", "signal_1d", "tabular_numeric", 1))
609 .unwrap();
610 registry
611 .register_adapter(adapter("b.flatten", "signal_1d", "tabular_numeric", 1))
612 .unwrap();
613
614 let resolution = registry.find_path(
615 &tid("dense_signal"),
616 &rid("signal_1d"),
617 &tid("table"),
618 &rid("tabular_numeric"),
619 &PlanningPolicy::default(),
620 );
621
622 assert!(resolution.path.is_none());
623 assert!(resolution.requires_user_choice);
624 assert_eq!(resolution.issues[0].code, "ambiguous_path");
625 assert_eq!(resolution.issues[0].choices.len(), 2);
626 }
627}