pdf-quiz-generator/backend/app/services/vector_service.py

import chromadb
from chromadb import EmbeddingFunction, Embeddings

from app.config import settings

_client = None


class LiteLLMEmbeddingFunction(EmbeddingFunction):
    """ChromaDB embedding function backed by LiteLLM."""

    def __call__(self, input: list[str]) -> Embeddings:
        import litellm
        kwargs = {
            "model": settings.LITELLM_EMBEDDING_MODEL,
            "input": input,
        }
        if settings.LITELLM_API_KEY:
            kwargs["api_key"] = settings.LITELLM_API_KEY
        if settings.LITELLM_API_BASE:
            kwargs["api_base"] = settings.LITELLM_API_BASE
        response = litellm.embedding(**kwargs)
        return [item["embedding"] for item in response.data]


def get_client() -> chromadb.PersistentClient:
    global _client
    if _client is None:
        _client = chromadb.PersistentClient(path=settings.CHROMA_PERSIST_DIR)
    return _client


def get_or_create_collection(document_id: int):
    client = get_client()
    # Only use custom embedding if model is configured and has a provider prefix
    use_ef = bool(settings.LITELLM_EMBEDDING_MODEL and "/" in settings.LITELLM_EMBEDDING_MODEL)
    ef = LiteLLMEmbeddingFunction() if use_ef else None
    kwargs = {"name": f"doc_{document_id}"}
    if ef:
        kwargs["embedding_function"] = ef
    return client.get_or_create_collection(**kwargs)


def delete_collection(document_id: int):
    client = get_client()
    try:
        client.delete_collection(name=f"doc_{document_id}")
    except Exception:
        pass


def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 200) -> list[str]:
    """Split text into overlapping chunks."""
    if len(text) <= chunk_size:
        return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap
    return chunks


def store_pages(document_id: int, pages: dict[int, str]):
    """Store page text as chunked embeddings in ChromaDB."""
    collection = get_or_create_collection(document_id)

    all_ids = []
    all_docs = []
    all_metadatas = []

    for page_num, text in pages.items():
        chunks = chunk_text(text)
        for i, chunk in enumerate(chunks):
            doc_id = f"doc_{document_id}_page_{page_num}_chunk_{i}"
            all_ids.append(doc_id)
            all_docs.append(chunk)
            all_metadatas.append({"page_num": page_num, "document_id": document_id})

    # ChromaDB has a batch limit; add in batches of 500
    batch_size = 500
    for i in range(0, len(all_ids), batch_size):
        collection.add(
            ids=all_ids[i:i + batch_size],
            documents=all_docs[i:i + batch_size],
            metadatas=all_metadatas[i:i + batch_size],
        )


def query_pages(
    document_id: int,
    query: str,
    start_page: int | None = None,
    end_page: int | None = None,
    n_results: int = 20,
) -> list[dict]:
    """Query vectorized content with optional page range filter."""
    collection = get_or_create_collection(document_id)

    where_filter = None
    if start_page is not None and end_page is not None:
        where_filter = {
            "$and": [
                {"page_num": {"$gte": start_page}},
                {"page_num": {"$lte": end_page}},
            ]
        }

    results = collection.query(
        query_texts=[query],
        n_results=n_results,
        where=where_filter,
    )

    docs = []
    if results and results["documents"]:
        for i, doc in enumerate(results["documents"][0]):
            meta = results["metadatas"][0][i] if results["metadatas"] else {}
            docs.append({"text": doc, "page_num": meta.get("page_num")})
    return docs


def get_pages_text(document_id: int, start_page: int, end_page: int) -> str:
    """Retrieve all stored text for a page range, ordered by page number."""
    collection = get_or_create_collection(document_id)

    where_filter = {
        "$and": [
            {"page_num": {"$gte": start_page}},
            {"page_num": {"$lte": end_page}},
        ]
    }

    # Get all documents in the range
    results = collection.get(
        where=where_filter,
        include=["documents", "metadatas"],
    )

    if not results or not results["documents"]:
        return ""

    # Sort by page number and chunk order
    paired = list(zip(results["documents"], results["metadatas"], results["ids"]))
    paired.sort(key=lambda x: (x[1].get("page_num", 0), x[2]))

    # Deduplicate overlapping chunks per page
    seen_pages = {}
    for doc, meta, doc_id in paired:
        page = meta.get("page_num", 0)
        if page not in seen_pages:
            seen_pages[page] = []
        seen_pages[page].append(doc)

    texts = []
    for page in sorted(seen_pages.keys()):
        # Join chunks for each page, removing overlap duplicates
        page_text = seen_pages[page][0]
        for chunk in seen_pages[page][1:]:
            # Find overlap and append only new content
            overlap_len = 200
            if len(chunk) > overlap_len:
                page_text += chunk[overlap_len:]
            else:
                page_text += chunk
        texts.append(f"--- Page {page} ---\n{page_text}")

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