src/rag/vectorstore.py

"""
Vector Store for Kalim Chatbot
Uses ChromaDB for storing and retrieving major embeddings
"""
import os
from typing import List, Dict, Any, Optional
from pathlib import Path

try:
    import chromadb
    from chromadb.config import Settings
    HAS_CHROMADB = True
except ImportError:
    HAS_CHROMADB = False

from config.settings import VECTORSTORE_DIR, TOP_K_RESULTS
from data.loader import load_majors
from rag.embeddings import (
    get_embedding_model,
    create_major_document,
    create_major_metadata
)


class KalimVectorStore:
    """Vector store for Kalim major recommendations."""

    COLLECTION_NAME = "kalim_majors"

    def __init__(self, persist_directory: Optional[Path] = None):
        """Initialize the vector store.

        Args:
            persist_directory: Directory to persist the database
        """
        if not HAS_CHROMADB:
            raise ImportError(
                "chromadb is required. "
                "Install with: pip install chromadb"
            )

        self.persist_directory = persist_directory or VECTORSTORE_DIR
        self.persist_directory.mkdir(parents=True, exist_ok=True)

        # Initialize ChromaDB client with persistence
        self._client = chromadb.PersistentClient(
            path=str(self.persist_directory),
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )

        self._collection = None
        self._embedding_model = None

    @property
    def embedding_model(self):
        """Lazy load embedding model."""
        if self._embedding_model is None:
            self._embedding_model = get_embedding_model()
        return self._embedding_model

    @property
    def collection(self):
        """Get or create the collection."""
        if self._collection is None:
            self._collection = self._client.get_or_create_collection(
                name=self.COLLECTION_NAME,
                metadata={"description": "Lebanese University majors for Kalim chatbot"}
            )
        return self._collection

    def build_index(self, majors: Optional[List[Dict]] = None, force_rebuild: bool = False):
        """Build the vector index from majors data.

        Args:
            majors: List of major dictionaries. If None, loads from JSON.
            force_rebuild: If True, rebuilds even if collection exists
        """
        if majors is None:
            majors = load_majors()

        # Check if already built
        if self.collection.count() > 0 and not force_rebuild:
            print(f"Collection already has {self.collection.count()} documents. "
                  "Use force_rebuild=True to rebuild.")
            return

        # Clear existing data if rebuilding
        if force_rebuild and self.collection.count() > 0:
            print("Clearing existing collection...")
            self._client.delete_collection(self.COLLECTION_NAME)
            self._collection = self._client.create_collection(
                name=self.COLLECTION_NAME,
                metadata={"description": "Lebanese University majors for Kalim chatbot"}
            )

        print(f"Building index for {len(majors)} majors...")

        # Prepare documents, embeddings, and metadata
        documents = []
        embeddings = []
        metadatas = []
        ids = []

        for major in majors:
            doc = create_major_document(major)
            metadata = create_major_metadata(major)
            major_id = str(major.get('id', len(ids) + 1))

            documents.append(doc)
            metadatas.append(metadata)
            ids.append(major_id)

        # Generate embeddings in batch
        print("Generating embeddings...")
        embeddings = self.embedding_model.embed_texts(documents)

        # Add to collection
        print("Adding to collection...")
        self.collection.add(
            documents=documents,
            embeddings=embeddings,
            metadatas=metadatas,
            ids=ids
        )

        print(f"Index built successfully with {self.collection.count()} documents")

    def search(
        self,
        query: str,
        n_results: int = TOP_K_RESULTS,
        where: Optional[Dict] = None,
        where_document: Optional[Dict] = None
    ) -> List[Dict[str, Any]]:
        """Search for similar majors.

        Args:
            query: Search query text
            n_results: Number of results to return
            where: Metadata filter conditions
            where_document: Document content filter

        Returns:
            List of matching majors with scores
        """
        # Generate query embedding
        query_embedding = self.embedding_model.embed_query(query)

        # Search
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
            where_document=where_document,
            include=["documents", "metadatas", "distances"]
        )

        # Format results
        formatted_results = []
        if results and results['ids'] and results['ids'][0]:
            for i, doc_id in enumerate(results['ids'][0]):
                result = {
                    'id': doc_id,
                    'document': results['documents'][0][i] if results['documents'] else '',
                    'metadata': results['metadatas'][0][i] if results['metadatas'] else {},
                    'distance': results['distances'][0][i] if results['distances'] else 0.0,
                    'score': 1 - results['distances'][0][i] if results['distances'] else 1.0
                }
                formatted_results.append(result)

        return formatted_results

    def search_by_holland(
        self,
        holland_codes: List[str],
        n_results: int = TOP_K_RESULTS
    ) -> List[Dict[str, Any]]:
        """Search for majors matching Holland codes.

        Args:
            holland_codes: List of Holland codes (e.g., ['E', 'C', 'S'])
            n_results: Number of results

        Returns:
            List of matching majors
        """
        if not holland_codes:
            return []

        # Build query from Holland descriptions
        from config.settings import HOLLAND_CODES
        query_parts = []
        for code in holland_codes:
            if code in HOLLAND_CODES:
                query_parts.append(HOLLAND_CODES[code]['description_ar'])

        query = " ".join(query_parts) if query_parts else " ".join(holland_codes)

        # Build filter for primary Holland code
        primary_code = holland_codes[0] if holland_codes else None
        where_filter = None
        if primary_code:
            where_filter = {"holland_primary": primary_code}

        return self.search(
            query=query,
            n_results=n_results,
            where=where_filter
        )

    def get_all_majors(self) -> List[Dict[str, Any]]:
        """Get all majors from the collection.

        Returns:
            List of all major metadata
        """
        results = self.collection.get(
            include=["documents", "metadatas"]
        )

        formatted = []
        if results and results['ids']:
            for i, doc_id in enumerate(results['ids']):
                formatted.append({
                    'id': doc_id,
                    'document': results['documents'][i] if results['documents'] else '',
                    'metadata': results['metadatas'][i] if results['metadatas'] else {}
                })

        return formatted

    def get_major_by_id(self, major_id: str) -> Optional[Dict[str, Any]]:
        """Get a specific major by ID.

        Args:
            major_id: Major ID string

        Returns:
            Major data or None
        """
        results = self.collection.get(
            ids=[major_id],
            include=["documents", "metadatas"]
        )

        if results and results['ids']:
            return {
                'id': results['ids'][0],
                'document': results['documents'][0] if results['documents'] else '',
                'metadata': results['metadatas'][0] if results['metadatas'] else {}
            }

        return None

    def reset(self):
        """Reset the collection (delete all data)."""
        self._client.delete_collection(self.COLLECTION_NAME)
        self._collection = None
        print("Collection reset successfully")


def get_vectorstore() -> KalimVectorStore:
    """Get the vector store instance."""
    return KalimVectorStore()


def build_vectorstore(force_rebuild: bool = False) -> KalimVectorStore:
    """Build and return the vector store.

    Args:
        force_rebuild: If True, rebuilds the index

    Returns:
        Built vector store
    """
    store = get_vectorstore()
    store.build_index(force_rebuild=force_rebuild)
    return store


if __name__ == "__main__":
    # Build the vector store when run directly
    print("Building Kalim Vector Store...")
    store = build_vectorstore(force_rebuild=True)
    print(f"Vector store built with {store.collection.count()} documents")

    # Test search
    print("\nTesting search for 'هندسة البرمجيات'...")
    results = store.search("هندسة البرمجيات", n_results=5)
    for r in results:
        print(f"  - {r['metadata'].get('name_ar')} (score: {r['score']:.3f})")