Capstone 2 — Domain A: Clinical Policy Q&A System

{
  "documents": [
    {
      "doc_id": "POLICY-AETNA-ORTHO-2024",
      "title": "Aetna Clinical Policy: Knee Arthroplasty",
      "payer": "Aetna",
      "category": "Orthopedic Surgery",
      "effective_date": "2024-01-01",
      "sections": [
        {
          "section_id": "1.0",
          "title": "Policy Statement",
          "content": "Aetna considers total knee arthroplasty (TKA, CPT 27447) medically necessary when ALL of the following criteria are met. Unicompartmental knee arthroplasty (UKA, CPT 27446) is considered medically necessary when criteria 1-3 below are met and the disease is limited to a single compartment."
        },
        {
          "section_id": "1.1",
          "title": "Clinical Criteria",
          "content": "1. Documented diagnosis of severe osteoarthritis (ICD-10: M17.11, M17.12) OR rheumatoid arthritis (M05.x, M06.x) confirmed by weight-bearing radiographs showing Kellgren-Lawrence Grade III or IV.\n2. Minimum 6 months of documented conservative treatment including: (a) physical therapy (8+ sessions), (b) NSAIDs or analgesics, and (c) at least one corticosteroid injection.\n3. Functional impairment documented by validated outcome measure (WOMAC score > 50 or equivalent).\n4. BMI below 40 (relative contraindication above 40; requires additional documentation)."
        },
        {
          "section_id": "2.0",
          "title": "Covered CPT Codes",
          "content": "27447 - Total Knee Arthroplasty (TKA)\n27446 - Partial Knee Arthroplasty (Unicompartmental, UKA)\n27486 - Revision of Total Knee Arthroplasty, one component\n27487 - Revision of Total Knee Arthroplasty, all components (see separate revision policy)"
        },
        {
          "section_id": "3.0",
          "title": "Exclusions",
          "content": "This policy does not cover: (a) revision arthroplasty for all components (CPT 27487) — see POLICY-AETNA-ORTHO-REV-2024; (b) bilateral simultaneous TKA without separate medical justification for each knee; (c) knee arthroplasty solely for pain management when functional criteria are not met."
        }
      ]
    },
    {
      "doc_id": "POLICY-UHC-ORTHO-2024",
      "title": "UnitedHealthcare Clinical Policy: Knee Replacement Surgery",
      "payer": "UnitedHealthcare",
      "category": "Orthopedic Surgery",
      "effective_date": "2024-02-15",
      "sections": [
        {
          "section_id": "1.0",
          "title": "Policy Statement",
          "content": "UnitedHealthcare requires prior authorization for total knee replacement (CPT 27447) and partial knee replacement (CPT 27446). Authorization is granted when clinical criteria demonstrate medical necessity as defined below."
        },
        {
          "section_id": "1.1",
          "title": "Medical Necessity Criteria",
          "content": "All of the following must be documented:\n1. Diagnosis of degenerative joint disease (M17.x) or inflammatory arthritis (M05.x, M06.x) with radiographic confirmation.\n2. Failure of conservative management for at least 3 months including: physical therapy, oral anti-inflammatories, and at least one intra-articular injection.\n3. Significant functional limitation documented by standardized assessment.\n4. Patient has been evaluated and deemed an appropriate surgical candidate by the operating surgeon."
        },
        {
          "section_id": "2.0",
          "title": "Documentation Requirements",
          "content": "The following must be submitted with the prior authorization request: (a) office notes documenting conservative treatment history, (b) radiographic reports dated within 6 months, (c) completed functional assessment score, (d) operative plan including implant type."
        }
      ]
    },
    {
      "doc_id": "POLICY-AETNA-IMAGING-2024",
      "title": "Aetna Clinical Policy: Advanced Imaging (MRI/CT)",
      "payer": "Aetna",
      "category": "Radiology",
      "effective_date": "2024-01-01",
      "sections": [
        {
          "section_id": "1.0",
          "title": "Policy Statement",
          "content": "Aetna requires prior authorization for advanced imaging studies including MRI (CPT 70553, 73721, 73723) and CT (CPT 70551, 72131). Authorization is managed through the AIM Specialty Health program."
        },
        {
          "section_id": "1.1",
          "title": "MRI Brain Criteria",
          "content": "MRI Brain with and without contrast (CPT 70553) is considered medically necessary for: (a) new-onset severe headache with neurological deficit, (b) suspected intracranial mass or lesion, (c) follow-up of known brain tumor, (d) evaluation of multiple sclerosis, (e) pre-surgical planning. Not covered for: routine headache evaluation without red flags."
        }
      ]
    }
  ]
}

