p6_agentic_rag.py

"""Pattern 6: agentic RAG. Retrieval is a tool the agent calls as often as it likes; the answer must cite sources."""

import asyncio
import re
from typing import Annotated, Any

from claude_agent_sdk import (AssistantMessage, ClaudeAgentOptions, ResultMessage, ToolAnnotations, ToolUseBlock,
                              create_sdk_mcp_server, query, tool)

# A toy "knowledge base": in production this is an embedding index. Here: keyword overlap scoring, no dependencies.
DOCS = {
    "kb-101": ("Refund policy", "Refunds are issued within 5 business days to the original payment method. Annual plans are refundable within 30 days of purchase; monthly plans are not refundable after the billing date."),
    "kb-102": ("Double charges", "If a customer sees two charges in one month, the second is usually an authorization hold that drops off in 3-7 days. If both charges settle, support should issue a refund of the duplicate and file ticket type BILL-DUP."),
    "kb-203": ("Login outage runbook", "When login fails for all users, check the auth service health endpoint first, then the token signing key rotation job, which runs at 02:00 UTC."),
    "kb-204": ("Password reset", "Password reset emails expire after 15 minutes. Resending generates a new token and invalidates the old one."),
    "kb-305": ("Data retention", "Account data is retained for 90 days after cancellation, then permanently deleted. Backups are purged within 30 further days."),
}


def score(query_text: str, text: str) -> int:
    q = set(re.findall(r"\w+", query_text.lower()))
    return len(q & set(re.findall(r"\w+", text.lower())))


@tool("search_docs",
      "Keyword search over the support knowledge base. Returns the best-matching passages with their document IDs. "
      "Call again with different wording if results look off-topic.",
      {"query": Annotated[str, "A short search query"], "k": Annotated[int, "How many results, 1-5"]},
      annotations=ToolAnnotations(readOnlyHint=True))
async def search_docs(args: dict[str, Any]):
    ranked = sorted(DOCS.items(), key=lambda kv: -score(args["query"], kv[1][0] + " " + kv[1][1]))[: max(1, min(args["k"], 5))]
    hits = [f"[{doc_id}] {title}: {text[:120]}..." for doc_id, (title, text) in ranked if score(args["query"], title + " " + text) > 0]
    return {"content": [{"type": "text", "text": "\n".join(hits) or "No passages matched. Try different wording."}]}


@tool("get_doc", "Fetch the full text of one knowledge-base document by ID", {"doc_id": str},
      annotations=ToolAnnotations(readOnlyHint=True))
async def get_doc(args: dict[str, Any]):
    if args["doc_id"] not in DOCS:
        return {"content": [{"type": "text", "text": f"No document with id {args['doc_id']}. IDs look like kb-101."}], "is_error": True}
    title, text = DOCS[args["doc_id"]]
    return {"content": [{"type": "text", "text": f"[{args['doc_id']}] {title}\n{text}"}]}


ANSWER_SCHEMA = {"type": "object",
                 "properties": {"answer": {"type": "string"},
                                "citations": {"type": "array", "items": {"type": "string"}},
                                "confidence": {"type": "string", "enum": ["low", "medium", "high"]}},
                 "required": ["answer", "citations", "confidence"]}

kb = create_sdk_mcp_server(name="kb", version="1.0.0", tools=[search_docs, get_doc])


async def ask(question: str):
    calls = []
    async for m in query(
        prompt=question,
        options=ClaudeAgentOptions(
            mcp_servers={"kb": kb}, allowed_tools=["mcp__kb__*"], tools=[], setting_sources=[], max_turns=8,
            system_prompt="Answer ONLY from the knowledge base tools. Cite document IDs. If the KB has nothing relevant, say so with confidence=low.",
            output_format={"type": "json_schema", "schema": ANSWER_SCHEMA},
        ),
    ):
        if isinstance(m, AssistantMessage):
            for b in m.content:
                if isinstance(b, ToolUseBlock):
                    calls.append(f"{b.name.split('__')[-1]}({b.input.get('query') or b.input.get('doc_id')})")
        elif isinstance(m, ResultMessage):
            print(f"\nQ: {question}\n  retrieval calls: {calls}\n  A: {m.structured_output}\n  cost=${m.total_cost_usd:.4f}")


async def main():
    await ask("A customer on a monthly plan was charged twice this month and wants their money back. What should I do?")
    await ask("How long do we keep a cancelled account's data?")
    await ask("What is our policy on employee parking?")


asyncio.run(main())