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 the LiteLLM proxy (direct httpx)."""

    def __call__(self, input: list[str]) -> Embeddings:
        import httpx, json as _json
        from app.services.embedding_service import _get_embedding_model
        model = _get_embedding_model() or settings.LITELLM_EMBEDDING_MODEL
        api_base = (settings.LITELLM_API_BASE or "").rstrip("/")
        body = {"model": model, "input": input, "dimensions": settings.EMBEDDING_DIMENSIONS}
        resp = httpx.post(
            f"{api_base}/v1/embeddings",
            headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}", "Content-Type": "application/json"},
            content=_json.dumps(body),
            timeout=60,
        )
        resp.raise_for_status()
        return [item["embedding"] for item in resp.json()["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()
    from app.services.embedding_service import _get_embedding_model
    active_model = _get_embedding_model() or settings.LITELLM_EMBEDDING_MODEL
    # Only use custom embedding if model and API base are configured
    use_ef = bool(active_model and settings.LITELLM_API_BASE)
    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)