synwire 0.1.0

Convenience re-exports and reference implementations for Synwire
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

//! Cache-backed embeddings that memoize results using moka.

use std::sync::Arc;

use moka::future::Cache;
use synwire_core::BoxFuture;
use synwire_core::embeddings::Embeddings;
use synwire_core::error::SynwireError;

/// Wraps an [`Embeddings`] implementation with an in-memory moka cache.
///
/// Repeated calls with the same text return cached vectors instead of
/// re-invoking the underlying embeddings model.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use synwire::cache::CacheBackedEmbeddings;
/// use synwire_core::embeddings::FakeEmbeddings;
///
/// let inner = Arc::new(FakeEmbeddings::new(64));
/// let cached = CacheBackedEmbeddings::new(inner, 1000);
/// ```
pub struct CacheBackedEmbeddings {
    inner: Arc<dyn Embeddings>,
    cache: Cache<String, Vec<f32>>,
}

impl CacheBackedEmbeddings {
    /// Creates a new cache-backed embeddings wrapper.
    ///
    /// `max_capacity` controls how many embedding vectors to cache.
    pub fn new(inner: Arc<dyn Embeddings>, max_capacity: u64) -> Self {
        Self {
            inner,
            cache: Cache::new(max_capacity),
        }
    }
}

impl Embeddings for CacheBackedEmbeddings {
    fn embed_documents<'a>(
        &'a self,
        texts: &'a [String],
    ) -> BoxFuture<'a, Result<Vec<Vec<f32>>, SynwireError>> {
        Box::pin(async move {
            let mut results = Vec::with_capacity(texts.len());
            let mut uncached_indices = Vec::new();
            let mut uncached_texts = Vec::new();

            // Check cache for each text
            for (i, text) in texts.iter().enumerate() {
                if let Some(cached) = self.cache.get(text).await {
                    results.push(Some(cached));
                } else {
                    results.push(None);
                    uncached_indices.push(i);
                    uncached_texts.push(text.clone());
                }
            }

            // Embed uncached texts
            if !uncached_texts.is_empty() {
                let embedded = self.inner.embed_documents(&uncached_texts).await?;
                for (idx, vec) in uncached_indices.into_iter().zip(embedded) {
                    self.cache.insert(texts[idx].clone(), vec.clone()).await;
                    results[idx] = Some(vec);
                }
            }

            // All entries should be Some now; collect them
            Ok(results.into_iter().map(Option::unwrap_or_default).collect())
        })
    }

    fn embed_query<'a>(&'a self, text: &'a str) -> BoxFuture<'a, Result<Vec<f32>, SynwireError>> {
        Box::pin(async move {
            if let Some(cached) = self.cache.get(text).await {
                return Ok(cached);
            }
            let vec = self.inner.embed_query(text).await?;
            self.cache.insert(text.to_owned(), vec.clone()).await;
            Ok(vec)
        })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use synwire_core::embeddings::FakeEmbeddings;

    #[tokio::test]
    async fn cache_hit_returns_same_vector() {
        let inner = Arc::new(FakeEmbeddings::new(16));
        let cached = CacheBackedEmbeddings::new(inner, 100);

        let v1 = cached.embed_query("hello").await.unwrap();
        let v2 = cached.embed_query("hello").await.unwrap();
        assert_eq!(v1, v2);
    }

    #[tokio::test]
    async fn cache_miss_invokes_inner() {
        let inner = Arc::new(FakeEmbeddings::new(16));
        let cached = CacheBackedEmbeddings::new(inner, 100);

        let v1 = cached.embed_query("alpha").await.unwrap();
        let v2 = cached.embed_query("beta").await.unwrap();
        assert_ne!(v1, v2);
    }

    #[tokio::test]
    async fn embed_documents_caches_per_text() {
        let inner = Arc::new(FakeEmbeddings::new(8));
        let cached = CacheBackedEmbeddings::new(inner, 100);

        let texts = vec!["foo".into(), "bar".into()];
        let first = cached.embed_documents(&texts).await.unwrap();

        // Second call should hit cache
        let second = cached.embed_documents(&texts).await.unwrap();
        assert_eq!(first, second);
    }

    #[tokio::test]
    async fn embed_documents_partial_cache() {
        let inner = Arc::new(FakeEmbeddings::new(8));
        let cached = CacheBackedEmbeddings::new(inner, 100);

        // Pre-populate cache with "foo"
        let _v = cached.embed_query("foo").await.unwrap();

        // Now embed ["foo", "bar"] -- "foo" from cache, "bar" fresh
        let texts = vec!["foo".into(), "bar".into()];
        let result = cached.embed_documents(&texts).await.unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].len(), 8);
        assert_eq!(result[1].len(), 8);
    }
}