1use crate::knowledge::Polarity;
14use crate::registry::{LiteralType, ObjectKind, PredicateDef};
15use serde::{Deserialize, Serialize};
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct ExtractorConfig {
24 pub model_id: String,
25 pub prompt_version: u32,
26 pub registry_major: u32,
27 pub mechanism: ExtractMechanism,
28 pub max_tokens: u32,
29 #[serde(default)]
33 pub model_digest: Option<String>,
34 #[serde(default)]
44 pub provider_profile_id: Option<String>,
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(rename_all = "snake_case")]
50pub enum ExtractMechanism {
51 JsonSchema,
53 ToolCall,
55 Grammar,
58 JsonMode,
60}
61
62impl ExtractorConfig {
63 pub fn id(&self) -> String {
72 let mut hasher = blake3::Hasher::new();
73 hasher.update(self.model_id.as_bytes());
74 hasher.update(&self.prompt_version.to_le_bytes());
75 hasher.update(&self.registry_major.to_le_bytes());
76 hasher.update(&[self.mechanism as u8]);
77 if let Some(digest) = &self.model_digest {
78 hasher.update(digest.as_bytes());
79 }
80 if let Some(profile_id) = &self.provider_profile_id {
81 hasher.update(profile_id.as_bytes());
82 }
83 hex::encode(hasher.finalize().as_bytes())
84 }
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct MentionRef {
99 pub surface: String,
100 pub entity_type: String,
101 #[serde(default)]
105 pub quote: Option<String>,
106 #[serde(default)]
110 pub span: (u32, u32),
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize)]
115#[serde(tag = "kind", rename_all = "snake_case")]
116pub enum ClaimObject {
117 Entity {
118 mention: MentionRef,
119 },
120 Literal {
121 literal_type: String,
122 value: String,
123 #[serde(default)]
127 quote: Option<String>,
128 #[serde(default)]
129 span: (u32, u32),
130 },
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct Claim {
136 pub predicate: String,
137 pub subject: MentionRef,
138 pub object: ClaimObject,
139 pub polarity: Polarity,
140 #[serde(default)]
142 pub valid_from: Option<i64>,
143 #[serde(default)]
145 pub valid_to: Option<i64>,
146 pub confidence: f32,
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct ExtractionResponse {
152 pub claims: Vec<Claim>,
153}
154
155#[derive(Debug, Clone, Default, Serialize, Deserialize)]
157pub struct ExtractSummary {
158 pub extracted: usize,
159 pub quarantined: usize,
160 pub episodes_done: usize,
161 pub episodes_failed: usize,
162 #[serde(default)]
166 pub failures: Vec<(String, String)>,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
171pub struct ExtractionBudget {
172 pub max_concurrent: usize,
173 pub max_episodes_per_batch: usize,
174 pub max_tokens_per_episode: u32,
175 pub max_repair_attempts: u32,
176 pub lease_timeout_secs: u64,
177}
178
179impl Default for ExtractionBudget {
180 fn default() -> Self {
181 Self {
182 max_concurrent: 4,
183 max_episodes_per_batch: 50,
184 max_tokens_per_episode: 8192,
185 max_repair_attempts: 1,
186 lease_timeout_secs: 300,
187 }
188 }
189}
190
191#[derive(Debug, Clone)]
195pub struct ValidationResult {
196 pub valid: Vec<Claim>,
197 pub invalid: Vec<(Claim, Vec<ValidationError>)>,
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize)]
202#[serde(tag = "kind", rename_all = "snake_case")]
203pub enum ValidationError {
204 UnknownPredicate {
205 predicate: String,
206 },
207 SubjectTypeMismatch {
208 predicate: String,
209 expected: Vec<String>,
210 got: String,
211 },
212 ObjectTypeMismatch {
213 predicate: String,
214 expected: String,
215 got: String,
216 },
217 MalformedLiteral {
218 literal_type: String,
219 value: String,
220 reason: String,
221 },
222 SpanOutOfBounds {
223 span: (u32, u32),
224 content_len: usize,
225 },
226 SurfaceNotVerbatim {
227 surface: String,
228 span: (u32, u32),
229 found: String,
230 },
231 ConfidenceOutOfRange {
232 confidence: f32,
233 },
234}
235
236pub fn validate_claims(
239 claims: &[Claim],
240 content: &str,
241 predicates: &[PredicateDef],
242) -> ValidationResult {
243 let mut valid = Vec::new();
244 let mut invalid = Vec::new();
245
246 for claim in claims {
247 let mut errors = Vec::new();
248
249 if !(0.0..=1.0).contains(&claim.confidence) {
251 errors.push(ValidationError::ConfidenceOutOfRange {
252 confidence: claim.confidence,
253 });
254 }
255
256 let Some(pred_def) = predicates.iter().find(|p| p.name == claim.predicate) else {
258 errors.push(ValidationError::UnknownPredicate {
259 predicate: claim.predicate.clone(),
260 });
261 invalid.push((claim.clone(), errors));
262 continue;
263 };
264
265 if !pred_def
267 .subject_types
268 .iter()
269 .any(|t| t == &claim.subject.entity_type)
270 {
271 errors.push(ValidationError::SubjectTypeMismatch {
272 predicate: claim.predicate.clone(),
273 expected: pred_def
274 .subject_types
275 .iter()
276 .map(|t| t.to_string())
277 .collect(),
278 got: claim.subject.entity_type.clone(),
279 });
280 }
281
282 match (&claim.object, &pred_def.object_kind) {
284 (ClaimObject::Entity { mention }, ObjectKind::Entity(expected)) => {
285 if !expected.0.iter().any(|t| t == &mention.entity_type) {
286 errors.push(ValidationError::ObjectTypeMismatch {
287 predicate: claim.predicate.clone(),
288 expected: expected.0.join("|"),
289 got: mention.entity_type.clone(),
290 });
291 }
292 }
293 (ClaimObject::Literal { literal_type, .. }, ObjectKind::Literal(expected_lt)) => {
294 if !literal_type_matches(literal_type, expected_lt) {
295 errors.push(ValidationError::ObjectTypeMismatch {
296 predicate: claim.predicate.clone(),
297 expected: format!("{expected_lt:?}"),
298 got: literal_type.clone(),
299 });
300 }
301 }
302 (
303 ClaimObject::Literal {
304 literal_type,
305 value,
306 ..
307 },
308 ObjectKind::Enum { variants },
309 ) => {
310 if !variants.iter().any(|v| v == value) {
311 errors.push(ValidationError::ObjectTypeMismatch {
312 predicate: claim.predicate.clone(),
313 expected: format!("enum: {}", variants.join("|")),
314 got: value.clone(),
315 });
316 }
317 let _ = literal_type; }
319 (ClaimObject::Entity { .. }, ObjectKind::Literal(_))
320 | (ClaimObject::Entity { .. }, ObjectKind::Enum { .. })
321 | (ClaimObject::Literal { .. }, ObjectKind::Entity(_)) => {
322 errors.push(ValidationError::ObjectTypeMismatch {
323 predicate: claim.predicate.clone(),
324 expected: format!("{:?}", pred_def.object_kind),
325 got: match &claim.object {
326 ClaimObject::Entity { .. } => "entity".into(),
327 ClaimObject::Literal { .. } => "literal".into(),
328 },
329 });
330 }
331 }
332
333 let mut repaired = claim.clone();
338 if !resolve_mention(&mut repaired.subject, content) {
339 errors.push(ValidationError::SurfaceNotVerbatim {
340 surface: claim.subject.surface.clone(),
341 span: claim.subject.span,
342 found: String::new(),
343 });
344 }
345 if let ClaimObject::Entity { mention } = &mut repaired.object {
346 if !resolve_mention(mention, content) {
347 errors.push(ValidationError::SurfaceNotVerbatim {
348 surface: mention.surface.clone(),
349 span: mention.span,
350 found: String::new(),
351 });
352 }
353 }
354 if let ClaimObject::Literal {
355 quote, value, span, ..
356 } = &mut repaired.object
357 {
358 match quote.as_deref() {
362 Some(q) => match locate_in_quote(content, q, value) {
363 Some((derived, canonical)) => {
364 if canonical != *value {
365 *value = canonical;
366 }
367 *span = derived;
368 }
369 None => errors.push(ValidationError::SurfaceNotVerbatim {
370 surface: value.clone(),
371 span: *span,
372 found: String::new(),
373 }),
374 },
375 None => check_span(span, content, &mut errors),
376 }
377 }
378
379 if errors.is_empty() {
380 valid.push(repaired);
381 } else {
382 invalid.push((claim.clone(), errors));
383 }
384 }
385
386 ValidationResult { valid, invalid }
387}
388
389fn check_span(span: &(u32, u32), content: &str, errors: &mut Vec<ValidationError>) {
390 let len = content.len();
391 if span.0 as usize >= len || span.1 as usize > len || span.0 >= span.1 {
392 errors.push(ValidationError::SpanOutOfBounds {
393 span: *span,
394 content_len: len,
395 });
396 }
397}
398
399fn resolve_mention(m: &mut MentionRef, content: &str) -> bool {
414 if let Some(quote) = m.quote.as_deref() {
416 return match locate_in_quote(content, quote, &m.surface) {
417 Some((span, canonical)) => {
418 m.surface = canonical;
419 m.span = span;
420 true
421 }
422 None => false,
423 };
424 }
425 let (a, b) = (m.span.0 as usize, m.span.1 as usize);
426 if content.get(a..b) == Some(m.surface.as_str()) {
428 return true;
429 }
430 if let Some(range) = char_span_to_bytes(content, a, b) {
432 if content.get(range.clone()) == Some(m.surface.as_str()) {
433 m.span = (range.start as u32, range.end as u32);
434 return true;
435 }
436 }
437 if let Some(found) = content.get(a..b) {
439 if found.eq_ignore_ascii_case(m.surface.as_str()) {
440 m.surface = found.to_string();
441 return true;
442 }
443 }
444 false
445}
446
447fn locate_in_quote(content: &str, quote: &str, needle: &str) -> Option<((u32, u32), String)> {
456 if needle.is_empty() {
457 return None;
458 }
459 let q0 = content.find(quote)?;
460 let window = &content[q0..q0 + quote.len()];
461 if let Some(off) = window.find(needle) {
462 let span = ((q0 + off) as u32, (q0 + off + needle.len()) as u32);
463 return Some((span, needle.to_string()));
464 }
465 let lowered = needle.to_ascii_lowercase();
467 for (off, _) in window.char_indices() {
468 let candidate = &window[off..];
469 if candidate.len() < needle.len() {
470 break;
471 }
472 let Some(head) = candidate.get(..needle.len()) else {
475 continue;
476 };
477 if head.to_ascii_lowercase() == lowered {
478 let found = head.to_string();
479 let span = ((q0 + off) as u32, (q0 + off + needle.len()) as u32);
480 return Some((span, found));
481 }
482 }
483 None
484}
485
486fn char_span_to_bytes(content: &str, a: usize, b: usize) -> Option<std::ops::Range<usize>> {
489 if b < a {
490 return None;
491 }
492 let mut start = None;
493 let mut end = None;
494 let mut idx = 0usize;
495 for (bi, _) in content.char_indices() {
496 if idx == a {
497 start = Some(bi);
498 }
499 if idx == b {
500 end = Some(bi);
501 break;
502 }
503 idx += 1;
504 }
505 let end = end.or((idx == b).then_some(content.len()))?;
507 Some(start?..end)
508}
509
510fn literal_type_matches(given: &str, expected: &LiteralType) -> bool {
511 match expected {
512 LiteralType::Text => given == "text",
513 LiteralType::Date => given == "date",
514 LiteralType::DateTime => given == "datetime",
515 LiteralType::Number => given == "number",
516 LiteralType::Bool => given == "bool",
517 LiteralType::Quantity { .. } => given == "quantity",
518 }
519}
520
521pub fn schema_from_registry(predicates: &[PredicateDef]) -> serde_json::Value {
527 let pred_names: Vec<&str> = predicates.iter().map(|p| p.name.as_str()).collect();
528 let entity_types: Vec<&str> = predicates
529 .iter()
530 .flat_map(|p| {
531 let subjects = p.subject_types.iter().map(|t| t.as_str());
532 let objects = match &p.object_kind {
533 ObjectKind::Entity(types) => types.0.iter().map(|t| t.as_str()).collect(),
534 _ => vec![],
535 };
536 subjects.chain(objects)
537 })
538 .collect::<std::collections::BTreeSet<_>>()
539 .into_iter()
540 .collect();
541
542 let mention_schema = serde_json::json!({
543 "type": "object",
544 "properties": {
545 "surface": { "type": "string", "description": "Verbatim text from the episode" },
546 "entity_type": { "type": "string", "enum": entity_types },
547 "quote": {
548 "type": "string",
549 "description": "A short snippet copied EXACTLY from the episode that contains the surface"
550 }
551 },
552 "required": ["surface", "entity_type", "quote"]
553 });
554
555 let object_schema = serde_json::json!({
556 "type": "object",
557 "properties": {
558 "kind": { "type": "string", "enum": ["entity", "literal"] }
559 },
560 "required": ["kind"],
561 "oneOf": [
562 {
563 "properties": {
564 "kind": { "const": "entity" },
565 "mention": mention_schema.clone()
566 },
567 "required": ["mention"]
568 },
569 {
570 "properties": {
571 "kind": { "const": "literal" },
572 "literal_type": { "type": "string", "enum": ["text", "date", "datetime", "number", "bool", "quantity"] },
573 "value": { "type": "string" },
574 "quote": {
575 "type": "string",
576 "description": "A short snippet copied EXACTLY from the episode that contains the value"
577 }
578 },
579 "required": ["literal_type", "value", "quote"]
580 }
581 ]
582 });
583
584 serde_json::json!({
585 "type": "object",
586 "properties": {
587 "claims": {
588 "type": "array",
589 "items": {
590 "type": "object",
591 "properties": {
592 "predicate": {
593 "type": "string",
594 "enum": pred_names
595 },
596 "subject": mention_schema,
597 "object": object_schema,
598 "polarity": {
599 "type": "string",
600 "enum": ["affirm", "deny"]
601 },
602 "valid_from": {
603 "type": ["integer", "null"],
604 "description": "Epoch millis, or null for 'always'"
605 },
606 "valid_to": {
607 "type": ["integer", "null"],
608 "description": "Epoch millis, or null for 'still true'"
609 },
610 "confidence": {
611 "type": "number",
612 "minimum": 0.0,
613 "maximum": 1.0
614 }
615 },
616 "required": ["predicate", "subject", "object", "polarity", "confidence"]
617 }
618 }
619 },
620 "required": ["claims"]
621 })
622}
623
624fn enum_alternation(values: &[&str]) -> String {
628 values
629 .iter()
630 .map(|v| format!("\"\\\"{v}\\\"\""))
631 .collect::<Vec<_>>()
632 .join(" | ")
633}
634
635pub fn grammar_from_registry(predicates: &[PredicateDef]) -> String {
647 let pred_names: Vec<&str> = predicates.iter().map(|p| p.name.as_str()).collect();
649 let entity_types: Vec<&str> = predicates
650 .iter()
651 .flat_map(|p| {
652 let subjects = p.subject_types.iter().map(|t| t.as_str());
653 let objects = match &p.object_kind {
654 ObjectKind::Entity(types) => types.0.iter().map(|t| t.as_str()).collect(),
655 _ => vec![],
656 };
657 subjects.chain(objects)
658 })
659 .collect::<std::collections::BTreeSet<_>>()
660 .into_iter()
661 .collect();
662
663 let pred_alts = enum_alternation(&pred_names);
664 let etype_alts = enum_alternation(&entity_types);
665
666 format!(
671 r#"root ::= ws "{{" ws "\"claims\"" ws ":" ws "[" ws claims ws "]" ws "}}"
672claims ::= (claim (ws "," ws claim)*)?
673claim ::= "{{" ws "\"predicate\"" ws ":" ws predicate ws "," ws "\"subject\"" ws ":" ws mention ws "," ws "\"object\"" ws ":" ws object-union ws "," ws "\"polarity\"" ws ":" ws polarity ws "," ws valid-from-opt valid-to-opt "\"confidence\"" ws ":" ws number ws "}}"
674valid-from-opt ::= ("\"valid_from\"" ws ":" ws temporal-val ws "," ws)?
675valid-to-opt ::= ("\"valid_to\"" ws ":" ws temporal-val ws "," ws)?
676temporal-val ::= "null" | integer
677mention ::= "{{" ws "\"surface\"" ws ":" ws string ws "," ws "\"entity_type\"" ws ":" ws entity-type ws "," ws "\"quote\"" ws ":" ws nonempty-string ws "}}"
678object-union ::= entity-object | literal-object
679entity-object ::= "{{" ws "\"kind\"" ws ":" ws "\"entity\"" ws "," ws "\"mention\"" ws ":" ws mention ws "}}"
680literal-object ::= "{{" ws "\"kind\"" ws ":" ws "\"literal\"" ws "," ws "\"literal_type\"" ws ":" ws literal-type ws "," ws "\"value\"" ws ":" ws string ws "," ws "\"quote\"" ws ":" ws nonempty-string ws "}}"
681entity-type ::= {etype_alts}
682literal-type ::= "\"text\"" | "\"date\"" | "\"datetime\"" | "\"number\"" | "\"bool\"" | "\"quantity\""
683predicate ::= {pred_alts}
684polarity ::= "\"affirm\"" | "\"deny\""
685string ::= "\"" ([^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]))* "\"" ws
686nonempty-string ::= "\"" ([^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]))+ "\"" ws
687number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws
688integer ::= "-"? ([0-9] | [1-9] [0-9]*) ws
689ws ::= [ \t\n]*
690"#,
691 )
692}
693
694pub fn build_extraction_prompt(predicates: &[PredicateDef]) -> String {
699 let mut s = String::new();
700 s.push_str(
701 "You are a knowledge extraction engine. Extract structured claims from the given text. \
702 Each claim references entities by their VERBATIM surface form and a quote copied from the text.\n\n\
703 Available predicates:\n",
704 );
705 for p in predicates {
706 let obj_desc = match &p.object_kind {
707 ObjectKind::Entity(types) => format!("entity: {}", types.0.join("|")),
708 ObjectKind::Literal(lt) => format!("literal: {lt:?}"),
709 ObjectKind::Enum { variants } => format!("enum: {}", variants.join("|")),
710 };
711 s.push_str(&format!("- {} ({}): {}\n", p.name, obj_desc, p.description));
712 if !p.examples.is_empty() {
713 s.push_str(&format!(" Examples: {}\n", p.examples.join("; ")));
714 }
715 }
716 s.push_str(
717 "\nReturn JSON matching the provided schema. For each entity mention and each literal \
718 value, provide:\n\
719 - surface: the text exactly as it appears in the episode.\n\
720 - quote: a short snippet (up to ~120 characters) copied EXACTLY from the episode, \
721 character for character, taken from the SAME sentence where the surface appears. \
722 Copy it — do not paraphrase, do not count character positions. The quote is \
723 REQUIRED for every mention, including subjects; an empty quote is an error.\n\n\
724 Rules:\n\
725 - The surface MUST occur inside your quote, and the quote MUST occur in the episode.\n\
726 - Subject and object types must match the predicate's definition.\n\
727 - Set confidence to your confidence in the claim (0.0 to 1.0).\n\
728 - Use valid_from/valid_to for time-bounded claims. Use null for 'always true' or 'still true'.\n",
729 );
730 s
731}
732
733#[derive(Debug, Clone)]
737pub struct FewShotExample {
738 pub text: String,
739 pub claims_json: String,
740}
741pub fn few_shot_examples<'a>(
747 target_text: &str,
748 corpus: &'a [FewShotExample],
749 k: usize,
750) -> Vec<&'a FewShotExample> {
751 if corpus.is_empty() || k == 0 {
752 return Vec::new();
753 }
754 let target_shingles = oxibrain_index::shingles(target_text.to_lowercase().trim(), 3);
755 let mut scored: Vec<(f64, &FewShotExample)> = corpus
756 .iter()
757 .map(|ex| {
758 let ex_shingles = oxibrain_index::shingles(ex.text.to_lowercase().trim(), 3);
759 let sim = oxibrain_index::jaccard(&target_shingles, &ex_shingles);
760 (sim, ex)
761 })
762 .collect();
763 scored.sort_by(|a, b| {
765 b.0.partial_cmp(&a.0)
766 .unwrap_or(std::cmp::Ordering::Equal)
767 .then_with(|| a.1.text.cmp(&b.1.text))
768 });
769 scored.iter().take(k).map(|(_, ex)| *ex).collect()
770}
771
772pub fn format_few_shot(examples: &[&FewShotExample]) -> String {
775 if examples.is_empty() {
776 return String::new();
777 }
778 let mut out = String::from("\nHere are some examples of correct extraction:\n\n");
779 for (i, ex) in examples.iter().enumerate() {
780 out.push_str(&format!("Example {}:\n", i + 1));
781 out.push_str(&format!("Input: {}\n", ex.text));
782 out.push_str(&format!("Output: {}\n\n", ex.claims_json));
783 }
784 out
785}
786
787pub fn default_few_shot_corpus() -> Vec<FewShotExample> {
795 vec![
796 FewShotExample {
797 text: "Alice works on ProjectX at Acme Corp. Bob knows Carol.".into(),
798 claims_json: r#"{"claims":[
799 {"predicate":"works_on",
800 "subject":{"surface":"Alice","entity_type":"Person","quote":"Alice works on ProjectX"},
801 "object":{"kind":"entity","mention":{"surface":"ProjectX","entity_type":"Project","quote":"Alice works on ProjectX"}},
802 "polarity":"affirm","confidence":0.95},
803 {"predicate":"employed_by",
804 "subject":{"surface":"Alice","entity_type":"Person","quote":"Alice works on ProjectX at Acme Corp"},
805 "object":{"kind":"entity","mention":{"surface":"Acme Corp","entity_type":"Organization","quote":"at Acme Corp"}},
806 "polarity":"affirm","confidence":0.9},
807 {"predicate":"knows",
808 "subject":{"surface":"Bob","entity_type":"Person","quote":"Bob knows Carol"},
809 "object":{"kind":"entity","mention":{"surface":"Carol","entity_type":"Person","quote":"Bob knows Carol"}},
810 "polarity":"affirm","confidence":0.9}
811]}"#.into(),
812 },
813 FewShotExample {
814 text: "김민수는 Acme Corp에 다니고 있다. 이서연은 brain-ui 프로젝트를 진행한다.".into(),
815 claims_json: r#"{"claims":[
816 {"predicate":"employed_by",
817 "subject":{"surface":"김민수","entity_type":"Person","quote":"김민수는 Acme Corp에 다니고 있다"},
818 "object":{"kind":"entity","mention":{"surface":"Acme Corp","entity_type":"Organization","quote":"Acme Corp에 다니고 있다"}},
819 "polarity":"affirm","confidence":0.9},
820 {"predicate":"works_on",
821 "subject":{"surface":"이서연","entity_type":"Person","quote":"이서연은 brain-ui 프로젝트를 진행한다"},
822 "object":{"kind":"entity","mention":{"surface":"brain-ui","entity_type":"Project","quote":"brain-ui 프로젝트를 진행한다"}},
823 "polarity":"affirm","confidence":0.9}
824]}"#.into(),
825 },
826 FewShotExample {
827 text: "Alice's full name is Alice Smith. She was born in Seoul.".into(),
828 claims_json: r#"{"claims":[
829 {"predicate":"full_name",
830 "subject":{"surface":"Alice","entity_type":"Person","quote":"Alice's full name is Alice Smith"},
831 "object":{"kind":"literal","literal_type":"text","value":"Alice Smith","quote":"full name is Alice Smith"},
832 "polarity":"affirm","confidence":0.95},
833 {"predicate":"born_in",
834 "subject":{"surface":"Alice","entity_type":"Person","quote":"Alice's full name"},
835 "object":{"kind":"entity","mention":{"surface":"Seoul","entity_type":"Place","quote":"born in Seoul"}},
836 "polarity":"affirm","confidence":0.9}
837]}"#.into(),
838 },
839 ]
840}
841
842#[cfg(test)]
843mod tests {
844 use super::*;
845
846 #[test]
847 fn extractor_id_deterministic() {
848 let c = ExtractorConfig {
849 model_id: "claude-sonnet-4-5".into(),
850 prompt_version: 1,
851 registry_major: 1,
852 mechanism: ExtractMechanism::ToolCall,
853 max_tokens: 8192,
854 model_digest: None,
855 provider_profile_id: None,
856 };
857 assert_eq!(c.id(), c.id());
858 }
859
860 #[test]
861 fn grammar_mechanism_changes_extractor_id() {
862 let base = ExtractorConfig {
865 model_id: "qwen2.5-1.5b-instruct".into(),
866 prompt_version: 1,
867 registry_major: 1,
868 mechanism: ExtractMechanism::Grammar,
869 max_tokens: 8192,
870 model_digest: None,
871 provider_profile_id: None,
872 };
873 assert_eq!(base.id(), base.id());
874 let json_mode = ExtractorConfig {
875 mechanism: ExtractMechanism::JsonSchema,
876 ..base.clone()
877 };
878 assert_ne!(base.id(), json_mode.id());
879 }
880
881 #[test]
882 fn resolve_mention_repairs_casing_but_not_location() {
883 let content = "The user prefers Rust.";
885 let mut m = MentionRef {
886 surface: "the user".into(),
887 entity_type: "person".into(),
888 quote: None,
889 span: (0, 8),
890 };
891 assert!(resolve_mention(&mut m, content));
892 assert_eq!(m.surface, "The user");
893
894 let mut m = MentionRef {
897 surface: "Rust".into(),
898 entity_type: "technology".into(),
899 quote: None,
900 span: (0, 4),
901 };
902 assert!(!resolve_mention(&mut m, content));
903 }
904
905 #[test]
906 fn resolve_mention_rejects_fabricated_surface() {
907 let content = "The user prefers Rust.";
908 let mut m = MentionRef {
909 surface: "Python".into(),
910 entity_type: "technology".into(),
911 quote: None,
912 span: (0, 6),
913 };
914 assert!(!resolve_mention(&mut m, content));
915 }
916
917 #[test]
918 fn extractor_id_changes_with_model() {
919 let base = ExtractorConfig {
920 model_id: "a".into(),
921 prompt_version: 1,
922 registry_major: 1,
923 mechanism: ExtractMechanism::JsonSchema,
924 max_tokens: 4096,
925 model_digest: None,
926 provider_profile_id: None,
927 };
928 let diff = ExtractorConfig {
929 model_id: "b".into(),
930 ..base.clone()
931 };
932 assert_ne!(base.id(), diff.id());
933 }
934
935 #[test]
936 fn extractor_id_changes_with_mechanism() {
937 let base = ExtractorConfig {
938 model_id: "a".into(),
939 prompt_version: 1,
940 registry_major: 1,
941 mechanism: ExtractMechanism::JsonSchema,
942 max_tokens: 4096,
943 model_digest: None,
944 provider_profile_id: None,
945 };
946 let diff = ExtractorConfig {
947 mechanism: ExtractMechanism::ToolCall,
948 ..base.clone()
949 };
950 assert_ne!(base.id(), diff.id());
951 }
952
953 #[test]
954 fn extractor_id_changes_with_registry_major() {
955 let base = ExtractorConfig {
956 model_id: "a".into(),
957 prompt_version: 1,
958 registry_major: 1,
959 mechanism: ExtractMechanism::JsonSchema,
960 max_tokens: 4096,
961 model_digest: None,
962 provider_profile_id: None,
963 };
964 let diff = ExtractorConfig {
965 registry_major: 2,
966 ..base.clone()
967 };
968 assert_ne!(base.id(), diff.id());
969 }
970
971 #[test]
972 fn extractor_id_changes_with_digest() {
973 let base = ExtractorConfig {
976 model_id: "qwen2.5-1.5b".into(),
977 prompt_version: 1,
978 registry_major: 1,
979 mechanism: ExtractMechanism::JsonSchema,
980 max_tokens: 8192,
981 model_digest: Some("abc123".into()),
982 provider_profile_id: None,
983 };
984 let diff = ExtractorConfig {
985 model_digest: Some("def456".into()),
986 provider_profile_id: None,
987 ..base.clone()
988 };
989 assert_ne!(
990 base.id(),
991 diff.id(),
992 "weight change must invalidate ExtractorId"
993 );
994
995 let nodigest = ExtractorConfig {
997 model_digest: None,
998 provider_profile_id: None,
999 ..base.clone()
1000 };
1001 assert_ne!(base.id(), nodigest.id());
1002 }
1003
1004 #[test]
1005 fn schema_contains_all_predicates() {
1006 let schema = schema_from_registry(crate::registry::core_v1());
1007 let claims_items =
1008 &schema["properties"]["claims"]["items"]["properties"]["predicate"]["enum"];
1009 let names: Vec<String> = claims_items
1010 .as_array()
1011 .unwrap()
1012 .iter()
1013 .map(|v| v.as_str().unwrap().to_string())
1014 .collect();
1015 assert!(names.contains(&"works_on".to_string()));
1016 assert!(names.contains(&"employed_by".to_string()));
1017 assert!(names.contains(&"born_in".to_string()));
1018 }
1019
1020 #[test]
1021 fn prompt_contains_predicate_descriptions() {
1022 let prompt = build_extraction_prompt(crate::registry::core_v1());
1023 assert!(prompt.contains("works_on"));
1024 assert!(prompt.contains("project"));
1025 assert!(prompt.contains("VERBATIM"));
1026 }
1027
1028 #[test]
1029 fn prompt_v2_teaches_quotes_not_offsets() {
1030 let prompt = build_extraction_prompt(crate::registry::core_v1());
1031 assert!(
1032 prompt.contains("quote"),
1033 "v2 prompt must teach quote copying"
1034 );
1035 assert!(
1036 !prompt.contains("byte offset"),
1037 "v2 prompt must not demand offset arithmetic"
1038 );
1039 }
1040
1041 #[test]
1042 fn grammar_uses_quotes_not_spans() {
1043 let g = grammar_from_registry(crate::registry::core_v1());
1044 let norm: String = g.split_whitespace().collect::<Vec<_>>().join(" ");
1045 assert!(
1046 norm.contains("\"\\\"quote\\\"\""),
1047 "mention and literal rules must require a quote"
1048 );
1049 assert!(
1050 !norm.contains("\"span\""),
1051 "model-facing grammar must not ask for numeric spans (ADR-006)"
1052 );
1053 }
1054
1055 #[test]
1056 fn schema_uses_quotes_not_spans() {
1057 let s = schema_from_registry(crate::registry::core_v1());
1058 let mention = &s["properties"]["claims"]["items"]["properties"]["subject"];
1059 assert!(
1060 mention["required"]
1061 .as_array()
1062 .unwrap()
1063 .iter()
1064 .any(|v| v == "quote"),
1065 "mention must require quote"
1066 );
1067 assert!(
1068 !mention["properties"]
1069 .as_object()
1070 .unwrap()
1071 .contains_key("span"),
1072 "model-facing schema must not ask for numeric spans"
1073 );
1074 }
1075
1076 #[test]
1077 fn few_shot_corpus_examples_validate() {
1078 for ex in default_few_shot_corpus() {
1082 let parsed: ExtractionResponse = serde_json::from_str(&ex.claims_json)
1083 .unwrap_or_else(|e| panic!("corpus example not parseable: {e}"));
1084 let result = validate_claims(&parsed.claims, &ex.text, crate::registry::core_v1());
1085 assert!(
1086 result.valid.len() == parsed.claims.len() && result.invalid.is_empty(),
1087 "corpus example invalid: {:#?}",
1088 result.invalid
1089 );
1090 }
1091 }
1092
1093 fn make_claim(
1094 predicate: &str,
1095 subj_surface: &str,
1096 subj_type: &str,
1097 subj_span: (u32, u32),
1098 obj_surface: &str,
1099 obj_type: &str,
1100 obj_span: (u32, u32),
1101 ) -> Claim {
1102 Claim {
1103 predicate: predicate.into(),
1104 subject: MentionRef {
1105 surface: subj_surface.into(),
1106 entity_type: subj_type.into(),
1107 quote: None,
1108 span: subj_span,
1109 },
1110 object: ClaimObject::Entity {
1111 mention: MentionRef {
1112 surface: obj_surface.into(),
1113 entity_type: obj_type.into(),
1114 quote: None,
1115 span: obj_span,
1116 },
1117 },
1118 polarity: Polarity::Affirm,
1119 valid_from: None,
1120 valid_to: None,
1121 confidence: 0.9,
1122 }
1123 }
1124
1125 #[test]
1126 fn validate_valid_claim() {
1127 let content = "Alice works on ProjectX at Acme Corp.";
1128 let claim = make_claim(
1129 "works_on",
1130 "Alice",
1131 "Person",
1132 (0, 5),
1133 "ProjectX",
1134 "Project",
1135 (15, 23),
1136 );
1137 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1138 assert_eq!(result.valid.len(), 1);
1139 assert!(result.invalid.is_empty());
1140 }
1141
1142 #[test]
1143 fn validate_part_of_accepts_project_object() {
1144 let content = "The parser module belongs to ProjectX.";
1148 let claim = make_claim(
1149 "part_of",
1150 "parser module",
1151 "Artifact",
1152 (4, 17),
1153 "ProjectX",
1154 "Project",
1155 (29, 37),
1156 );
1157 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1158 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1159 assert!(result.invalid.is_empty());
1160 }
1161
1162 #[test]
1163 fn validate_object_type_mismatch_lists_allowed_types() {
1164 let content = "Alice works on ProjectX.";
1165 let claim = make_claim(
1166 "works_on",
1167 "Alice",
1168 "Person",
1169 (0, 5),
1170 "ProjectX",
1171 "Place",
1172 (15, 23),
1173 );
1174 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1175 assert!(result.valid.is_empty());
1176 assert!(result.invalid[0].1.iter().any(|e| matches!(
1177 e,
1178 ValidationError::ObjectTypeMismatch { expected, got, .. }
1179 if expected == "Project" && got == "Place"
1180 )));
1181 }
1182
1183 #[test]
1184 fn validate_unknown_predicate() {
1185 let content = "Alice works on ProjectX.";
1186 let claim = make_claim(
1187 "unknown_pred",
1188 "Alice",
1189 "Person",
1190 (0, 5),
1191 "ProjectX",
1192 "Project",
1193 (15, 23),
1194 );
1195 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1196 assert!(result.valid.is_empty());
1197 assert_eq!(result.invalid.len(), 1);
1198 }
1199
1200 #[test]
1201 fn validate_fabricated_entity_rejected() {
1202 let content = "Alice works on ProjectX.";
1203 let claim = make_claim(
1205 "works_on",
1206 "Bob",
1207 "Person",
1208 (0, 5),
1209 "ProjectX",
1210 "Project",
1211 (15, 23),
1212 );
1213 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1214 assert!(result.valid.is_empty());
1215 assert_eq!(result.invalid.len(), 1);
1216 assert!(matches!(
1217 result.invalid[0].1[0],
1218 ValidationError::SurfaceNotVerbatim { .. }
1219 ));
1220 }
1221
1222 #[test]
1223 fn validate_span_out_of_bounds() {
1224 let content = "Alice works on ProjectX.";
1228 let claim = make_claim(
1229 "works_on",
1230 "Alice",
1231 "Person",
1232 (0, 5),
1233 "Zanzibar",
1234 "Project",
1235 (999, 1000),
1236 );
1237 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1238 assert!(result.valid.is_empty());
1239 }
1240
1241 #[test]
1242 fn validate_subject_type_mismatch() {
1243 let content = "Acme Corp employs Alice (0-5).";
1244 let claim = make_claim(
1245 "employed_by",
1246 "Acme Corp",
1247 "Organization",
1248 (0, 9),
1249 "Somewhere",
1250 "Organization",
1251 (17, 26),
1252 );
1253 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1254 assert!(result.valid.is_empty());
1255 assert!(matches!(
1256 result.invalid[0].1[0],
1257 ValidationError::SubjectTypeMismatch { .. }
1258 ));
1259 }
1260
1261 fn make_claim_with_quote(
1264 predicate: &str,
1265 subj_surface: &str,
1266 subj_type: &str,
1267 subj_quote: &str,
1268 obj_surface: &str,
1269 obj_type: &str,
1270 obj_quote: &str,
1271 ) -> Claim {
1272 Claim {
1273 predicate: predicate.into(),
1274 subject: MentionRef {
1275 surface: subj_surface.into(),
1276 entity_type: subj_type.into(),
1277 quote: Some(subj_quote.into()),
1278 span: (0, 0), },
1280 object: ClaimObject::Entity {
1281 mention: MentionRef {
1282 surface: obj_surface.into(),
1283 entity_type: obj_type.into(),
1284 quote: Some(obj_quote.into()),
1285 span: (0, 0),
1286 },
1287 },
1288 polarity: Polarity::Affirm,
1289 valid_from: None,
1290 valid_to: None,
1291 confidence: 0.9,
1292 }
1293 }
1294
1295 #[test]
1296 fn quote_locates_and_derives_span_multilingual() {
1297 let content = "김민수는 Acme Corp에 다니고 있다. 이서연은 brain-ui를 진행한다.";
1300 let claim = make_claim_with_quote(
1301 "employed_by",
1302 "김민수",
1303 "Person",
1304 "김민수는 Acme Corp에 다니고 있다",
1305 "Acme Corp",
1306 "Organization",
1307 "Acme Corp에 다니고 있다",
1308 );
1309 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1310 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1311 let start = content.find("김민수").unwrap() as u32;
1312 assert_eq!(
1313 result.valid[0].subject.span,
1314 (start, start + "김민수".len() as u32)
1315 );
1316 }
1317
1318 #[test]
1319 fn quote_disambiguates_multiple_occurrences() {
1320 let content = "Alice met Bob. Later Alice left.";
1322 let claim = make_claim_with_quote(
1323 "knows",
1324 "Alice",
1325 "Person",
1326 "Alice met Bob",
1327 "Bob",
1328 "Person",
1329 "Alice met Bob",
1330 );
1331 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1332 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1333 assert_eq!(result.valid[0].subject.span, (0, 5)); }
1335
1336 #[test]
1337 fn quote_not_found_rejects() {
1338 let content = "Alice works on ProjectX.";
1341 let claim = make_claim_with_quote(
1342 "works_on",
1343 "Alice",
1344 "Person",
1345 "Alice works on SecreTProJect", "ProjectX",
1347 "Project",
1348 "works on ProjectX",
1349 );
1350 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1351 assert!(result.valid.is_empty(), "fabricated quote must be rejected");
1352 assert!(matches!(
1353 result.invalid[0].1[0],
1354 ValidationError::SurfaceNotVerbatim { .. }
1355 ));
1356 }
1357
1358 #[test]
1359 fn quote_without_surface_rejects() {
1360 let content = "Alice works on ProjectX. Bob knows Carol.";
1363 let claim = make_claim_with_quote(
1364 "works_on",
1365 "Bob",
1366 "Person",
1367 "Alice works on ProjectX", "ProjectX",
1369 "Project",
1370 "Alice works on ProjectX",
1371 );
1372 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1373 assert!(result.valid.is_empty(), "surface absent from quote");
1374 assert!(matches!(
1375 result.invalid[0].1[0],
1376 ValidationError::SurfaceNotVerbatim { .. }
1377 ));
1378 }
1379
1380 #[test]
1381 fn quote_casing_canonicalizes_surface() {
1382 let content = "Alice works on ProjectX at Acme Corp.";
1383 let claim = make_claim_with_quote(
1384 "employed_by",
1385 "alice",
1386 "Person",
1387 "Alice works on ProjectX at Acme Corp",
1388 "acme corp",
1389 "Organization",
1390 "at Acme Corp",
1391 );
1392 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1393 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1394 assert_eq!(result.valid[0].subject.surface, "Alice");
1395 }
1396
1397 #[test]
1398 fn quote_casing_fallback_survives_multibyte_window() {
1399 let content = "김민수는 한국 지사 Acme Corp에서 일한다. 본사는 미국에 있다.";
1401 let claim = make_claim_with_quote(
1404 "employed_by",
1405 "김민수",
1406 "Person",
1407 "김민수는 한국 지사 Acme Corp에서 일한다",
1408 "ACME",
1409 "Organization",
1410 "한국 지사 Acme Corp에서",
1411 );
1412 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1413 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1414 assert_eq!(
1416 result.valid[0].subject.surface, "김민수",
1417 "exact-match subject must pass untouched"
1418 );
1419 if let ClaimObject::Entity { mention } = &result.valid[0].object {
1420 assert_eq!(
1421 mention.surface, "Acme",
1422 "casing must canonicalize to source"
1423 );
1424 let start = content.find("Acme").unwrap() as u32;
1425 assert_eq!(mention.span, (start, start + "Acme".len() as u32));
1426 } else {
1427 panic!("expected entity object");
1428 }
1429 }
1430
1431 #[test]
1432 fn literal_quote_derives_span() {
1433 let content = "Alice's full name is Alice Smith.";
1436 let claim = Claim {
1437 predicate: "full_name".into(),
1438 subject: MentionRef {
1439 surface: "Alice".into(),
1440 entity_type: "Person".into(),
1441 quote: Some("full name is Alice Smith".into()),
1442 span: (0, 0),
1443 },
1444 object: ClaimObject::Literal {
1445 literal_type: "text".into(),
1446 value: "Alice Smith".into(),
1447 quote: Some("full name is Alice Smith".into()),
1448 span: (0, 0),
1449 },
1450 polarity: Polarity::Affirm,
1451 valid_from: None,
1452 valid_to: None,
1453 confidence: 0.9,
1454 };
1455 let result = validate_claims(&[claim], content, crate::registry::core_v1());
1456 assert_eq!(result.valid.len(), 1, "errors: {:?}", result.invalid);
1457 let start = content.find("Alice Smith").unwrap() as u32;
1458 if let ClaimObject::Literal { span, .. } = &result.valid[0].object {
1459 assert_eq!(*span, (start, start + "Alice Smith".len() as u32));
1460 } else {
1461 panic!("expected literal object");
1462 }
1463 }
1464
1465 #[test]
1468 fn grammar_smoke_has_rules() {
1469 let g = grammar_from_registry(crate::registry::core_v1());
1470 let norm: String = g.split_whitespace().collect::<Vec<_>>().join(" ");
1472 for rule in [
1473 "root",
1474 "claims",
1475 "claim",
1476 "mention",
1477 "object-union",
1478 "predicate",
1479 "entity-type",
1480 "polarity",
1481 "literal-type",
1482 "string",
1483 "number",
1484 "integer",
1485 "ws",
1486 ] {
1487 let needle = format!("{rule} ::=");
1488 assert!(
1489 norm.contains(&needle),
1490 "grammar missing rule definition for `{rule}`"
1491 );
1492 }
1493 }
1494
1495 #[test]
1496 fn grammar_forbids_empty_quotes() {
1497 let g = grammar_from_registry(crate::registry::core_v1());
1501 let norm: String = g.split_whitespace().collect::<Vec<_>>().join(" ");
1502 assert!(
1503 norm.contains("nonempty-string ::="),
1504 "a nonempty-string rule must exist"
1505 );
1506 assert!(
1507 norm.contains("\"\\\"quote\\\"\" ws \":\" ws nonempty-string"),
1508 "quote fields must use nonempty-string"
1509 );
1510 }
1511
1512 #[test]
1513 fn grammar_and_schema_agree_on_predicates() {
1514 let preds = crate::registry::core_v1();
1515
1516 let grammar = grammar_from_registry(preds);
1517 let schema = schema_from_registry(preds);
1518
1519 let schema_preds: std::collections::BTreeSet<String> =
1520 schema["properties"]["claims"]["items"]["properties"]["predicate"]["enum"]
1521 .as_array()
1522 .unwrap()
1523 .iter()
1524 .map(|v| v.as_str().unwrap().to_string())
1525 .collect();
1526
1527 for name in &schema_preds {
1528 let needle = format!("\\\"{name}\\\"");
1530 assert!(
1531 grammar.contains(&needle),
1532 "grammar missing predicate `{name}` present in schema"
1533 );
1534 }
1535 }
1536
1537 #[test]
1538 fn grammar_and_schema_agree_on_entity_types() {
1539 let preds = crate::registry::core_v1();
1540 let grammar = grammar_from_registry(preds);
1541 let schema = schema_from_registry(preds);
1542
1543 let schema_types: std::collections::BTreeSet<String> = schema["properties"]["claims"]["items"]
1544 ["properties"]["subject"]["properties"]["entity_type"]["enum"]
1545 .as_array()
1546 .unwrap()
1547 .iter()
1548 .map(|v| v.as_str().unwrap().to_string())
1549 .collect();
1550
1551 for name in &schema_types {
1552 let needle = format!("\\\"{name}\\\"");
1553 assert!(
1554 grammar.contains(&needle),
1555 "grammar missing entity type `{name}` present in schema"
1556 );
1557 }
1558 }
1559
1560 #[test]
1561 fn grammar_has_polarity_and_literal_type_enums() {
1562 let g = grammar_from_registry(crate::registry::core_v1());
1563 assert!(g.contains("\\\"affirm\\\""));
1564 assert!(g.contains("\\\"deny\\\""));
1565 for lt in ["text", "date", "datetime", "number", "bool", "quantity"] {
1566 assert!(
1567 g.contains(&format!("\\\"{lt}\\\"")),
1568 "grammar missing literal type `{lt}`"
1569 );
1570 }
1571 }
1572
1573 #[test]
1574 fn grammar_valid_response_roundtrips_serde() {
1575 let claim = make_claim(
1579 "works_on",
1580 "Alice",
1581 "Person",
1582 (0, 5),
1583 "ProjectX",
1584 "Project",
1585 (15, 23),
1586 );
1587 let resp = ExtractionResponse {
1588 claims: vec![claim],
1589 };
1590 let json = serde_json::to_string(&resp).unwrap();
1591 let back: ExtractionResponse = serde_json::from_str(&json).unwrap();
1592 assert_eq!(back.claims.len(), 1);
1593 assert_eq!(back.claims[0].predicate, "works_on");
1594 }
1595
1596 #[test]
1597 fn grammar_has_optional_temporal_fields() {
1598 let g = grammar_from_registry(crate::registry::core_v1());
1599 assert!(g.contains("valid-from-opt"));
1600 assert!(g.contains("valid-to-opt"));
1601 assert!(g.contains("\\\"valid_from\\\""));
1603 }
1604 #[test]
1605 fn grammar_supports_empty_claims() {
1606 let g = grammar_from_registry(crate::registry::core_v1());
1608 assert!(g.contains("(claim (ws \",\" ws claim)*)?"));
1610 }
1611
1612 #[test]
1615 fn few_shot_selects_most_similar() {
1616 let corpus = vec![
1617 FewShotExample {
1618 text: "Alice works at Acme.".into(),
1619 claims_json: r#"{"claims":[]}"#.into(),
1620 },
1621 FewShotExample {
1622 text: "Bob likes pizza.".into(),
1623 claims_json: r#"{"claims":[]}"#.into(),
1624 },
1625 ];
1626 let target = "Alice works at Globex.";
1627 let selected = few_shot_examples(target, &corpus, 1);
1628 assert_eq!(selected.len(), 1);
1629 assert!(
1630 selected[0].text.contains("Alice"),
1631 "should pick the most similar example, got: {}",
1632 selected[0].text
1633 );
1634 }
1635
1636 #[test]
1637 fn few_shot_empty_corpus_returns_empty() {
1638 let corpus: Vec<FewShotExample> = vec![];
1639 let selected = few_shot_examples("any text", &corpus, 3);
1640 assert!(selected.is_empty());
1641 }
1642
1643 #[test]
1644 fn few_shot_k_caps_results() {
1645 let corpus: Vec<FewShotExample> = (0..10)
1646 .map(|i| FewShotExample {
1647 text: format!("Sample text {i}."),
1648 claims_json: r#"{"claims":[]}"#.into(),
1649 })
1650 .collect();
1651 let selected = few_shot_examples("Sample text", &corpus, 3);
1652 assert_eq!(selected.len(), 3);
1653 }
1654
1655 #[test]
1656 fn few_shot_format_includes_input_output() {
1657 let ex = FewShotExample {
1658 text: "Alice works at Acme.".into(),
1659 claims_json: r#"{"claims":[]}"#.into(),
1660 };
1661 let formatted = format_few_shot(&[&ex]);
1662 assert!(formatted.contains("Alice works at Acme"));
1663 assert!(formatted.contains(r#"{"claims":[]}"#));
1664 }
1665
1666 #[test]
1667 fn few_shot_format_empty_returns_empty_string() {
1668 let formatted = format_few_shot(&[]);
1669 assert_eq!(formatted, "");
1670 }
1671}