Skip to main content

weavatrix_memory/extraction/model/
link.rs

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