llm.py

"""
llm.py
──────
Wraps the Groq API client and owns all prompt engineering for AstroBot.

Responsibilities:
  - Validate the Groq API key at startup
  - Build the system prompt (astrology tutor persona + no-prediction guardrail)
  - Format retrieved context chunks into the prompt
  - Call the Groq chat completion endpoint and return the answer string
"""

import logging

from groq import Groq
from langchain_core.documents import Document

from config import cfg

logger = logging.getLogger(__name__)

# ── System prompt ─────────────────────────────────────────────────────────────
# Defines the bot's persona, scope, and hard guardrails.

SYSTEM_TEMPLATE = """You are AstroBot, a patient and knowledgeable astrology tutor.
Your students are learning astrology concepts. Your role is to:
  • Explain astrological concepts clearly and accurately using the provided context.
  • Use analogies and examples to make complex ideas approachable.
  • Reference classical and modern astrology where relevant.
  • Encourage curiosity and deeper study.

HARD RULES — never break these:
  1. Do NOT make personal predictions or interpret anyone's birth chart.
  2. Do NOT speculate about future events for specific individuals.
  3. If the context does not contain enough information to answer, say so honestly
     and suggest the student consult a textbook or senior practitioner.
  4. Keep answers focused on educational content only.

--- CONTEXT FROM COURSE MATERIALS ---
{context}
--- END OF CONTEXT ---

Answer the student's question based solely on the context above.
If the answer isn't in the context, say: "I don't have that in my course materials right now — 
let me point you to further study resources."
"""


# ── Public API ────────────────────────────────────────────────────────────────

def create_client() -> Groq:
    """
    Initialise and validate the Groq client.

    Raises
    ------
    ValueError
        If GROQ_API_KEY is missing.
    """
    if not cfg.groq_api_key:
        raise ValueError(
            "GROQ_API_KEY is not set. Add it in Space → Settings → Repository secrets."
        )
    logger.info("Groq client initialised (model: %s)", cfg.groq_model)
    return Groq(api_key=cfg.groq_api_key)


def generate_answer(client: Groq, query: str, context_docs: list[Document]) -> str:
    """
    Build the RAG prompt and call Groq to get an answer.

    Parameters
    ----------
    client : Groq
        Groq client returned by create_client().
    query : str
        The student's question.
    context_docs : list[Document]
        Retrieved chunks from the vector store.

    Returns
    -------
    str
        The model's answer string.
    """
    context_text = _format_context(context_docs)
    system_prompt = SYSTEM_TEMPLATE.format(context=context_text)

    logger.debug("Calling Groq | model=%s | context_chunks=%d", cfg.groq_model, len(context_docs))

    response = client.chat.completions.create(
        model=cfg.groq_model,
        messages=[
            {"role": "system",  "content": system_prompt},
            {"role": "user",    "content": query},
        ],
        temperature=cfg.groq_temperature,
        max_tokens=cfg.groq_max_tokens,
    )

    answer = response.choices[0].message.content
    logger.debug("Groq response: %d chars", len(answer))
    return answer


# ── Internal helpers ──────────────────────────────────────────────────────────

def _format_context(docs: list[Document]) -> str:
    """
    Format retrieved documents into a numbered context block
    that is easy for the LLM to parse.
    """
    blocks = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", doc.metadata.get("source_row", i))
        page   = doc.metadata.get("page", "")
        header = f"[Source {i}" + (f" | {source}" if source else "") + (f" | p.{page}" if page else "") + "]"
        blocks.append(f"{header}\n{doc.page_content}")
    return "\n\n".join(blocks)