# requirements.txt
anthropic>=0.30.0
pytest>=7.4.0

mkdir clinical-policy-rag && cd clinical-policy-rag
python -m venv venv && source venv/bin/activate   # Windows: venv\Scripts\activate
pip install "anthropic>=0.30.0" pytest
export ANTHROPIC_API_KEY=your-key-here             # Windows: set ANTHROPIC_API_KEY=your-key-here

mkdir -p data    # Windows: mkdir data
# Paste the policy JSON into data/policies.json

python -c "import json; d=json.load(open('data/policies.json')); print(f'{len(d[\"documents\"])} documents loaded')"

python -c "from vector_store import MockVectorStore; s=MockVectorStore(); s.add('t1','hello world',{}); print(s.search('hello',1))"

python -c "from agent import ingest_policies; count=ingest_policies(); print(f'Ingested {count} chunks')"

python agent.py

python -m pytest test_rag.py -v

"""vector_store.py — Mock vector store with TF-IDF similarity.

WHAT: In-memory vector store for RAG retrieval.
WHY:  Lets you build the full RAG pipeline without external
      dependencies. Same API as a real vector DB.
GOTCHA: TF-IDF is a bag-of-words approximation. Real embedding
        models capture semantic meaning far better. This is a
        learning scaffold, not production code.
"""

import math
import re
from collections import Counter
from typing import Optional


class MockVectorStore:
    """In-memory vector store using TF-IDF + cosine similarity."""

    def __init__(self):
        self.documents = []  # List of {chunk_id, content, metadata, vector}
        self.vocab = {}      # term -> document frequency
        self.total_docs = 0

    def _tokenize(self, text: str) -> list[str]:
        """Split text into lowercase tokens."""
        return re.findall(r'\b[a-z0-9]+(?:[\-\.][a-z0-9]+)*\b', text.lower())

    def _build_tfidf(self, tokens: list[str]) -> dict[str, float]:
        """Build a TF-IDF vector from tokens."""
        tf = Counter(tokens)
        total = len(tokens) or 1
        vector = {}
        for term, count in tf.items():
            tf_val = count / total
            idf_val = math.log((self.total_docs + 1) / (self.vocab.get(term, 0) + 1))
            vector[term] = tf_val * idf_val
        return vector

    def _cosine_sim(self, v1: dict, v2: dict) -> float:
        """Cosine similarity between two sparse vectors."""
        common = set(v1) & set(v2)
        if not common:
            return 0.0
        dot = sum(v1[k] * v2[k] for k in common)
        mag1 = math.sqrt(sum(v ** 2 for v in v1.values()))
        mag2 = math.sqrt(sum(v ** 2 for v in v2.values()))
        if mag1 == 0 or mag2 == 0:
            return 0.0
        return dot / (mag1 * mag2)

    def add(self, chunk_id: str, content: str, metadata: dict) -> None:
        """Add a document chunk to the store."""
        tokens = self._tokenize(content)
        # Update document frequency counts
        unique_terms = set(tokens)
        for term in unique_terms:
            self.vocab[term] = self.vocab.get(term, 0) + 1
        self.total_docs += 1
        vector = self._build_tfidf(tokens)
        self.documents.append({
            "chunk_id": chunk_id,
            "content": content,
            "metadata": metadata,
            "vector": vector,
        })

    def search(
        self,
        query: str,
        top_k: int = 5,
        filters: Optional[dict] = None,
    ) -> list[dict]:
        """Search for relevant chunks.

        WHAT: Computes cosine similarity between query and all chunks,
              applies optional metadata filters, returns top_k results.
        WHY:  This is the retrieval step of RAG — the quality of these
              results directly determines Claude's answer quality.
        """
        if not query.strip():
            return [{"error": "EMPTY_QUERY", "message": "Query cannot be empty."}]

        query_tokens = self._tokenize(query)
        query_vector = self._build_tfidf(query_tokens)

        results = []
        for doc in self.documents:
            # Apply metadata filters
            if filters:
                skip = False
                for key, val in filters.items():
                    if key in doc["metadata"]:
                        doc_val = doc["metadata"][key]
                        if isinstance(doc_val, list):
                            if val not in doc_val:
                                skip = True
                        elif doc_val.lower() != val.lower():
                            skip = True
                if skip:
                    continue

            score = self._cosine_sim(query_vector, doc["vector"])
            results.append({
                "chunk_id": doc["chunk_id"],
                "doc_id": doc["metadata"].get("doc_id", ""),
                "section_id": doc["metadata"].get("section_id", ""),
                "content": doc["content"],
                "similarity_score": round(score, 4),
                "metadata": doc["metadata"],
            })

        results.sort(key=lambda x: x["similarity_score"], reverse=True)
        return results[:top_k]

    def clear(self) -> None:
        """Reset the store."""
        self.documents.clear()
        self.vocab.clear()
        self.total_docs = 0

