Skip to main content

weavatrix_memory/context/
model.rs

1use crate::{
2    domain::MemoryView,
3    error::{MemoryError, Result},
4    id::EntityId,
5    time::Timestamp,
6};
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeSet;
9use weavatrix_graph::Graph;
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct ContextRequest {
13    pub seeds: Vec<EntityId>,
14    pub valid_at: Timestamp,
15    pub known_at: Timestamp,
16    pub token_budget: usize,
17    pub max_depth: usize,
18    #[serde(default)]
19    pub relations: BTreeSet<String>,
20    #[serde(default)]
21    pub repositories: BTreeSet<String>,
22    #[serde(default)]
23    pub branches: BTreeSet<String>,
24}
25
26impl ContextRequest {
27    /// Creates a bounded context request.
28    ///
29    /// # Errors
30    ///
31    /// Requires at least one seed and a non-zero token budget.
32    pub fn new(
33        seeds: Vec<EntityId>,
34        valid_at: Timestamp,
35        known_at: Timestamp,
36        token_budget: usize,
37    ) -> Result<Self> {
38        if seeds.is_empty() {
39            return Err(MemoryError::InvalidValue {
40                field: "context.seeds",
41                reason: "at least one seed is required",
42            });
43        }
44        if token_budget == 0 {
45            return Err(MemoryError::InvalidValue {
46                field: "context.token_budget",
47                reason: "must be greater than zero",
48            });
49        }
50        Ok(Self {
51            seeds,
52            valid_at,
53            known_at,
54            token_budget,
55            max_depth: 2,
56            relations: BTreeSet::new(),
57            repositories: BTreeSet::new(),
58            branches: BTreeSet::new(),
59        })
60    }
61
62    /// Creates a seedless template for `ContextCompiler::compile_with_retrieval`.
63    ///
64    /// # Errors
65    ///
66    /// Rejects a zero token budget.
67    pub fn for_retrieval(
68        valid_at: Timestamp,
69        known_at: Timestamp,
70        token_budget: usize,
71    ) -> Result<Self> {
72        if token_budget == 0 {
73            return Err(MemoryError::InvalidValue {
74                field: "context.token_budget",
75                reason: "must be greater than zero",
76            });
77        }
78        Ok(Self {
79            seeds: Vec::new(),
80            valid_at,
81            known_at,
82            token_budget,
83            max_depth: 2,
84            relations: BTreeSet::new(),
85            repositories: BTreeSet::new(),
86            branches: BTreeSet::new(),
87        })
88    }
89
90    #[must_use]
91    pub const fn with_max_depth(mut self, max_depth: usize) -> Self {
92        self.max_depth = max_depth;
93        self
94    }
95
96    #[must_use]
97    pub fn include_relation(mut self, relation: impl Into<String>) -> Self {
98        self.relations.insert(relation.into());
99        self
100    }
101
102    #[must_use]
103    pub fn in_repository(mut self, repository: impl Into<String>) -> Self {
104        self.repositories.insert(repository.into());
105        self
106    }
107
108    #[must_use]
109    pub fn on_branch(mut self, branch: impl Into<String>) -> Self {
110        self.branches.insert(branch.into());
111        self
112    }
113}
114
115#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
116pub struct ContextReceipt {
117    pub valid_at: Timestamp,
118    pub known_at: Timestamp,
119    pub source_position: Option<u64>,
120    pub estimator: String,
121    pub token_budget: usize,
122    pub estimated_tokens: usize,
123    pub examined_facts: usize,
124    pub selected_facts: usize,
125    pub omitted_by_budget: usize,
126    pub excluded_by_scope: usize,
127}
128
129#[derive(Debug, Clone)]
130pub struct ContextBundle {
131    pub view: MemoryView,
132    pub graph: Graph,
133    pub receipt: ContextReceipt,
134}