From f0699d91bf07ff9de43e7f28d64e5768d4d3e619 Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 20:45:39 +0800
Subject: [PATCH 1/6] feat: add dynamic callback serialization via
 _callback_type marker
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Implement to_dict/from_dict on Callback base class using a
"module:ClassName" type marker (_callback_type) for dynamic dispatch.
Any Callback subclass — built-in or third-party — can now round-trip
through JSON without a hardcoded registry.

- Callback.to_dict() emits _callback_type with fully-qualified class path
- Callback.from_dict() dynamically imports and instantiates the correct
  subclass, delegating to overridden from_dict when present
- CostModel.to_dict() injects _callback_type alongside existing _type
- CostModel.from_dict() strips both markers before construction
- MlflowCallback gains to_dict/from_dict (replaces NotImplementedError stubs)
- _RunConfig.save()/load() now serialize and restore callbacks
- Replace old Callback base tests with serialization round-trip tests
---
 llmeter/callbacks/base.py               | 179 ++++++++++++++++++++----
 llmeter/callbacks/cost/model.py         |  19 ++-
 llmeter/callbacks/mlflow.py             |  31 ++--
 llmeter/runner.py                       |   7 +
 tests/unit/callbacks/cost/test_model.py |   1 +
 tests/unit/callbacks/test_base.py       |  92 +++++++-----
 6 files changed, 249 insertions(+), 80 deletions(-)

diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py
index f4c4dba..5212c71 100644
--- a/llmeter/callbacks/base.py
+++ b/llmeter/callbacks/base.py
@@ -4,14 +4,21 @@
 
 from __future__ import annotations
 
+import importlib
+import json
+import logging
 from abc import ABC
-from typing import final
+from typing import Any, final
 
 from upath.types import ReadablePathLike, WritablePathLike
 
 from ..endpoints.base import InvocationResponse
+from ..json_utils import llmeter_default_serializer
 from ..results import Result
 from ..runner import _RunConfig
+from ..utils import ensure_path
+
+logger = logging.getLogger(__name__)
 
 
 class Callback(ABC):
@@ -22,8 +29,17 @@ class Callback(ABC):
     associated with test runs or individual model invocations.
 
     A Callback object may implement multiple of the defined lifecycle hooks (such as
-    `before_invoke`, `after_run`, etc). Callbacks must support serializing their configuration to
-    a file (by implementing `save_to_file`), and loading back (via `load_from_file`).
+    `before_invoke`, `after_run`, etc). Callbacks support serializing their configuration via
+    ``to_dict()`` / ``from_dict()`` (and the convenience wrappers ``save_to_file()`` /
+    ``load_from_file()``).
+
+    Serialization uses a ``_callback_type`` marker (``"module:ClassName"``) so that
+    ``Callback.from_dict()`` can dynamically import and reconstruct the correct subclass
+    without a hardcoded registry. This means third-party callbacks round-trip through JSON
+    automatically, as long as the defining module is importable at load time.
+
+    Subclasses with complex nested state (like ``CostModel``) can override ``to_dict()`` and
+    ``from_dict()`` while preserving the type marker by calling ``super()``.
     """
 
     async def before_invoke(self, payload: dict) -> None:
@@ -71,46 +87,149 @@ async def after_run(self, result: Result) -> None:
         """
         pass
 
-    def save_to_file(self, path: WritablePathLike) -> None:
-        """Save this Callback to file
+    # -- Serialization -----------------------------------------------------------------
 
-        Individual Callbacks implement this method to save their configuration to a file that will
-        be re-loadable with the equivalent `_load_from_file()` method.
+    def to_dict(self) -> dict:
+        """Serialize this callback's configuration to a JSON-safe dict.
 
