07_prompting.py

"""
07_prompting.py
-----------------
Stage 7 of the RAG pipeline: PROMPTING + GENERATION.

Builds the final prompt from the user's question and the retrieved context
chunks (stage 6), then sends it to an LLM through the OpenRouter API to
generate a grounded answer with source citations.

This module is imported by streamlit_app.py as the "rag" module. It is the
single place that owns:
    - OPENROUTER_API_KEY
    - OPENROUTER_MODEL
so that streamlit_app.py can fill them in from Streamlit secrets at deploy
time, exactly as required by the project instructions.

API key handling
-----------------
- No real API key is ever hardcoded in this file.
- Locally: set the OPENROUTER_API_KEY environment variable (e.g. via a
  local, un-committed .env file loaded with python-dotenv).
- On Streamlit Cloud: streamlit_app.py reads it from st.secrets and injects
  it into this module at runtime.

Run directly to try a full question -> answer flow against the existing
Chroma store (requires OPENROUTER_API_KEY to be set in your environment):
    python 07_prompting.py
"""

from __future__ import annotations

import importlib.util
import sys
import os
from dataclasses import dataclass
from typing import List, Optional

import requests

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))


def _import_module_from_file(module_name: str, filename: str):
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


_retrieval_module = _import_module_from_file("retrieval_stage", "06_retrieve_context.py")
RetrievedChunk = _retrieval_module.RetrievedChunk

# ---------------------------------------------------------------------------
# Configuration. These two module-level variables are intentionally mutable
# so streamlit_app.py can set them from Streamlit secrets at deploy time:
#
#   rag.OPENROUTER_API_KEY = st.secrets.get("OPENROUTER_API_KEY", "")
#   rag.OPENROUTER_MODEL = st.secrets.get("OPENROUTER_MODEL", rag.OPENROUTER_MODEL)
# ---------------------------------------------------------------------------
OPENROUTER_API_KEY: str = os.environ.get("OPENROUTER_API_KEY", "")
OPENROUTER_MODEL: str = os.environ.get("OPENROUTER_MODEL", "openai/gpt-4o-mini")
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"

NO_CONTEXT_MESSAGE = (
    "I don't have enough information in the provided documents to answer that "
    "question. Please try rephrasing, or add relevant documents to the knowledge base."
)

SYSTEM_PROMPT = (
    "You are a careful, honest assistant that answers questions strictly using "
    "the CONTEXT provided below. Rules:\n"
    "1. Only use information present in the CONTEXT to answer.\n"
    "2. If the CONTEXT does not contain the answer, say so explicitly instead "
    "of guessing or using outside knowledge.\n"
    "3. Always cite which source(s) you used, referencing them as [Source: <filename>].\n"
    "4. Be concise and accurate."
)


@dataclass
class RAGAnswer:
    """The final structured result returned to the Streamlit UI."""

    answer: str
    sources: List[str]
    retrieved_chunks: List[RetrievedChunk]
    used_context: bool


def build_prompt(question: str, chunks: List[RetrievedChunk]) -> str:
    """Assemble the user-facing prompt from retrieved context + the question."""
    context_blocks = []
    for i, chunk in enumerate(chunks, start=1):
        source = chunk.metadata.get("filename", chunk.metadata.get("doc_id", "unknown"))
        context_blocks.append(f"[Context {i} | Source: {source}]\n{chunk.text}")

    context_text = "\n\n".join(context_blocks)

    return (
        f"CONTEXT:\n{context_text}\n\n"
        f"QUESTION:\n{question}\n\n"
        "Answer the question using only the CONTEXT above, and cite your sources."
    )


def call_openrouter(
    prompt: str,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
    temperature: float = 0.2,
) -> str:
    """
    Send a chat completion request to OpenRouter and return the model's reply text.

    Raises
    ------
    RuntimeError
        If no API key is configured, or the API call fails.
    """
    resolved_key = api_key or OPENROUTER_API_KEY
    resolved_model = model or OPENROUTER_MODEL

    if not resolved_key:
        raise RuntimeError(
            "No OpenRouter API key configured. Set the OPENROUTER_API_KEY "
            "environment variable locally, or configure it in Streamlit secrets "
            "when deployed."
        )

    headers = {
        "Authorization": f"Bearer {resolved_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": resolved_model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        "temperature": temperature,
    }

    try:
        response = requests.post(
            OPENROUTER_API_URL, headers=headers, json=payload, timeout=60
        )
        response.raise_for_status()
    except requests.exceptions.RequestException as exc:
        raise RuntimeError(f"OpenRouter API request failed: {exc}") from exc

    data = response.json()

    try:
        return data["choices"][0]["message"]["content"].strip()
    except (KeyError, IndexError) as exc:
        raise RuntimeError(f"Unexpected OpenRouter API response format: {data}") from exc


def answer_question(
    question: str,
    top_k: int = _retrieval_module.DEFAULT_TOP_K,
) -> RAGAnswer:
    """
    Full end-to-end RAG answer generation for a single question.

    Retrieves context first; if retrieval returns nothing, the model is
    NEVER called and a clear "no context" message is returned instead of
    letting the LLM answer from its own knowledge.
    """
    retrieved_chunks = _retrieval_module.retrieve_context(question, top_k=top_k)

    if not retrieved_chunks:
        return RAGAnswer(
            answer=NO_CONTEXT_MESSAGE,
            sources=[],
            retrieved_chunks=[],
            used_context=False,
        )

    prompt = build_prompt(question, retrieved_chunks)
    answer_text = call_openrouter(prompt)

    sources = sorted(
        {
            chunk.metadata.get("filename", chunk.metadata.get("doc_id", "unknown"))
            for chunk in retrieved_chunks
        }
    )

    return RAGAnswer(
        answer=answer_text,
        sources=sources,
        retrieved_chunks=retrieved_chunks,
        used_context=True,
    )


if __name__ == "__main__":
    try:
        from dotenv import load_dotenv

        load_dotenv()  # loads a local, un-committed .env file if present
        OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", OPENROUTER_API_KEY)
    except ImportError:
        pass

    test_question = "What is this project about?"
    result = answer_question(test_question)

    print(f"Q: {test_question}\n")
    print(f"A: {result.answer}\n")
    if result.sources:
        print(f"Sources: {', '.join(result.sources)}")
    else:
        print("Sources: none (no context was retrieved)")