"""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))