src/vector_store.py

"""Embedding and ChromaDB vector-store functions."""

from functools import lru_cache
import hashlib
import os

os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")

import chromadb
from sentence_transformers import SentenceTransformer

from src.config import CHROMA_DIR, COLLECTION_NAME, EMBEDDING_MODEL_NAME, TOP_K_RESULTS


@lru_cache(maxsize=1)
def load_embedding_model() -> SentenceTransformer:
    """Load the embedding model once so repeated questions are faster."""
    return SentenceTransformer(EMBEDDING_MODEL_NAME)


@lru_cache(maxsize=1)
def get_chroma_client():
    """Create one persistent ChromaDB client for the running server."""
    CHROMA_DIR.mkdir(parents=True, exist_ok=True)
    return chromadb.PersistentClient(path=str(CHROMA_DIR))


def get_chroma_collection():
    """Create or open the ChromaDB collection used by this app."""
    client = get_chroma_client()
    return client.get_or_create_collection(name=COLLECTION_NAME)


def clear_vector_database() -> None:
    """Clear old notes by resetting the ChromaDB collection."""
    client = get_chroma_client()

    try:
        client.delete_collection(name=COLLECTION_NAME)
    except Exception:
        pass

    client.get_or_create_collection(name=COLLECTION_NAME)


def get_vector_count() -> int:
    """Return how many chunks are stored in ChromaDB."""
    collection = get_chroma_collection()
    return collection.count()


def store_chunks_in_vector_database(chunks: list[dict]) -> int:
    """Convert chunks into embeddings and save them in ChromaDB."""
    if not chunks:
        raise ValueError("There are no chunks to store in the vector database.")

    model = load_embedding_model()
    collection = get_chroma_collection()

    texts = [chunk["text"] for chunk in chunks]
    embeddings = model.encode(texts, show_progress_bar=False).tolist()

    ids = [
        hashlib.sha1(
            f"{chunk['source_file']}|{chunk['page_number']}|{chunk['chunk_id']}|{chunk['text']}".encode("utf-8")
        ).hexdigest()
        for chunk in chunks
    ]
    metadata = [
        {
            "chunk_id": chunk["chunk_id"],
            "page_number": chunk["page_number"],
            "source_file": chunk["source_file"],
            "extraction_method": chunk.get("extraction_method", "embedded_text"),
            "ocr_mode": chunk.get("ocr_mode", "not_used"),
            "ocr_engine": chunk.get("ocr_engine", "not_used"),
        }
        for chunk in chunks
    ]

    collection.add(ids=ids, documents=texts, embeddings=embeddings, metadatas=metadata)
    return len(chunks)


def search_relevant_chunks(question: str, top_k: int = TOP_K_RESULTS) -> list[dict]:
    """Find chunks that are most similar to the user's question."""
    if not question.strip():
        raise ValueError("Please enter a question first.")

    model = load_embedding_model()
    collection = get_chroma_collection()

    if collection.count() == 0:
        raise ValueError("Please upload and process notes before asking a question.")

    question_embedding = model.encode([question], show_progress_bar=False).tolist()[0]
    results = collection.query(
        query_embeddings=[question_embedding],
        n_results=top_k,
        include=["documents", "metadatas", "distances"],
    )

    relevant_chunks = []
    documents = results.get("documents", [[]])[0]
    metadatas = results.get("metadatas", [[]])[0]
    distances = results.get("distances", [[]])[0]

    for document, metadata, distance in zip(documents, metadatas, distances):
        relevant_chunks.append(
            {
                "text": document,
                "page_number": metadata["page_number"],
                "source_file": metadata["source_file"],
                "chunk_id": metadata["chunk_id"],
                "extraction_method": metadata.get("extraction_method", "embedded_text"),
                "ocr_mode": metadata.get("ocr_mode", "not_used"),
                "ocr_engine": metadata.get("ocr_engine", "not_used"),
                "distance": distance,
            }
        )

    return relevant_chunks


def get_all_stored_chunks(max_characters: int = 12000) -> str:
    """Read stored chunks and combine them for summary-style tasks."""
    collection = get_chroma_collection()

    if collection.count() == 0:
        raise ValueError("Please upload and process notes first.")

    results = collection.get(include=["documents", "metadatas"])
    combined_parts = []
    current_length = 0

    for document, metadata in zip(results["documents"], results["metadatas"]):
        source_label = f"[{metadata['source_file']} - Page {metadata['page_number']}]"
        part = f"{source_label}\n{document}"

        if current_length + len(part) > max_characters:
            break

        combined_parts.append(part)
        current_length += len(part)

    return "\n\n".join(combined_parts)