litellm_rs/storage/vector/
backend.rs

1//! Vector store backend enum with dispatch methods
2
3use crate::config::models::file_storage::VectorDbConfig;
4use crate::utils::error::gateway_error::{GatewayError, Result};
5use tracing::info;
6
7use super::pinecone::PineconeStore;
8use super::qdrant::QdrantStore;
9use super::types::{SearchResult, VectorPoint};
10use super::weaviate::WeaviateStore;
11
12/// Vector store backend enum
13#[derive(Debug, Clone)]
14pub enum VectorStoreBackend {
15    /// Qdrant vector database
16    Qdrant(QdrantStore),
17    /// Weaviate vector database
18    Weaviate(WeaviateStore),
19    /// Pinecone vector database
20    Pinecone(PineconeStore),
21}
22
23impl VectorStoreBackend {
24    /// Create a new vector store instance
25    pub async fn new(config: &VectorDbConfig) -> Result<Self> {
26        info!("Initializing vector database: {}", config.db_type);
27
28        match config.db_type.as_str() {
29            "qdrant" => Ok(VectorStoreBackend::Qdrant(QdrantStore::new(config).await?)),
30            "weaviate" | "pinecone" => Err(GatewayError::Config(format!(
31                "Vector DB type '{}' is declared but not implemented yet. Implemented types: [\"qdrant\"]",
32                config.db_type
33            ))),
34            _ => Err(GatewayError::Config(format!(
35                "Unsupported vector DB type: {}",
36                config.db_type
37            ))),
38        }
39    }
40
41    /// Store a vector with metadata
42    pub async fn store(
43        &self,
44        id: &str,
45        vector: &[f32],
46        metadata: Option<serde_json::Value>,
47    ) -> Result<()> {
48        match self {
49            VectorStoreBackend::Qdrant(store) => store.store(id, vector, metadata).await,
50            VectorStoreBackend::Weaviate(store) => store.store(id, vector, metadata).await,
51            VectorStoreBackend::Pinecone(store) => store.store(id, vector, metadata).await,
52        }
53    }
54
55    /// Search for similar vectors
56    pub async fn search(
57        &self,
58        query_vector: &[f32],
59        limit: usize,
60        threshold: Option<f32>,
61    ) -> Result<Vec<SearchResult>> {
62        match self {
63            VectorStoreBackend::Qdrant(store) => store.search(query_vector, limit, threshold).await,
64            VectorStoreBackend::Weaviate(store) => {
65                store.search(query_vector, limit, threshold).await
66            }
67            VectorStoreBackend::Pinecone(store) => {
68                store.search(query_vector, limit, threshold).await
69            }
70        }
71    }
72
73    /// Delete a vector by ID
74    pub async fn delete(&self, id: &str) -> Result<()> {
75        match self {
76            VectorStoreBackend::Qdrant(store) => store.delete(id).await,
77            VectorStoreBackend::Weaviate(store) => store.delete(id).await,
78            VectorStoreBackend::Pinecone(store) => store.delete(id).await,
79        }
80    }
81
82    /// Get a vector by ID
83    pub async fn get(&self, id: &str) -> Result<Option<VectorPoint>> {
84        match self {
85            VectorStoreBackend::Qdrant(store) => store.get(id).await,
86            VectorStoreBackend::Weaviate(store) => store.get(id).await,
87            VectorStoreBackend::Pinecone(store) => store.get(id).await,
88        }
89    }
90
91    /// Health check
92    pub async fn health_check(&self) -> Result<()> {
93        match self {
94            VectorStoreBackend::Qdrant(store) => store.health_check().await,
95            VectorStoreBackend::Weaviate(store) => store.health_check().await,
96            VectorStoreBackend::Pinecone(store) => store.health_check().await,
97        }
98    }
99
100    /// Close connections
101    pub async fn close(&self) -> Result<()> {
102        match self {
103            VectorStoreBackend::Qdrant(_store) => Ok(()), // No explicit close needed for HTTP clients
104            VectorStoreBackend::Weaviate(_store) => Ok(()),
105            VectorStoreBackend::Pinecone(_store) => Ok(()),
106        }
107    }
108
109    /// Batch store vectors
110    pub async fn batch_store(&self, points: &[VectorPoint]) -> Result<()> {
111        match self {
112            VectorStoreBackend::Qdrant(store) => store.batch_store(points).await,
113            VectorStoreBackend::Weaviate(store) => store.batch_store(points).await,
114            VectorStoreBackend::Pinecone(store) => store.batch_store(points).await,
115        }
116    }
117
118    /// Count vectors in collection
119    pub async fn count(&self) -> Result<u64> {
120        match self {
121            VectorStoreBackend::Qdrant(store) => store.count().await,
122            VectorStoreBackend::Weaviate(store) => store.count().await,
123            VectorStoreBackend::Pinecone(store) => store.count().await,
124        }
125    }
126}
litellm_rs/storage/vector/backend.rs

litellm_rs/storage/vector/
backend.rs
