generation.py

"""
generation.py
-------------
Takes a user query + retrieved chunks and generates a grounded answer
using Groq's llama-3.3-70b-versatile model.

The system prompt explicitly forbids the model from using knowledge
outside the provided chunks. If the chunks don't contain the answer,
the model says so.

Usage (standalone test):
    python generation.py "Who is the best CS 367 professor for exams?"
"""

import os
import sys
from dotenv import load_dotenv
from groq import Groq

from retrieval import retrieve

load_dotenv()

# ── Config ────────────────────────────────────────────────────────────────────
GROQ_MODEL = "llama-3.3-70b-versatile"
MAX_TOKENS = 1024
# ─────────────────────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are a helpful assistant for GMU students researching CS 367 (Computer Systems and Programming) professors and course experiences.

CRITICAL RULES — follow these exactly:
1. Answer ONLY using information from the provided student review excerpts below.
2. Do NOT use your general training knowledge about professors, courses, or GMU.
3. If the provided excerpts do not contain enough information to answer the question, respond with exactly: "I don't have enough information in the provided reviews to answer that question."
4. Always cite which source document(s) your answer draws from.
5. Be specific and quote or paraphrase student language where helpful.
6. Keep answers concise (3–6 sentences) unless the question requires more detail."""


def build_context(chunks: list[dict]) -> str:
    """Format retrieved chunks into a numbered context block."""
    parts = []
    for i, chunk in enumerate(chunks, 1):
        parts.append(
            f"[Excerpt {i} — Source: {chunk['source']}]\n{chunk['text']}"
        )
    return "\n\n".join(parts)


def ask(query: str, top_k: int = 5) -> dict:
    """
    End-to-end: retrieve → generate → return.

    Returns:
        {
            "answer":  str,          # LLM-generated grounded answer
            "sources": list[str],    # unique source filenames cited
            "chunks":  list[dict],   # raw retrieved chunks for transparency
        }
    """
    # 1. Retrieve relevant chunks
    chunks = retrieve(query, top_k=top_k)

    # 2. Build context string
    context = build_context(chunks)

    # 3. Call Groq
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "GROQ_API_KEY not found. Add it to your .env file."
        )

    client = Groq(api_key=api_key)

    user_message = (
        f"Here are student review excerpts about CS 367 at GMU:\n\n"
        f"{context}\n\n"
        f"Question: {query}"
    )

    response = client.chat.completions.create(
        model      = GROQ_MODEL,
        max_tokens = MAX_TOKENS,
        messages   = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": user_message},
        ],
    )

    answer = response.choices[0].message.content.strip()

    # 4. Collect unique sources
    sources = list(dict.fromkeys(c["source"] for c in chunks))

    return {
        "answer":  answer,
        "sources": sources,
        "chunks":  chunks,
    }


# ── Standalone test ───────────────────────────────────────────────────────────
if __name__ == "__main__":
    query = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else \
            "What do students say about Professor Zhong's exams?"

    print(f"\nQuery: {query!r}\n{'='*60}")
    result = ask(query)

    print(f"\nAnswer:\n{result['answer']}")
    print(f"\nSources: {', '.join(result['sources'])}")
    print(f"\nRetrieved chunks ({len(result['chunks'])}):")
    for i, c in enumerate(result["chunks"], 1):
        print(f"  [{i}] {c['source']} (dist={c['distance']}) — {c['text'][:120]}…")