app/services/vector_store.py

"""Author RAG Chatbot SaaS — Vector Retriever.

Retrieves relevant text chunks from ChromaDB using semantic search.
RULE: Always filter by author_id in metadata — no cross-tenant leakage.
RULE: Run search for all query variations, then deduplicate by chunk ID.

Collection architecture: one ChromaDB collection PER BOOK (not per author).
This means no where= metadata filter is needed at query time — collection
isolation already ensures no cross-book contamination.
Phase 2B confirmed: B7 (stale chunks) already handled in embed_and_store().
"""

from dataclasses import dataclass

import structlog

from app.config import get_settings
from app.services.embeddings import _get_chroma, get_collection_name

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


@dataclass
class RetrievedChunk:
    """A single retrieved text chunk from ChromaDB."""

    chunk_id: str
    text: str
    book_id: str
    book_title: str
    chunk_index: int
    score: float     # Initial cosine similarity score (0 to 1)
    rerank_score: float = 0.0   # Updated by re-ranker


async def retrieve_chunks(
    queries: list[str],
    author_id: str,
    book_id: str | None,
    top_k: int | None = None,
) -> list[RetrievedChunk]:
    """Retrieve relevant chunks from ChromaDB for a list of query variations.

    Searches each query variation and deduplicates results by chunk ID.

    Args:
        queries: List of query strings (original + rewritten variations).
        author_id: UUID of the author (enforces tenant isolation).
        book_id: UUID of the selected book, or None for cross-book search.
        top_k: Number of results to retrieve per query variation.

    Returns:
        Deduplicated list of RetrievedChunk objects (not yet re-ranked).
    """
    top_k = top_k or cfg.RAG_RETRIEVAL_TOP_K
    chroma = _get_chroma()

    # Get all collections to search
    collections_to_search = await _get_target_collections(chroma, author_id, book_id)
    if not collections_to_search:
        logger.warning("No collections found for author", author_id=author_id)
        return []

    # Embed all query variations at once
    query_embeddings = await _embed_queries(queries)

    # Search each collection with each query embedding
    seen_ids: set[str] = set()
    all_chunks: list[RetrievedChunk] = []

    for collection_name, book_meta in collections_to_search:
        try:
            collection = chroma.get_collection(collection_name)
        except Exception:
            logger.warning("Collection not found", name=collection_name)
            continue

        for embedding in query_embeddings:
            results = collection.query(
                query_embeddings=[embedding],
                n_results=min(top_k, collection.count()),
                include=["documents", "metadatas", "distances"],
            )
            if not results["ids"] or not results["ids"][0]:
                continue

            for chunk_id, doc, meta, distance in zip(
                results["ids"][0],
                results["documents"][0],
                results["metadatas"][0],
                results["distances"][0],
            ):
                if chunk_id in seen_ids:
                    continue
                seen_ids.add(chunk_id)

                # ChromaDB returns L2 distance — convert to similarity (lower = more similar)
                similarity = max(0.0, 1.0 - (distance / 2.0))

                all_chunks.append(RetrievedChunk(
                    chunk_id=chunk_id,
                    text=doc,
                    book_id=meta.get("book_id", ""),
                    book_title=meta.get("book_title", "Unknown"),
                    chunk_index=int(meta.get("chunk_index", 0)),
                    score=similarity,
                ))

    # Sort by initial similarity score
    all_chunks.sort(key=lambda c: c.score, reverse=True)
    logger.debug("Retrieved chunks", count=len(all_chunks), queries=len(queries))
    return all_chunks


async def _get_target_collections(
    chroma,
    author_id: str,
    book_id: str | None,
) -> list[tuple[str, dict]]:
    """Identify which ChromaDB collections to search.

    Args:
        chroma: ChromaDB client.
        author_id: UUID of the author.
        book_id: Specific book UUID or None (all books).

    Returns:
        List of (collection_name, metadata) tuples.
    """
    try:
        all_collections = chroma.list_collections()
    except Exception as e:
        logger.error("Failed to list ChromaDB collections", error=str(e))
        return []

    # SEC-4 fix: use the full 32-char hex UUID instead of [:12] prefix —
    # the prefix could collide between authors sharing similar UUID prefixes
    # (common with time-based UUIDs). Full UUID guarantees tenant isolation.
    author_prefix = author_id.replace("-", "")
    author_tag = f"a{author_prefix}"

    targets = []
    for col in all_collections:
        if not col.name.startswith(author_tag):
            continue  # Skip other authors' collections
        if book_id is None:
            targets.append((col.name, col.metadata or {}))
        else:
            expected_name = get_collection_name(author_id, book_id)
            if col.name == expected_name:
                targets.append((col.name, col.metadata or {}))
                break

    return targets


async def _embed_queries(queries: list[str]) -> list[list[float]]:
    """Embed query strings using OpenAI embeddings.

    Phase 2B fix: was creating AsyncOpenAI() on every call (same B3 bug as call_llm).
    Now uses the shared singleton from pipeline/helpers.py.

    Args:
        queries: List of query strings.

    Returns:
        List of embedding vectors.
    """
    from app.services.pipeline.helpers import _get_openai_client
    client = _get_openai_client()
    response = await client.embeddings.create(
        model=cfg.OPENAI_EMBEDDING_MODEL,
        input=queries,
    )
    return [item.embedding for item in response.data]