rag.py

"""
Document RAG: chunking text a user uploads, and retrieving the most
relevant chunks for their current message via MongoDB Atlas Vector Search.

Requires a Vector Search index (see DEPLOYMENT_GUIDE.md for the one-time
Atlas setup) named settings.VECTOR_INDEX_NAME on the `documents` collection,
covering the `embedding` field.
"""
from config import settings
from database import documents_collection
from embeddings import embed_text


def chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
    """Simple fixed-size character chunking with overlap - good enough for
    short personal notes/documents without pulling in a heavier NLP splitter."""
    text = text.strip()
    if not text:
        return []
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start = end - overlap
    return chunks


async def retrieve_relevant_chunks(email: str, query: str, top_k: int = 4) -> list[str]:
    query_vector = await embed_text(query)

    pipeline = [
        {
            "$vectorSearch": {
                "index": settings.VECTOR_INDEX_NAME,
                "path": "embedding",
                "queryVector": query_vector,
                "numCandidates": 100,
                "limit": top_k,
                "filter": {"email": email},
            }
        },
        {"$project": {"text": 1, "_id": 0}},
    ]
    try:
        cursor = documents_collection().aggregate(pipeline)
        results = await cursor.to_list(length=top_k)
    except Exception:
        # The vector index may not exist yet (e.g. not set up in Atlas), or
        # the query briefly failed - fail soft so chat still works without
        # document context rather than breaking the conversation.
        return []
    return [r["text"] for r in results]