skm_select/
llm.rs

1//! LLM-based selection strategy (s latency).
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7
8use skm_core::SkillMetadata;
9
10use crate::error::{LlmError, SelectError};
11use crate::strategy::{
12    Confidence, LatencyClass, SelectionContext, SelectionResult, SelectionStrategy,
13};
14
15/// LLM client trait — users implement for their LLM provider.
16///
17/// Both methods are async (all LLM calls are network IO).
18#[async_trait]
19pub trait LlmClient: Send + Sync {
20    /// Send a prompt, receive a text response.
21    async fn complete(&self, prompt: &str, max_tokens: usize) -> Result<String, LlmError>;
22
23    /// Send a prompt with structured output (JSON mode).
24    /// Default implementation calls complete() and parses JSON.
25    async fn complete_structured(
26        &self,
27        prompt: &str,
28        _schema: &serde_json::Value,
29        max_tokens: usize,
30    ) -> Result<serde_json::Value, LlmError> {
31        let text = self.complete(prompt, max_tokens).await?;
32        serde_json::from_str(&text).map_err(|e| LlmError::ParseError(e.to_string()))
33    }
34}
35
36/// Configuration for LLM strategy.
37#[derive(Debug, Clone)]
38pub struct LlmStrategyConfig {
39    /// System prompt for the LLM.
40    pub system_prompt: String,
41
42    /// Include few-shot examples from skill metadata.
43    pub include_few_shot: bool,
44
45    /// Max skills to include in prompt.
46    pub max_candidates: usize,
47
48    /// Temperature for generation.
49    pub temperature: f32,
50}
51
52impl Default for LlmStrategyConfig {
53    fn default() -> Self {
54        Self {
55            system_prompt: DEFAULT_SYSTEM_PROMPT.to_string(),
56            include_few_shot: true,
57            max_candidates: 20,
58            temperature: 0.0,
59        }
60    }
61}
62
63const DEFAULT_SYSTEM_PROMPT: &str = r#"You are a skill selector. Given a user query and a list of available skills, determine which skill(s) should handle the query.
64
65Respond ONLY with valid JSON in this format:
66{"skills": [{"name": "skill-name", "confidence": 0.0-1.0, "reasoning": "brief explanation"}]}
67
68If no skill is appropriate, return: {"skills": []}
69"#;
70
71/// LLM response structure.
72#[derive(Debug, Deserialize)]
73struct LlmResponse {
74    skills: Vec<SkillSelection>,
75}
76
77#[derive(Debug, Deserialize)]
78struct SkillSelection {
79    name: String,
80    confidence: f32,
81    #[serde(default)]
82    reasoning: Option<String>,
83}
84
85/// LLM-based intent classification for ambiguous queries.
86pub struct LlmStrategy {
87    client: Arc<dyn LlmClient>,
88    config: LlmStrategyConfig,
89}
90
91impl LlmStrategy {
92    /// Create a new LLM strategy.
93    pub fn new(client: Arc<dyn LlmClient>, config: LlmStrategyConfig) -> Self {
94        Self { client, config }
95    }
96
97    /// Build the prompt for the LLM.
98    fn build_prompt(&self, query: &str, candidates: &[&SkillMetadata]) -> String {
99        let mut prompt = self.config.system_prompt.clone();
100        prompt.push_str("\n\nAvailable skills:\n");
101
102        for (i, skill) in candidates.iter().take(self.config.max_candidates).enumerate() {
103            prompt.push_str(&format!("{}. {}: {}\n", i + 1, skill.name, skill.description));
104        }
105
106        prompt.push_str(&format!("\nUser query: \"{}\"\n", query));
107        prompt.push_str("\nWhich skill(s) should handle this query? Respond with JSON:");
108
109        prompt
110    }
111
112    /// Parse the LLM response.
113    fn parse_response(
114        &self,
115        response: &str,
116        candidates: &[&SkillMetadata],
117    ) -> Result<Vec<SelectionResult>, SelectError> {
118        // Try to find JSON in the response
119        let json_str = if let Some(start) = response.find('{') {
120            if let Some(end) = response.rfind('}') {
121                &response[start..=end]
122            } else {
123                response
124            }
125        } else {
126            response
127        };
128
129        let parsed: LlmResponse = serde_json::from_str(json_str)
130            .map_err(|e| SelectError::Llm(LlmError::ParseError(e.to_string())))?;
131
132        // Validate and convert
133        let candidate_names: std::collections::HashSet<_> =
134            candidates.iter().map(|c| c.name.as_str()).collect();
135
136        let results: Vec<SelectionResult> = parsed
137            .skills
138            .into_iter()
139            .filter(|s| candidate_names.contains(s.name.as_str()))
140            .map(|s| {
141                let confidence = Confidence::from_score(s.confidence);
142                let skill_name = skm_core::SkillName::new(&s.name).unwrap_or_else(|_| {
143                    skm_core::SkillName::new("unknown").unwrap()
144                });
145                
146                let mut result = SelectionResult::new(skill_name, s.confidence, confidence, "llm");
147                if let Some(reasoning) = s.reasoning {
148                    result = result.with_reasoning(reasoning);
149                }
150                result
151            })
152            .collect();
153
154        Ok(results)
155    }
156}
157
158#[async_trait]
159impl SelectionStrategy for LlmStrategy {
160    async fn select(
161        &self,
162        query: &str,
163        candidates: &[&SkillMetadata],
164        _ctx: &SelectionContext,
165    ) -> Result<Vec<SelectionResult>, SelectError> {
166        if candidates.is_empty() {
167            return Ok(Vec::new());
168        }
169
170        let prompt = self.build_prompt(query, candidates);
171        let response = self.client.complete(&prompt, 500).await?;
172        self.parse_response(&response, candidates)
173    }
174
175    fn name(&self) -> &str {
176        "llm"
177    }
178
179    fn latency_class(&self) -> LatencyClass {
180        LatencyClass::Seconds
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use std::path::PathBuf;
188
189    struct MockLlmClient {
190        response: String,
191    }
192
193    impl MockLlmClient {
194        fn new(response: &str) -> Self {
195            Self {
196                response: response.to_string(),
197            }
198        }
199    }
200
201    #[async_trait]
202    impl LlmClient for MockLlmClient {
203        async fn complete(&self, _prompt: &str, _max_tokens: usize) -> Result<String, LlmError> {
204            Ok(self.response.clone())
205        }
206    }
207
208    fn make_metadata(name: &str, description: &str) -> SkillMetadata {
209        SkillMetadata {
210            name: skm_core::SkillName::new(name).unwrap(),
211            description: description.to_string(),
212            tags: Vec::new(),
213            triggers: Vec::new(),
214            source_path: PathBuf::new(),
215            content_hash: 0,
216            estimated_tokens: 100,
217        }
218    }
219
220    #[tokio::test]
221    async fn test_llm_strategy() {
222        let response = r#"{"skills": [{"name": "pdf-skill", "confidence": 0.9, "reasoning": "Query mentions PDF"}]}"#;
223        let client = Arc::new(MockLlmClient::new(response));
224
225        let strategy = LlmStrategy::new(client, LlmStrategyConfig::default());
226
227        let skills = vec![
228            make_metadata("pdf-skill", "Extract text from PDFs"),
229            make_metadata("weather-skill", "Get weather info"),
230        ];
231        let refs: Vec<_> = skills.iter().collect();
232        let ctx = SelectionContext::new();
233
234        let results = strategy.select("extract pdf text", &refs, &ctx).await.unwrap();
235
236        assert_eq!(results.len(), 1);
237        assert_eq!(results[0].skill.as_str(), "pdf-skill");
238        assert_eq!(results[0].score, 0.9);
239        assert!(results[0].reasoning.is_some());
240    }
241
242    #[tokio::test]
243    async fn test_llm_strategy_no_match() {
244        let response = r#"{"skills": []}"#;
245        let client = Arc::new(MockLlmClient::new(response));
246
247        let strategy = LlmStrategy::new(client, LlmStrategyConfig::default());
248
249        let skills = vec![make_metadata("pdf-skill", "Extract text from PDFs")];
250        let refs: Vec<_> = skills.iter().collect();
251        let ctx = SelectionContext::new();
252
253        let results = strategy.select("play music", &refs, &ctx).await.unwrap();
254
255        assert!(results.is_empty());
256    }
257
258    #[tokio::test]
259    async fn test_llm_strategy_invalid_skill() {
260        let response = r#"{"skills": [{"name": "nonexistent-skill", "confidence": 0.9}]}"#;
261        let client = Arc::new(MockLlmClient::new(response));
262
263        let strategy = LlmStrategy::new(client, LlmStrategyConfig::default());
264
265        let skills = vec![make_metadata("pdf-skill", "Extract text from PDFs")];
266        let refs: Vec<_> = skills.iter().collect();
267        let ctx = SelectionContext::new();
268
269        let results = strategy.select("query", &refs, &ctx).await.unwrap();
270
271        // Should filter out the nonexistent skill
272        assert!(results.is_empty());
273    }
274
275    #[test]
276    fn test_build_prompt() {
277        let client = Arc::new(MockLlmClient::new(""));
278        let strategy = LlmStrategy::new(client, LlmStrategyConfig::default());
279
280        let skills = vec![
281            make_metadata("pdf-skill", "Extract text from PDFs"),
282            make_metadata("weather-skill", "Get weather info"),
283        ];
284        let refs: Vec<_> = skills.iter().collect();
285
286        let prompt = strategy.build_prompt("test query", &refs);
287
288        assert!(prompt.contains("pdf-skill"));
289        assert!(prompt.contains("weather-skill"));
290        assert!(prompt.contains("test query"));
291    }
292}
skm_select/llm.rs

skm_select/
llm.rs
