Skip to main content

weavatrix_memory/context/
model.rs

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