MCPcopy Create free account
hub / github.com/atomicdotdev/atomic / search_embeddings

Method search_embeddings

atomic-core/src/pristine/txn/write/embeddings.rs:115–167  ·  view source on GitHub ↗
(
        &self,
        query_vector: &[f32],
        top_k: usize,
    )

Source from the content-addressed store, hash-verified

113 }
114
115 fn search_embeddings(
116 &self,
117 query_vector: &[f32],
118 top_k: usize,
119 ) -> PristineResult<Vec<SearchResult>> {
120 let table = match self.txn.open_table(EMBEDDINGS) {
121 Ok(table) => table,
122 Err(redb::TableError::TableDoesNotExist(_)) => return Ok(Vec::new()),
123 Err(e) => return Err(PristineError::from(e)),
124 };
125
126 // Brute-force cosine similarity search over all embeddings
127 let mut scored: Vec<SearchResult> = Vec::new();
128
129 for item in table.iter()? {
130 let (key_guard, value_guard) = item?;
131 let key_str = key_guard.value();
132
133 let (path, chunk_idx) = match decode_embedding_key(key_str) {
134 Some(decoded) => decoded,
135 None => continue,
136 };
137
138 let bytes = value_guard.value();
139 let record: EmbeddingRecord =
140 postcard::from_bytes(bytes).map_err(|e| PristineError::Serialization {
141 message: format!(
142 "failed to deserialize EmbeddingRecord at '{}': {}",
143 key_str, e
144 ),
145 })?;
146
147 let score = cosine_similarity(query_vector, &record.vector);
148
149 scored.push(SearchResult {
150 path: path.to_string(),
151 chunk_idx,
152 score,
153 preview: record.preview,
154 });
155 }
156
157 // Sort by descending score
158 scored.sort_by(|a, b| {
159 b.score
160 .partial_cmp(&a.score)
161 .unwrap_or(std::cmp::Ordering::Equal)
162 });
163
164 // Return top-k
165 scored.truncate(top_k);
166 Ok(scored)
167 }
168}
169
170impl<'a> EmbeddingsMutTxnT for WriteTxn<'a> {

Callers 2

test_embedding_searchFunction · 0.45

Calls 6

decode_embedding_keyFunction · 0.85
partial_cmpMethod · 0.80
truncateMethod · 0.80
cosine_similarityFunction · 0.70
iterMethod · 0.45
pushMethod · 0.45

Tested by 2

test_embedding_searchFunction · 0.36