api/db.py

import os
from typing import List, Dict, Any
from dotenv import load_dotenv
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct, VectorParams, Distance
from sentence_transformers import SentenceTransformer

# Load environment variables from a .env file if it exists
load_dotenv()

# --- Qdrant Configuration ---
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333").strip()
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", "").strip()
COLLECTION = "reports"

# --- Model Loading (from local files within the Docker image) ---
# This path corresponds to where the Dockerfile copies the model.
MODEL_PATH = "./models/all-MiniLM-L6-v2" 

# Initialize the embedding model from the local path.
try:
    print(f"Loading sentence-transformer model from local path: {MODEL_PATH}")
    embedding_model = SentenceTransformer(MODEL_PATH)
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ FATAL: Could not load the embedding model from {MODEL_PATH}.")
    print("This indicates an issue with the Docker build or the file path in db.py.")
    raise e

# --- Qdrant Client and Collection Setup ---
def _make_client() -> QdrantClient:
    """Creates a Qdrant client based on environment variables."""
    if QDRANT_API_KEY:
        return QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=30.0, check_compatibility=False)
    else:
        return QdrantClient(url=QDRANT_URL, timeout=30.0, check_compatibility=False)

qdrant = _make_client()

def _ensure_collection():
    """Ensures the Qdrant collection exists, creating it if necessary."""
    try:
        qdrant.get_collection(collection_name=COLLECTION)
        print(f"✅ Collection '{COLLECTION}' already exists.")
    except Exception:
        print(f"🔧 Collection '{COLLECTION}' not found. Creating it...")
        qdrant.create_collection(
            collection_name=COLLECTION,
            vectors_config=VectorParams(size=384, distance=Distance.COSINE),
        )
        print("✅ Collection created.")

_ensure_collection()

# --- Database Functions ---
def save_report(report_id: str, text: str, title: str):
    """Encodes and saves a report to Qdrant."""
    vector = embedding_model.encode(text).tolist()
    qdrant.upsert(
        collection_name=COLLECTION,
        points=[PointStruct(id=report_id, vector=vector, payload={"text": text, "title": title})],
    )

def list_reports() -> List[Dict[str, Any]]:
    """Lists recent reports, including their titles."""
    hits, _ = qdrant.scroll(collection_name=COLLECTION, limit=50)
    return [{"id": h.id, "title": h.payload.get("title", "(untitled)"), "text": h.payload.get("text", "")} for h in hits]

def search_reports(query: str) -> List[Dict[str, Any]]:
    """Performs semantic search and returns reports with titles."""
    vector = embedding_model.encode(query).tolist()
    hits = qdrant.search(collection_name=COLLECTION, query_vector=vector, limit=5)
    return [{"id": hit.id, "score": float(hit.score), "title": hit.payload.get("title", "(untitled)"), "text": hit.payload.get("text", "")} for hit in hits]