stowken 0.7.0

Compressed storage and retrieval of LLM token sequences
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

//! Custom backend: implement StorageBackend for your own persistence layer.
//!
//! This example shows a minimal in-memory backend backed by a HashMap,
//! demonstrating the interface you need to implement for any custom store.

use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use tokio::sync::RwLock;

use stowken::{
    storage::{StorageBackend, StorageError, StorageResult},
    types::{
        AnalyticsQuery, Conversation, ConversationManifest, Message, MessageContent,
        SegmentHash, StoredSegment, StowkenConfig,
    },
    Stowken,
};

#[derive(Default, Clone)]
struct MyBackend {
    segments: Arc<RwLock<HashMap<String, StoredSegment>>>,
    manifests: Arc<RwLock<HashMap<String, ConversationManifest>>>,
}

#[async_trait]
impl StorageBackend for MyBackend {
    async fn put_segment(&self, segment: &StoredSegment) -> StorageResult<()> {
        let mut map = self.segments.write().await;
        map.entry(segment.hash.0.clone())
            .and_modify(|s| s.ref_count += 1)
            .or_insert_with(|| segment.clone());
        Ok(())
    }

    async fn get_segment(&self, hash: &SegmentHash) -> StorageResult<StoredSegment> {
        self.segments
            .read()
            .await
            .get(&hash.0)
            .cloned()
            .ok_or_else(|| StorageError::SegmentNotFound(hash.0.clone()))
    }

    async fn has_segment(&self, hash: &SegmentHash) -> StorageResult<bool> {
        Ok(self.segments.read().await.contains_key(&hash.0))
    }

    async fn increment_ref(&self, hash: &SegmentHash) -> StorageResult<()> {
        let mut map = self.segments.write().await;
        if let Some(seg) = map.get_mut(&hash.0) {
            seg.ref_count += 1;
            Ok(())
        } else {
            Err(StorageError::SegmentNotFound(hash.0.clone()))
        }
    }

    async fn replace_segment_data(
        &self,
        hash: &SegmentHash,
        new_data: Vec<u8>,
    ) -> StorageResult<()> {
        let mut map = self.segments.write().await;
        if let Some(seg) = map.get_mut(&hash.0) {
            seg.compressed_size = new_data.len() as u32;
            seg.compressed_data = new_data;
            Ok(())
        } else {
            Err(StorageError::SegmentNotFound(hash.0.clone()))
        }
    }

    async fn decrement_ref(&self, hash: &SegmentHash) -> StorageResult<bool> {
        let mut map = self.segments.write().await;
        if let Some(seg) = map.get_mut(&hash.0) {
            seg.ref_count = seg.ref_count.saturating_sub(1);
            Ok(seg.ref_count == 0)
        } else {
            Err(StorageError::SegmentNotFound(hash.0.clone()))
        }
    }

    async fn delete_segment(&self, hash: &SegmentHash) -> StorageResult<()> {
        self.segments.write().await.remove(&hash.0);
        Ok(())
    }

    async fn put_manifest(&self, manifest: &ConversationManifest) -> StorageResult<()> {
        self.manifests.write().await.insert(manifest.id.clone(), manifest.clone());
        Ok(())
    }

    async fn get_manifest(&self, id: &str) -> StorageResult<ConversationManifest> {
        self.manifests
            .read()
            .await
            .get(id)
            .cloned()
            .ok_or_else(|| StorageError::ConversationNotFound(id.to_owned()))
    }

    async fn delete_manifest(&self, id: &str) -> StorageResult<()> {
        self.manifests.write().await.remove(id);
        Ok(())
    }

    async fn list_conversations(
        &self,
        _query: &AnalyticsQuery,
        limit: u64,
        offset: u64,
    ) -> StorageResult<Vec<String>> {
        let ids: Vec<String> = self.manifests.read().await.keys().cloned().collect();
        Ok(ids.into_iter().skip(offset as usize).take(limit as usize).collect())
    }

    async fn garbage_collect(&self) -> StorageResult<u64> {
        let mut segs = self.segments.write().await;
        let before = segs.len() as u64;
        segs.retain(|_, seg| seg.ref_count > 0);
        Ok(before - segs.len() as u64)
    }

    async fn storage_size_bytes(&self) -> StorageResult<u64> {
        Ok(self.segments.read().await.values().map(|s| s.compressed_size as u64).sum())
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let vault = Stowken::new(MyBackend::default(), StowkenConfig::default()).await?;

    let conv = Conversation {
        id: None,
        model: "custom-model".into(),
        tokenizer: "cl100k_base".into(),
        application: Some("custom-backend-demo".into()),
        metadata: None,
        messages: vec![Message {
            role: "user".into(),
            content: MessageContent::Text("Stored in my custom backend!".into()),
            name: None,
            tool_call_id: None,
        }],
    };

    let result = vault.store(conv).await?;
    println!("Stored via custom backend: {} segments ({} new)", result.total_segments, result.new_segments);

    let stats = vault.stats().await?;
    println!("Storage size: {} bytes", stats.storage_bytes);

    Ok(())
}