weavatrix_memory/extraction/model/
output.rs1use super::TextSpan;
2use crate::{
3 domain::{Confidence, validate_text},
4 error::Result,
5 time::Timestamp,
6};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10pub struct ExtractedRelation {
11 pub local_id: String,
12 pub source: String,
13 pub relation: String,
14 pub target: String,
15 pub confidence: Confidence,
16 pub valid_from: Option<Timestamp>,
17 pub valid_until: Option<Timestamp>,
18 pub span: Option<TextSpan>,
19}
20
21impl ExtractedRelation {
22 pub fn new(
28 local_id: impl Into<String>,
29 source: impl Into<String>,
30 relation: impl Into<String>,
31 target: impl Into<String>,
32 confidence: Confidence,
33 ) -> Result<Self> {
34 let relation = Self {
35 local_id: local_id.into(),
36 source: source.into(),
37 relation: relation.into(),
38 target: target.into(),
39 confidence,
40 valid_from: None,
41 valid_until: None,
42 span: None,
43 };
44 relation.validate()?;
45 Ok(relation)
46 }
47
48 #[must_use]
49 pub const fn valid_from(mut self, value: Timestamp) -> Self {
50 self.valid_from = Some(value);
51 self
52 }
53
54 #[must_use]
55 pub const fn valid_until(mut self, value: Timestamp) -> Self {
56 self.valid_until = Some(value);
57 self
58 }
59
60 #[must_use]
61 pub const fn with_span(mut self, span: TextSpan) -> Self {
62 self.span = Some(span);
63 self
64 }
65
66 pub(crate) fn validate(&self) -> Result<()> {
67 validate_text("extracted_relation.local_id", &self.local_id)?;
68 validate_text("extracted_relation.source", &self.source)?;
69 validate_text("extracted_relation.relation", &self.relation)?;
70 validate_text("extracted_relation.target", &self.target)
71 }
72}
73
74#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
75pub struct ExtractionOutput {
76 pub entities: Vec<super::ExtractedEntity>,
77 pub relations: Vec<ExtractedRelation>,
78}