Skip to main content

weavatrix_memory/extraction/model/
output.rs

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