|
|
""" |
|
|
Vector Database Service implementation for Qdrant |
|
|
""" |
|
|
|
|
|
from typing import List, Dict, Any, Optional |
|
|
from qdrant_client import QdrantClient |
|
|
from qdrant_client.models import PointStruct, VectorParams, Distance, Record |
|
|
|
|
|
|
|
|
|
|
|
class VectorDatabaseClient: |
|
|
"""Client for interacting with Qdrant vector database""" |
|
|
|
|
|
def __init__(self, url: str, api_key: str, collection_name: str, embedding_size: int): |
|
|
"""Initialize Qdrant client and collection settings |
|
|
|
|
|
Args: |
|
|
url: Qdrant server URL |
|
|
api_key: API key for Qdrant |
|
|
collection_name: Name of the collection to use |
|
|
embedding_size: Size of embedding vectors |
|
|
""" |
|
|
self.client = QdrantClient(url=url, api_key=api_key) |
|
|
self.collection_name = collection_name |
|
|
self.embedding_size = embedding_size |
|
|
|
|
|
def ensure_collection_exists(self): |
|
|
"""Ensure the collection exists, create it if it doesn't""" |
|
|
collections = [c.name for c in self.client.get_collections().collections] |
|
|
|
|
|
if self.collection_name not in collections: |
|
|
self.client.create_collection( |
|
|
collection_name=self.collection_name, |
|
|
vectors_config=VectorParams( |
|
|
size=self.embedding_size, |
|
|
distance=Distance.COSINE |
|
|
) |
|
|
) |
|
|
print(f"✅ Collection '{self.collection_name}' created.") |
|
|
else: |
|
|
print(f"ℹ️ Collection '{self.collection_name}' already exists.") |
|
|
|
|
|
def add_embedding(self, id: str, embedding: List[float], filename: str, metadata: Optional[str] = None) -> str: |
|
|
"""Add an embedding to the collection |
|
|
|
|
|
Args: |
|
|
id: Unique ID for the point |
|
|
embedding: Vector embedding |
|
|
filename: Original filename |
|
|
metadata: Optional metadata as JSON string |
|
|
|
|
|
Returns: |
|
|
ID of the added point |
|
|
""" |
|
|
payload = {"filename": filename} |
|
|
if metadata: |
|
|
payload["metadata"] = metadata |
|
|
|
|
|
self.client.upsert( |
|
|
collection_name=self.collection_name, |
|
|
points=[ |
|
|
PointStruct( |
|
|
id=id, |
|
|
vector=embedding, |
|
|
payload=payload |
|
|
) |
|
|
] |
|
|
) |
|
|
return id |
|
|
|
|
|
def add_embedding_with_payload(self, id: str, embedding: List[float], payload: Dict[str, Any]) -> str: |
|
|
"""Add an embedding with a custom payload |
|
|
|
|
|
Args: |
|
|
id: Unique ID for the point |
|
|
embedding: Vector embedding |
|
|
payload: Dictionary of metadata to store |
|
|
|
|
|
Returns: |
|
|
ID of the added point |
|
|
""" |
|
|
self.client.upsert( |
|
|
collection_name=self.collection_name, |
|
|
points=[ |
|
|
PointStruct( |
|
|
id=id, |
|
|
vector=embedding, |
|
|
payload=payload |
|
|
) |
|
|
] |
|
|
) |
|
|
return id |
|
|
|
|
|
def search_by_embedding(self, embedding: List[float], limit: int = 5) -> List[Record]: |
|
|
"""Search for similar vectors |
|
|
|
|
|
Args: |
|
|
embedding: Query vector |
|
|
limit: Maximum number of results |
|
|
|
|
|
Returns: |
|
|
List of search results |
|
|
""" |
|
|
results = self.client.search( |
|
|
collection_name=self.collection_name, |
|
|
query_vector=embedding, |
|
|
limit=limit |
|
|
) |
|
|
return results |
|
|
|
|
|
def search_by_id(self, id: str, limit: int = 1) -> List[Record]: |
|
|
"""Search for similar vectors using an existing vector as query |
|
|
|
|
|
Args: |
|
|
id: ID of the existing vector to use as query |
|
|
limit: Maximum number of results |
|
|
|
|
|
Returns: |
|
|
List of search results |
|
|
""" |
|
|
|
|
|
vector = self.client.retrieve( |
|
|
collection_name=self.collection_name, |
|
|
ids=[id] |
|
|
) |
|
|
|
|
|
if not vector or len(vector) == 0: |
|
|
return [] |
|
|
|
|
|
|
|
|
return self.search_by_embedding(vector[0].vector, limit) |
|
|
|
|
|
def delete_embedding(self, id: str) -> bool: |
|
|
"""Delete an embedding from the collection |
|
|
|
|
|
Args: |
|
|
id: ID of the embedding to delete |
|
|
|
|
|
Returns: |
|
|
True if deleted, False if not found |
|
|
""" |
|
|
self.client.delete( |
|
|
collection_name=self.collection_name, |
|
|
points_selector=[id] |
|
|
) |
|
|
return True |
|
|
|
|
|
def list_collections(self) -> List[str]: |
|
|
"""List all collections in the database |
|
|
|
|
|
Returns: |
|
|
List of collection names |
|
|
""" |
|
|
return [c.name for c in self.client.get_collections().collections] |
