skm_select/
semantic.rs

1//! Semantic embedding-based selection strategy (ms latency).
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use tokio::sync::RwLock;
7
8use skm_core::SkillMetadata;
9use skm_embed::{ComponentWeights, EmbeddingIndex, EmbeddingProvider};
10
11use crate::error::SelectError;
12use crate::strategy::{
13    Confidence, LatencyClass, SelectionContext, SelectionResult, SelectionStrategy,
14};
15
16/// Configuration for semantic selection.
17#[derive(Debug, Clone)]
18pub struct SemanticConfig {
19    /// Maximum candidates to return.
20    pub top_k: usize,
21
22    /// Minimum similarity threshold.
23    pub min_score: f32,
24
25    /// Score gap for adaptive cutoff.
26    pub gap_threshold: f32,
27
28    /// Component weights for multi-component embeddings.
29    pub component_weights: ComponentWeights,
30
31    /// Enable gap-based adaptive selection.
32    pub use_adaptive_k: bool,
33}
34
35impl Default for SemanticConfig {
36    fn default() -> Self {
37        Self {
38            top_k: 5,
39            min_score: 0.3,
40            gap_threshold: 0.15,
41            component_weights: ComponentWeights::default(),
42            use_adaptive_k: true,
43        }
44    }
45}
46
47impl SemanticConfig {
48    /// Create with custom top_k.
49    pub fn with_top_k(mut self, top_k: usize) -> Self {
50        self.top_k = top_k;
51        self
52    }
53
54    /// Create with custom minimum score.
55    pub fn with_min_score(mut self, min_score: f32) -> Self {
56        self.min_score = min_score;
57        self
58    }
59
60    /// Create with custom gap threshold.
61    pub fn with_gap_threshold(mut self, gap: f32) -> Self {
62        self.gap_threshold = gap;
63        self
64    }
65
66    /// Disable adaptive k (always return top_k).
67    pub fn without_adaptive_k(mut self) -> Self {
68        self.use_adaptive_k = false;
69        self
70    }
71}
72
73/// Embedding-based semantic similarity search.
74pub struct SemanticStrategy {
75    provider: Arc<dyn EmbeddingProvider>,
76    index: Arc<RwLock<EmbeddingIndex>>,
77    config: SemanticConfig,
78}
79
80impl SemanticStrategy {
81    /// Create a new semantic strategy.
82    pub fn new(
83        provider: Arc<dyn EmbeddingProvider>,
84        index: EmbeddingIndex,
85        config: SemanticConfig,
86    ) -> Self {
87        Self {
88            provider,
89            index: Arc::new(RwLock::new(index)),
90            config,
91        }
92    }
93
94    /// Update the index.
95    pub async fn update_index(&self, index: EmbeddingIndex) {
96        let mut guard = self.index.write().await;
97        *guard = index;
98    }
99
100    /// Get the current index.
101    pub async fn index(&self) -> EmbeddingIndex {
102        self.index.read().await.clone()
103    }
104}
105
106#[async_trait]
107impl SelectionStrategy for SemanticStrategy {
108    async fn select(
109        &self,
110        query: &str,
111        candidates: &[&SkillMetadata],
112        _ctx: &SelectionContext,
113    ) -> Result<Vec<SelectionResult>, SelectError> {
114        // Embed the query
115        let query_embedding = self.provider.embed_one(query).await?;
116
117        // Get index
118        let index = self.index.read().await;
119
120        // Query the index
121        let scored = if self.config.use_adaptive_k {
122            index.query_adaptive(
123                &query_embedding,
124                self.config.min_score,
125                self.config.top_k,
126                self.config.gap_threshold,
127            )
128        } else {
129            index.query(&query_embedding, self.config.top_k)
130        };
131
132        // Filter to only candidates
133        let candidate_names: std::collections::HashSet<_> =
134            candidates.iter().map(|c| &c.name).collect();
135
136        let results: Vec<SelectionResult> = scored
137            .into_iter()
138            .filter(|s| candidate_names.contains(&s.name))
139            .filter(|s| s.score >= self.config.min_score)
140            .map(|s| {
141                let confidence = Confidence::from_score(s.score);
142                SelectionResult::new(s.name, s.score, confidence, "semantic").with_reasoning(
143                    format!(
144                        "desc={:.2}, trig={:.2}, tags={:.2}, ex={:.2}",
145                        s.component_scores.description,
146                        s.component_scores.triggers,
147                        s.component_scores.tags,
148                        s.component_scores.examples
149                    ),
150                )
151            })
152            .collect();
153
154        Ok(results)
155    }
156
157    fn name(&self) -> &str {
158        "semantic"
159    }
160
161    fn latency_class(&self) -> LatencyClass {
162        LatencyClass::Milliseconds
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn test_semantic_config() {
172        let config = SemanticConfig::default()
173            .with_top_k(3)
174            .with_min_score(0.5)
175            .with_gap_threshold(0.2)
176            .without_adaptive_k();
177
178        assert_eq!(config.top_k, 3);
179        assert_eq!(config.min_score, 0.5);
180        assert_eq!(config.gap_threshold, 0.2);
181        assert!(!config.use_adaptive_k);
182    }
183
184    #[test]
185    fn test_semantic_config_defaults() {
186        let config = SemanticConfig::default();
187        assert_eq!(config.top_k, 5);
188        assert_eq!(config.min_score, 0.3);
189        assert!(config.use_adaptive_k);
190    }
191}
skm_select/semantic.rs

skm_select/
semantic.rs
