Spaces:
Running
Running
File size: 2,602 Bytes
2210481 fc0b415 2210481 126e90e fc0b415 126e90e fc0b415 126e90e fc0b415 126e90e 2210481 fc0b415 2210481 fc0b415 2210481 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
from typing import List, Dict, Optional
import os
class VectorDatabase:
def __init__(
self,
url=os.getenv("QDRANT_URL"),
api_key=os.getenv("QDRANT_API_KEY"),
collection_name="testing_col",
# embedding_model_name: str = "BAAI/bge-small-en", # Default model
):
"""
Initialize the Qdrant client, FastEmbed, and collection.
Args:
host (str): Host address of the Qdrant server.
port (int): Port of the Qdrant server.
collection_name (str): Name of the collection to use or create.
embedding_model_name (str): Name of the FastEmbed model to use.
"""
self.client = QdrantClient(url=url, api_key=api_key)
self.collection_name = collection_name
def upsert_documents(self, texts: List[str]):
# Insert into Qdrant
self.client.add(
collection_name=self.collection_name,
documents=texts,
)
print(
f"Inserted {len(texts)} documents into collection '{self.collection_name}'."
)
def _delete_collection(self):
"""
Delete a collection from the Qdrant database.
Args:
collection_name (str): Name of the collection to delete.
"""
# Check if the collection exists
collections = self.client.get_collections()
collection_names = [collection.name for collection in collections.collections]
print(f"Existing collections: {collection_names}")
if self.collection_name in collection_names:
# Delete the collection
self.client.delete_collection(self.collection_name)
print(f"Collection '{self.collection_name}' deleted.")
else:
print(f"Collection '{self.collection_name}' does not exist.")
def search_similar(self, query_text: str):
search_result = self.client.query(
collection_name=self.collection_name,
query_text=query_text,
limit=4,
)
documents = [item.document for item in search_result]
result = " ".join(document for document in documents if document)
return result
def list_collections(self):
"""
List all collections in the Qdrant database.
Returns:
List[str]: List of collection names.
"""
collections = self.client.get_collections().collections
return [collection.name for collection in collections]
|