vector_store.py

"""
Vector store
============

One ChromaDB collection per role. Role-scoped retrieval is the crux of
the RBAC wall: a finance user's query literally cannot pull engineering
documents because we only query the `finance_docs` collection.

Embeddings run locally via sentence-transformers (`all-MiniLM-L6-v2` by
default). No API key needed for the embedding step. The only external
call in the whole pipeline is the LLM generation in rag_chat.py.
"""

from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import Dict, List

import chromadb
from chromadb.errors import NotFoundError
from chromadb.utils import embedding_functions

from config import CHROMA_PERSIST_DIR, DOCUMENTS_DIR, EMBEDDING_MODEL, ROLES

logger = logging.getLogger(__name__)


class VectorDocumentStore:
    def __init__(
        self,
        db_path: Path = CHROMA_PERSIST_DIR,
        docs_dir: Path = DOCUMENTS_DIR,
        embedding_model: str = EMBEDDING_MODEL,
    ) -> None:
        self.db_path = db_path
        self.docs_dir = docs_dir

        self.docs_dir.mkdir(parents=True, exist_ok=True)
        for role in ROLES:
            (self.docs_dir / role).mkdir(parents=True, exist_ok=True)

        self.chroma_client = chromadb.PersistentClient(path=str(db_path))
        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=embedding_model
        )

        self.collections: Dict[str, chromadb.Collection] = {}
        self._init_collections()

    def _init_collections(self) -> None:
        for role in ROLES:
            name = f"{role}_docs"
            try:
                self.collections[role] = self.chroma_client.get_collection(
                    name=name,
                    embedding_function=self.embedding_function,
                )
                logger.info("Loaded existing collection %s", name)
            except (NotFoundError, ValueError):
                self.collections[role] = self.chroma_client.create_collection(
                    name=name,
                    embedding_function=self.embedding_function,
                )
                logger.info("Created new collection %s", name)

    def add_document(self, role: str, title: str, content: str) -> bool:
        if role not in ROLES:
            logger.warning("Invalid role: %s", role)
            return False

        try:
            safe_title = title.replace(" ", "_").replace("/", "_").replace("\\", "_")
            doc_path = self.docs_dir / role / f"{safe_title}.txt"
            with open(doc_path, "w") as f:
                f.write(content)

            doc_id = f"{role}_{safe_title}_{int(time.time())}"

            existing = self.collections[role].get(where={"title": title})
            if existing and existing.get("ids"):
                self.collections[role].delete(ids=existing["ids"])
                logger.info("Replaced existing document '%s' for role '%s'", title, role)

            self.collections[role].add(
                documents=[content],
                metadatas=[{"title": title, "path": str(doc_path)}],
                ids=[doc_id],
            )
            logger.info("Added document '%s' for role '%s'", title, role)
            return True
        except Exception as e:
            logger.error("Error adding document: %s", e)
            return False

    def list_documents(self, role: str) -> List[str]:
        if role not in ROLES:
            return []
        try:
            result = self.collections[role].get()
            metadatas = result.get("metadatas") or []
            return [m.get("title", "Untitled") for m in metadatas]
        except Exception as e:
            logger.error("Error listing documents: %s", e)
            return []

    def search_documents(self, role: str, query: str, top_k: int = 3) -> List[Dict]:
        if role not in ROLES:
            return []

        try:
            results = self.collections[role].query(query_texts=[query], n_results=top_k)
            documents = results.get("documents") or [[]]
            metadatas = results.get("metadatas") or [[]]
            distances = results.get("distances") or [[]]

            if not documents or not documents[0]:
                return []

            formatted: List[Dict] = []
            for i, doc in enumerate(documents[0]):
                meta = metadatas[0][i] if i < len(metadatas[0]) else {}
                distance = distances[0][i] if distances and i < len(distances[0]) else 1.0
                formatted.append(
                    {
                        "title": meta.get("title", "Untitled"),
                        "content": doc,
                        "score": 1.0 - min(distance, 1.0),
                    }
                )
            return formatted
        except Exception as e:
            logger.error("Error searching documents: %s", e)
            return []