Review fixes and cassette-backed provider tests

README.md

@@ -144,7 +144,7 @@ TypeSafe response:
 - By default, it requests probabilities, but there is also a mode that supports getting discrete values
   - it will then map the discrete value to a probability distribution of all 0s, except one element which is 1.0
 - Telemetry
-  - `input_tokens` and `output_tokens` aggregate all PydanticAI requests in the successful run, including malformed-structure retries
+  - `input_tokens` and `output_tokens` aggregate every PydanticAI request made during the call, including malformed-structure retries and attempts that later failed and were retried, so a retried call is never under-billed
   - `n_retries` counts retries of transient provider failures
   - `n_retries_malformed_structure` counts PydanticAI corrective retries for malformed output
   - `latency` is end-to-end request latency in seconds, including retries
@@ -159,8 +159,11 @@ TypeSafe response:
   - `retry` configures retries for transient connection, timeout, and retryable HTTP failures
   - `n_retry_malformed_structure` configures corrective retries for output that fails structural validation
   - authentication, context-window, and malformed-structure errors are not retried by `retry`
-- Tests make real LLM and TypeSafe API calls, but also caches the calls to make testing deterministic (see json_cache.py)
-  - the json_cache.json files should live in the test directory next to the tests
+- Tests make real LLM and TypeSafe API calls, recorded as HTTP cassettes so replay is deterministic (vcrpy via pytest-recording)
+  - cassettes are JSON and live in `tests/cassettes`, next to the tests
+  - replay is the default and needs no credentials or network; the whole client stack runs against recorded provider traffic
+  - re-record with `uv run pytest tests/test_live_apis.py --record-mode=rewrite`, which makes real billable calls and needs `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, and `TYPESAFE_API_KEY`
+  - request/response credentials are stripped at record time (see `tests/conftest.py`); matching includes the request body because every call posts to the same endpoint
 - Compatibility scope
   - `OpenTypeSafeClient` subclasses `TypeSafeClient`
   - supports synchronous and asynchronous `system_one`, context management, `close`, and `aclose`

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
-name = "opentypesafe2"
+name = "open-typesafe-client"
 version = "0.1.0"
-description = "Add your description here"
+description = "Drop-in TypeSafeClient replacement backed by LLM APIs via PydanticAI"
 readme = "README.md"
 authors = [
     { name = "Erik Gafni", email = "erik@typesafe.ai" }
@@ -17,6 +17,7 @@ dependencies = [
 [dependency-groups]
 dev = [
     "pytest>=9.0.2",
+    "pytest-recording>=0.13.4",
     "ruff>=0.14.14",
 ]
 
@@ -27,9 +28,29 @@ url = "https://pypi.typesafe.ai/"
 [tool.uv.build-backend]
 module-name = "open_typesafe_client"
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+# Fail rather than silently making a live call when a cassette is missing. Passing
+# --record-mode (anything but "none") lifts this for cassette-backed tests.
+addopts = "--block-network"
+
 [tool.ruff]
 target-version = "py311"
 
+[tool.ruff.lint]
+select = [
+    "E",    # pycodestyle
+    "W",    # pycodestyle warnings
+    "F",    # pyflakes
+    "I",    # import sorting
+    "UP",   # pyupgrade
+    "B",    # bugbear
+    "BLE",  # blind except
+    "C4",   # comprehensions
+    "SIM",  # simplify
+    "RUF",  # ruff-specific
+]
+
 [build-system]
 requires = ["uv_build>=0.11.7,<0.12.0"]
 build-backend = "uv_build"

src/open_typesafe_client/client.py

@@ -4,12 +4,14 @@
 
 import json
 import time
-from collections.abc import Mapping
+from dataclasses import dataclass, field
 from types import TracebackType
 from typing import Any, Self, cast
 
 from pydantic import BaseModel
+from pydantic_ai import Agent
 from pydantic_ai.models import Model
+from pydantic_ai.usage import RunUsage
 from typesafe_client import RetryConfig, TypeSafeClient
 from typesafe_client.api.models import (
     ChoiceAnswer,
@@ -38,8 +40,9 @@
 from open_typesafe_client.utils.probability_normalization import (
     AnswerMode,
     ProbabilityNormalization,
-    ProbabilityNormalizationStats,
     normalize_probabilities,
+    probability_usage_data,
+    to_distribution,
 )
 from open_typesafe_client.utils.pydantic_utils import (
     Question,
@@ -82,7 +85,13 @@ def _answer(
             should_normalize_probabilities,
         )
         probabilities = probability_normalization.probabilities
-        score = sum(index * probabilities[str(index)] for index in range(len(labels)))
+        # The score is an expected value, so it is only meaningful over a distribution
+        # summing to 1. Normalize explicitly here: the reported probabilities are left
+        # untouched when ``normalize_probabilities`` is disabled.
+        score_distribution = to_distribution(probabilities)
+        score = sum(
+            index * score_distribution[str(index)] for index in range(len(labels))
+        )
         answer = ScoreAnswer(
             type=QuestionType.Score,
             score=score,
@@ -109,41 +118,74 @@ def _answer(
     return answer, probability_normalization
 
 
-def _response(
-    model: str,
-    questions: Mapping[str, Question],
-    output: BaseModel,
-    input_tokens: int,
-    output_tokens: int,
-    n_retries: int,
-    n_retries_malformed_structure: int,
-    latency: float,
-    llm_answer_mode: AnswerMode,
-    should_normalize_probabilities: bool,
-) -> SystemOneResponse:
-    raw_answers = output.model_dump(mode="python", by_alias=True)["answers"]
-    answers: dict[str, Answer] = {}
-    probability_normalization_stats = ProbabilityNormalizationStats()
-    for question_id, question in questions.items():
-        answer, probability_normalization = _answer(
-            question,
-            raw_answers[question_id],
-            llm_answer_mode,
-            should_normalize_probabilities,
-        )
-        answers[question_id] = answer
-        probability_normalization_stats.add(question_id, probability_normalization)
+@dataclass
+class _Evaluation:
+    """State shared by the synchronous and asynchronous ``system_one`` paths.
 
-    usage_data: dict[str, Any] = {
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "n_retries": n_retries,
-        "n_retries_malformed_structure": n_retries_malformed_structure,
-        "latency": latency,
-    }
-    usage_data.update(probability_normalization_stats.usage_data())
-    usage = Usage(**usage_data)
-    return SystemOneResponse(model=model, answers=answers, usage=usage)
+    ``run_usage`` is threaded through every PydanticAI attempt so that tokens spent on
+    attempts which later failed are still billed to the caller.
+    """
+
+    model_name: str
+    questions: dict[str, Question]
+    pydantic_agent: Agent
+    llm_answer_mode: AnswerMode
+    should_normalize_probabilities: bool
+    run_usage: RunUsage = field(default_factory=RunUsage)
+    started_at: float = field(default_factory=time.perf_counter)
+    _requests_before_attempt: int | None = None
+    _n_retries_malformed_structure: int = 0
+
+    def begin_attempt(self) -> None:
+        """Finish accounting for the prior attempt and start the next one."""
+        if self._requests_before_attempt is not None:
+            self._n_retries_malformed_structure += (
+                self.run_usage.requests - self._requests_before_attempt
+            )
+        self._requests_before_attempt = self.run_usage.requests
+
+    def response(self, output: BaseModel, n_retries: int) -> SystemOneResponse:
+        """Build the response from a successful attempt.
+
+        :param output: Validated LLM output model.
+        :param n_retries: Transient-failure retries performed.
+        :return: TypeSafe-shaped response.
+        """
+        latency = time.perf_counter() - self.started_at
+        raw_answers = output.model_dump(mode="python", by_alias=True)["answers"]
+        answers: dict[str, Answer] = {}
+        probability_normalizations: dict[
+            str,
+            ProbabilityNormalization | None,
+        ] = {}
+        for question_id, question in self.questions.items():
+            answer, probability_normalization = _answer(
+                question,
+                raw_answers[question_id],
+                self.llm_answer_mode,
+                self.should_normalize_probabilities,
+            )
+            answers[question_id] = answer
+            probability_normalizations[question_id] = probability_normalization
+
+        assert self._requests_before_attempt is not None
+        n_retries_malformed_structure = self._n_retries_malformed_structure + max(
+            0,
+            self.run_usage.requests - self._requests_before_attempt - 1,
+        )
+        usage_data: dict[str, Any] = {
+            "input_tokens": self.run_usage.input_tokens,
+            "output_tokens": self.run_usage.output_tokens,
+            "n_retries": n_retries,
+            "n_retries_malformed_structure": n_retries_malformed_structure,
+            "latency": latency,
+        }
+        usage_data.update(probability_usage_data(probability_normalizations))
+        return SystemOneResponse(
+            model=self.model_name,
+            answers=answers,
+            usage=Usage(**usage_data),
+        )
 
 
 class OpenTypeSafeClient(TypeSafeClient):
@@ -164,6 +206,10 @@ def __init__(
         n_retry_malformed_structure: int = 0,
         retry: RetryConfig = NoRetries(),  # noqa: B008 - reference-compatible signature
     ) -> None:
+        # ``TypeSafeClient.__init__`` is deliberately not called: it requires a TypeSafe
+        # API key and builds ``self._api_client``, neither of which this client uses.
+        # Every inherited method that touches ``self._api_client`` is overridden
+        # below, so a new one must be overridden here too or it raises AttributeError.
         if llm_answer_mode not in ("probabilities", "discrete"):
             raise ValueError("llm_answer_mode must be 'probabilities' or 'discrete'")
         if n_retry_malformed_structure < 0:
@@ -175,13 +221,12 @@ def __init__(
         self.n_retry_malformed_structure = n_retry_malformed_structure
         self.retry = retry
 
-    def system_one(
+    def _evaluation(
         self,
         model: str | Model,
-        document: InstructionValue,
         questions: QuestionCollectionType,
-    ) -> SystemOneResponse:
-        """Synchronously evaluate ``questions`` against one ``document``."""
+    ) -> _Evaluation:
+        """Prepare the questions, output model, and agent for one evaluation."""
         prepared_questions = prepare_questions(questions)
         output_model = create_llm_output_model(
             prepared_questions,
@@ -194,64 +239,51 @@ def system_one(
             self.n_retry_malformed_structure,
             _SYSTEM_PROMPT,
         )
-        started_at = time.perf_counter()
-
-        result, n_retries = run_with_retries(
-            lambda: pydantic_agent.run_sync(_prompt(document)),
-            self.retry,
-        )
-        latency = time.perf_counter() - started_at
-        return _response(
-            get_model_name(model),
-            prepared_questions,
-            cast(BaseModel, result.output),
-            result.usage.input_tokens,
-            result.usage.output_tokens,
-            n_retries,
-            max(0, result.usage.requests - 1),
-            latency,
-            self.llm_answer_mode,
-            self.normalize_probabilities,
+        return _Evaluation(
+            model_name=get_model_name(model),
+            questions=prepared_questions,
+            pydantic_agent=pydantic_agent,
+            llm_answer_mode=self.llm_answer_mode,
+            should_normalize_probabilities=self.normalize_probabilities,
         )
 
+    def system_one(
+        self,
+        model: str | Model,
+        document: InstructionValue,
+        questions: QuestionCollectionType,
+    ) -> SystemOneResponse:
+        """Synchronously evaluate ``questions`` against one ``document``."""
+        evaluation = self._evaluation(model, questions)
+
+        def attempt() -> Any:
+            evaluation.begin_attempt()
+            return evaluation.pydantic_agent.run_sync(
+                _prompt(document),
+                usage=evaluation.run_usage,
+            )
+
+        result, n_retries = run_with_retries(attempt, self.retry)
+        return evaluation.response(cast(BaseModel, result.output), n_retries)
+
     async def system_one_async(
         self,
         model: str | Model,
         document: InstructionValue,
         questions: QuestionCollectionType,
     ) -> SystemOneResponse:
         """Asynchronously evaluate ``questions`` against one ``document``."""
-        prepared_questions = prepare_questions(questions)
-        output_model = create_llm_output_model(
-            prepared_questions,
-            self.llm_answer_mode,
-        )
-        pydantic_agent = create_pydantic_ai_agent(
-            model,
-            output_model,
-            self.structured_outputs,
-            self.n_retry_malformed_structure,
-            _SYSTEM_PROMPT,
-        )
-        started_at = time.perf_counter()
+        evaluation = self._evaluation(model, questions)
 
-        result, n_retries = await run_with_retries_async(
-            lambda: pydantic_agent.run(_prompt(document)),
-            self.retry,
-        )
-        latency = time.perf_counter() - started_at
-        return _response(
-            get_model_name(model),
-            prepared_questions,
-            cast(BaseModel, result.output),
-            result.usage.input_tokens,
-            result.usage.output_tokens,
-            n_retries,
-            max(0, result.usage.requests - 1),
-            latency,
-            self.llm_answer_mode,
-            self.normalize_probabilities,
-        )
+        def attempt() -> Any:
+            evaluation.begin_attempt()
+            return evaluation.pydantic_agent.run(
+                _prompt(document),
+                usage=evaluation.run_usage,
+            )
+
+        result, n_retries = await run_with_retries_async(attempt, self.retry)
+        return evaluation.response(cast(BaseModel, result.output), n_retries)
 
     def close(self) -> None:
         """Close the client.