lc_vector_stores/
chromadb.rs

1// lc-vector-stores/src/chromadb.rs
2//! ChromaDB vector store implementation (HTTP API)
3//!
4//! Uses ChromaDB's REST API for vector storage and retrieval.
5//! Supports connecting to a remote ChromaDB service (docker run -p 8000:8000 chromadb/chroma).
6
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9use serde_json::json;
10use std::collections::HashMap;
11
12use crate::{Document, FilterOp, MetadataFilter, SearchResult, VectorStore, VectorStoreError};
13
14/// ChromaDB configuration
15#[derive(Debug, Clone)]
16pub struct ChromaDBConfig {
17    /// ChromaDB service URL, default http://localhost:8000
18    pub host: String,
19    /// Collection name
20    pub collection_name: String,
21    /// Vector dimension
22    pub vector_size: usize,
23    /// Collection metadata (optional)
24    pub metadata: Option<HashMap<String, String>>,
25}
26
27impl Default for ChromaDBConfig {
28    fn default() -> Self {
29        Self {
30            host: "http://localhost:8000".to_string(),
31            collection_name: "langchainrust".to_string(),
32            vector_size: 1536,
33            metadata: None,
34        }
35    }
36}
37
38impl ChromaDBConfig {
39    /// Creates a new ChromaDB configuration
40    pub fn new(
41        host: impl Into<String>,
42        collection_name: impl Into<String>,
43        vector_size: usize,
44    ) -> Self {
45        Self {
46            host: host.into(),
47            collection_name: collection_name.into(),
48            vector_size,
49            metadata: None,
50        }
51    }
52}
53
54/// ChromaDB collection info (parsed from the API response)
55#[derive(Debug, Deserialize)]
56#[allow(dead_code)]
57struct ChromaCollection {
58    id: String,
59    name: String,
60    #[serde(default)]
61    metadata: Option<serde_json::Value>,
62}
63
64/// ChromaDB add request body
65#[derive(Debug, Serialize)]
66struct ChromaAddRequest {
67    ids: Vec<String>,
68    embeddings: Vec<Vec<f32>>,
69    documents: Vec<String>,
70    #[serde(skip_serializing_if = "Option::is_none")]
71    metadatas: Option<Vec<HashMap<String, serde_json::Value>>>,
72}
73
74/// ChromaDB query request body
75#[derive(Debug, Serialize)]
76struct ChromaQueryRequest {
77    query_embeddings: Vec<Vec<f32>>,
78    n_results: usize,
79    #[serde(skip_serializing_if = "Option::is_none")]
80    include: Option<Vec<String>>,
81    /// Chroma `where` filter dict (see [`filter_to_chroma`])
82    #[serde(rename = "where", skip_serializing_if = "Option::is_none")]
83    where_filter: Option<serde_json::Value>,
84}
85
86/// ChromaDB query response
87#[derive(Debug, Deserialize)]
88struct ChromaQueryResponse {
89    ids: Vec<Vec<String>>,
90    distances: Vec<Vec<f64>>,
91    documents: Vec<Vec<String>>,
92    #[serde(default)]
93    metadatas: Vec<Vec<Option<HashMap<String, serde_json::Value>>>>,
94}
95
96/// ChromaDB get response
97#[derive(Debug, Deserialize)]
98struct ChromaGetResponse {
99    ids: Vec<String>,
100    documents: Vec<Option<String>>,
101    #[serde(default)]
102    metadatas: Vec<Option<HashMap<String, serde_json::Value>>>,
103    embeddings: Option<Vec<Vec<f32>>>,
104}
105
106/// ChromaDB vector store
107///
108/// Connects to a ChromaDB service via HTTP API.
109///
110/// # Example
111/// ```ignore
112/// use lc_vector_stores::ChromaDBVectorStore;
113///
114/// let store = ChromaDBVectorStore::new(
115///     ChromaDBConfig::new("http://localhost:8000", "my_collection", 384)
116/// ).await?;
117/// ```
118pub struct ChromaDBVectorStore {
119    config: ChromaDBConfig,
120    client: reqwest::Client,
121    collection_id: Option<String>,
122}
123
124impl ChromaDBVectorStore {
125    /// Creates a ChromaDB vector store and initializes the collection automatically
126    pub async fn new(config: ChromaDBConfig) -> Result<Self, VectorStoreError> {
127        let client = reqwest::Client::new();
128        let mut store = Self {
129            config,
130            client,
131            collection_id: None,
132        };
133        store.init_collection().await?;
134        Ok(store)
135    }
136
137    /// Initializes or fetches the collection
138    async fn init_collection(&mut self) -> Result<(), VectorStoreError> {
139        // try to fetch the existing collection
140        let url = format!(
141            "{}/api/v1/collections/{}",
142            self.config.host, self.config.collection_name
143        );
144        let response = self
145            .client
146            .get(&url)
147            .send()
148            .await
149            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
150
151        if response.status().is_success() {
152            let collection: ChromaCollection = response.json().await.map_err(|e| {
153                VectorStoreError::StorageError(format!("failed to parse collection info: {}", e))
154            })?;
155            self.collection_id = Some(collection.id);
156            return Ok(());
157        }
158
159        // collection does not exist, create a new one
160        let create_url = format!("{}/api/v1/collections", self.config.host);
161        let mut body = json!({
162            "name": self.config.collection_name,
163        });
164
165        if let Some(ref meta) = self.config.metadata {
166            body["metadata"] = serde_json::to_value(meta).unwrap_or(json!({}));
167        }
168
169        let response = self
170            .client
171            .post(&create_url)
172            .json(&body)
173            .send()
174            .await
175            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
176
177        if response.status().is_success() {
178            let collection: ChromaCollection = response.json().await.map_err(|e| {
179                VectorStoreError::StorageError(format!(
180                    "failed to parse new collection info: {}",
181                    e
182                ))
183            })?;
184            self.collection_id = Some(collection.id);
185            Ok(())
186        } else {
187            let text = response.text().await.unwrap_or_default();
188            Err(VectorStoreError::StorageError(format!(
189                "failed to create collection: {}",
190                text
191            )))
192        }
193    }
194
195    /// Gets the collection ID
196    fn get_collection_id(&self) -> Result<&str, VectorStoreError> {
197        self.collection_id.as_deref().ok_or_else(|| {
198            VectorStoreError::StorageError("collection is not initialized".to_string())
199        })
200    }
201
202    /// Builds the collection API base URL
203    fn collection_url(&self, endpoint: &str) -> Result<String, VectorStoreError> {
204        let cid = self.get_collection_id()?;
205        Ok(format!(
206            "{}/api/v1/collections/{}/{}",
207            self.config.host, cid, endpoint
208        ))
209    }
210
211    /// Builds a Chroma query request body (pure function, convenient for testing).
212    fn query_request(
213        query_embedding: &[f32],
214        k: usize,
215        filter: Option<&MetadataFilter>,
216    ) -> ChromaQueryRequest {
217        ChromaQueryRequest {
218            query_embeddings: vec![query_embedding.to_vec()],
219            n_results: k,
220            include: Some(vec![
221                "documents".to_string(),
222                "distances".to_string(),
223                "metadatas".to_string(),
224            ]),
225            where_filter: filter.map(filter_to_chroma),
226        }
227    }
228
229    /// POSTs to `/query` and parses the result (shared by plain and filtered retrieval).
230    async fn query_impl(
231        &self,
232        request: ChromaQueryRequest,
233    ) -> Result<Vec<SearchResult>, VectorStoreError> {
234        let url = self.collection_url("query")?;
235        let response = self
236            .client
237            .post(&url)
238            .json(&request)
239            .send()
240            .await
241            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
242
243        if !response.status().is_success() {
244            let text = response.text().await.unwrap_or_default();
245            return Err(VectorStoreError::StorageError(format!(
246                "query failed: {}",
247                text
248            )));
249        }
250
251        let query_result: ChromaQueryResponse = response.json().await.map_err(|e| {
252            VectorStoreError::StorageError(format!("failed to parse query results: {}", e))
253        })?;
254
255        let mut results = Vec::new();
256
257        // ChromaDB returns nested arrays (one result set per query)
258        if let Some(doc_list) = query_result.documents.into_iter().next() {
259            let dist_list = query_result
260                .distances
261                .into_iter()
262                .next()
263                .unwrap_or_default();
264            let meta_list = query_result
265                .metadatas
266                .into_iter()
267                .next()
268                .unwrap_or_default();
269            let id_list = query_result.ids.into_iter().next().unwrap_or_default();
270
271            for (i, content) in doc_list.into_iter().enumerate() {
272                let score = dist_list.get(i).copied().unwrap_or(0.0);
273                // ChromaDB returns L2 distance; convert to a similarity score (1 / (1 + dist))
274                let similarity = 1.0 / (1.0 + score);
275                let metadata = meta_list
276                    .get(i)
277                    .unwrap_or(&None)
278                    .clone()
279                    .unwrap_or_default();
280                let doc_id = id_list.get(i).cloned();
281
282                results.push(SearchResult {
283                    document: Document {
284                        content,
285                        metadata,
286                        id: doc_id,
287                    },
288                    score: similarity as f32,
289                });
290            }
291        }
292
293        // sort by similarity descending
294        results.sort_by(|a, b| {
295            b.score
296                .partial_cmp(&a.score)
297                .unwrap_or(std::cmp::Ordering::Equal)
298        });
299        Ok(results)
300    }
301}
302
303#[async_trait]
304impl VectorStore for ChromaDBVectorStore {
305    async fn add_documents(
306        &self,
307        documents: Vec<Document>,
308        embeddings: Vec<Vec<f32>>,
309    ) -> Result<Vec<String>, VectorStoreError> {
310        if documents.is_empty() {
311            return Ok(Vec::new());
312        }
313
314        let count = documents.len();
315        let ids: Vec<String> = (0..count)
316            .map(|i| {
317                documents[i]
318                    .id
319                    .clone()
320                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
321            })
322            .collect();
323
324        let contents: Vec<String> = documents.iter().map(|d| d.content.clone()).collect();
325        let metadatas: Vec<HashMap<String, serde_json::Value>> =
326            documents.iter().map(|d| d.metadata.clone()).collect();
327        let has_metadata = metadatas.iter().any(|m| !m.is_empty());
328
329        let request = ChromaAddRequest {
330            ids: ids.clone(),
331            embeddings,
332            documents: contents,
333            metadatas: if has_metadata { Some(metadatas) } else { None },
334        };
335
336        let url = self.collection_url("add")?;
337        let response = self
338            .client
339            .post(&url)
340            .json(&request)
341            .send()
342            .await
343            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
344
345        if !response.status().is_success() {
346            let text = response.text().await.unwrap_or_default();
347            return Err(VectorStoreError::StorageError(format!(
348                "failed to add documents: {}",
349                text
350            )));
351        }
352
353        Ok(ids)
354    }
355
356    async fn similarity_search(
357        &self,
358        query_embedding: &[f32],
359        k: usize,
360    ) -> Result<Vec<SearchResult>, VectorStoreError> {
361        let request = Self::query_request(query_embedding, k, None);
362        self.query_impl(request).await
363    }
364
365    /// S3: similarity search with metadata filtering — filtering is delegated to the server (Chroma `where` syntax).
366    async fn similarity_search_with_filter(
367        &self,
368        query_embedding: &[f32],
369        k: usize,
370        filter: Option<&MetadataFilter>,
371    ) -> Result<Vec<SearchResult>, VectorStoreError> {
372        let request = Self::query_request(query_embedding, k, filter);
373        self.query_impl(request).await
374    }
375
376    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
377        let url = self.collection_url("get")?;
378        let body = json!({
379            "ids": [id],
380            "include": ["documents", "metadatas"]
381        });
382
383        let response = self
384            .client
385            .post(&url)
386            .json(&body)
387            .send()
388            .await
389            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
390
391        if !response.status().is_success() {
392            return Ok(None);
393        }
394
395        let get_result: ChromaGetResponse = response.json().await.map_err(|e| {
396            VectorStoreError::StorageError(format!("failed to parse document: {}", e))
397        })?;
398
399        if get_result.ids.is_empty() {
400            return Ok(None);
401        }
402
403        let content = get_result
404            .documents
405            .into_iter()
406            .next()
407            .flatten()
408            .unwrap_or_default();
409        let metadata = get_result
410            .metadatas
411            .into_iter()
412            .next()
413            .flatten()
414            .unwrap_or_default();
415
416        Ok(Some(Document {
417            content,
418            metadata,
419            id: Some(id.to_string()),
420        }))
421    }
422
423    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
424        let url = self.collection_url("get")?;
425        let body = json!({
426            "ids": [id],
427            "include": ["embeddings"]
428        });
429
430        let response = self
431            .client
432            .post(&url)
433            .json(&body)
434            .send()
435            .await
436            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
437
438        if !response.status().is_success() {
439            return Ok(None);
440        }
441
442        let get_result: ChromaGetResponse = response.json().await.map_err(|e| {
443            VectorStoreError::StorageError(format!("failed to parse document: {}", e))
444        })?;
445
446        if let Some(embeddings) = get_result.embeddings {
447            Ok(embeddings.into_iter().next())
448        } else {
449            Ok(None)
450        }
451    }
452
453    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
454        let url = self.collection_url("delete")?;
455        let body = json!({
456            "ids": [id]
457        });
458
459        let response = self
460            .client
461            .post(&url)
462            .json(&body)
463            .send()
464            .await
465            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
466
467        if !response.status().is_success() {
468            let text = response.text().await.unwrap_or_default();
469            return Err(VectorStoreError::StorageError(format!(
470                "failed to delete document: {}",
471                text
472            )));
473        }
474
475        Ok(())
476    }
477
478    async fn count(&self) -> usize {
479        let url = match self.collection_url("count") {
480            Ok(u) => u,
481            Err(e) => {
482                log::warn!("ChromaDB count() failed to build URL: {}", e);
483                return 0;
484            }
485        };
486
487        let response = self.client.post(&url).send().await;
488        match response {
489            Ok(resp) => {
490                if resp.status().is_success() {
491                    match resp.json::<usize>().await {
492                        Ok(count) => count,
493                        Err(e) => {
494                            log::warn!("ChromaDB count() failed to parse response: {}", e);
495                            0
496                        }
497                    }
498                } else {
499                    log::warn!("ChromaDB count() request failed with non-success status");
500                    0
501                }
502            }
503            Err(e) => {
504                log::warn!("ChromaDB count() request error: {}", e);
505                0
506            }
507        }
508    }
509
510    async fn clear(&self) -> Result<(), VectorStoreError> {
511        // fetch all document IDs, then delete in bulk
512        let get_url = self.collection_url("get")?;
513        let body = json!({
514            "include": []
515        });
516
517        let response = self
518            .client
519            .post(&get_url)
520            .json(&body)
521            .send()
522            .await
523            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
524
525        if !response.status().is_success() {
526            let text = response.text().await.unwrap_or_default();
527            return Err(VectorStoreError::StorageError(format!(
528                "failed to fetch document list: {}",
529                text
530            )));
531        }
532
533        let get_result: ChromaGetResponse = response.json().await.map_err(|e| {
534            VectorStoreError::StorageError(format!("failed to parse document list: {}", e))
535        })?;
536
537        if get_result.ids.is_empty() {
538            return Ok(());
539        }
540
541        // delete in bulk
542        let del_url = self.collection_url("delete")?;
543        let del_body = json!({
544            "ids": get_result.ids
545        });
546
547        let response = self
548            .client
549            .post(&del_url)
550            .json(&del_body)
551            .send()
552            .await
553            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
554
555        if !response.status().is_success() {
556            let text = response.text().await.unwrap_or_default();
557            return Err(VectorStoreError::StorageError(format!(
558                "failed to clear collection: {}",
559                text
560            )));
561        }
562
563        Ok(())
564    }
565}
566
567/// S3: translates [`MetadataFilter`] → Chroma `where` filter dict.
568///
569/// A single-field condition becomes `{ key: { "$op": value } }` (Chroma v2 supports
570/// `$eq $ne $gt $gte $lt $lte $in $nin`); AND/OR combinations become
571/// `{ "$and": [...] }` / `{ "$or": [...] }`. Isomorphic to the Pinecone translation,
572/// but maintained per backend independently; the semantics are fully delegated to the server.
573pub fn filter_to_chroma(filter: &MetadataFilter) -> serde_json::Value {
574    fn op_str(op: FilterOp) -> &'static str {
575        match op {
576            FilterOp::Eq => "$eq",
577            FilterOp::Ne => "$ne",
578            FilterOp::Gt => "$gt",
579            FilterOp::Gte => "$gte",
580            FilterOp::Lt => "$lt",
581            FilterOp::Lte => "$lte",
582            FilterOp::In => "$in",
583            FilterOp::Nin => "$nin",
584        }
585    }
586    match filter {
587        MetadataFilter::Field { key, op, value } => {
588            serde_json::json!({ key.clone(): { op_str(*op): value.clone() } })
589        }
590        MetadataFilter::And(filters) => {
591            let items: Vec<serde_json::Value> = filters.iter().map(filter_to_chroma).collect();
592            serde_json::json!({ "$and": items })
593        }
594        MetadataFilter::Or(filters) => {
595            let items: Vec<serde_json::Value> = filters.iter().map(filter_to_chroma).collect();
596            serde_json::json!({ "$or": items })
597        }
598    }
599}
600
601#[cfg(test)]
602mod tests {
603    use super::*;
604
605    /// S3: single-field condition → Chroma `where` dict.
606    #[test]
607    fn test_filter_to_chroma_field() {
608        assert_eq!(
609            filter_to_chroma(&MetadataFilter::field("lang", FilterOp::Eq, "rust")),
610            serde_json::json!({ "lang": { "$eq": "rust" } })
611        );
612        assert_eq!(
613            filter_to_chroma(&MetadataFilter::field("year", FilterOp::Lt, 2020)),
614            serde_json::json!({ "year": { "$lt": 2020 } })
615        );
616    }
617
618    /// S3: AND/OR combination → nested `$and`/`$or`.
619    #[test]
620    fn test_filter_to_chroma_and_or() {
621        let f = MetadataFilter::or(vec![
622            MetadataFilter::field("lang", FilterOp::Eq, "python"),
623            MetadataFilter::and(vec![
624                MetadataFilter::field("lang", FilterOp::Eq, "rust"),
625                MetadataFilter::field("tags", FilterOp::In, vec!["ml"]),
626            ]),
627        ]);
628        assert_eq!(
629            filter_to_chroma(&f),
630            serde_json::json!({
631                "$or": [
632                    { "lang": { "$eq": "python" } },
633                    { "$and": [
634                        { "lang": { "$eq": "rust" } },
635                        { "tags": { "$in": ["ml"] } }
636                    ]}
637                ]
638            })
639        );
640    }
641
642    /// S3: without a filter, `where_filter` is None and the `where` field is not serialized.
643    #[test]
644    fn test_query_request_no_filter() {
645        let req = ChromaDBVectorStore::query_request(&[1.0, 2.0], 3, None);
646        assert!(req.where_filter.is_none());
647        let v = serde_json::to_value(&req).unwrap();
648        assert!(v.get("where").is_none());
649        assert_eq!(v["n_results"], 3);
650    }
651
652    /// S3: with a filter, the `where` field serializes to a Chroma dict.
653    #[test]
654    fn test_query_request_with_filter() {
655        let f = MetadataFilter::field("lang", FilterOp::Eq, "rust");
656        let req = ChromaDBVectorStore::query_request(&[1.0, 2.0], 3, Some(&f));
657        let v = serde_json::to_value(&req).unwrap();
658        assert_eq!(v["where"], serde_json::json!({ "lang": { "$eq": "rust" } }));
659    }
660}
lc_vector_stores/chromadb.rs

lc_vector_stores/
chromadb.rs
