Skip to main content

weavatrix_memory/extraction/model/
link.rs

1use crate::{Confidence, EntityId, MemoryEvent, NewEvent, Result};
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
5pub struct LinkPolicy {
6    pub minimum_score: u16,
7    pub minimum_margin: u16,
8    pub create_unmatched: bool,
9}
10
11impl Default for LinkPolicy {
12    fn default() -> Self {
13        Self {
14            minimum_score: 8_000,
15            minimum_margin: 500,
16            create_unmatched: true,
17        }
18    }
19}
20
21impl LinkPolicy {
22    /// Creates a strict entity-linking policy in basis points.
23    ///
24    /// # Errors
25    ///
26    /// Rejects scores or margins above 10,000.
27    pub fn new(minimum_score: u16, minimum_margin: u16, create_unmatched: bool) -> Result<Self> {
28        if minimum_score > 10_000 || minimum_margin > 10_000 {
29            return Err(crate::MemoryError::InvalidValue {
30                field: "link_policy",
31                reason: "scores must be between 0 and 10,000 basis points",
32            });
33        }
34        Ok(Self {
35            minimum_score,
36            minimum_margin,
37            create_unmatched,
38        })
39    }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
43#[serde(rename_all = "snake_case")]
44pub enum LinkMethod {
45    StableId,
46    ExternalId,
47    ProviderHint,
48    ScopedLabel,
49    Label,
50    Alias,
51    Created,
52    Ambiguous,
53    Unresolved,
54}
55
56#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
57pub struct LinkCandidate {
58    pub entity_id: EntityId,
59    pub score: Confidence,
60    pub method: LinkMethod,
61}
62
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64pub struct LinkDecision {
65    pub mention_id: String,
66    pub entity_id: Option<EntityId>,
67    pub score: Confidence,
68    pub method: LinkMethod,
69    pub candidates: Vec<LinkCandidate>,
70}
71
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73pub struct RejectedRelation {
74    pub relation_id: String,
75    pub source_mention: String,
76    pub target_mention: String,
77    pub reason: String,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81pub struct ExtractionPlan {
82    pub provider: String,
83    pub source: String,
84    pub events: Vec<NewEvent<MemoryEvent>>,
85    pub links: Vec<LinkDecision>,
86    pub rejected_relations: Vec<RejectedRelation>,
87}
88
89impl ExtractionPlan {
90    #[must_use]
91    pub fn node_event_count(&self) -> usize {
92        self.events
93            .iter()
94            .filter(|event| matches!(event.payload, MemoryEvent::NodeUpserted { .. }))
95            .count()
96    }
97
98    #[must_use]
99    pub fn fact_event_count(&self) -> usize {
100        self.events
101            .iter()
102            .filter(|event| matches!(event.payload, MemoryEvent::FactRecorded { .. }))
103            .count()
104    }
105}