app/services/embeddings.py

"""Author RAG Chatbot SaaS — Document Embedder.

Generates vector embeddings using OpenAI text-embedding-3-small.
Batches chunks to minimize API calls. Stores vectors in ChromaDB.
RULE: Always batch embeddings — never embed one chunk at a time.
RULE: Always namespace ChromaDB collections by author_id + book_id.
"""

import asyncio
from typing import Any

import structlog
from openai import AsyncOpenAI

from app.config import get_settings
from app.core.chroma_client import get_chroma_client
from app.services.chunker import TextChunk

logger = structlog.get_logger(__name__)
cfg = get_settings()

_openai_client: AsyncOpenAI | None = None


def _get_openai() -> AsyncOpenAI:
    """Lazily create and cache OpenAI async client."""
    global _openai_client
    if _openai_client is None:
        _openai_client = AsyncOpenAI(api_key=cfg.OPENAI_API_KEY)
    return _openai_client


def _get_chroma():
    """Return the shared ChromaDB client."""
    return get_chroma_client()


def get_collection_name(author_id: str, book_id: str) -> str:
    """Build the ChromaDB collection name for an author's book.

    Format: author_{short_id}_book_{short_id} (ChromaDB name limits apply).

    Args:
        author_id: UUID of the author.
        book_id: UUID of the book.

    Returns:
        Collection name string.
    """
    # SEC-4 fix: use the full 32-char hex UUID for both segments instead of [:12]
    # prefix. Truncation to 12 chars risked cross-tenant prefix collision between
    # authors whose UUIDs share the same time-based prefix. Full UUID guarantees
    # tenant isolation. Both creator (embed_and_store) and reader (vector_store.py)
    # call this function, so one change fixes the entire naming contract.
    full_author = author_id.replace("-", "")
    full_book = book_id.replace("-", "")
    return f"a{full_author}_b{full_book}"


async def embed_and_store(
    chunks: list[TextChunk],
    author_id: str,
    book_id: str,
    book_title: str,
) -> str:
    """Generate embeddings for all chunks and store them in ChromaDB.

    Processes chunks in batches of EMBEDDING_BATCH_SIZE.

    Args:
        chunks: List of TextChunk objects from the chunker.
        author_id: UUID of the author (for namespacing).
        book_id: UUID of the book (for collection naming).
        book_title: Title of the book (stored as metadata).

    Returns:
        ChromaDB collection name (stored on the book record).
    """
    collection_name = get_collection_name(author_id, book_id)
    chroma = _get_chroma()

    # Create or get collection
    collection = chroma.get_or_create_collection(
        name=collection_name,
        metadata={"author_id": author_id, "book_id": book_id, "book_title": book_title},
    )

    # Delete any existing embeddings (re-processing case)
    existing = collection.count()
    if existing > 0:
        collection.delete(where={"book_id": {"$eq": book_id}})
        logger.info("Cleared existing embeddings for re-processing", collection=collection_name, count=existing)

    # Process in batches
    batch_size = cfg.EMBEDDING_BATCH_SIZE
    total_embedded = 0

    for batch_start in range(0, len(chunks), batch_size):
        batch = chunks[batch_start: batch_start + batch_size]
        texts = [chunk.text for chunk in batch]

        # Generate embeddings
        embeddings = await _generate_embeddings(texts)

        # Prepare ChromaDB documents
        ids = [f"{book_id}_chunk_{chunk.chunk_index}" for chunk in batch]
        metadatas = [
            {
                "author_id": author_id,
                "book_id": book_id,
                "book_title": book_title,
                "chunk_index": chunk.chunk_index,
                "char_start": chunk.char_start,
                "char_end": chunk.char_end,
                "token_count": chunk.token_count,
            }
            for chunk in batch
        ]

        collection.add(
            ids=ids,
            embeddings=embeddings,
            documents=texts,
            metadatas=metadatas,
        )
        total_embedded += len(batch)
        logger.debug("Embedded batch", batch_size=len(batch), total=total_embedded)

    logger.info("Embedding complete", collection=collection_name, total_chunks=total_embedded)
    return collection_name


async def _generate_embeddings(texts: list[str]) -> list[list[float]]:
    """Call OpenAI Embeddings API for a batch of texts.

    Args:
        texts: List of strings to embed.

    Returns:
        List of embedding vectors (floats).
    """
    client = _get_openai()
    response = await client.embeddings.create(
        model=cfg.OPENAI_EMBEDDING_MODEL,
        input=texts,
    )
    return [item.embedding for item in response.data]


def delete_book_embeddings(author_id: str, book_id: str) -> None:
    """Delete all embeddings for a book from ChromaDB.

    Args:
        author_id: UUID of the author.
        book_id: UUID of the book.
    """
    collection_name = get_collection_name(author_id, book_id)
    chroma = _get_chroma()
    try:
        chroma.delete_collection(collection_name)
        logger.info("Deleted ChromaDB collection", collection=collection_name)
    except Exception as e:
        logger.warning("Could not delete collection (may not exist)", collection=collection_name, error=str(e))