Spaces:
Sleeping
Sleeping
File size: 4,080 Bytes
aabd1d8 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import os
from qdrant_client import QdrantClient
from qdrant_client.http import models
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
class VectorDB:
def __init__(self, collection_name: str = "pro_rag_container"):
self.collection_name = collection_name
# --- 1. CLOUD vs LOCAL LOGIC ---
qdrant_url = os.getenv("QDRANT_URL")
qdrant_key = os.getenv("QDRANT_API_KEY")
if qdrant_url and qdrant_key:
print("☁️ Connecting to Qdrant Cloud...")
self.client = QdrantClient(url=qdrant_url, api_key=qdrant_key)
else:
print("🏠 Connecting to Local Docker...")
self.client = QdrantClient(url="http://localhost:6333")
# --- 2. THE MISSING FUNCTION ---
def create_collection(self, vector_size: int = 3072):
"""
Creates the collection if it doesn't exist.
Using 3072 dimensions for OpenAI text-embedding-3-large.
"""
# Check if collection exists
if self.client.collection_exists(collection_name=self.collection_name):
print(f"ℹ️ Collection '{self.collection_name}' already exists. Skipping creation.")
return
print(f"⚙️ Creating collection '{self.collection_name}' with size {vector_size}...")
# Create Collection with Cosine Similarity
self.client.create_collection(
collection_name=self.collection_name,
vectors_config=models.VectorParams(
size=vector_size,
distance=models.Distance.COSINE
)
)
print(f"✅ Collection '{self.collection_name}' created successfully!")
def reset_database(self):
"""
Deletes the collection.
"""
self.client.delete_collection(collection_name=self.collection_name)
print(f"⚠️ Collection '{self.collection_name}' has been DELETED.")
# import os
# from qdrant_client import QdrantClient
# from qdrant_client.http import models
# from dotenv import load_dotenv
# # Load environment variables (API Keys, etc.)
# load_dotenv()
# class VectorDB:
# def __init__(self, collection_name: str = "pro_rag_v1"):
# """
# Initialize connection to Qdrant (Docker).
# """
# self.collection_name = collection_name
# self.client = QdrantClient(url="http://localhost:6333")
# # Verify connection immediately
# try:
# self.client.get_collections()
# print(f"✅ Connected to Qdrant Database at http://localhost:6333")
# except Exception as e:
# print(f"❌ Could not connect to Qdrant. Is Docker running? Error: {e}")
# def create_collection(self, vector_size: int = 3072):
# """
# Creates the collection if it doesn't exist.
# Using 3072 dimensions for OpenAI text-embedding-3-large.
# """
# # Check if collection exists
# if self.client.collection_exists(collection_name=self.collection_name):
# print(f"ℹ️ Collection '{self.collection_name}' already exists. Skipping creation.")
# return
# print(f"⚙️ Creating collection '{self.collection_name}' with size {vector_size}...")
# # Create Collection with Cosine Similarity
# self.client.create_collection(
# collection_name=self.collection_name,
# vectors_config=models.VectorParams(
# size=vector_size,
# distance=models.Distance.COSINE
# )
# )
# print(f"✅ Collection '{self.collection_name}' created successfully!")
# def reset_database(self):
# """
# DANGEROUS: Deletes the collection. Used for restarting the POC.
# """
# self.client.delete_collection(collection_name=self.collection_name)
# print(f"⚠️ Collection '{self.collection_name}' has been DELETED.") |