"""agent.py — Clinical Policy RAG Agent (Capstone 2-A)

A RAG-powered conversational agent that answers questions about
payer clinical policies with cited references.

Usage:
  export ANTHROPIC_API_KEY=your-key-here
  python agent.py
"""

import json
import re
import anthropic
from vector_store import MockVectorStore

# ── WHAT: Initialize client and vector store ───────────────────
client = anthropic.Anthropic()
MODEL = "claude-sonnet-4-6"
store = MockVectorStore()

# ── WHAT: Ingest the policy corpus ─────────────────────────────
# WHY: Section-aware chunking preserves the logical structure of
#      clinical policies. Each section becomes one chunk with
#      rich metadata for filtering.
def ingest_policies(corpus_path: str = "data/policies.json") -> int:
    """Load policies, chunk by section, add to vector store."""
    with open(corpus_path) as f:
        corpus = json.load(f)

    chunk_count = 0
    for doc in corpus["documents"]:
        doc_id = doc["doc_id"]
        payer = doc["payer"]
        category = doc["category"]
        effective_date = doc["effective_date"]

        for section in doc["sections"]:
            # ── WHAT: Extract CPT/ICD codes from content ───────
            # WHY: Storing codes in metadata enables exact-match
            #      filtering alongside semantic search.
            cpt_codes = re.findall(r'\b\d{5}\b', section["content"])
            icd_codes = re.findall(r'\b[A-Z]\d{2}\.\w+\b', section["content"])

            chunk_id = f"{doc_id}:{section['section_id']}"
            content = f"{section['title']}: {section['content']}"
            metadata = {
                "doc_id": doc_id,
                "doc_title": doc["title"],
                "section_id": section["section_id"],
                "section_title": section["title"],
                "payer": payer,
                "category": category,
                "effective_date": effective_date,
                "cpt_codes": cpt_codes,
                "icd_codes": icd_codes,
            }

            store.add(chunk_id, content, metadata)
            chunk_count += 1

    return chunk_count

# ── WHAT: Define the search tool ───────────────────────────────
TOOLS = [
    {
        "name": "search_policy_knowledge_base",
        "description": (
            "Search the clinical policy knowledge base for relevant "
            "policy sections. Returns ranked chunks with similarity scores "
            "and metadata (payer, category, effective date, CPT/ICD codes). "
            "Use this to answer questions about prior authorization "
            "requirements, clinical criteria, covered procedures, and "
            "payer-specific policies."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Natural language search query about clinical policies",
                },
                "top_k": {
                    "type": "integer",
                    "description": "Number of results to return (default: 5)",
                },
                "filters": {
                    "type": "object",
                    "description": "Optional filters: payer, category, cpt_code",
                    "properties": {
                        "payer": {"type": "string", "description": "Filter by payer name (e.g., Aetna)"},
                        "category": {"type": "string", "description": "Filter by category (e.g., Orthopedic Surgery)"},
                        "cpt_code": {"type": "string", "description": "Filter by CPT code (e.g., 27447)"},
                    },
                },
            },
            "required": ["query"],
        },
    },
]

SYSTEM_PROMPT = """You are a clinical policy reference assistant. You help \
healthcare provider staff determine prior authorization requirements by \
searching payer clinical policy documents.

Rules:
- ALWAYS cite your sources using the format [Source: DOC_ID, Section X.Y].
- Every factual claim must reference a specific policy section.
- If the knowledge base does not contain relevant information, say so clearly. \
  NEVER make up criteria or policy details.
- You provide INFORMATION ONLY — you do not approve or deny authorizations.
- If a query is vague, ask for clarification (which payer? which procedure?).
- Note the effective date of cited policies so users know about currency.
- When comparing policies across payers, clearly label which criteria belong \
  to which payer.
- Do not provide medical advice. Relay policy criteria only."""

def handle_search(args: dict) -> str:
    """Execute the search tool against the vector store."""
    query = args.get("query", "")
    top_k = args.get("top_k", 5)
    filters = args.get("filters")

    # Convert cpt_code filter to cpt_codes list filter
    if filters and "cpt_code" in filters:
        filters["cpt_codes"] = filters.pop("cpt_code")

    results = store.search(query, top_k=top_k, filters=filters)
    return json.dumps(results, indent=2)


