weavatrix_memory/context/
model.rs1use 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 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 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}