`5ee947d`

Add token usage tracking to OpenAI and Anthropic teachers

Authored by mfwolffe <wolffemf@dukes.jmu.edu> 2 weeks ago

SHA: 5ee947d5d21792013fc7e6843286b6dd344834a7
Parents: a016dea
Tree: 2712546

3 changed files

Status	File	+
M	`src/dlm/synth/__init__.py`	2
M	`src/dlm/synth/teachers.py`	59
M	`tests/unit/synth/test_teachers.py`	61

src/dlm/synth/__init__.pymodified

      SynthTeacher,
      TeacherKind,
      TeacherRef,
 +    TeacherUsage,
      VllmServerTeacher,
      build_teacher,
      parse_teacher_ref,
      "TeacherKind",
      "TeacherInvocationError",
      "TeacherRef",
 +    "TeacherUsage",
      "TeacherUnavailableError",
      "VllmServerTeacher",
      "apply_plan",

src/dlm/synth/teachers.pymodified

  import importlib
  import json
 +import logging
  import os
  import urllib.error
  import urllib.request
      TeacherUnavailableError,
+ )
 +_log = logging.getLogger(__name__)
++
  TeacherKind = Literal["self", "hf", "openai", "anthropic", "vllm-server"]
  _DEFAULT_MAX_NEW_TOKENS = 512
  _ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
 +@dataclass
 +class TeacherUsage:
 +    """Accumulated token usage from API-backed teachers."""
++
 +    prompt_tokens: int = 0
 +    completion_tokens: int = 0
 +    requests: int = 0
++
 +    @property
 +    def total_tokens(self) -> int:
 +        return self.prompt_tokens + self.completion_tokens
++
 +    def log_summary(self, teacher_name: str) -> None:
 +        if self.requests == 0:
 +            return
 +        _log.info(
 +            "teacher %s usage: %d requests, %d prompt tokens, "
 +            "%d completion tokens, %d total tokens",
 +            teacher_name,
 +            self.requests,
 +            self.prompt_tokens,
 +            self.completion_tokens,
 +            self.total_tokens,
 +        )
++
++
  @dataclass(frozen=True)
  class TeacherRef:
      """Parsed `--teacher` selector from the CLI."""
      client_factory: OpenAiClientFactory | None = field(default=None, repr=False, compare=False)
      api_key_env: str = field(default=_OPENAI_API_KEY_ENV, repr=False, compare=False)
      name: str = field(init=False)
 +    usage: TeacherUsage = field(default_factory=TeacherUsage, init=False, repr=False, compare=False)
      _client: Any = field(default=None, init=False, repr=False, compare=False)
      def __post_init__(self) -> None:
              response = client.chat.completions.create(**payload)
          except Exception as exc:
              raise TeacherInvocationError(f"{self.name} request failed: {exc}") from exc
 +        _accumulate_openai_usage(self.usage, response)
          return _require_non_empty_teacher_output(
              _extract_openai_message_text(response),
              teacher=self.name,
      client_factory: AnthropicClientFactory | None = field(default=None, repr=False, compare=False)
      api_key_env: str = field(default=_ANTHROPIC_API_KEY_ENV, repr=False, compare=False)
      name: str = field(init=False)
 +    usage: TeacherUsage = field(default_factory=TeacherUsage, init=False, repr=False, compare=False)
      _client: Any = field(default=None, init=False, repr=False, compare=False)
      def __post_init__(self) -> None:
              response = client.messages.create(**payload)
          except Exception as exc:
              raise TeacherInvocationError(f"{self.name} request failed: {exc}") from exc
 +        _accumulate_anthropic_usage(self.usage, response)
          return _require_non_empty_teacher_output(
              _extract_anthropic_text(response),
              teacher=self.name,
      return getattr(obj, name, None)
 +def _accumulate_openai_usage(usage: TeacherUsage, response: Any) -> None:
 +    usage.requests += 1
 +    u = _obj_get(response, "usage")
 +    if u is None:
 +        return
 +    pt = _obj_get(u, "prompt_tokens")
 +    ct = _obj_get(u, "completion_tokens")
 +    if isinstance(pt, int):
 +        usage.prompt_tokens += pt
 +    if isinstance(ct, int):
 +        usage.completion_tokens += ct
++
++
 +def _accumulate_anthropic_usage(usage: TeacherUsage, response: Any) -> None:
 +    usage.requests += 1
 +    u = _obj_get(response, "usage")
 +    if u is None:
 +        return
 +    pt = _obj_get(u, "input_tokens")
 +    ct = _obj_get(u, "output_tokens")
 +    if isinstance(pt, int):
 +        usage.prompt_tokens += pt
 +    if isinstance(ct, int):
 +        usage.completion_tokens += ct
++
++
  def _normalize_openai_compat_base_url(url: str) -> str:
      stripped = url.rstrip("/")
      if stripped.endswith("/v1/chat/completions"):

tests/unit/synth/test_teachers.pymodified

      SelfTeacher,
      TeacherInvocationError,
      TeacherUnavailableError,
 +    TeacherUsage,
      VllmServerTeacher,
      build_teacher,
      parse_teacher_ref,
          assert payloads[0]["seed"] == 5
          assert factories == ["secret"]
 +    def test_openai_teacher_accumulates_usage(self, monkeypatch: pytest.MonkeyPatch) -> None:
 +        monkeypatch.setenv("OPENAI_API_KEY", "secret")
++
 +        def _create(**_kwargs: Any) -> Any:
 +            return SimpleNamespace(
 +                choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))],
 +                usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
 +            )
