Skip to main content

weavatrix_memory/extraction/model/
input.rs

1use crate::{
2    AgentId, Confidence, EntityId, MemoryError, Result, SessionId, Timestamp,
3    domain::{validate_optional_text, validate_text},
4};
5use serde::{Deserialize, Serialize};
6use std::collections::BTreeMap;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
9pub struct TextSpan {
10    pub start: usize,
11    pub end: usize,
12}
13
14impl TextSpan {
15    /// Creates a half-open UTF-8 byte span.
16    ///
17    /// # Errors
18    ///
19    /// Rejects empty or reversed spans.
20    pub fn new(start: usize, end: usize) -> Result<Self> {
21        if start >= end {
22            return Err(MemoryError::InvalidValue {
23                field: "text_span",
24                reason: "start must be before end",
25            });
26        }
27        Ok(Self { start, end })
28    }
29}
30
31#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
32pub struct ExtractionInput {
33    pub source: String,
34    pub content: String,
35    pub occurred_at: Timestamp,
36    pub recorded_at: Timestamp,
37    pub agent_id: AgentId,
38    pub session_id: SessionId,
39    pub repository: Option<String>,
40    pub branch: Option<String>,
41    pub locator: Option<String>,
42    pub digest: Option<String>,
43}
44
45impl ExtractionInput {
46    /// Creates source material with caller-controlled temporal provenance.
47    ///
48    /// # Errors
49    ///
50    /// Rejects empty content, invalid source text, or future occurrence time.
51    pub fn new(
52        source: impl Into<String>,
53        content: impl Into<String>,
54        occurred_at: Timestamp,
55        recorded_at: Timestamp,
56        agent_id: AgentId,
57        session_id: SessionId,
58    ) -> Result<Self> {
59        let input = Self {
60            source: source.into(),
61            content: content.into(),
62            occurred_at,
63            recorded_at,
64            agent_id,
65            session_id,
66            repository: None,
67            branch: None,
68            locator: None,
69            digest: None,
70        };
71        input.validate()?;
72        Ok(input)
73    }
74
75    #[must_use]
76    pub fn in_repository(mut self, repository: impl Into<String>) -> Self {
77        self.repository = Some(repository.into());
78        self
79    }
80
81    #[must_use]
82    pub fn on_branch(mut self, branch: impl Into<String>) -> Self {
83        self.branch = Some(branch.into());
84        self
85    }
86
87    #[must_use]
88    pub fn with_locator(mut self, locator: impl Into<String>) -> Self {
89        self.locator = Some(locator.into());
90        self
91    }
92
93    #[must_use]
94    pub fn with_digest(mut self, digest: impl Into<String>) -> Self {
95        self.digest = Some(digest.into());
96        self
97    }
98
99    pub(crate) fn validate(&self) -> Result<()> {
100        validate_text("extraction.source", &self.source)?;
101        if self.content.trim().is_empty() {
102            return Err(MemoryError::InvalidValue {
103                field: "extraction.content",
104                reason: "must contain non-whitespace text",
105            });
106        }
107        if self.occurred_at > self.recorded_at {
108            return Err(MemoryError::InvalidValue {
109                field: "extraction.occurred_at",
110                reason: "must not be later than recorded_at",
111            });
112        }
113        validate_optional_text("extraction.repository", self.repository.as_deref())?;
114        validate_optional_text("extraction.branch", self.branch.as_deref())?;
115        validate_optional_text("extraction.locator", self.locator.as_deref())?;
116        validate_optional_text("extraction.digest", self.digest.as_deref())
117    }
118}
119
120#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
121pub struct EntityHint {
122    pub entity_id: EntityId,
123    pub confidence: Confidence,
124    pub source: String,
125}
126
127impl EntityHint {
128    /// Creates a provider-supplied link candidate.
129    ///
130    /// # Errors
131    ///
132    /// Rejects an empty or whitespace-padded hint source.
133    pub fn new(
134        entity_id: EntityId,
135        confidence: Confidence,
136        source: impl Into<String>,
137    ) -> Result<Self> {
138        let hint = Self {
139            entity_id,
140            confidence,
141            source: source.into(),
142        };
143        validate_text("entity_hint.source", &hint.source)?;
144        Ok(hint)
145    }
146}
147
148#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
149pub struct ExtractedEntity {
150    pub local_id: String,
151    pub kind: String,
152    pub label: String,
153    pub confidence: Confidence,
154    pub stable_id: Option<EntityId>,
155    pub aliases: Vec<String>,
156    pub attributes: BTreeMap<String, String>,
157    pub hints: Vec<EntityHint>,
158    pub span: Option<TextSpan>,
159}
160
161impl ExtractedEntity {
162    /// Creates a provider-local entity mention.
163    ///
164    /// # Errors
165    ///
166    /// Rejects invalid local identifier, kind, or label text.
167    pub fn new(
168        local_id: impl Into<String>,
169        kind: impl Into<String>,
170        label: impl Into<String>,
171        confidence: Confidence,
172    ) -> Result<Self> {
173        let entity = Self {
174            local_id: local_id.into(),
175            kind: kind.into(),
176            label: label.into(),
177            confidence,
178            stable_id: None,
179            aliases: Vec::new(),
180            attributes: BTreeMap::new(),
181            hints: Vec::new(),
182            span: None,
183        };
184        entity.validate()?;
185        Ok(entity)
186    }
187
188    #[must_use]
189    pub fn with_stable_id(mut self, id: EntityId) -> Self {
190        self.stable_id = Some(id);
191        self
192    }
193
194    #[must_use]
195    pub fn with_alias(mut self, alias: impl Into<String>) -> Self {
196        self.aliases.push(alias.into());
197        self
198    }
199
200    #[must_use]
201    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
202        self.attributes.insert(key.into(), value.into());
203        self
204    }
205
206    #[must_use]
207    pub fn with_hint(mut self, hint: EntityHint) -> Self {
208        self.hints.push(hint);
209        self
210    }
211
212    #[must_use]
213    pub const fn with_span(mut self, span: TextSpan) -> Self {
214        self.span = Some(span);
215        self
216    }
217
218    pub(crate) fn validate(&self) -> Result<()> {
219        validate_text("extracted_entity.local_id", &self.local_id)?;
220        validate_text("extracted_entity.kind", &self.kind)?;
221        validate_text("extracted_entity.label", &self.label)?;
222        for alias in &self.aliases {
223            validate_text("extracted_entity.alias", alias)?;
224        }
225        for key in self.attributes.keys() {
226            validate_text("extracted_entity.attribute.key", key)?;
227        }
228        for (key, value) in &self.attributes {
229            if key.starts_with("external_id.") {
230                validate_text("extracted_entity.external_id", value)?;
231            }
232        }
233        for hint in &self.hints {
234            validate_text("entity_hint.source", &hint.source)?;
235        }
236        Ok(())
237    }
238}