senti/backend/database/vector/client.py

"""
Senti AI — Qdrant Vector Store Client.

Production vector store using Qdrant.
Collections:
    senti_knowledge → laws, regulations, finance docs
    senti_memory    → user memory embeddings (future)

Uses `query_points()` API (qdrant-client >= 1.12).
"""

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue,
)
import os
import uuid
from typing import Optional

QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333")
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "")

_qdrant_instance = None

def get_qdrant_client() -> QdrantClient:
    global _qdrant_instance
    if _qdrant_instance is not None:
        return _qdrant_instance
    
    import os
    # Use in-memory for test environment
    if os.environ.get('SENTI_ENV') == 'test' or \
       os.environ.get('PYTEST_CURRENT_TEST'):
        _qdrant_instance = QdrantClient(location=":memory:")
        return _qdrant_instance
    
    # Use persistent storage for dev/prod
    storage_path = os.path.join(
        os.path.dirname(__file__),
        '..', '..', 'storage', 'vector_store'
    )
    os.makedirs(storage_path, exist_ok=True)
    
    try:
        _qdrant_instance = QdrantClient(path=storage_path)
    except Exception:
        # If lock exists, fall back to in-memory
        _qdrant_instance = QdrantClient(location=":memory:")
    
    return _qdrant_instance

class SentiVectorStore:
    """
    Production vector store using Qdrant.
    Collections:
        senti_knowledge → laws, regulations, finance docs
        senti_memory    → user memory embeddings (future)
    """

    COLLECTIONS = {
        "knowledge": "senti_knowledge",
        "memory":    "senti_memory",
    }

    EMBEDDING_DIM = 1024  # BGE-M3 dimension

    def __init__(self):
        self.client = get_qdrant_client()

    def create_collections(self) -> None:
        """Create all required collections if they don't exist."""

        existing = [
            c.name for c in self.client.get_collections().collections
        ]

        for name, col_name in self.COLLECTIONS.items():
            if col_name not in existing:
                self.client.create_collection(
                    collection_name=col_name,
                    vectors_config=VectorParams(
                        size=self.EMBEDDING_DIM,
                        distance=Distance.COSINE,
                    ),
                )
                print(f"[QDRANT] Created collection: {col_name}")
            else:
                count = self.client.get_collection(col_name).points_count
                print(f"[QDRANT] Collection {col_name}: {count} documents")

    def add_documents(
        self,
        collection: str,
        texts: list[str],
        embeddings: list[list[float]],
        metadata: list[dict],
    ) -> int:
        """Add documents with embeddings to a collection."""

        col_name = self.COLLECTIONS.get(collection, collection)

        points = []
        for text, embedding, meta in zip(texts, embeddings, metadata):
            points.append(
                PointStruct(
                    id=str(uuid.uuid4()),
                    vector=embedding,
                    payload={
                        "text": text,
                        "source": meta.get("source", "unknown"),
                        "category": meta.get("category", "general"),
                        "jurisdiction": meta.get("jurisdiction", "KE"),
                        "effective_date": meta.get("effective_date", ""),
                        "chunk_type": meta.get("chunk_type", "content"),
                    },
                )
            )

        self.client.upsert(
            collection_name=col_name,
            points=points,
        )

        return len(points)

    def search(
        self,
        collection: str,
        query_embedding: list[float],
        limit: int = 5,
        filters: Optional[dict] = None,
    ) -> list[dict]:
        """
        Search for similar documents.
        Returns list with text, score, and metadata.

        Uses `query_points()` (qdrant-client >= 1.12).
        Falls back to keyword-only if collection is empty.
        """

        col_name = self.COLLECTIONS.get(collection, collection)

        qdrant_filter = None
        if filters:
            conditions = []
            for key, value in filters.items():
                conditions.append(
                    FieldCondition(key=key, match=MatchValue(value=value))
                )
            if conditions:
                qdrant_filter = Filter(must=conditions)

        try:
            response = self.client.query_points(
                collection_name=col_name,
                query=query_embedding,
                limit=limit,
                query_filter=qdrant_filter,
                with_payload=True,
            )

            return [
                {
                    "text": pt.payload.get("text", ""),
                    "score": pt.score if hasattr(pt, "score") else 0.0,
                    "source": pt.payload.get("source", ""),
                    "category": pt.payload.get("category", ""),
                    "jurisdiction": pt.payload.get("jurisdiction", "KE"),
                    "effective_date": pt.payload.get("effective_date", ""),
                }
                for pt in response.points
            ]
        except Exception as e:
            print(f"[QDRANT] Search failed: {e}")
            return []

    def get_count(self, collection: str) -> int:
        col_name = self.COLLECTIONS.get(collection, collection)
        return self.client.get_collection(col_name).points_count

    def health_check(self) -> bool:
        try:
            self.client.get_collections()
            return True
        except Exception:
            return False


vector_store = SentiVectorStore()