++
 +        client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=_create)))
 +        teacher = OpenAiTeacher("gpt-4o-mini", client_factory=lambda _k: client)
 +        teacher.generate("sys", "usr")
 +        teacher.generate("sys", "usr")
 +        assert teacher.usage.requests == 2
 +        assert teacher.usage.prompt_tokens == 20
 +        assert teacher.usage.completion_tokens == 10
 +        assert teacher.usage.total_tokens == 30
++
      def test_openai_teacher_wraps_request_failures(self, monkeypatch: pytest.MonkeyPatch) -> None:
          monkeypatch.setenv("OPENAI_API_KEY", "secret")
          assert captured["payload"]["model"] == "claude-3-5-haiku-latest"
          assert factories == ["secret"]
 +    def test_anthropic_teacher_accumulates_usage(self, monkeypatch: pytest.MonkeyPatch) -> None:
 +        monkeypatch.setenv("ANTHROPIC_API_KEY", "secret")
++
 +        class _Messages:
 +            @staticmethod
 +            def create(**_kwargs: Any) -> Any:
 +                return SimpleNamespace(
 +                    content=[SimpleNamespace(type="text", text="ok")],
 +                    usage=SimpleNamespace(input_tokens=8, output_tokens=3),
 +                )
++
 +        teacher = AnthropicTeacher(
 +            "claude-3-5-haiku-latest",
 +            client_factory=lambda _k: SimpleNamespace(messages=_Messages()),
 +        )
 +        teacher.generate("sys", "usr")
 +        teacher.generate("sys", "usr")
 +        assert teacher.usage.requests == 2
 +        assert teacher.usage.prompt_tokens == 16
 +        assert teacher.usage.completion_tokens == 6
 +        assert teacher.usage.total_tokens == 22
++
      def test_anthropic_teacher_wraps_request_failures(
          self, monkeypatch: pytest.MonkeyPatch
      ) -> None:
          assert completion_calls[0][3:] == (29, 0.4, 0.75, 9, 30.0)
 +class TestTeacherUsage:
 +    def test_total_tokens(self) -> None:
 +        u = TeacherUsage(prompt_tokens=10, completion_tokens=5, requests=1)
 +        assert u.total_tokens == 15
++
 +    def test_log_summary_skips_zero_requests(self, caplog: pytest.LogCaptureFixture) -> None:
 +        u = TeacherUsage()
 +        u.log_summary("test")
 +        assert "test" not in caplog.text
++
 +    def test_log_summary_emits_on_nonzero(self, caplog: pytest.LogCaptureFixture) -> None:
 +        import logging
++
 +        u = TeacherUsage(prompt_tokens=100, completion_tokens=50, requests=3)
 +        with caplog.at_level(logging.INFO, logger="dlm.synth.teachers"):
 +            u.log_summary("openai:gpt-4o")
 +        assert "openai:gpt-4o" in caplog.text
 +        assert "150 total tokens" in caplog.text
++
++
  class TestTeacherHelpers:
      def test_flatten_teacher_prompt_handles_partial_inputs(self) -> None:
          assert teachers_mod._flatten_teacher_prompt("system", "user").startswith("System:\n")