app/rag.py

"""RAG pipeline: retrieve relevant chunks, then generate a grounded answer."""
from __future__ import annotations

from functools import lru_cache
from typing import List, Optional

from .config import get_settings
from .embeddings import embed_query
from .models import ChatResponse, Source
from . import vector_store

SYSTEM_PROMPT = (
    "You are a document assistant. Answer the user's question using ONLY the provided "
    "context excerpts from their uploaded documents. Cite the supporting sources inline "
    "as [1], [2], etc., matching the numbered context blocks. If the answer is not "
    "contained in the context, say you could not find it in the documents. Be concise "
    "and accurate."
)


@lru_cache
def _llm():
    from openai import OpenAI

    settings = get_settings()
    if not settings.openrouter_api_key:
        raise RuntimeError("OPENROUTER_API_KEY is not set in the environment / .env file.")
    return OpenAI(
        base_url=settings.openrouter_base_url,
        api_key=settings.openrouter_api_key,
    )


def _build_context(hits) -> str:
    blocks = []
    for i, hit in enumerate(hits, start=1):
        payload = hit.payload or {}
        blocks.append(
            f"[{i}] (source: {payload.get('filename', '?')}, "
            f"chunk {payload.get('chunk_index', '?')})\n{payload.get('text', '')}"
        )
    return "\n\n".join(blocks)


def answer_question(
    question: str,
    document_ids: Optional[List[str]] = None,
    top_k: Optional[int] = None,
) -> ChatResponse:
    settings = get_settings()
    k = top_k or settings.top_k

    query_vector = embed_query(question)
    hits = vector_store.search(query_vector, top_k=k, document_ids=document_ids)

    if not hits:
        return ChatResponse(
            answer="I couldn't find anything relevant in your uploaded documents.",
            sources=[],
        )

    context = _build_context(hits)
    user_message = (
        f"Context excerpts:\n\n{context}\n\n"
        f"Question: {question}\n\n"
        "Answer using only the context above and cite sources as [n]."
    )

    client = _llm()
    response = client.chat.completions.create(
        model=settings.openrouter_model,
        max_tokens=1024,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
    )
    answer_text = (response.choices[0].message.content or "").strip()

    sources = [
        Source(
            document_id=(hit.payload or {}).get("document_id", ""),
            filename=(hit.payload or {}).get("filename", "?"),
            chunk_index=(hit.payload or {}).get("chunk_index", -1),
            score=float(hit.score),
            snippet=((hit.payload or {}).get("text", "")[:280]),
        )
        for hit in hits
    ]

    return ChatResponse(answer=answer_text, sources=sources)