weavatrix_memory/extraction/model/
input.rs1use 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 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 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 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 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}