feat(server): add qdrant vector db integration
This commit is contained in:
1 parent
def473c856
commit
61c0e88ad8
3 files changed
+497
No files matched your search
@@ -0,0 +1,85 @@
|
||||
use qdrant_client::qdrant::{CreateCollectionBuilder, Distance, PointStruct, VectorParamsBuilder, UpsertPointsBuilder};
|
||||
use qdrant_client::Qdrant;
|
||||
use std::sync::Arc;
|
||||
use std::error::Error;
|
||||
use uuid::Uuid;
|
||||
use log::{info, error};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct VectorDB {
|
||||
client: Arc<Qdrant>,
|
||||
collection_name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct VectorSearchResult {
|
||||
pub id: String,
|
||||
pub doc_type: String,
|
||||
pub text: String,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
impl VectorDB {
|
||||
pub async fn new(url: &str, collection_name: &str) -> Result<Self, Box<dyn Error + Send + Sync>> {
|
||||
let client = Qdrant::from_url(url).build()?;
|
||||
|
||||
let db = Self {
|
||||
client: Arc::new(client),
|
||||
collection_name: collection_name.to_string(),
|
||||
};
|
||||
|
||||
db.init_collection().await?;
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
async fn init_collection(&self) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
// Fastembed AllMiniLML6V2 uses 384 dimensions
|
||||
let vector_params = VectorParamsBuilder::new(384, Distance::Cosine).build();
|
||||
|
||||
let collection_exists = self.client.collection_exists(&self.collection_name).await?;
|
||||
if !collection_exists {
|
||||
self.client
|
||||
.create_collection(
|
||||
CreateCollectionBuilder::new(&self.collection_name)
|
||||
.vectors_config(vector_params)
|
||||
)
|
||||
.await?;
|
||||
info!("Created Qdrant collection: {}", self.collection_name);
|
||||
} else {
|
||||
info!("Qdrant collection {} already exists", self.collection_name);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn index_document(
|
||||
&self,
|
||||
id: &str,
|
||||
doc_type: &str,
|
||||
text: &str,
|
||||
vector: Vec<f32>,
|
||||
) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let point_id = match Uuid::parse_str(id) {
|
||||
Ok(u) => u.to_string(),
|
||||
Err(_) => {
|
||||
// If it's not a valid UUID, let's create a deterministic UUID based on the string
|
||||
let namespace = Uuid::NAMESPACE_OID;
|
||||
Uuid::new_v5(&namespace, id.as_bytes()).to_string()
|
||||
}
|
||||
};
|
||||
|
||||
let mut payload = std::collections::HashMap::new();
|
||||
payload.insert("doc_type".to_string(), doc_type.into());
|
||||
payload.insert("text".to_string(), text.into());
|
||||
payload.insert("original_id".to_string(), id.into());
|
||||
|
||||
let point = PointStruct::new(point_id, vector, payload);
|
||||
|
||||
self.client
|
||||
.upsert_points(UpsertPointsBuilder::new(&self.collection_name, vec![point]))
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user