weavatrix_memory/context/
token.rs

1use crate::{MemoryError, MemoryFact, MemoryNode, Result};
2
3pub trait TokenEstimator {
4    fn estimate(&self, value: &str) -> usize;
5    fn name(&self) -> &'static str;
6}
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub struct BytesTokenEstimator {
10    bytes_per_token: usize,
11}
12
13impl BytesTokenEstimator {
14    /// Creates a deterministic byte-based fallback estimator.
15    ///
16    /// # Errors
17    ///
18    /// Rejects zero bytes per token.
19    pub fn new(bytes_per_token: usize) -> Result<Self> {
20        if bytes_per_token == 0 {
21            return Err(MemoryError::InvalidValue {
22                field: "bytes_per_token",
23                reason: "must be greater than zero",
24            });
25        }
26        Ok(Self { bytes_per_token })
27    }
28}
29
30impl Default for BytesTokenEstimator {
31    fn default() -> Self {
32        Self { bytes_per_token: 4 }
33    }
34}
35
36impl TokenEstimator for BytesTokenEstimator {
37    fn estimate(&self, value: &str) -> usize {
38        value.len().div_ceil(self.bytes_per_token).max(1)
39    }
40
41    fn name(&self) -> &'static str {
42        "utf8_bytes"
43    }
44}
45
46pub(super) fn node_tokens(estimator: &impl TokenEstimator, node: &MemoryNode) -> usize {
47    estimator.estimate(node.id.as_str())
48        + estimator.estimate(&node.kind)
49        + estimator.estimate(&node.label)
50        + node
51            .repository
52            .iter()
53            .chain(node.branch.iter())
54            .map(|value| estimator.estimate(value))
55            .sum::<usize>()
56        + node
57            .attributes
58            .iter()
59            .map(|(key, value)| estimator.estimate(key) + estimator.estimate(value))
60            .sum::<usize>()
61        + 6
62}
63
64pub(super) fn fact_tokens(estimator: &impl TokenEstimator, fact: &MemoryFact) -> usize {
65    estimator.estimate(fact.id.as_str())
66        + estimator.estimate(fact.source.as_str())
67        + estimator.estimate(&fact.relation)
68        + estimator.estimate(fact.target.as_str())
69        + fact
70            .evidence
71            .iter()
72            .map(|item| {
73                estimator.estimate(&item.kind)
74                    + estimator.estimate(&item.source)
75                    + item
76                        .locator
77                        .iter()
78                        .chain(item.digest.iter())
79                        .map(|value| estimator.estimate(value))
80                        .sum::<usize>()
81            })
82            .sum::<usize>()
83        + 16
84}
weavatrix_memory/context/token.rs

weavatrix_memory/context/
token.rs