-        Args:
-            path: (Local or Cloud) path where the callback is saved
+        The returned dict includes a ``_callback_type`` key with the fully-qualified
+        class path (``"module:ClassName"``), enabling ``Callback.from_dict`` to
+        reconstruct the correct subclass without a hardcoded registry.
+
+        By default, all public (non-underscore-prefixed) instance attributes are
+        included. Subclasses with richer state should override this method and call
+        ``super().to_dict()`` to preserve the type marker.
+
+        Returns:
+            dict: A JSON-serializable dictionary representation of this callback.
+
+        Example::
+
+            >>> from llmeter.callbacks import CostModel
+            >>> from llmeter.callbacks.cost.dimensions import InputTokens
+            >>> model = CostModel(request_dims=[InputTokens(price_per_million=3.0)])
+            >>> d = model.to_dict()
+            >>> d["_callback_type"]
+            'llmeter.callbacks.cost.model:CostModel'
         """
-        raise NotImplementedError("TODO: Callback.save_to_file is not yet implemented!")
+        cls = self.__class__
+        data: dict[str, Any] = {
+            "_callback_type": f"{cls.__module__}:{cls.__qualname__}",
+        }
+        data.update({k: v for k, v in vars(self).items() if not k.startswith("_")})
+        return data
 
-    @staticmethod
-    @final
-    def load_from_file(path: ReadablePathLike) -> Callback:
-        """Load (any type of) Callback from file
+    @classmethod
+    def from_dict(cls, raw: dict, **kwargs: Any) -> Callback:
+        """Reconstruct a Callback from a dict produced by ``to_dict()``.
 
-        `Callback.load_from_file()` attempts to detect the type of Callback saved in a given file,
-        and use the relevant implementation's `_load_from_file` method to load it.
+        Uses the ``_callback_type`` field to dynamically import and instantiate
+        the correct subclass. If called on a concrete subclass (e.g.
+        ``CostModel.from_dict(...)``), the ``_callback_type`` is still respected
+        so that the dict always controls which class is created.
+
+        Args:
+            raw: A dictionary previously produced by ``to_dict()`` (or loaded from
+                JSON). Must contain a ``_callback_type`` key.
+            **kwargs: Extra keyword arguments forwarded to the resolved class
+                constructor (or its own ``from_dict`` if it overrides this method).
+
+        Returns:
+            Callback: An instance of the appropriate Callback subclass.
+
+        Raises:
+            ValueError: If ``_callback_type`` is missing from *raw*.
+            ImportError: If the module referenced by ``_callback_type`` cannot be
+                imported.
+            AttributeError: If the class name cannot be found in the referenced
+                module.
+
+        Example::
+
+            >>> from llmeter.callbacks.base import Callback
+            >>> d = {
+            ...     "_callback_type": "llmeter.callbacks.mlflow:MlflowCallback",
+            ...     "step": 1,
+            ...     "nested": False,
+            ... }
+            >>> cb = Callback.from_dict(d)  # returns an MlflowCallback instance
+        """
+        raw = dict(raw)  # shallow copy — don't mutate caller's dict
+        callback_type = raw.pop("_callback_type", None)
+        if callback_type is None:
+            raise ValueError(
+                "Cannot deserialize Callback: dict is missing '_callback_type' key. "
+                f"Got keys: {list(raw.keys())}"
+            )
+
+        module_path, class_name = callback_type.rsplit(":", 1)
+        module = importlib.import_module(module_path)
+        callback_cls = getattr(module, class_name)
+
+        # If the resolved class has its own from_dict (e.g. CostModel), delegate to it
+        # so that subclass-specific deserialization logic is honoured.
+        if callback_cls is not cls and "from_dict" in callback_cls.__dict__:
+            # Re-inject _callback_type so the subclass from_dict can pop it if needed
+            return callback_cls.from_dict(raw, **kwargs)
+
+        # Remove any keys the constructor doesn't expect (e.g. _type from JSONableBase)
+        raw.pop("_type", None)
+        return callback_cls(**raw, **kwargs)
+
+    def to_json(self, **kwargs: Any) -> str:
+        """Serialize this callback to a JSON string.
 
         Args:
-            path: (Local or Cloud) path where the callback is saved
+            **kwargs: Extra keyword arguments forwarded to ``json.dumps``
+                (e.g. ``indent``).
+
         Returns:
-            callback: A loaded Callback - for example an `MlflowCallback`.
+            str: JSON representation of this callback.
         """
-        raise NotImplementedError(
-            "TODO: Callback.load_from_file is not yet implemented!"
-        )
+        kwargs.setdefault("default", llmeter_default_serializer)
+        return json.dumps(self.to_dict(), **kwargs)
 
     @classmethod
-    def _load_from_file(cls, path: ReadablePathLike) -> Callback:
-        """Load this Callback from file
+    def from_json(cls, json_string: str, **kwargs: Any) -> Callback:
+        """Reconstruct a Callback from a JSON string produced by ``to_json()``.
+
+        Args:
+            json_string: A valid JSON string.
+            **kwargs: Extra keyword arguments forwarded to ``from_dict``.
+
+        Returns:
+            Callback: An instance of the appropriate Callback subclass.
+        """
+        return cls.from_dict(json.loads(json_string), **kwargs)
+
+    def save_to_file(self, path: WritablePathLike) -> None:
+        """Save this Callback's configuration to a JSON file.
 
