Switching between model providers

pip install pytest

"""Two simulated LLM providers, standing in for two real ones, behind one
shared interface -- so the calling code never needs to know which is behind it.
"""
import random
import time
from dataclasses import dataclass


@dataclass
class NormalizedResponse:
    text: str
    tokens_used: int
    cost_usd: float
    latency_ms: float
    provider: str


def _fake_tokenize_provider_a(text: str) -> int:
    """Provider A: counts tokens roughly as whitespace-split words."""
    return max(1, len(text.split()))


def _fake_tokenize_provider_b(text: str) -> int:
    """Provider B: counts tokens roughly as groups of 4 characters -- a
    genuinely different tokenizer, the way real providers differ.
    """
    return max(1, len(text) // 4)


class ProviderA:
    name = "provider-a"
    price_per_1k_tokens = 0.002

    def generate(self, prompt: str, seed: int | None = None) -> NormalizedResponse:
        time.sleep(0.01)
        reply_text = f"[A] Here is an answer about {prompt[:30]}, in a few steps."
        tokens = _fake_tokenize_provider_a(prompt) + _fake_tokenize_provider_a(reply_text)
        return NormalizedResponse(
            text=reply_text, tokens_used=tokens,
            cost_usd=round(tokens / 1000 * self.price_per_1k_tokens, 6),
            latency_ms=10.0, provider=self.name,
        )


class ProviderB:
    name = "provider-b"
    price_per_1k_tokens = 0.0009

    def generate(self, prompt: str, seed: int | None = None) -> NormalizedResponse:
        time.sleep(0.025)
        reply_text = f"Sure, regarding {prompt[:30]}: here's a short answer."
        tokens = _fake_tokenize_provider_b(prompt) + _fake_tokenize_provider_b(reply_text)
        return NormalizedResponse(
            text=reply_text, tokens_used=tokens,
            cost_usd=round(tokens / 1000 * self.price_per_1k_tokens, 6),
            latency_ms=25.0, provider=self.name,
        )


prompt = "how do I reset my account password"
print("provider A:", ProviderA().generate(prompt))
print("provider B:", ProviderB().generate(prompt))

"""Runs the same real test prompts through both providers and reports the
differences that actually matter before switching: cost, latency, and
whether the new provider's answers still pass the same structural checks.
"""
from providers import ProviderA, ProviderB

TEST_PROMPTS = [
    "how do I reset my account password",
    "what is your refund policy",
    "how do I cancel my subscription",
    "why was my payment declined",
    "how do I update my delivery address",
]


def structural_checks_pass(response) -> bool:
    return (
        len(response.text) > 10
        and len(response.text) < 400
        and response.tokens_used > 0
    )


current, candidate = ProviderA(), ProviderB()
total_cost_current = total_cost_candidate = 0.0
total_latency_current = total_latency_candidate = 0.0

for prompt in TEST_PROMPTS:
    r_current = current.generate(prompt)
    r_candidate = candidate.generate(prompt)
    total_cost_current += r_current.cost_usd
    total_cost_candidate += r_candidate.cost_usd
    total_latency_current += r_current.latency_ms
    total_latency_candidate += r_candidate.latency_ms
    print(f"{prompt!r}: current {r_current.tokens_used}tok/${r_current.cost_usd:.6f}  "
          f"candidate {r_candidate.tokens_used}tok/${r_candidate.cost_usd:.6f}")

n = len(TEST_PROMPTS)
print(f"total cost -- current: ${total_cost_current:.6f}   candidate: ${total_cost_candidate:.6f}")
print(f"cost change: {(total_cost_candidate/total_cost_current - 1) * 100:+.1f}%")
print(f"avg latency -- current: {total_latency_current/n:.1f} ms   candidate: {total_latency_candidate/n:.1f} ms")

from compare_providers import structural_checks_pass
from providers import NormalizedResponse, ProviderA, ProviderB


def test_both_providers_return_the_same_normalized_shape():
    a = ProviderA().generate("test prompt")
    b = ProviderB().generate("test prompt")
    assert isinstance(a, NormalizedResponse)
    assert isinstance(b, NormalizedResponse)
    assert set(vars(a).keys()) == set(vars(b).keys())


def test_both_providers_pass_the_same_structural_checks():
    a = ProviderA().generate("how do I reset my password")
    b = ProviderB().generate("how do I reset my password")
    assert structural_checks_pass(a)
    assert structural_checks_pass(b)


def test_cost_is_always_a_small_positive_number():
    for provider in (ProviderA(), ProviderB()):
        response = provider.generate("some prompt")
        assert 0 < response.cost_usd < 0.01

pytest test_providers.py -q