1use std::collections::HashMap;
40use std::path::{Path, PathBuf};
41use std::sync::OnceLock;
42
43use compact_str::CompactString;
44use serde::Deserialize;
45use tracing::warn;
46use vtcode_commons::VtCodePaths;
47use vtcode_commons::tokens::estimate_tokens;
48
49use super::resource_cache::{ResourceCache, canonical_cache_path, fingerprint_markdown_directories};
50
51const EXAMPLES_DIR: &str = "examples";
52const PROMPTS_PARENT: &str = ".vtcode/prompts";
53const MIN_TOKEN_BUDGET: usize = 16;
54
55fn user_examples_dir(home: &Path) -> PathBuf {
56 if dirs::home_dir().as_deref() == Some(home)
57 && let Ok(paths) = VtCodePaths::resolve()
58 && let Ok(prompt_dir) = paths.config_path("prompts")
59 {
60 return prompt_dir.join(EXAMPLES_DIR);
61 }
62 home.join(PROMPTS_PARENT).join(EXAMPLES_DIR)
63}
64
65pub const DEFAULT_FEW_SHOT_BUDGET_TOKENS: usize = 800;
69
70pub const FEW_SHOT_SECTION_HEADER: &str = "[Few-Shot Examples]";
73
74#[derive(Debug, Clone, PartialEq, Eq)]
81pub struct FewShotExample {
82 pub id: String,
83 pub tags: Vec<String>,
84 pub summary: String,
85 pub body: String,
86 pub token_count: usize,
87 pub source_path: PathBuf,
88}
89
90#[derive(Debug, Clone, Default, Deserialize)]
93struct FewShotFrontmatter {
94 #[serde(default)]
95 id: Option<String>,
96 #[serde(default)]
97 tags: Option<Vec<String>>,
98 #[serde(default)]
99 summary: Option<String>,
100}
101
102#[derive(Debug, Clone, Default)]
109pub struct FewShotStore {
110 examples: Vec<FewShotExample>,
111 normalized_tags: Vec<Vec<CompactString>>,
112}
113
114impl FewShotStore {
115 pub fn load(workspace_root: Option<&Path>, home_dir: Option<&Path>) -> Self {
123 let key = FewShotCacheKey::new(workspace_root, home_dir);
124 load_cached_sync(key, workspace_root.map(Path::to_path_buf), home_dir.map(Path::to_path_buf))
125 }
126
127 pub async fn load_async(workspace_root: Option<&Path>, home_dir: Option<&Path>) -> Self {
131 let workspace_root = workspace_root.map(Path::to_path_buf);
132 let home_dir = home_dir.map(Path::to_path_buf);
133 match tokio::task::spawn_blocking(move || {
134 let key = FewShotCacheKey::new(workspace_root.as_deref(), home_dir.as_deref());
137 load_cached_sync(key, workspace_root, home_dir)
138 })
139 .await
140 {
141 Ok(store) => store,
142 Err(error) => {
143 warn!(%error, "few-shot resource load task failed");
144 Self::default()
145 }
146 }
147 }
148
149 pub fn from_examples(examples: Vec<FewShotExample>) -> Self {
152 let mut sorted = examples;
153 sorted.sort_by(|a, b| a.id.cmp(&b.id));
154 let normalized_tags = sorted.iter().map(|example| normalize_tags(&example.tags)).collect();
155 Self { examples: sorted, normalized_tags }
156 }
157
158 pub fn len(&self) -> usize {
160 self.examples.len()
161 }
162
163 pub fn is_empty(&self) -> bool {
164 self.examples.is_empty()
165 }
166
167 pub fn iter(&self) -> impl Iterator<Item = &FewShotExample> {
169 self.examples.iter()
170 }
171
172 pub fn select(&self, query: &str, budget_tokens: usize) -> Vec<&FewShotExample> {
184 if self.examples.is_empty() || budget_tokens < MIN_TOKEN_BUDGET {
185 return Vec::new();
186 }
187
188 let query_words = tokenize(query);
189 if query_words.is_empty() {
190 return Vec::new();
191 }
192 let query_lower = query.to_ascii_lowercase();
193
194 let mut scored: Vec<(f64, &FewShotExample)> = self
195 .examples
196 .iter()
197 .zip(&self.normalized_tags)
198 .filter_map(|(example, normalized_tags)| {
199 let score = score_example(normalized_tags, &query_words, &query_lower);
200 if score <= 0.0 {
201 return None;
202 }
203 Some((score, example))
204 })
205 .collect();
206
207 scored.sort_by(|a, b| {
209 b.0.partial_cmp(&a.0)
210 .unwrap_or(std::cmp::Ordering::Equal)
211 .then_with(|| a.1.id.cmp(&b.1.id))
212 });
213
214 let mut chosen: Vec<&FewShotExample> = Vec::new();
215 let mut used = 0usize;
216 for (_, example) in scored {
217 if example.token_count == 0 {
218 continue;
220 }
221 if used.saturating_add(example.token_count) > budget_tokens {
222 continue;
223 }
224 chosen.push(example);
225 used = used.saturating_add(example.token_count);
226 }
227
228 chosen
229 }
230}
231
232pub fn render_few_shot_section(examples: &[&FewShotExample]) -> String {
236 if examples.is_empty() {
237 return String::new();
238 }
239
240 let mut out = String::new();
241 out.push_str(FEW_SHOT_SECTION_HEADER);
242 out.push_str("\nThe following examples illustrate expected behavior for similar tasks.\n\n");
243 for example in examples {
244 let _ = std::fmt::Write::write_fmt(&mut out, format_args!("### {}\n", example.id));
245 if !example.summary.is_empty() {
246 let _ = std::fmt::Write::write_fmt(&mut out, format_args!("_{}_\n\n", example.summary));
247 }
248 let _ = std::fmt::Write::write_fmt(&mut out, format_args!("{}\n\n", example.body.trim_end()));
249 }
250 out
251}
252
253#[derive(Debug, Clone, PartialEq, Eq, Hash)]
258struct FewShotCacheKey {
259 workspace_examples_dir: Option<PathBuf>,
260 home_examples_dir: Option<PathBuf>,
261}
262
263impl FewShotCacheKey {
264 fn new(workspace_root: Option<&Path>, home_dir: Option<&Path>) -> Self {
265 Self {
266 workspace_examples_dir: workspace_root
267 .map(|root| canonical_cache_path(&root.join(PROMPTS_PARENT).join(EXAMPLES_DIR))),
268 home_examples_dir: home_dir.map(|home| canonical_cache_path(&user_examples_dir(home))),
269 }
270 }
271}
272
273static FEW_SHOT_CACHE: OnceLock<ResourceCache<FewShotCacheKey, FewShotStore>> = OnceLock::new();
274
275fn few_shot_cache() -> &'static ResourceCache<FewShotCacheKey, FewShotStore> {
276 FEW_SHOT_CACHE.get_or_init(ResourceCache::default)
277}
278
279fn load_cached_sync(key: FewShotCacheKey, workspace_root: Option<PathBuf>, home_dir: Option<PathBuf>) -> FewShotStore {
280 if let Some(cached) = few_shot_cache().fast_get(&key) {
281 return (*cached).clone();
282 }
283
284 few_shot_cache().with_load_gate(|| {
285 if let Some(cached) = few_shot_cache().fast_get(&key) {
286 return (*cached).clone();
287 }
288
289 let source_dirs = key
290 .workspace_examples_dir
291 .iter()
292 .chain(key.home_examples_dir.iter())
293 .map(PathBuf::as_path)
294 .collect::<Vec<_>>();
295 let fingerprint = fingerprint_markdown_directories(&source_dirs);
296 if let Some(cached) = few_shot_cache().get_if_unchanged(&key, &fingerprint) {
297 return (*cached).clone();
298 }
299
300 let store = load_uncached(workspace_root.as_deref(), home_dir.as_deref());
301 few_shot_cache().insert(key, store.clone(), fingerprint);
302 store
303 })
304}
305
306fn load_uncached(workspace_root: Option<&Path>, home_dir: Option<&Path>) -> FewShotStore {
307 let mut by_id: HashMap<String, FewShotExample> = HashMap::new();
308
309 if let Some(workspace) = workspace_root {
311 merge_from_dir(&mut by_id, &workspace.join(PROMPTS_PARENT).join(EXAMPLES_DIR));
312 }
313
314 if let Some(home) = home_dir {
315 merge_from_dir(&mut by_id, &user_examples_dir(home));
316 }
317
318 FewShotStore::from_examples(by_id.into_values().collect())
319}
320
321fn merge_from_dir(by_id: &mut HashMap<String, FewShotExample>, dir: &Path) {
322 let entries = match std::fs::read_dir(dir) {
323 Ok(entries) => entries,
324 Err(err) if err.kind() == std::io::ErrorKind::NotFound => return,
325 Err(err) => {
326 warn!("few_shot: failed to read directory {}: {err}", dir.display());
327 return;
328 }
329 };
330
331 let mut paths: Vec<PathBuf> = entries
332 .filter_map(Result::ok)
333 .map(|entry| entry.path())
334 .filter(|path| {
335 path.is_file()
336 && path
337 .extension()
338 .and_then(|ext| ext.to_str())
339 .map(|ext| ext.eq_ignore_ascii_case("md"))
340 .unwrap_or(false)
341 })
342 .collect();
343 paths.sort();
344
345 for path in paths {
346 let Some(stem) = path
347 .file_stem()
348 .and_then(|s| s.to_str())
349 .map(str::trim)
350 .filter(|s| !s.is_empty())
351 else {
352 continue;
353 };
354
355 let Ok(raw) = std::fs::read_to_string(&path) else {
356 warn!("few_shot: could not read {}", path.display());
357 continue;
358 };
359
360 if let Some(existing) = by_id.get(stem) {
361 warn!(
362 "few_shot: duplicate id '{stem}' ({}); keeping earlier {}",
363 path.display(),
364 existing.source_path.display()
365 );
366 continue;
367 }
368
369 match parse_example(stem, &raw, &path) {
370 Ok(example) => {
371 by_id.insert(stem.to_string(), example);
372 }
373 Err(err) => {
374 warn!("few_shot: failed to parse {}: {err}", path.display());
375 }
376 }
377 }
378}
379
380fn parse_example(fallback_id: &str, raw: &str, source_path: &Path) -> Result<FewShotExample, String> {
381 let (frontmatter, body) = parse_frontmatter(raw);
382 let id = frontmatter
383 .id
384 .map(|value| value.trim().to_string())
385 .filter(|value| !value.is_empty())
386 .unwrap_or_else(|| fallback_id.to_string());
387 let tags = frontmatter.tags.unwrap_or_default();
388 let summary = frontmatter.summary.map(|value| value.trim().to_string()).unwrap_or_default();
389 let body = body.trim().to_string();
390 let token_count = estimate_tokens(&body);
391
392 Ok(FewShotExample {
393 id,
394 tags,
395 summary,
396 body,
397 token_count,
398 source_path: source_path.to_path_buf(),
399 })
400}
401
402fn parse_frontmatter(raw: &str) -> (FewShotFrontmatter, String) {
403 let trimmed = raw.trim_start_matches('\u{feff}');
404 if !trimmed.starts_with("---\n") {
405 return (FewShotFrontmatter::default(), raw.to_string());
406 }
407 let Some(end_rel) = trimmed[4..].find("\n---") else {
408 return (FewShotFrontmatter::default(), raw.to_string());
409 };
410 let yaml = &trimmed[4..4 + end_rel];
411 let body_start = 4 + end_rel + "\n---".len();
412 let body = trimmed[body_start..].trim_start_matches('\n').to_string();
413
414 let frontmatter = match serde_saphyr::from_str::<FewShotFrontmatter>(yaml) {
415 Ok(value) => value,
416 Err(err) => {
417 warn!("few_shot: frontmatter parse failed: {err}");
418 FewShotFrontmatter::default()
419 }
420 };
421
422 (frontmatter, body)
423}
424
425fn tokenize(query: &str) -> Vec<String> {
426 query
427 .split(|c: char| !c.is_alphanumeric() && c != '_' && c != '-')
428 .filter(|token| !token.is_empty())
429 .map(|token| token.to_ascii_lowercase())
430 .collect()
431}
432
433fn normalize_tags(tags: &[String]) -> Vec<CompactString> {
434 tags.iter().map(|tag| CompactString::from(tag.to_ascii_lowercase())).collect()
435}
436
437fn score_example(normalized_tags: &[CompactString], query_words: &[String], query_lower: &str) -> f64 {
438 let mut score = 0.0;
439 for tag in normalized_tags {
440 if query_words.iter().any(|word| word == tag.as_str()) {
441 score += 1.0;
442 continue;
443 }
444 if query_lower.contains(tag.as_str()) {
445 score += 0.5;
446 }
447 }
448 score
449}
450
451#[cfg(test)]
456mod tests {
457 use super::*;
458 use std::path::PathBuf;
459
460 fn example(id: &str, tags: &[&str], body: &str) -> FewShotExample {
461 let token_count = estimate_tokens(body);
462 FewShotExample {
463 id: id.to_string(),
464 tags: tags.iter().map(|s| s.to_string()).collect(),
465 summary: String::new(),
466 body: body.to_string(),
467 token_count,
468 source_path: PathBuf::from(format!("/tmp/{id}.md")),
469 }
470 }
471
472 #[test]
473 fn empty_store_returns_empty_selection() {
474 let store = FewShotStore::default();
475 assert!(store.select("anything", 1000).is_empty());
476 }
477
478 #[test]
479 fn zero_budget_returns_empty() {
480 let store = FewShotStore::from_examples(vec![example("edit", &["edit"], "do an edit")]);
481 assert!(store.select("edit a file", 0).is_empty());
482 assert!(store.select("edit a file", 8).is_empty());
483 }
484
485 #[test]
486 fn matching_tag_is_selected() {
487 let store = FewShotStore::from_examples(vec![
488 example("edit", &["edit"], "an edit example"),
489 example("search", &["search"], "a search example"),
490 ]);
491 let chosen = store.select("please edit foo.rs", 1000);
492 assert_eq!(chosen.len(), 1);
493 assert_eq!(chosen[0].id, "edit");
494 }
495
496 #[test]
497 fn no_overlap_returns_empty() {
498 let store = FewShotStore::from_examples(vec![example("deploy", &["deploy", "kubernetes"], "a deploy example")]);
499 let chosen = store.select("read a local file", 1000);
500 assert!(chosen.is_empty());
501 }
502
503 #[test]
504 fn example_exceeding_budget_is_skipped() {
505 let store = FewShotStore::from_examples(vec![
507 example("a", &["foo"], &"foo ".repeat(50)),
508 example("b", &["foo"], &"foo ".repeat(50)),
509 ]);
510 let chosen = store.select("foo bar", 60);
511 assert_eq!(chosen.len(), 1);
512 assert_eq!(chosen[0].id, "a"); }
514
515 #[test]
516 fn multiple_matches_sorted_by_score_then_id() {
517 let store = FewShotStore::from_examples(vec![
518 example("git-commit", &["git", "commit"], "git commit example"),
520 example("git-only", &["git"], "git only example"),
522 example("zzz-other", &["git"], "git other example"),
524 ]);
525 let chosen = store.select("git commit message", 5000);
526 assert_eq!(chosen[0].id, "git-commit");
528 let rest: Vec<&str> = chosen.iter().skip(1).map(|e| e.id.as_str()).collect();
530 assert_eq!(rest, vec!["git-only", "zzz-other"]);
531 }
532
533 #[test]
534 fn parse_frontmatter_extracts_metadata() {
535 let raw = "---\nid: read-then-edit\ntags: [read, edit]\nsummary: read first\n---\nbody here\n";
536 let (fm, body) = parse_frontmatter(raw);
537 assert_eq!(fm.id.as_deref(), Some("read-then-edit"));
538 assert_eq!(fm.tags, Some(vec!["read".to_string(), "edit".to_string()]));
539 assert_eq!(fm.summary.as_deref(), Some("read first"));
540 assert_eq!(body, "body here\n");
543 }
544
545 #[test]
546 fn parse_frontmatter_without_metadata_returns_body() {
547 let raw = "no frontmatter here\nbody text\n";
548 let (fm, body) = parse_frontmatter(raw);
549 assert!(fm.id.is_none());
550 assert_eq!(body, raw);
551 }
552
553 #[test]
554 fn render_section_is_empty_when_no_examples() {
555 assert!(render_few_shot_section(&[]).is_empty());
556 }
557
558 #[test]
559 fn render_section_includes_id_and_body() {
560 let store =
561 FewShotStore::from_examples(vec![example("demo", &["hi", "demo"], "user said hi\nassistant said hello")]);
562 let chosen = store.select("hi there", 1000);
563 let rendered = render_few_shot_section(&chosen);
564 assert!(rendered.starts_with(FEW_SHOT_SECTION_HEADER));
565 assert!(rendered.contains("[Few-Shot Examples]\nThe following examples illustrate"));
566 assert!(rendered.contains("### demo"));
567 assert!(rendered.contains("user said hi"));
568 }
569
570 #[test]
571 fn load_discovers_examples_from_disk() {
572 let workspace = tempfile::tempdir().expect("workspace");
573 let examples_dir = workspace.path().join(".vtcode/prompts/examples");
574 std::fs::create_dir_all(&examples_dir).expect("mkdir");
575 std::fs::write(examples_dir.join("demo.md"), "---\ntags: [demo]\nsummary: a demo\n---\nDemo body\n")
576 .expect("write");
577
578 let store = FewShotStore::load(Some(workspace.path()), None);
579 assert_eq!(store.len(), 1);
580 let only = &store.examples[0];
581 assert_eq!(only.id, "demo");
582 assert_eq!(only.tags, vec!["demo".to_string()]);
583 assert_eq!(only.summary, "a demo");
584 assert!(only.body.contains("Demo body"));
585 assert!(only.token_count > 0);
586 }
587
588 #[test]
589 fn load_silently_skips_missing_directories() {
590 let store = FewShotStore::load(Some(Path::new("/does/not/exist/anywhere")), None);
592 assert!(store.is_empty());
593 }
594
595 #[test]
596 #[serial_test::serial]
597 fn load_preserves_workspace_precedence_and_refreshes_edits() {
598 few_shot_cache().clear();
599 let workspace = tempfile::tempdir().expect("workspace");
600 let home = tempfile::tempdir().expect("home");
601 let workspace_dir = workspace.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
602 let home_dir = home.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
603 std::fs::create_dir_all(&workspace_dir).expect("workspace examples");
604 std::fs::create_dir_all(&home_dir).expect("home examples");
605 std::fs::write(home_dir.join("review.md"), "---\ntags: [review]\n---\nhome\n").expect("home example");
606 std::fs::write(workspace_dir.join("review.md"), "---\ntags: [review]\n---\nworkspace\n")
607 .expect("workspace example");
608
609 let first = FewShotStore::load(Some(workspace.path()), Some(home.path()));
610 assert_eq!(first.examples[0].body, "workspace");
611
612 std::fs::write(workspace_dir.join("review.md"), "---\ntags: [review]\n---\nupdated\n").expect("edit example");
613 few_shot_cache().force_metadata_poll();
614 let second = FewShotStore::load(Some(workspace.path()), Some(home.path()));
615 assert_eq!(second.examples[0].body, "updated");
616 }
617
618 #[test]
619 #[serial_test::serial]
620 fn cache_detects_additions_deletions_and_missing_sources() {
621 few_shot_cache().clear();
622 let workspace = tempfile::tempdir().expect("workspace");
623 let examples_dir = workspace.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
624
625 assert!(FewShotStore::load(Some(workspace.path()), None).is_empty());
626 std::fs::create_dir_all(&examples_dir).expect("examples directory");
627 std::fs::write(examples_dir.join("one.md"), "---\ntags: [one]\n---\none\n").expect("one example");
628 few_shot_cache().force_metadata_poll();
629 assert_eq!(FewShotStore::load(Some(workspace.path()), None).len(), 1);
630
631 std::fs::remove_file(examples_dir.join("one.md")).expect("remove example");
632 few_shot_cache().force_metadata_poll();
633 assert!(FewShotStore::load(Some(workspace.path()), None).is_empty());
634 }
635
636 #[test]
637 #[serial_test::serial]
638 fn expired_cache_entry_reparses_even_when_metadata_is_unchanged() {
639 few_shot_cache().clear();
640 let workspace = tempfile::tempdir().expect("workspace");
641 let examples_dir = workspace.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
642 std::fs::create_dir_all(&examples_dir).expect("examples directory");
643 std::fs::write(examples_dir.join("ttl.md"), "---\ntags: [ttl]\n---\nold\n").expect("old example");
644
645 let first = FewShotStore::load(Some(workspace.path()), None);
646 assert_eq!(first.examples[0].body, "old");
647 std::fs::write(examples_dir.join("ttl.md"), "---\ntags: [ttl]\n---\nnew\n").expect("new example");
648 few_shot_cache().force_expiration();
649 let second = FewShotStore::load(Some(workspace.path()), None);
650 assert_eq!(second.examples[0].body, "new");
651 }
652
653 #[test]
654 #[serial_test::serial]
655 fn cache_hit_preserves_rendered_few_shot_prompt() {
656 few_shot_cache().clear();
657 let workspace = tempfile::tempdir().expect("workspace");
658 let examples_dir = workspace.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
659 std::fs::create_dir_all(&examples_dir).expect("examples directory");
660 std::fs::write(
661 examples_dir.join("render.md"),
662 "---\ntags: [render]\nsummary: Stable output\n---\nRender this example.\n",
663 )
664 .expect("render example");
665
666 let cold = FewShotStore::load(Some(workspace.path()), None);
667 let cold_selected = cold.select("render", DEFAULT_FEW_SHOT_BUDGET_TOKENS);
668 let cold_rendered = render_few_shot_section(&cold_selected);
669
670 let warm = FewShotStore::load(Some(workspace.path()), None);
671 let warm_selected = warm.select("render", DEFAULT_FEW_SHOT_BUDGET_TOKENS);
672 let warm_rendered = render_few_shot_section(&warm_selected);
673 assert_eq!(cold_rendered, warm_rendered);
674 }
675
676 #[tokio::test(flavor = "current_thread")]
677 #[serial_test::serial]
678 async fn concurrent_cache_misses_share_the_parsed_store() {
679 few_shot_cache().clear();
680 let workspace = tempfile::tempdir().expect("workspace");
681 let examples_dir = workspace.path().join(PROMPTS_PARENT).join(EXAMPLES_DIR);
682 std::fs::create_dir_all(&examples_dir).expect("examples directory");
683 std::fs::write(examples_dir.join("concurrent.md"), "---\ntags: [concurrent]\n---\nbody\n").expect("example");
684
685 let (first, second) = tokio::join!(
686 FewShotStore::load_async(Some(workspace.path()), None),
687 FewShotStore::load_async(Some(workspace.path()), None),
688 );
689 assert_eq!(first.examples, second.examples);
690 }
691}