test1_retrieval.py

"""
Test 1 — Deterministic Data Retrieval: JSON vs simulated Vector DB.

Hypothesis: Structured JSON retrieval is deterministically correct; a simulated
Vector DB has a measurable error rate due to probabilistic top-k retrieval.
"""

import random
from mock_data import SESSIONS

VECTOR_DB_ERROR_RATE = 0.15   # 15% chance of top-k mismatch

QUERIES = [
    (1,  "weight_kg",      "What was the patient's weight in session 1?"),
    (2,  "dyspnea_score",  "What was the dyspnea score in session 2?"),
    (3,  "bp_mmhg",        "What was the blood pressure in session 3?"),
    (4,  "edema",          "Was there edema in session 4?"),
    (5,  "weight_kg",      "What was the patient's weight in session 5?"),
    (6,  "dyspnea_score",  "What was the dyspnea score in session 6?"),
    (7,  "bp_mmhg",        "What was the blood pressure in session 7?"),
    (8,  "edema",          "Was there edema in session 8?"),
    (9,  "weight_kg",      "What was the patient's weight in session 9?"),
    (10, "dyspnea_score",  "What was the dyspnea score in session 10?"),
    (3,  "weight_kg",      "How much did the patient weigh in session 3?"),
    (7,  "weight_kg",      "How much did the patient weigh in session 7?"),
    (1,  "edema",          "Did the patient have edema in session 1?"),
    (5,  "edema",          "Did the patient have edema in session 5?"),
    (10, "edema",          "Did the patient have edema in session 10?"),
    (2,  "bp_mmhg",        "What was the blood pressure in session 2?"),
    (6,  "bp_mmhg",        "What was the blood pressure in session 6?"),
    (4,  "dyspnea_score",  "How severe was the dyspnea in session 4?"),
    (8,  "dyspnea_score",  "How severe was the dyspnea in session 8?"),
    (9,  "bp_mmhg",        "What was the blood pressure in session 9?"),
]


def json_retrieve(session_id: int, field: str):
    """Exact field lookup in the structured JSON schema."""
    session = next(s for s in SESSIONS if s["session_id"] == session_id)
    return session[field]


def vector_db_retrieve(session_id: int, field: str):
    """
    Simulated Vector DB: with probability VECTOR_DB_ERROR_RATE the wrong
    session is returned (top-k mismatch — adjacent session ±1).
    """
    if random.random() < VECTOR_DB_ERROR_RATE:
        wrong_id = session_id + random.choice([-1, 1])
        wrong_id = max(1, min(10, wrong_id))
        session = next(s for s in SESSIONS if s["session_id"] == wrong_id)
    else:
        session = next(s for s in SESSIONS if s["session_id"] == session_id)
    return session[field]


def run_test1() -> dict:
    print("\n" + "=" * 60)
    print("TEST 1 — Deterministic Data Retrieval")
    print("=" * 60)

    json_correct = 0
    vector_correct = 0
    results = []

    for session_id, field, query in QUERIES:
        expected = json_retrieve(session_id, field)
        json_answer = json_retrieve(session_id, field)
        vector_answer = vector_db_retrieve(session_id, field)

        j_ok = json_answer == expected
        v_ok = vector_answer == expected

        json_correct += j_ok
        vector_correct += v_ok

        results.append({
            "query": query,
            "session_id": session_id,
            "field": field,
            "expected": str(expected),
            "json_ok": j_ok,
            "vector_ok": v_ok,
        })
        status = "OK" if j_ok else "FAIL"
        vstatus = "OK" if v_ok else "FAIL (mismatch)"
        print(f"  [{status}|{vstatus}] Session {session_id} / {field}: expected={expected}, vector={vector_answer}")

    total = len(QUERIES)
    print(f"\n  JSON method:       {json_correct}/{total} correct ({json_correct/total*100:.0f}%)")
    print(f"  Vector DB (sim):   {vector_correct}/{total} correct ({vector_correct/total*100:.0f}%)")

    return {
        "total": total,
        "json_correct": json_correct,
        "vector_correct": vector_correct,
        "results": results,
    }