lance_context_core/
context.rs

1use std::collections::HashMap;
2
3#[derive(Debug, Clone)]
4pub struct Context {
5    uri: String,
6    branch: String,
7    next_id: u64,
8    entries: Vec<ContextEntry>,
9    snapshots: HashMap<String, Snapshot>,
10}
11
12impl Context {
13    pub fn new(uri: impl Into<String>) -> Self {
14        Self {
15            uri: uri.into(),
16            branch: "main".to_string(),
17            next_id: 1,
18            entries: Vec::new(),
19            snapshots: HashMap::new(),
20        }
21    }
22
23    pub fn uri(&self) -> &str {
24        &self.uri
25    }
26
27    pub fn branch(&self) -> &str {
28        &self.branch
29    }
30
31    pub fn entries(&self) -> u64 {
32        self.entries.len() as u64
33    }
34
35    pub fn add(&mut self, role: &str, content: &str, data_type: Option<&str>) -> u64 {
36        let entry_id = self.next_id;
37        self.next_id += 1;
38        self.entries.push(ContextEntry {
39            id: entry_id,
40            role: role.to_string(),
41            data_type: data_type.map(|value| value.to_string()),
42            content: content.to_string(),
43        });
44        entry_id
45    }
46
47    pub fn snapshot(&mut self, label: Option<&str>) -> String {
48        let id = match label {
49            Some(label) if !label.is_empty() => label.to_string(),
50            _ => format!("snapshot-{}", self.entries.len()),
51        };
52        let snapshot = Snapshot {
53            id: id.clone(),
54            label: label.map(|value| value.to_string()),
55            entry_count: self.entries.len() as u64,
56            branch: self.branch.clone(),
57        };
58        self.snapshots.insert(id.clone(), snapshot);
59        id
60    }
61
62    pub fn fork(&self, branch_name: impl Into<String>) -> Self {
63        Self {
64            uri: self.uri.clone(),
65            branch: branch_name.into(),
66            next_id: self.next_id,
67            entries: self.entries.clone(),
68            snapshots: self.snapshots.clone(),
69        }
70    }
71
72    pub fn checkout(&mut self, snapshot_id: &str) {
73        if let Some(snapshot) = self.snapshots.get(snapshot_id) {
74            self.entries.truncate(snapshot.entry_count as usize);
75            self.next_id = self.entries.len() as u64 + 1;
76        }
77    }
78}
79
80#[derive(Debug, Clone)]
81pub struct ContextEntry {
82    pub id: u64,
83    pub role: String,
84    pub data_type: Option<String>,
85    pub content: String,
86}
87
88#[derive(Debug, Clone)]
89pub struct Snapshot {
90    pub id: String,
91    pub label: Option<String>,
92    pub entry_count: u64,
93    pub branch: String,
94}
lance_context_core/context.rs

lance_context_core/
context.rs