def chat(user_message: str, conversation_history: list) -> str:
    """Send a message through the RAG agent loop."""
    conversation_history.append({"role": "user", "content": user_message})

    while True:
        response = client.messages.create(
            model=MODEL,
            max_tokens=1500,
            system=SYSTEM_PROMPT,
            tools=TOOLS,
            messages=conversation_history,
        )

        if response.stop_reason == "tool_use":
            conversation_history.append({
                "role": "assistant",
                "content": response.content,
            })

            tool_results = []
            for block in response.content:
                if block.type == "tool_use":
                    try:
                        result = handle_search(block.input)
                    except Exception as e:
                        result = json.dumps({"error": "INDEX_UNAVAILABLE", "message": str(e)})
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result,
                    })

            conversation_history.append({"role": "user", "content": tool_results})
            continue

        conversation_history.append({
            "role": "assistant",
            "content": response.content,
        })
        return "\n".join(
            b.text for b in response.content if hasattr(b, "text")
        )


def main():
    """Run the interactive RAG agent."""
    print("Loading clinical policy knowledge base...")
    count = ingest_policies()
    print(f"Ingested {count} policy chunks.")
    print("=" * 60)
    print("  Clinical Policy Q&A — Capstone 2-A")
    print("  Ask about prior auth requirements, clinical criteria, etc.")
    print("  Type 'quit' to exit.")
    print("=" * 60)

    history = []
    while True:
        user_input = input("\nYou: ").strip()
        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "q"):
            print("Goodbye!")
            break
        try:
            response = chat(user_input, history)
            print(f"\nAgent: {response}")
        except anthropic.APIError as e:
            print(f"\n[API Error] {e.message}")
        except Exception as e:
            print(f"\n[Error] {e}")


if __name__ == "__main__":
    main()

# verify.py — End-to-end smoke test
from agent import ingest_policies, chat

ingest_policies()

test_queries = [
    ("Is prior auth required for CPT 27447 under Aetna?",
     ["POLICY-AETNA-ORTHO-2024"]),
    ("What conservative treatments does UHC require before knee replacement?",
     ["POLICY-UHC-ORTHO-2024"]),
    ("Does Aetna require prior auth for brain MRI?",
     ["POLICY-AETNA-IMAGING-2024"]),
]

print("=== End-to-End Verification ===\n")
passed = 0
for query, expected_citations in test_queries:
    history = []
    response = chat(query, history)
    found = [cit for cit in expected_citations if cit in response]
    status = "PASS" if len(found) == len(expected_citations) else "FAIL"
    if status == "PASS":
        passed += 1
    print(f"[{status}] Query: {query}")
    print(f"       Expected citations: {expected_citations}")
    print(f"       Found: {found}\n")

print(f"Result: {passed}/{len(test_queries)} tests passed.")
if passed == len(test_queries):
    print("All tests passed — your RAG agent is working correctly!")

python verify.py

"""test_rag.py — Tests for the Clinical Policy RAG pipeline."""

import json
import pytest
from vector_store import MockVectorStore

@pytest.fixture
def loaded_store():
    """Create a vector store with test policy data."""
    store = MockVectorStore()
    with open("data/policies.json") as f:
        corpus = json.load(f)
    for doc in corpus["documents"]:
        for section in doc["sections"]:
            store.add(
                f"{doc['doc_id']}:{section['section_id']}",
                f"{section['title']}: {section['content']}",
                {"doc_id": doc["doc_id"], "payer": doc["payer"],
                 "section_id": section["section_id"],
                 "category": doc["category"]},
            )
    return store

class TestVectorStore:
    def test_search_returns_results(self, loaded_store):
        results = loaded_store.search("knee arthroplasty prior auth", top_k=3)
        assert len(results) > 0
        assert results[0]["similarity_score"] > 0

    def test_payer_filter(self, loaded_store):
        results = loaded_store.search("knee replacement", filters={"payer": "Aetna"})
        for r in results:
            assert r["metadata"]["payer"] == "Aetna"

    def test_empty_query_error(self, loaded_store):
        results = loaded_store.search("")
        assert results[0].get("error") == "EMPTY_QUERY"

    def test_no_match_returns_empty_or_low_score(self, loaded_store):
        results = loaded_store.search("quantum physics")
        if results:
            assert results[0]["similarity_score"] < 0.3

class TestChunking:
    def test_section_metadata_preserved(self, loaded_store):
        results = loaded_store.search("clinical criteria knee")
        for r in results:
            assert "doc_id" in r["metadata"]
            assert "section_id" in r["metadata"]

if __name__ == "__main__":
    pytest.main([__file__, "-v"])