-        Individual Callbacks implement this method to define how they can be loaded from files
-        created by the equivalent `save_to_file()` method.
+        The file can be loaded back with ``Callback.load_from_file(path)``.
 
         Args:
-            path: (Local or Cloud) path where the callback is saved
+            path: (Local or Cloud) path where the callback should be saved.
+        """
+        path = ensure_path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("w") as f:
+            f.write(self.to_json(indent=4))
+
+    @staticmethod
+    @final
+    def load_from_file(path: ReadablePathLike) -> Callback:
+        """Load (any type of) Callback from a JSON file.
+
+        The ``_callback_type`` field inside the file determines which subclass is
+        instantiated, so callers don't need to know the concrete type in advance.
+
+        Args:
+            path: (Local or Cloud) path to a JSON file previously created by
+                ``save_to_file()``.
+
         Returns:
-            callback: The loaded Callback object
+            Callback: The deserialized callback instance.
         """
-        raise NotImplementedError(
-            "TODO: Callback._load_from_file is not yet implemented!"
-        )
+        path = ensure_path(path)
+        with path.open("r") as f:
+            return Callback.from_dict(json.load(f))
diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py
index dd867f4..782654b 100644
--- a/llmeter/callbacks/cost/model.py
+++ b/llmeter/callbacks/cost/model.py
@@ -15,7 +15,7 @@
 from ..base import Callback
 from .dimensions import IRequestCostDimension, IRunCostDimension
 from .results import CalculatedCostWithDimensions
-from .serde import JSONableBase, from_dict_with_class_map
+from .serde import JSONableBase, from_dict_with_class, from_dict_with_class_map
 
 
 @dataclass
@@ -203,6 +203,18 @@ async def after_run(self, result: Result) -> None:
             result, recalculate_request_costs=False, save=True
         )
 
+    def to_dict(self, **kwargs) -> dict:
+        """Serialize the cost model to a JSON-safe dict.
+
+        Injects the ``_callback_type`` marker from ``Callback.to_dict()`` into the
+        dict produced by ``JSONableBase.to_dict()``, so that
+        ``Callback.from_dict()`` can reconstruct this ``CostModel`` dynamically.
+        """
+        data = JSONableBase.to_dict(self, **kwargs)
+        cls = self.__class__
+        data["_callback_type"] = f"{cls.__module__}:{cls.__qualname__}"
+        return data
+
     def save_to_file(self, path: WritablePathLike) -> None:
         """Save the cost model (including all dimensions) to a JSON file"""
         path = ensure_path(path)
@@ -217,13 +229,16 @@ def from_dict(cls, raw: dict, alt_classes: dict = {}, **kwargs) -> "CostModel":
             **alt_classes,
         }
         raw_args = {**raw}
+        # Strip callback/serde type markers — they're not constructor args
+        raw_args.pop("_callback_type", None)
+        raw_args.pop("_type", None)
         for key in ("request_dims", "run_dims"):
             if key in raw_args:
                 raw_args[key] = {
                     name: from_dict_with_class_map(d, class_map=dim_classes)
                     for name, d in raw_args[key].items()
                 }
-        return super().from_dict(raw_args, alt_classes=alt_classes, **kwargs)
+        return from_dict_with_class(raw=raw_args, cls=cls, **kwargs)
 
     @classmethod
     def _load_from_file(cls, path: ReadablePathLike):
diff --git a/llmeter/callbacks/mlflow.py b/llmeter/callbacks/mlflow.py
index bac5ba1..bfb10de 100644
--- a/llmeter/callbacks/mlflow.py
+++ b/llmeter/callbacks/mlflow.py
@@ -1,8 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from upath.types import ReadablePathLike, WritablePathLike
-
 from ..results import Result
 from ..utils import DeferredError
 from .base import Callback
@@ -68,16 +66,29 @@ def __init__(self, step: int | None = None, nested: bool = False) -> None:
         # Check MLFlow is installed by polling any attribute on the module to trigger DeferredError
         mlflow.__version__
 
+    def to_dict(self) -> dict:
+        """Serialize this callback's configuration to a JSON-safe dict.
+
+        Returns:
+            dict: Contains ``_callback_type``, ``step``, and ``nested``.
+        """
+        data = super().to_dict()
+        # Exclude the mlflow version check side-effect attribute if present
+        return data
+
     @classmethod
-    def _load_from_file(cls, path: ReadablePathLike):
-        raise NotImplementedError(
-            "TODO: MlflowCallback does not yet support loading from file"
-        )
+    def from_dict(cls, raw: dict, **kwargs) -> "MlflowCallback":
+        """Reconstruct an MlflowCallback from a dict produced by ``to_dict()``.
 
