feat: RAG answers with sutta citations (M1)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
marcuspaico
2026-08-14 17:26:26 -07:00
parent c14b314107
commit 37caa87b49
3 changed files with 42 additions and 0 deletions

25
src/buddhagpt/rag.py Normal file
View File

@@ -0,0 +1,25 @@
from pathlib import Path
from buddhagpt.index import search
SYSTEM = (
"You are a thoughtful guide grounded in the Pali Canon. Answer with warmth and "
"precision. Base doctrinal claims on the provided passages and cite them by uid "
"(e.g. [mn21]). If the passages do not cover the question, say so plainly."
)
def build_prompt(question: str, passages: list[dict]) -> list[dict]:
ctx = "\n\n".join(f"[{p['uid']}] {p['title']}:\n{p['chunk']}" for p in passages)
return [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": f"Passages:\n{ctx}\n\nQuestion: {question}"},
]
def answer(question: str, model_path: str, db_path: Path) -> dict:
from mlx_lm import load, generate
passages = search(db_path, question, k=4)
model, tokenizer = load(model_path)
prompt = tokenizer.apply_chat_template(
build_prompt(question, passages), tokenize=False, add_generation_prompt=True
)
text = generate(model, tokenizer, prompt=prompt, max_tokens=500)
return {"answer": text, "citations": [{"uid": p["uid"], "title": p["title"]} for p in passages]}