Skip to main content

weavatrix_memory/extraction/model/
output.rs

1use 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    /// Creates a provider-local relation between two mentions.
23    ///
24    /// # Errors
25    ///
26    /// Rejects invalid identifiers or relation text.
27    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}