-    def save_to_file(self, path: WritablePathLike) -> None:
-        raise NotImplementedError(
-            "TODO: MlflowCallback does not yet support saving to file"
-        )
+        Args:
+            raw: Dictionary with ``step`` and ``nested`` keys.
+
+        Returns:
+            MlflowCallback: The deserialized callback.
+        """
+        raw = dict(raw)
+        raw.pop("_callback_type", None)
+        return cls(**raw, **kwargs)
 
     async def _log_llmeter_run(self, result: Result):
         """Log parameters and metrics from an LLMeter run to MLflow.
diff --git a/llmeter/runner.py b/llmeter/runner.py
index 7101b5f..9471dd8 100644
--- a/llmeter/runner.py
+++ b/llmeter/runner.py
@@ -187,6 +187,9 @@ def save(
         if not isinstance(self.tokenizer, dict):
             config_copy.tokenizer = Tokenizer.to_dict(self.tokenizer)
 
+        if self.callbacks:
+            config_copy.callbacks = [cb.to_dict() for cb in self.callbacks]
+
         with run_config_path.open("w") as f:
             f.write(
                 json.dumps(
@@ -207,6 +210,10 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"):
             config = json.load(f)
         config["endpoint"] = Endpoint.load(config["endpoint"])
         config["tokenizer"] = Tokenizer.load(config["tokenizer"])
+        if config.get("callbacks"):
+            from .callbacks.base import Callback  # deferred: callbacks.base imports _RunConfig
+
+            config["callbacks"] = [Callback.from_dict(cb) for cb in config["callbacks"]]
         return cls(**config)
 
 
diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py
index b1ea1a8..6eed5bd 100644
--- a/tests/unit/callbacks/cost/test_model.py
+++ b/tests/unit/callbacks/cost/test_model.py
@@ -25,6 +25,7 @@ def test_cost_model_serialization():
     assert model.run_dims["ComputeSeconds"].price_per_hour == 50
     assert model.to_dict() == {
         "_type": "CostModel",
+        "_callback_type": "llmeter.callbacks.cost.model:CostModel",
         "request_dims": {
             "TokensIn": {
                 "_type": "InputTokens",
diff --git a/tests/unit/callbacks/test_base.py b/tests/unit/callbacks/test_base.py
index b3aa824..c236350 100644
--- a/tests/unit/callbacks/test_base.py
+++ b/tests/unit/callbacks/test_base.py
@@ -1,51 +1,67 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 import pytest
 
 from llmeter.callbacks.base import Callback
 
 
-class TestBase:
-    def test__load_from_file_not_implemented(self):
-        """
-        Test that _load_from_file raises NotImplementedError.
-        """
-        with pytest.raises(NotImplementedError):
-            Callback._load_from_file("valid_path.json")
-
-    def test_load_from_file_not_implemented(self):
-        """
-        Test that load_from_file raises a NotImplementedError as it's not yet implemented.
-        """
-        with pytest.raises(NotImplementedError):
-            Callback.load_from_file("valid_path.txt")
-
-    def test_load_from_file_raises_not_implemented_error(self):
-        """
-        Test that Callback.load_from_file raises a NotImplementedError.
-
-        This test verifies that calling the static method load_from_file
-        on the Callback class raises a NotImplementedError, as the method
-        is not yet implemented.
-        """
-        with pytest.raises(NotImplementedError) as excinfo:
-            Callback.load_from_file("dummy_path")
+class _DummyCallback(Callback):
+    """A minimal concrete callback for testing serialization."""
 
-        assert (
-            str(excinfo.value)
-            == "TODO: Callback.load_from_file is not yet implemented!"
-        )
+    def __init__(self, alpha: int = 1, beta: str = "hello"):
+        self.alpha = alpha
+        self.beta = beta
 
-    def test_load_from_file_raises_not_implemented_error_2(self):
-        """
-        Test that _load_from_file raises NotImplementedError when called.
-        This method is not yet implemented in the base Callback class.
-        """
-        with pytest.raises(NotImplementedError) as excinfo:
-            Callback._load_from_file("dummy_path")
 
+class TestCallbackSerialization:
+    def test_to_dict_includes_callback_type(self):
+        cb = _DummyCallback(alpha=42, beta="world")
+        d = cb.to_dict()
         assert (
-            str(excinfo.value)
-            == "TODO: Callback._load_from_file is not yet implemented!"
+            d["_callback_type"]
+            == f"{_DummyCallback.__module__}:{_DummyCallback.__qualname__}"
         )
+        assert d["alpha"] == 42
+        assert d["beta"] == "world"
+
+    def test_to_dict_excludes_private_attrs(self):
+        cb = _DummyCallback()
+        cb._internal = "secret"
+        d = cb.to_dict()
+        assert "_internal" not in d
+
+    def test_from_dict_round_trip(self):
+        cb = _DummyCallback(alpha=7, beta="test")
+        d = cb.to_dict()
+        restored = Callback.from_dict(d)
+        assert isinstance(restored, _DummyCallback)
+        assert restored.alpha == 7
+        assert restored.beta == "test"
+
+    def test_from_dict_missing_callback_type_raises(self):
+        with pytest.raises(ValueError, match="_callback_type"):
+            Callback.from_dict({"alpha": 1})
+
+    def test_to_json_round_trip(self):
+        cb = _DummyCallback(alpha=99)
+        json_str = cb.to_json()
+        restored = Callback.from_json(json_str)
+        assert isinstance(restored, _DummyCallback)
+        assert restored.alpha == 99
+
+    def test_save_and_load_from_file(self, tmp_path):
+        cb = _DummyCallback(alpha=5, beta="file_test")
+        file_path = tmp_path / "callback.json"
+        cb.save_to_file(file_path)
+
+        loaded = Callback.load_from_file(file_path)
+        assert isinstance(loaded, _DummyCallback)
+        assert loaded.alpha == 5
+        assert loaded.beta == "file_test"
+
+        # Verify the file is valid JSON with the type marker
+        with open(file_path) as f:
+            data = json.load(f)
+        assert "_callback_type" in data

From b0e57b143853111980e4956006de3afacb1472ff Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 20:45:40 +0800
Subject: [PATCH 2/6] fix: Python 3.10 compatibility for multi-version testing

- Add uv constraint-dependencies to cap onnxruntime<1.24 and gevent<25.5
  on Python 3.10 (newer versions dropped cp310 wheels)
- Set default-groups to exclude heavy dev group (sagemaker) from test runs
- Rewrite multiline f-string in heatmap.py to 3.10-compatible syntax
- Replace standalone mock package imports with stdlib unittest.mock
---
 llmeter/plotting/heatmap.py                 | 6 +++---
 pyproject.toml                              | 9 +++++++++
 tests/unit/endpoints/test_bedrock.py        | 2 +-
 tests/unit/endpoints/test_bedrock_invoke.py | 3 ++-
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/llmeter/plotting/heatmap.py b/llmeter/plotting/heatmap.py
index 8995fc5..f382d12 100644
--- a/llmeter/plotting/heatmap.py
+++ b/llmeter/plotting/heatmap.py
@@ -47,9 +47,9 @@ def __str__(self) -> str:
         Returns:
             str: String in interval notation, e.g. "[1,2]" or "(1,2)"
         """
-        return f"{'[' if self.closed in ['left', 'both'] else '('}{self.left}, {
-            self.right
-        }{']' if self.closed in ['right', 'both'] else ')'}"
+        left_bracket = "[" if self.closed in ["left", "both"] else "("
+        right_bracket = "]" if self.closed in ["right", "both"] else ")"
+        return f"{left_bracket}{self.left}, {self.right}{right_bracket}"
 
     @property
     def mid(self):
diff --git a/pyproject.toml b/pyproject.toml
index 886929b..3f4d658 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -64,6 +64,15 @@ test = [
     "llmeter[all]",
 ]
 
+[tool.uv]
+# Exclude heavy dev group (sagemaker etc.) from default install/run
+default-groups = ["test", "docs"]
+# Transitive deps that dropped cp310 wheels
+constraint-dependencies = [
+    "onnxruntime<1.24; python_version == '3.10'",
+    "gevent<25.5; python_version == '3.10'",
+]
+
 [tool.pytest.ini_options]
 asyncio_default_fixture_loop_scope = "function"
 filterwarnings = ["ignore::DeprecationWarning"]
diff --git a/tests/unit/endpoints/test_bedrock.py b/tests/unit/endpoints/test_bedrock.py
index 1dd8066..8a0b1f7 100644
--- a/tests/unit/endpoints/test_bedrock.py
+++ b/tests/unit/endpoints/test_bedrock.py
@@ -6,7 +6,7 @@
 
 import pytest
 from botocore.exceptions import ClientError
-from mock import MagicMock
+from unittest.mock import MagicMock
 
 from llmeter.endpoints.bedrock import (
     BedrockBase,
diff --git a/tests/unit/endpoints/test_bedrock_invoke.py b/tests/unit/endpoints/test_bedrock_invoke.py
index af101bc..124c142 100644
--- a/tests/unit/endpoints/test_bedrock_invoke.py
+++ b/tests/unit/endpoints/test_bedrock_invoke.py
@@ -3,9 +3,10 @@
 from contextlib import contextmanager
 from io import BytesIO
 
+from unittest.mock import MagicMock, Mock
+
 import pytest
 from botocore.exceptions import ClientError
-from mock import MagicMock, Mock
 
 from llmeter.endpoints.bedrock_invoke import (
     BedrockInvoke,

From 8ac34f4883c27289a2864691fa0b8158b0fdd3da Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 21:01:19 +0800
Subject: [PATCH 3/6] ci: add multi-Python version test workflow driven by
 classifiers

Add CI workflow (.github/workflows/ci.yml) with lint and test jobs.
The test matrix is dynamically derived from pyproject.toml classifiers,
so adding/removing a Python version only requires updating one file.

Uses actions/checkout@v7 and astral-sh/setup-uv@v8.3.2 (Node.js 22+).
setup-uv handles Python installation directly, no setup-python needed.

Also add trove classifiers to pyproject.toml.
---
 .github/workflows/ci.yml | 82 ++++++++++++++++++++++++++++++++++++++++
 pyproject.toml           |  8 ++++
 2 files changed, 90 insertions(+)
 create mode 100644 .github/workflows/ci.yml

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..aa3e0bc
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,82 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+permissions:
+  contents: read
+
+jobs:
+  # Derive the Python version matrix from pyproject.toml classifiers
+  resolve-python-versions:
+    runs-on: ubuntu-latest
+    outputs:
+      versions: ${{ steps.parse.outputs.versions }}
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Extract Python versions from classifiers
+        id: parse
+        run: |
+          versions=$(python3 -c "
+          import tomllib, json
+          with open('pyproject.toml', 'rb') as f:
+              data = tomllib.load(f)
+          classifiers = data['project'].get('classifiers', [])
+          versions = []
+          for c in classifiers:
+              # Match 'Programming Language :: Python :: 3.X'
+              parts = [p.strip() for p in c.split('::')]
+              if len(parts) == 3 and parts[0] == 'Programming Language' and parts[1] == 'Python':
+                  v = parts[2]
+                  # Only include specific minor versions (e.g. '3.10'), not just '3'
+                  if '.' in v:
+                      versions.append(v)
+          print(json.dumps(versions))
+          ")
+          echo "versions=$versions" >> "$GITHUB_OUTPUT"
+          echo "Resolved Python versions: $versions"
+
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+
+      - uses: astral-sh/setup-uv@v8.3.2
+        with:
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Ruff check
+        run: uv run ruff check .
+
+      - name: Ruff format check
+        run: uv run ruff format --check .
+
+      - name: Bandit security scan
+        run: uv run bandit -c pyproject.toml -r llmeter
+
+  test:
+    needs: resolve-python-versions
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ${{ fromJson(needs.resolve-python-versions.outputs.versions) }}
+    steps:
+      - uses: actions/checkout@v7
+
+      - uses: astral-sh/setup-uv@v8.3.2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Run tests
+        run: uv run pytest --cov=llmeter -m "not integ"
diff --git a/pyproject.toml b/pyproject.toml
index 3f4d658..93fffc3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,6 +9,14 @@ keywords = ["llm", "genai", "testing", "performance"]
 license = { text = "Apache-2.0" }
 maintainers = [{ name = "llmeter-maintainers", email = "llmeter-maintainers@amazon.com" }]
 requires-python = ">=3.10,<4"
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
 dependencies = [
     "tqdm>=4.66.0",
     "jmespath>=0.7.1,<2.0.0",

From 3b6288adfeaa7648d1d61ba5f28bdf8e9fc8e932 Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 21:12:31 +0800
Subject: [PATCH 4/6] fix: CI test failures from missing AWS region and deps

- Set dummy AWS credentials and region in test step env
- Add ruff, bandit, and ipykernel to test dependency group
---
 .github/workflows/ci.yml | 4 ++++
 pyproject.toml           | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index aa3e0bc..978d1c6 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -80,3 +80,7 @@ jobs:
 
       - name: Run tests
         run: uv run pytest --cov=llmeter -m "not integ"
+        env:
+          AWS_DEFAULT_REGION: us-east-1
+          AWS_ACCESS_KEY_ID: testing
+          AWS_SECRET_ACCESS_KEY: testing
diff --git a/pyproject.toml b/pyproject.toml
index 93fffc3..6c557c1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -68,6 +68,9 @@ test = [
     "pytest-cov>=5.0.0",
     "pillow>=12.1.1",
     "aws-bedrock-token-generator>=1.1.0",
+    "ruff>=0.6.5",
+    "bandit>=1.7.10",
+    "ipykernel>=6.29.0",
     # Pull in all optional dependencies for comprehensive testing
     "llmeter[all]",
 ]

From c683fd43e51a5e1d0f6aa34a00bd7ae299eaa6c1 Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 21:18:17 +0800
Subject: [PATCH 5/6] fix: CI lint and test failures across Python versions

- Scope ruff/format to llmeter/ and tests/ (skip example notebooks)
- Add ruff, bandit, ipykernel to test dependency group
- Fix abstract class error message match for Python 3.10
- Add audio/wave MIME type support in OpenAI endpoint
- Handle missing task.cancelling() on Python 3.10
- Remove unused imports and variables in tests
---
 .github/workflows/ci.yml                     | 4 ++--
 llmeter/endpoints/openai.py                  | 3 ++-
 tests/unit/callbacks/cost/test_dimensions.py | 4 ++--
 tests/unit/endpoints/test_llmeter_invoke.py  | 2 +-
 tests/unit/test_interrupted_run.py           | 8 +++++++-
 tests/unit/test_runner_backlog_ui.py         | 3 ---
 6 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 978d1c6..9ec0375 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -53,10 +53,10 @@ jobs:
         run: uv sync
 
       - name: Ruff check
-        run: uv run ruff check .
+        run: uv run ruff check llmeter tests
 
       - name: Ruff format check
-        run: uv run ruff format --check .
+        run: uv run ruff format --check llmeter tests
 
       - name: Bandit security scan
         run: uv run bandit -c pyproject.toml -r llmeter
diff --git a/llmeter/endpoints/openai.py b/llmeter/endpoints/openai.py
index 79bcf72..1806909 100644
--- a/llmeter/endpoints/openai.py
+++ b/llmeter/endpoints/openai.py
@@ -34,13 +34,14 @@
 
 # MIME types supported by OpenAI, grouped by content part type
 _OPENAI_IMAGE_MIMES = {"image/jpeg", "image/png", "image/gif", "image/webp"}
-_OPENAI_AUDIO_MIMES = {"audio/mpeg", "audio/wav"}
+_OPENAI_AUDIO_MIMES = {"audio/mpeg", "audio/wav", "audio/wave"}
 _OPENAI_FILE_MIMES = {"application/pdf"}
 
 # Map MIME → OpenAI audio "format" field value
 _MIME_TO_OPENAI_AUDIO_FMT: dict[str, Literal["mp3", "wav"]] = {
     "audio/mpeg": "mp3",
     "audio/wav": "wav",
+    "audio/wave": "wav",
 }
 
 TOpenAICompletionBase = TypeVar(
diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py
index 68cbe1c..71c623f 100644
--- a/tests/unit/callbacks/cost/test_dimensions.py
+++ b/tests/unit/callbacks/cost/test_dimensions.py
@@ -17,13 +17,13 @@ def test_base_classes_require_implementing_calculate():
     class TestReqDim(RequestCostDimensionBase):
         pass
 
-    with pytest.raises(TypeError, match="without an implementation"):
+    with pytest.raises(TypeError, match="abstract method"):
         TestReqDim()
 
     class TestRunDim(RunCostDimensionBase):
         pass
 
-    with pytest.raises(TypeError, match="without an implementation"):
+    with pytest.raises(TypeError, match="abstract method"):
         TestRunDim()
 
 
diff --git a/tests/unit/endpoints/test_llmeter_invoke.py b/tests/unit/endpoints/test_llmeter_invoke.py
index 4fc9718..86500ca 100644
--- a/tests/unit/endpoints/test_llmeter_invoke.py
+++ b/tests/unit/endpoints/test_llmeter_invoke.py
@@ -9,7 +9,7 @@
 
 from datetime import datetime, timezone
 
-from llmeter.endpoints.base import Endpoint, InvocationResponse
+from llmeter.endpoints.base import Endpoint
 
 # ---------------------------------------------------------------------------
 # Minimal concrete endpoint for testing the decorator in isolation
diff --git a/tests/unit/test_interrupted_run.py b/tests/unit/test_interrupted_run.py
index e535d15..e70f23a 100644
--- a/tests/unit/test_interrupted_run.py
+++ b/tests/unit/test_interrupted_run.py
@@ -526,7 +526,13 @@ async def target():
         shutdown._task = task
 
         shutdown._handle(signal.SIGINT)
-        assert task.cancelling() > 0
+        # task.cancelling() added in 3.11; on 3.10 verify via await
+        if hasattr(task, "cancelling"):
+            assert task.cancelling() > 0
+        else:
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            return
 
         # Cleanup
         try:
diff --git a/tests/unit/test_runner_backlog_ui.py b/tests/unit/test_runner_backlog_ui.py
index c684793..6b2f0ba 100644
--- a/tests/unit/test_runner_backlog_ui.py
+++ b/tests/unit/test_runner_backlog_ui.py
@@ -382,9 +382,6 @@ def slow_invoke(payload):
 
         mock_endpoint.invoke.side_effect = slow_invoke
 
-        backlog_bar_created = [False]
-        original_run = run._run
-
         # Patch tqdm to track if a "Processing backlog" bar is created
         with patch("llmeter.runner.tqdm") as mock_tqdm:
             mock_tqdm.return_value = MagicMock()

From 2ef3f452e05929331477db3a800713c78221b574 Mon Sep 17 00:00:00 2001
From: Alessandro Cere <13648166+acere@users.noreply.github.com>
Date: Thu, 16 Jul 2026 21:27:55 +0800
Subject: [PATCH 6/6] ci: run full test matrix only on PRs, single version on
 push

- Full Python version matrix (from classifiers) on pull_request
- Lint + Python 3.12 tests only on push to main
- Skip CI entirely for docs-only changes (docs/**, *.md)
---
 .github/workflows/ci.yml | 64 +++++++++++++++++++++++++++++-----------
 1 file changed, 46 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9ec0375..13f96bf 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -3,15 +3,43 @@ name: CI
 on:
   push:
     branches: [main]
+    paths-ignore:
+      - "docs/**"
+      - "*.md"
   pull_request:
     branches: [main]
+    paths-ignore:
+      - "docs/**"
+      - "*.md"
 
 permissions:
   contents: read
 
 jobs:
-  # Derive the Python version matrix from pyproject.toml classifiers
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+
+      - uses: astral-sh/setup-uv@v8.3.2
+        with:
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Ruff check
+        run: uv run ruff check llmeter tests
+
+      - name: Ruff format check
+        run: uv run ruff format --check llmeter tests
+
+      - name: Bandit security scan
+        run: uv run bandit -c pyproject.toml -r llmeter
+
+  # Full matrix on PRs, single version on push to main
   resolve-python-versions:
+    if: github.event_name == 'pull_request'
     runs-on: ubuntu-latest
     outputs:
       versions: ${{ steps.parse.outputs.versions }}
@@ -40,40 +68,40 @@ jobs:
           echo "versions=$versions" >> "$GITHUB_OUTPUT"
           echo "Resolved Python versions: $versions"
 
-  lint:
+  test:
     runs-on: ubuntu-latest
+    needs: resolve-python-versions
+    if: github.event_name == 'pull_request'
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ${{ fromJson(needs.resolve-python-versions.outputs.versions) }}
     steps:
       - uses: actions/checkout@v7
 
       - uses: astral-sh/setup-uv@v8.3.2
         with:
-          python-version: "3.12"
+          python-version: ${{ matrix.python-version }}
 
       - name: Install dependencies
         run: uv sync
 
-      - name: Ruff check
-        run: uv run ruff check llmeter tests
-
-      - name: Ruff format check
-        run: uv run ruff format --check llmeter tests
-
-      - name: Bandit security scan
-        run: uv run bandit -c pyproject.toml -r llmeter
+      - name: Run tests
+        run: uv run pytest --cov=llmeter -m "not integ"
+        env:
+          AWS_DEFAULT_REGION: us-east-1
+          AWS_ACCESS_KEY_ID: testing
+          AWS_SECRET_ACCESS_KEY: testing
 
-  test:
-    needs: resolve-python-versions
+  test-push:
+    if: github.event_name == 'push'
     runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: ${{ fromJson(needs.resolve-python-versions.outputs.versions) }}
     steps:
       - uses: actions/checkout@v7
 
       - uses: astral-sh/setup-uv@v8.3.2
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: "3.12"
 
       - name: Install dependencies
         run: uv sync