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