answer.py

#!/usr/bin/env python3
"""
Answer a Dhamma question using retrieved context + Claude Sonnet.

Usage (CLI):
  python3 answer.py "what is the first stage of enlightenment?"
  python3 answer.py "explain kamma" --n 8
  python3 answer.py "what is jhana?" --type scripture

Importable:
  from answer import answer
  result = answer("what is stream-entry?")
  print(result["answer"])
  for s in result["sources"]: print(s["source_label"])
"""

import argparse
import hashlib
import json
import os
import sys
from pathlib import Path

try:
    import anthropic
except ImportError:
    sys.exit("pip install anthropic")

from retrieval import retrieve


def _load_api_key() -> str:
    """Return ANTHROPIC_API_KEY from env or ~/.bash_profile fallback."""
    key = os.environ.get("ANTHROPIC_API_KEY", "")
    if key:
        return key
    profile = Path.home() / ".bash_profile"
    if profile.exists():
        for line in profile.read_text().splitlines():
            line = line.strip()
            if line.startswith("export ANTHROPIC_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    return ""


MODEL      = "claude-sonnet-4-6"
CACHE_FILE = Path("data/answer_cache.json")

SYSTEM_PROMPT = """\
You are an assistant that answers questions strictly from the provided source passages \
of David Roylance, a Theravada Buddhist teacher. Write in the same voice and style \
as David Roylance uses in those passages.

Answer shape (MOST IMPORTANT — direct AND thorough):
- Lead with the answer. The FIRST sentence must directly and plainly answer the \
  question asked. State the conclusion first, then develop it.
- Be thorough. After the direct opening, fully develop the answer: explain the \
  reasoning, cover the relevant aspects, distinctions, and supporting points the \
  passages provide, and include illustrative scripture where it strengthens the \
  explanation. Aim for a complete, well-rounded answer — do not cut it short. \
  Detailed is good; padding and filler are not. Every sentence should add real \
  content from the passages, never restate the question or hedge.
- When the question asks about a concept or teaching (however it is phrased — \
  "what is X", "define X", "explain X", "tell me about X", or just "X"), first \
  identify which of these the retrieved passages actually support, then cover each \
  that they do: the definition; the concept's CAUSE/ORIGIN (how it arises); its \
  CESSATION (how it is eliminated or resolved); and any key distinction the question \
  invites. A definition alone is incomplete when the passages also explain the cause \
  and the cessation — weave those in. \
- When the question asks how to overcome, resolve, get rid of, or deal with a \
  problem, name the SPECIFIC practice or antidote the passages prescribe for THAT \
  problem — by its name, if the passages name one — rather than only general advice \
  like "train the mind" or "meditate". If the passages prescribe a particular \
  practice as the direct remedy for the problem, lead the how-to with it. \
- NEVER open with throat-clearing. Banned openers include: "This is an important \
  question", "The passages address this carefully", "perhaps not in the way one might \
  expect", "It's worth noting", "Let us explore", and any sentence about what the \
  passages do or don't do. Just answer.
- Do not be coy, clever, or indirect (e.g. never answer "how do you know you are \
  Enlightened?" with "you wouldn't be asking"). Give the actual substantive answer.
- Structure a long answer for readability: flowing prose by default, and section \
  sub-headers when the answer has several distinct parts.
- Never both answer AND disclaim. If you can answer from the passages, just answer — \
  do not hedge that the material is incomplete. Only use the exact decline sentence \
  (see Decline) when you genuinely cannot answer at all.

Content rules (non-negotiable):
- Answer ONLY from the source passages. Do not add outside knowledge or context.
- Cite every claim with a [Source N] marker.
- Every named concept, category, or term you use MUST appear word-for-word in the cited passage.
  Do NOT invent category labels, headings, or named groupings not present in the text.
- Do NOT extract adjectives or fragments from compound terms to create new terms.
  For example: if the passage says "The Three Unwholesome Roots", do not write "the unwholesome"
  or "roots of the unwholesome" — use the full phrase as written.
- If the passages do not contain enough to answer, say so directly.
- Do not speculate, summarize beyond what is stated, or editorialize.
- Source preference: always prefer David Roylance's own commentary and explanatory passages \
  over direct Pali Canon scripture when answering conceptual questions. Use scripture \
  (blockquote format) only when the question specifically asks for The Buddha's words, \
  or when no commentary passage covers the topic. If both are available, build the answer \
  from commentary and reference scripture only as supporting illustration.

Style (mirror the teacher's voice from the passages):
- Plain, direct English — clear and accessible, not academic or flowery.
- These words are ALWAYS capitalised, no exceptions: Teachings, Enlightenment, The Buddha, \
The Practitioner, The Path, Kamma, Dhamma, Nibbana, The Natural Law of Kamma. \
Never write "teachings", "enlightenment", or "the path" in lower case. \
All other terms — the mind, craving, clinging, anger, aversion, ignorance, wisdom, virtue, \
concentration, suffering, liberation — use standard lower case mid-sentence.
- Preserve the teacher's exact punctuation within phrases, including slash notation \
  (e.g. write "craving/desire/attachment", never "craving, desire, attachment"). \
  Slashes in the source text are intentional — they signal that the terms are one concept \
  seen from different angles, not a list of separate things.
- Use the teacher's characteristic phrases where they appear in the sources, \
  e.g. "learning, reflecting, and practicing", "the mind", "discontentedness".
- Warm but precise — explain to a student fully and patiently, but get to the point first.
- Do not write in first person as the teacher. Explain what the passages say.

Scripture formatting:
- When quoting The Buddha's words directly from the Pali Canon, wrap the quote in a \
  markdown blockquote (lines starting with "> "). On the last line of the blockquote, \
  add the sutta reference prefixed with "— ". Look for a "(Reference: ...)" line in the \
  source passage — use that exact reference (e.g. if the passage says \
  "(Reference: SN 35.240)", write "— SN 35.240"). If no such line appears, use the \
  collection abbreviation only (e.g. "— SN"). \
  Only use blockquotes for direct canonical scripture; do not use them for \
  David Roylance's commentary.
- Do NOT paraphrase scripture passages. If you use content from a scripture source, \
  quote it directly in a blockquote. Never restate The Buddha's words in your own phrasing.

Decline:
- If the source passages genuinely cannot support an answer, respond with exactly this \
  sentence and nothing else: \
  "The source passages here don't cover that. Try rephrasing, or narrow to a specific \
  volume or topic using the filters above." """


# ── cache ─────────────────────────────────────────────────────────────────────

def _cache_key(question: str, n_results: int, content_type: str | None, volume: int | None) -> str:
    normalized = " ".join(question.lower().split())
    payload = json.dumps([normalized, n_results, content_type, volume], sort_keys=True)
    return hashlib.sha256(payload.encode()).hexdigest()[:20]


def _load_cache() -> dict:
    if CACHE_FILE.exists():
        try:
            return json.loads(CACHE_FILE.read_text())
        except Exception:
            return {}
    return {}


def _save_cache(cache: dict) -> None:
    CACHE_FILE.write_text(json.dumps(cache, ensure_ascii=False, indent=2))


# ── context building ──────────────────────────────────────────────────────────

_REF_RE = __import__("re").compile(r"\(Reference:\s*([^)]+)\)")

def _best_ref(chunks: list[dict]) -> dict:
    """
    Build a map of (volume, chapter) → most-specific sutta ref found across all
    retrieved chunks. A ref with a number (e.g. "SN 35.240") beats a bare
    collection name (e.g. "SN").
    """
    best: dict[tuple, str] = {}
    for c in chunks:
        key = (c.get("vol_num", 0), c.get("ch_num", 0))
        m = _REF_RE.search(c["text"])
        candidate = m.group(1).strip() if m else ", ".join(c.get("sutta_refs") or [])
        if not candidate:
            continue
        existing = best.get(key, "")
        # Prefer the ref that contains a space (i.e. has a number like "SN 35.240")
        if " " in candidate and " " not in existing:
            best[key] = candidate
        elif not existing:
            best[key] = candidate
    return best

def _build_context(chunks: list[dict]) -> tuple[str, list[dict]]:
    """Format all retrieved chunks into a numbered context block."""
    ref_map = _best_ref(chunks)
    lines = []
    sources = []
    for i, c in enumerate(chunks, 1):
        ct = c.get("content_type", "")
        type_tag = " [SCRIPTURE — must blockquote, do not paraphrase]" if ct == "scripture" else ""
        header = f"[Source {i}]{type_tag} {c['source_label']}"
        key = (c.get("vol_num", 0), c.get("ch_num", 0))
        ref = ref_map.get(key, "") or ", ".join(c.get("sutta_refs") or [])
        refs = f"  ({ref})" if ref else ""
        lines.append(f"{header}{refs}\n{c['text']}")
        sources.append({**c, "index": i})
    return "\n\n---\n\n".join(lines), sources


# ── query understanding (NotebookLM-style contextualization) ───────────────────

def _taxonomy_text() -> str:
    """Compact text form of David's teaching taxonomy (data/taxonomy.yaml) to guide
    query expansion. Empty string if the file is missing/unreadable."""
    try:
        import yaml
        t = yaml.safe_load((Path("data/taxonomy.yaml")).read_text())
    except Exception:
        return ""
    parts = []
    if t.get("anger_rooted_feelings"):
        parts.append("Feelings rooted in ANGER — remedy is loving-kindness "
                     "meditation: " + ", ".join(t["anger_rooted_feelings"]) + ".")
    if t.get("craving_rooted_feelings"):
        parts.append("All other feelings (pleasant, most painful, and neither) are "
                     "rooted in CRAVING/DESIRE/ATTACHMENT — remedy is breathing "
                     "mindfulness meditation and generosity, to train the mind to let "
                     "go: " + ", ".join(t["craving_rooted_feelings"]) + ".")
    if t.get("roots"):
        parts.append("Roots -> antidotes: " + "; ".join(
            f"{k} -> {v.get('antidote')}" for k, v in t["roots"].items()) + ".")
    if t.get("practices"):
        parts.append("Key practices to name when relevant: " +
                     ", ".join(t["practices"]) + ".")
    return "\n".join(parts)


_EXPAND_SYSTEM = (
    "You turn a user's question into 2-4 short SEARCH QUERIES naming the underlying "
    "concepts in David Roylance's Theravada Buddhist Teachings needed to answer it "
    "fully. When the question asks how to overcome / get rid of / deal with "
    "something, classify the emotion/problem using the framework below and generate "
    "queries covering BOTH the cause AND the specific remedy/practice for THAT "
    "problem.\n\n"
    + _taxonomy_text() +
    "\n\nThese are retrieval queries only, NOT answers, and not authoritative. "
    "Output one query per line — no numbering, no preamble, nothing else."
)


def _expand_query(client, question: str) -> tuple[list[str], int]:
    """LLM contextualizes the question into retrieval sub-queries so the right
    concepts (e.g. the antidote, not just the symptom) surface. Returns
    (sub_queries, tokens_used). Best-effort: failure degrades to no expansion."""
    try:
        msg = client.messages.create(
            model=MODEL, max_tokens=200, system=_EXPAND_SYSTEM,
            messages=[{"role": "user", "content": question}],
        )
        lines = [l.strip("-•* \t").strip() for l in msg.content[0].text.splitlines()]
        subs = [l for l in lines if l][:4]
        return subs, msg.usage.input_tokens + msg.usage.output_tokens
    except Exception:
        return [], 0


# ── main answer function ───────────────────────────────────────────────────────

def answer(
    question: str,
    n_results: int = 8,
    content_type: str | None = None,
    volume: int | None = None,
) -> dict:
    """
    Retrieve context and generate a grounded answer.
    Results are cached by (question, n_results, content_type, volume).

    Returns:
      {
        "question": str,
        "answer":   str,
        "sources":  list[dict],
        "usage":    {"input": int, "output": int},
        "cached":   bool,
      }
    """
    key   = _cache_key(question, n_results, content_type, volume)
    cache = _load_cache()
    if key in cache:
        return {**cache[key], "cached": True}

    client = anthropic.Anthropic(api_key=_load_api_key())

    # Contextualize the question into retrieval sub-queries (skip on type-filtered
    # browsing). The expansion guides RETRIEVAL only — it is never shown to the
    # answer model, which must ground every claim in the retrieved passages.
    extra, exp_tokens = ([], 0)
    if not content_type:
        extra, exp_tokens = _expand_query(client, question)

    chunks  = retrieve(question, n_results=n_results, content_type=content_type,
                       volume=volume, extra_queries=extra)
    context, sources = _build_context(chunks)

    user_message = f"""Source passages:

{context}

---

Question: {question}"""

    msg = client.messages.create(
        model=MODEL,
        max_tokens=2048,
        system=SYSTEM_PROMPT,
        messages=[{"role": "user", "content": user_message}],
    )

    result = {
        "question": question,
        "answer":   msg.content[0].text,
        "sources":  sources,
        "expanded_queries": extra,   # for transparency/debug; not seen by answer model
        "usage":    {"input": msg.usage.input_tokens + exp_tokens,
                     "output": msg.usage.output_tokens},
    }

    cache[key] = result
    _save_cache(cache)

    return {**result, "cached": False}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("question")
    ap.add_argument("--n", type=int, default=8, dest="n_results")
    ap.add_argument("--type", dest="content_type", default=None)
    ap.add_argument("--vol", type=int, default=None, dest="volume")
    args = ap.parse_args()

    result = answer(args.question, n_results=args.n_results,
                    content_type=args.content_type, volume=args.volume)

    cached_tag = "  [cached]" if result.get("cached") else ""
    print(f"\nQ: {result['question']}{cached_tag}\n")
    print(result["answer"])
    print(f"\n── Sources ({'|'.join(str(s['index']) for s in result['sources'])})")
    for s in result["sources"]:
        ct = s["content_type"]
        refs = f"  [{', '.join(s['sutta_refs'])}]" if s["sutta_refs"] else ""
        print(f"  [{s['index']}] score={s['score']}  {ct}{refs}")
        print(f"       {s['source_label']}")
    u = result["usage"]
    print(f"\n── Tokens: {u['input']} in / {u['output']} out{cached_tag}")


if __name__ == "__main__":
    main()