test2_tokens.py

"""
Test 2 — Context Window / Token Efficiency.

Hypothesis: Loading profile + 2 recent sessions stays well under 4,000 tokens;
naively loading all 10 sessions exceeds that limit significantly.
"""

import json
import tiktoken
from mock_data import PATIENT_PROFILE, SESSIONS

TOKEN_LIMIT = 4000


def build_payload_optimized() -> str:
    """Profile + 2 most recent sessions (sessions 9 and 10)."""
    payload = {
        "patient_profile": PATIENT_PROFILE,
        "recent_sessions": [SESSIONS[8], SESSIONS[9]],  # index 8=session9, 9=session10
    }
    return json.dumps(payload, ensure_ascii=False, indent=2)


def build_payload_naive() -> str:
    """Profile + all 10 sessions loaded into the prompt."""
    payload = {
        "patient_profile": PATIENT_PROFILE,
        "all_sessions": SESSIONS,
    }
    return json.dumps(payload, ensure_ascii=False, indent=2)


def count_tokens(text: str, model: str = "gpt-4o") -> int:
    enc = tiktoken.encoding_for_model(model)
    return len(enc.encode(text))


def run_test2() -> dict:
    print("\n" + "=" * 60)
    print("TEST 2 — Context Window / Token Efficiency")
    print("=" * 60)

    opt_tokens = count_tokens(build_payload_optimized())
    naive_tokens = count_tokens(build_payload_naive())

    print(f"  Optimized payload (profile + 2 sessions):  {opt_tokens:>5} tokens")
    print(f"  Naive payload (profile + all 10 sessions): {naive_tokens:>5} tokens")
    print(f"  Limit:                                     {TOKEN_LIMIT:>5} tokens")
    print(f"  Savings:                                   {naive_tokens - opt_tokens:>5} tokens "
          f"({(1 - opt_tokens/naive_tokens)*100:.0f}% fewer)")

    assert opt_tokens < TOKEN_LIMIT, (
        f"ASSERT FAILED: optimized payload is {opt_tokens} tokens (limit: {TOKEN_LIMIT})"
    )
    print(f"\n  ASSERT passed: {opt_tokens} < {TOKEN_LIMIT} tokens")

    return {
        "opt_tokens": opt_tokens,
        "naive_tokens": naive_tokens,
        "limit": TOKEN_LIMIT,
        "assert_passed": opt_tokens < TOKEN_LIMIT,
    }