Spaces:
Sleeping
Sleeping
| 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.") |