code.py

"""Knowledge Retrieval (RAG) — answer from fetched documents, not from memory.

Retrieve the passages that look relevant, put them in the prompt, and require
the answer to come from them. The model stops guessing about things it was
never trained on - internal policies, this week's prices, your runbook.

The part people skip: telling the model to refuse when the documents don't
cover the question. Without that instruction RAG becomes a system that
retrieves three irrelevant passages and then confidently makes something up.
"""

from langchain_core.vectorstores import InMemoryVectorStore

from resources.agent import embeddings, llm
from resources.helper import show_response

# A tiny internal handbook. Nothing here is public knowledge, which is the
# point - the model cannot fall back on training data for any of it.
DOCS = [
    "Refund policy: orders may be refunded within 30 days of delivery. "
    "Opened software licences are non-refundable regardless of date.",
    "Shipping: standard delivery is 3-5 business days. Express is next-day "
    "if ordered before 14:00 GMT. We do not ship to PO boxes.",
    "Support hours: the helpdesk is staffed 08:00-18:00 GMT Monday to Friday. "
    "Priority customers get a 2-hour response SLA including weekends.",
    "Warranty: hardware carries a 24-month warranty. Accidental damage is "
    "excluded unless AccidentalCare was purchased at checkout.",
    "Account deletion requests are processed within 14 days. Deleted accounts "
    "cannot be restored; billing records are retained for 7 years by law.",
]

# Embeddings turn text into vectors so "can I get my money back" can match a
# passage that says "refund", with no shared keywords at all.
store = InMemoryVectorStore.from_texts(DOCS, embeddings)
retriever = store.as_retriever(search_kwargs={"k": 2})

# The refusal instruction is the load-bearing line. Everything else here is
# plumbing; this is what separates RAG from confident invention.
ANSWER_PROMPT = (
    "Answer the question using ONLY the context below.\n"
    "If the context does not contain the answer, reply exactly: "
    "I don't have that information.\n"
    "Do not use outside knowledge.\n\n"
    "Context:\n{context}\n\nQuestion: {question}"
)


def answer(question):
    """Retrieve, then generate. Prints what was retrieved so it's auditable."""
    hits = retriever.invoke(question)

    print("  retrieved:")
    for hit in hits:
        print(f"    - {hit.page_content[:70]}...")

    context = "\n\n".join(h.page_content for h in hits)
    return llm.invoke(ANSWER_PROMPT.format(context=context, question=question))


if __name__ == "__main__":
    questions = [
        # Answerable, and phrased without the words the document uses.
        "Can I get my money back on a laptop I received three weeks ago?",
        # Answerable, but needs the exclusion buried in the same passage.
        "I opened the software licence I bought last week - can I return it?",
        # Not in the handbook at all. The refusal instruction gets tested here.
        "What's your company's parental leave policy?",
    ]

    for question in questions:
        print(f"\n{'=' * 60}\nQUESTION: {question}\n{'=' * 60}")
        show_response(answer(question))