Skip to main content

weavatrix_memory/extraction/model/
input.rs

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