app/core/database.py

import os
import requests
from dotenv import load_dotenv

load_dotenv()

QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
COLLECTION_NAME = "physical_ai_textbook"

if not QDRANT_URL or not QDRANT_API_KEY:
    raise ValueError("QDRANT_URL and QDRANT_API_KEY must be set in the .env file")

# Ensure URL doesn't end with slash + handle if user put "https://" or not
if not QDRANT_URL.startswith("http"):
    QDRANT_URL = f"https://{QDRANT_URL}"
QDRANT_URL = QDRANT_URL.rstrip("/")

HEADERS = {
    "api-key": QDRANT_API_KEY,
    "Content-Type": "application/json"
}

def init_db():
    """
    Initializes the Qdrant collection via REST API.
    """
    # Check if collection exists
    check_url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
    response = requests.get(check_url, headers=HEADERS)

    if response.status_code == 200:
        print(f"Collection {COLLECTION_NAME} already exists.")
    else:
        print(f"Creating collection: {COLLECTION_NAME}")
        # Create collection
        create_url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}"
        payload = {
            "vectors": {
                "size": 768,
                "distance": "Cosine"
            }
        }
        resp = requests.put(create_url, headers=HEADERS, json=payload)
        if resp.status_code == 200:
            print("Collection created successfully.")
        else:
            print(f"Error creating collection: {resp.text}")

def search_points(vector, limit=5):
    url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points/search"
    payload = {
        "vector": vector,
        "limit": limit,
        "with_payload": True
    }
    response = requests.post(url, headers=HEADERS, json=payload)
    if response.status_code == 200:
        return response.json().get("result", [])
    else:
        print(f"Search Error: {response.text}")
        return []

def upsert_points(points):
    url = f"{QDRANT_URL}/collections/{COLLECTION_NAME}/points?wait=true"
    payload = {
        "points": points
    }
    response = requests.put(url, headers=HEADERS, json=payload)
    if response.status_code != 200:
        print(f"Upsert Error: {response.text}")