diff --git a/docs/reference/callbacks/cost/serde.md b/docs/reference/callbacks/cost/serde.md deleted file mode 100644 index 5ea2d0f..0000000 --- a/docs/reference/callbacks/cost/serde.md +++ /dev/null @@ -1 +0,0 @@ -::: llmeter.callbacks.cost.serde diff --git a/docs/reference/json_utils.md b/docs/reference/json_utils.md deleted file mode 100644 index 4fd15e6..0000000 --- a/docs/reference/json_utils.md +++ /dev/null @@ -1 +0,0 @@ -::: llmeter.json_utils diff --git a/docs/reference/serialization.md b/docs/reference/serialization.md new file mode 100644 index 0000000..c14cb1e --- /dev/null +++ b/docs/reference/serialization.md @@ -0,0 +1,4 @@ +::: llmeter.serialization + options: + filters: + - ".*" # Allow private methods for `Serializable._{get|set}_llmeter_state` diff --git a/docs/user_guide/callbacks/index.md b/docs/user_guide/callbacks/index.md index c1af74c..46e7e65 100644 --- a/docs/user_guide/callbacks/index.md +++ b/docs/user_guide/callbacks/index.md @@ -36,3 +36,16 @@ results = await runner.run( ``` Each callback will be processed in the **same order** as you provide them to the runner. This is important to remember and configure properly, if you're stacking multiple callbacks that access the same data (for example - transforming an invocation response *then* logging/exporting it somewhere, both using `after_invoke`). + +### Storing extra data on responses and results + +If your callback needs to attach **extra fields** on your `InvocationResponse`s (for example, a computed cost or a custom label), store them in the response's `annotations` dictionary: + +```python +async def after_invoke(self, response): + response.annotations["my_metric"] = compute_something(response) +``` + +Any extra fields added directly to the `InvocationResponse` object itself will *not* be preserved when the responses are saved to file. For legacy compatability, any unrecognized fields found in older responses.jsonl files are *currently* collected back to `annotations`, but this behaviour may be dropped in future. + +Likewise for extra **run-level** data, `after_run` can add numeric statistics to the `Result` via [`_update_contributed_stats`](../../reference/results/#llmeter.results.Result), which are persisted alongside the run's other stats. (This is how the built-in [cost modelling callback](cost.md) records run-level cost metrics.) diff --git a/docs/user_guide/serialization.md b/docs/user_guide/serialization.md new file mode 100644 index 0000000..a29557b --- /dev/null +++ b/docs/user_guide/serialization.md @@ -0,0 +1,199 @@ +# Save (and Reload) Results + +!!! warning "Important security warning" + Do not use LLMeter file load functions on data from untrusted sources! For more details on why, see the following section. + +LLMeter can save your test configurations and results to files - whether locally or on the Cloud - and load past runs back into Python for analysis later. For example: + +```python +from llmeter.results import Result +from llmeter.runner import Runner + +# Provide s3:// URI or local path: +base_output_path = f"s3://doc-example-bucket/llmeter/outputs/{endpoint.model_id}" + +runner = Runner( + endpoint, + # Configure the output path when creating a Runner... + output_path=base_output_path, +) +results = await runner.run( + payload=sample_payload, + n_requests=3, + clients=5, + # ...*Optionally* specify a specific path for an individual run: + # (Otherwise a subfolder will be created by run name, automatically) + output_path=f"{base_output_path}/my-cool-run" +) + +# At some later date, load your result back: +results = Result.load(f"{base_output_path}/my-cool-run") +``` + +!!! note "A note on performance" + While it's *possible* for LLMeter to write run outputs directly to Cloud object stores like + [Amazon S3](https://aws.amazon.com/s3/), remember it might reduce the maximum throughput you + can drive in *high-volume* tests, since it will consume network bandwidth. + +For the most part, generated files are JSON-based (or [JSON Lines](https://jsonlines.org/), for the +individual responses). However, LLMeter also handles some more complex data types including: + +- Binary data in request/response payloads (such as images) +- [Callbacks](./callbacks) configured on the test Run, including LLMeter built-ins as well as *your custom* callback classes + +This core de/serialization functionality is implemented in the [`llmeter.serialization`](../reference/serialization) module. + + +## How complex data types are represented and loaded + +Objects that are not natively JSON serializable (but support LLMeter's serialization protocol) are saved to JSON-based formats something like the below: + +```json +{ + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", + "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"} +} +``` + +- `datetime` objects are stored as ISO-8601 format strings in UTC timezone, like `2024-01-01T00:00:00Z`. +- `bytes` objects are serialized to base64 strings in a special `{"__llmeter_bytes__": ""}` wrapper. +- Objects implementing LLMeter's [`Serializable`](../reference/serialization.md#llmeter.serialization.Serializable) interface (including for example Endpoints and Callbacks) are represented as dicts with the `__llmeter_{class/state}__` properties as shown above + +When **loading** these objects back from file, LLMeter will try to import and instantiate **any** class path saved in the `__llmeter_class__` field. + +This is the same trust model as Python's native [pickle](https://docs.python.org/3/library/pickle.html) library: It is possible to construct malicious data which will run arbitrary code during loading, so **only load data that you trust**. + +## Under the hood: Serialization API components + +In many cases you'll be working with high-level classes like `Runner` and `Result` that already provide convenience methods to save to and load from file. However, building custom LLMeter extensions may require understanding how our [`llmeter.serialization`](../reference/serialization) components fit together: + +| Symbol | Purpose | +|--------|---------| +| [`Serializable`](../reference/serialization.md#llmeter.serialization.Serializable) | Mixin enabling your class to be serialized and loaded by LLMeter | +| [`dump_object`](../reference/serialization.md#llmeter.serialization.dump_object) / [`load_object`](../reference/serialization.md#llmeter.serialization.load_object) | Full round-trip persistence via a type-tagged envelope | +| [`json_default`](../reference/serialization.md#llmeter.serialization.json_default) | `json.dumps` fallback for bytes, datetime, PathLike | +| [`bytes_decoder`](../reference/serialization.md#llmeter.serialization.bytes_decoder) | `json.loads` object hook to restore `__llmeter_bytes__` markers | +| [`datetime_to_str`](../reference/serialization.md#llmeter.serialization.datetime_to_str) / [`str_to_datetime`](../reference/serialization.md#llmeter.serialization.str_to_datetime) | UTC ISO-8601 with `Z` suffix, both directions | + +```python +from llmeter.serialization import ( + dump_object, load_object, json_default, bytes_decoder, + datetime_to_str, str_to_datetime, +) +``` + +### Make custom classes serializable by LLMeter with the `Serializable` mixin + +The `Serializable` mixin provides two default methods: + +- `_get_llmeter_state`: Construct a JSON-ready dictionary of the state your class needs to be re-initialized + - The default implementation inspects the arguments of your `__init__` constructor and attempts to fetch those from the current object's fields, or with an `_` underscore prefix if the raw parameter name isn't present. +- `_set_llmeter_state`: Initialise an instance of your class, based on a state dictionary + - The default implementation calls your `__init__` with the arguments stored in the dictionary. + +Nested `Serializable` objects are recursively handled by default: `_get_llmeter_state` wraps them via `dump_object`, and `_set_llmeter_state` restores them via `load_object`. + +You'd only need to **override** these implementations if, for example: + +- An `__init__` parameter is consumed without being stored, or +- Reconstruction needs special logic beyond `__init__(**state)`, or +- You want to exclude large transient data from persistence + + +### Save and load your own objects + +Once your class inherits `Serializable`, it gets `save_to_file()` and `load_from_file()` for free - no extra code required: + +```python +from llmeter.callbacks.base import Callback +from llmeter.callbacks.mlflow import MlflowCallback + +cb = MlflowCallback(step=5, nested=True) +cb.save_to_file("/tmp/callback.json") + +# load_from_file is *polymorphic*: it reads the __llmeter_class__ recorded in the +# file and rebuilds the correct subclass, so you can call it on the base class. +restored = Callback.load_from_file("/tmp/callback.json") # -> MlflowCallback(step=5, nested=True) +``` + +Under the hood, those methods use two functions you can also call directly if you're managing the JSON yourself: + +- `dump_object(obj)` builds the type-tagged envelope (`{"__llmeter_class__": ..., "__llmeter_state__": ...}`). It reads the state from `obj._get_llmeter_state()`, or - for a plain `@dataclass` that *doesn't* inherit `Serializable` - from its fields. +- `load_object(data)` imports the class named in `__llmeter_class__`, creates a bare instance (via `__new__`, bypassing `__init__`), then repopulates it through `_set_llmeter_state()` - which by default re-runs `__init__` with the saved state. + +```python +from llmeter.serialization import dump_object, load_object + +data = dump_object(my_object) # -> plain dict +my_object = load_object(data) # -> reconstructed instance +``` + +Note the envelope produced by `dump_object` may still contain non-JSON values (like `bytes` or `datetime`) nested inside its state, so pair it with `json_default` when you actually write it out - see below. + +### Reading and writing the JSON yourself + +When you call `json.dump`/`json.dumps` directly on LLMeter data, pass `json_default` so the extra types are handled: + +```python +import json +from llmeter.serialization import json_default, bytes_decoder + +with open("my-file.json", "w") as f: + json.dump(my_data, f, default=json_default, indent=4) +``` + +`json_default` converts, in order: + +- **bytes** — wrapped in a `{"__llmeter_bytes__": ""}` marker +- **datetime** — UTC ISO-8601 string with a `Z` suffix +- **date / time** — `.isoformat()` +- **os.PathLike** — POSIX path string +- **anything else** — `str()` fallback + +To turn the `bytes` markers back into `bytes` on the way in, pass `bytes_decoder` as the `object_hook`: + +```python +with open("my-file.json") as f: + data = json.load(f, object_hook=bytes_decoder) +``` + +## How LLMeter's built-ins use this + +For everyday use you rarely touch the functions above directly, because the high-level classes wrap them for you. + +### Endpoints: config, not connections + +Endpoints are saved as *configuration only* - runtime state like a boto3 client is never written to file. When you load an endpoint back, `_set_llmeter_state` re-runs the constructor, which recreates that client for you: + +```python +from llmeter.serialization import dump_object, load_object + +data = dump_object(endpoint) +# → {"__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverse", +# "__llmeter_state__": {"model_id": "claude-3", "region": "us-west-2"}} + +restored = load_object(data) # boto3 client rebuilt via __init__ +``` + +### Runners: the whole run configuration + +When you give a `Runner` an `output_path`, it saves its full configuration - endpoint, tokenizer and callbacks included - as `run_config.json` at the start of each run. You can also trigger this yourself: + +```python +runner = Runner(endpoint=BedrockConverse(...), callbacks=[MlflowCallback(step=1)]) +runner.save(output_path="/tmp/run") # writes /tmp/run/run_config.json +``` + +If a callback (or any other configured object) can't be serialized because its class doesn't inherit `Serializable`, saving raises a `TypeError`. This is why custom callbacks and cost dimensions should subclass the relevant LLMeter base class - see [`Callback`](../reference/callbacks/base.md#llmeter.callbacks.base.Callback) and the [cost dimension base classes](../reference/callbacks/cost/dimensions.md). + +### Dataclasses work automatically + +Because the `Serializable` mixin introspects the constructor, any `@dataclass` that inherits it is serializable with no extra code - the generated `__init__` supplies the parameter names the mixin looks for. LLMeter's built-in cost dimensions are a good example: + +```python +from llmeter.callbacks.cost.dimensions import InputTokens +from llmeter.serialization import dump_object, load_object + +data = dump_object(InputTokens(price_per_million=3.0)) +restored = load_object(data) # -> InputTokens(price_per_million=3.0, granularity=1) +``` diff --git a/llmeter/callbacks/base.py b/llmeter/callbacks/base.py index f4c4dba..6ab8463 100644 --- a/llmeter/callbacks/base.py +++ b/llmeter/callbacks/base.py @@ -5,16 +5,14 @@ from __future__ import annotations from abc import ABC -from typing import final - -from upath.types import ReadablePathLike, WritablePathLike from ..endpoints.base import InvocationResponse from ..results import Result from ..runner import _RunConfig +from ..serialization import Serializable -class Callback(ABC): +class Callback(Serializable, ABC): """Base class for a callback in LLMeter Callbacks support extending LLMeter functionality by running additional code at defined points @@ -22,8 +20,11 @@ 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) - which have no-op implementations by default. Serialization + to/from file is inherited from `llmeter.serialization.Serializable`, and is necessary so your + callback(s) can be saved to file (and restored) as part of a Run configuration. Any custom + callback class that is *not* LLMeter-serializable will raise a `TypeError` when the run config + it belongs to is saved. """ async def before_invoke(self, payload: dict) -> None: @@ -44,7 +45,8 @@ async def after_invoke(self, response: InvocationResponse) -> None: timing and token counts) Returns: None: If you'd like to add information to the `response` logged in the Run, modify it - in-place. + in-place. To attach **extra custom fields** that you want preserved when responses + are saved to file, store them in the `response.annotations` map. """ pass @@ -70,47 +72,3 @@ async def after_run(self, result: Result) -> None: None: If you'd like to modify the run `result`, edit the argument in-place. """ pass - - def save_to_file(self, path: WritablePathLike) -> None: - """Save this Callback to file - - Individual Callbacks implement this method to save their configuration to a file that will - be re-loadable with the equivalent `_load_from_file()` method. - - Args: - path: (Local or Cloud) path where the callback is saved - """ - raise NotImplementedError("TODO: Callback.save_to_file is not yet implemented!") - - @staticmethod - @final - def load_from_file(path: ReadablePathLike) -> Callback: - """Load (any type of) Callback from file - - `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. - - Args: - path: (Local or Cloud) path where the callback is saved - Returns: - callback: A loaded Callback - for example an `MlflowCallback`. - """ - raise NotImplementedError( - "TODO: Callback.load_from_file is not yet implemented!" - ) - - @classmethod - def _load_from_file(cls, path: ReadablePathLike) -> Callback: - """Load this Callback from file - - Individual Callbacks implement this method to define how they can be loaded from files - created by the equivalent `save_to_file()` method. - - Args: - path: (Local or Cloud) path where the callback is saved - Returns: - callback: The loaded Callback object - """ - raise NotImplementedError( - "TODO: Callback._load_from_file is not yet implemented!" - ) diff --git a/llmeter/callbacks/cost/dimensions.py b/llmeter/callbacks/cost/dimensions.py index 60498be..466d7a3 100644 --- a/llmeter/callbacks/cost/dimensions.py +++ b/llmeter/callbacks/cost/dimensions.py @@ -15,57 +15,58 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from math import ceil -from typing import Optional +from typing import Protocol # Local Dependencies: from ...endpoints.base import InvocationResponse from ...results import Result from ...runner import _RunConfig -from .serde import ISerializable, JSONableBase +from ...serialization import Serializable -class IRequestCostDimension(ISerializable): +class IRequestCostDimension(Protocol): """Interface for one dimension of a per-request cost model Per-request cost components are calculated independently for each invocation in a test run, and can be used to model factors like cost-per-request, cost-per-input-tokens, cost-per-request-duration, etc. They're typically most relevant for serverless deployments like Amazon Bedrock, or estimating duration-based execution costs for AWS Lambda functions. - """ - - async def calculate(self, response: InvocationResponse) -> Optional[float]: - """Calculate (this component of) the cost for an individual request/response""" - ... - - -class RequestCostDimensionBase(ABC, JSONableBase): - """Base class for implementing per-request cost model dimensions - This class provides a default implementation of `ISerializable` and sets up an abstract method - for `calculate()`. It's fine if you don't want to derive from it directly - just be sure to - fully implement `IRequestCostDimension`! + !!! warning + This interface describes only the *calculation* behavior. It intentionally says nothing + about serialization: a dimension used purely in-memory need not be serializable. However, + to be usable inside a [`CostModel`][llmeter.callbacks.cost.model.CostModel] that gets + **saved** (e.g. as a run callback, or via ``save_to_file``), a dimension must + **additionally** be LLMeter-serializable. The supported way to get that is to subclass + [`RequestCostDimensionBase`][llmeter.callbacks.cost.dimensions.RequestCostDimensionBase] + (which mixes in [`Serializable`][llmeter.serialization.Serializable]; implementing this + Protocol alone is not sufficient for persistence. """ - @abstractmethod - async def calculate(self, response: InvocationResponse) -> Optional[float]: + async def calculate(self, response: InvocationResponse) -> float | None: """Calculate (this component of) the cost for an individual request/response""" - raise NotImplementedError( - "Children of RequestCostDimensionBase must implement `calculate()`! At: %s" - % (self.__class__,) - ) + ... -class IRunCostDimension(ISerializable): +class IRunCostDimension(Protocol): """Interface for one dimension of a per-Run cost model Per-run cost components are notified before the start of a test run via `before_run_start()`, and then requested to `calculate()` at the end of the run. They're most relevant for provisioned-infrastructure based deployments like Amazon SageMaker, where factors like a (request-independent) cost-per-hour are important. + + !!! warning + As with `IRequestCostDimension`, this interface covers only calculation behavior and not + serialization. To be usable inside a + [`CostModel`][llmeter.callbacks.cost.model.CostModel] that gets saved, subclass + [`RunCostDimensionBase`][llmeter.callbacks.cost.dimensions.RunCostDimensionBase] (which + mixes in [`Serializable`][llmeter.serialization.Serializable]) rather than implementing + this Protocol alone. """ async def before_run_start(self, run_config: _RunConfig) -> None: - """Function called to notify the cost component that a test run is about to start + """Notify the cost component that a test run is about to start This method is called before the test run starts, and can be used to perform any initialization or setup required for the cost component. In general, we assume a dimension @@ -75,7 +76,7 @@ async def before_run_start(self, run_config: _RunConfig) -> None: """ ... - async def calculate(self, result: Result) -> Optional[float]: + async def calculate(self, result: Result) -> float | None: """Calculate (this component of) the cost for a completed test run Dimensions that depend on `before_run_start` being called to return an accurate result @@ -85,20 +86,34 @@ async def calculate(self, result: Result) -> Optional[float]: ... -class RunCostDimensionBase(ABC, JSONableBase): +class RequestCostDimensionBase(Serializable, ABC): + """Base class for implementing per-request cost model dimensions + + This class provides a default implementation of serialization (via + [`Serializable`][llmeter.serialization.Serializable]) and sets up an abstract method for + `calculate()`. + """ + + @abstractmethod + async def calculate(self, response: InvocationResponse) -> float | None: + """Calculate (this component of) the cost for an individual request/response""" + raise NotImplementedError + + +class RunCostDimensionBase(Serializable, ABC): """Base class for implementing per-run cost model dimensions - This class provides a default implementation of `ISerializable`, a default empty - `before_run_start` implementation, and abstract methods for the other requirements of the - `IRunCostDimension` protocol. It's fine if you don't want to derive from it directly - just - make sure you fully implement `IRunCostDimension`! + This class provides a default implementation of serialization (via + [`Serializable`][llmeter.serialization.Serializable]), a default empty `before_run_start` + implementation, and abstract methods for the other requirements of the `IRunCostDimension` + protocol. """ async def before_run_start(self, run_config: _RunConfig) -> None: - """Function called to notify the cost component that a test run is about to start + """Notify the cost component that a test run is about to start This method is called before the test run starts, and can be used to perform any - initialization or setup required for the cost component. In general, we assume a dimension + initialization or setup required for the cost component. In general, we assume a dimension instance may be re-used for multiple test runs, but only one run at a time: Meaning `before_run_start()` should not be called again before `calculate()` is called for the previous run. @@ -108,17 +123,19 @@ async def before_run_start(self, run_config: _RunConfig) -> None: pass @abstractmethod - async def calculate(self, result: Result) -> Optional[float]: + async def calculate(self, result: Result) -> float | None: """Calculate (this component of) the cost for a completed test run Dimensions that depend on `before_run_start` being called to return an accurate result should throw an error if this was not done. Dimensions that only need `calculate()` should silently ignore if `before_run_start` was not called. """ - raise NotImplementedError( - "Children of RunCostDimensionBase must implement `calculate()`! At: %s" - % (self.__class__,) - ) + raise NotImplementedError + + +# --------------------------------------------------------------------------- +# Concrete dimension implementations +# --------------------------------------------------------------------------- @dataclass @@ -133,13 +150,11 @@ class InputTokens(RequestCostDimensionBase): price_per_million: float granularity: int = 1 - async def calculate(self, req: InvocationResponse) -> Optional[float]: + async def calculate(self, req: InvocationResponse) -> float | None: if req.num_tokens_input is None: return None - billable_tokens = ( - ceil(req.num_tokens_input / self.granularity) * self.granularity - ) - return billable_tokens * self.price_per_million / 1000000 + billable = ceil(req.num_tokens_input / self.granularity) * self.granularity + return billable * self.price_per_million / 1_000_000 @dataclass @@ -154,13 +169,11 @@ class OutputTokens(RequestCostDimensionBase): price_per_million: float granularity: int = 1 - async def calculate(self, req: InvocationResponse) -> Optional[float]: + async def calculate(self, req: InvocationResponse) -> float | None: if req.num_tokens_output is None: return None - billable_tokens = ( - ceil(req.num_tokens_output / self.granularity) * self.granularity - ) - return billable_tokens * self.price_per_million / 1000000 + billable = ceil(req.num_tokens_output / self.granularity) * self.granularity + return billable * self.price_per_million / 1_000_000 @dataclass @@ -175,10 +188,10 @@ class EndpointTime(RunCostDimensionBase): price_per_hour: float granularity_secs: float = 1 - async def calculate(self, result: Result) -> Optional[float]: + async def calculate(self, result: Result) -> float | None: if result.total_test_time is None: return None - billable_secs = ( + billable = ( ceil(result.total_test_time / self.granularity_secs) * self.granularity_secs ) - return billable_secs * self.price_per_hour / 3600 + return billable * self.price_per_hour / 3600 diff --git a/llmeter/callbacks/cost/model.py b/llmeter/callbacks/cost/model.py index dd867f4..a2a40f8 100644 --- a/llmeter/callbacks/cost/model.py +++ b/llmeter/callbacks/cost/model.py @@ -1,25 +1,16 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -# Python Built-Ins: -import importlib -from dataclasses import dataclass, field +"""Cost modelling callback for LLMeter test runs.""" -# Extrnal Dependencies: -from upath.types import ReadablePathLike, WritablePathLike - -# Local Dependencies: from ...endpoints.base import InvocationResponse from ...results import Result from ...runner import _RunConfig -from ...utils import ensure_path from ..base import Callback from .dimensions import IRequestCostDimension, IRunCostDimension from .results import CalculatedCostWithDimensions -from .serde import JSONableBase, from_dict_with_class_map -@dataclass -class CostModel(JSONableBase, Callback): +class CostModel(Callback): """Model costs of (test runs of) Foundation Models A Cost Model is composed of whatever pricing *"dimensions"* are relevant for your FM deployment @@ -33,15 +24,14 @@ class CostModel(JSONableBase, Callback): a Callback when running an LLMeter Run or Experiment, to annotate the results automatically. """ - request_dims: dict[str, IRequestCostDimension] = field(default_factory=dict) - run_dims: dict[str, IRunCostDimension] = field(default_factory=dict) - def __init__( self, - request_dims: dict[str, IRequestCostDimension] - | list[IRequestCostDimension] - | None = None, - run_dims: dict[str, IRunCostDimension] | list[IRunCostDimension] | None = None, + request_dims: ( + dict[str, IRequestCostDimension] | list[IRequestCostDimension] | None + ) = None, + run_dims: ( + dict[str, IRunCostDimension] | list[IRunCostDimension] | None + ) = None, ): """Create a CostModel @@ -58,21 +48,19 @@ def __init__( (including request_dims) have the same name. """ all_dims: dict[str, IRequestCostDimension | IRunCostDimension] = {} - self.request_dims = {} - self.run_dims = {} + self.request_dims: dict[str, IRequestCostDimension] = {} + self.run_dims: dict[str, IRunCostDimension] = {} if request_dims is not None: for name, dim in ( request_dims.items() if isinstance(request_dims, dict) - else zip( - (d.__class__.__name__ for d in request_dims), - request_dims, - ) + else ((d.__class__.__name__, d) for d in request_dims) ): if name in all_dims: raise ValueError( - f"Duplicate cost dimension name '{name}': Got both {dim} and {all_dims[name]}" + f"Duplicate cost dimension name '{name}': " + f"Got both {dim} and {all_dims[name]}" ) all_dims[name] = dim self.request_dims[name] = dim @@ -81,29 +69,30 @@ def __init__( for name, dim in ( run_dims.items() if isinstance(run_dims, dict) - else zip( - (d.__class__.__name__ for d in run_dims), - run_dims, - ) + else ((d.__class__.__name__, d) for d in run_dims) ): if name in all_dims: raise ValueError( - f"Duplicate cost dimension name '{name}': Got both {dim} and {all_dims[name]}" + f"Duplicate cost dimension name '{name}': " + f"Got both {dim} and {all_dims[name]}" ) all_dims[name] = dim self.run_dims[name] = dim + # ------------------------------------------------------------------ + # Callback lifecycle hooks + # ------------------------------------------------------------------ + async def calculate_request_cost( - self, - response: InvocationResponse, - save: bool = False, + self, response: InvocationResponse, save: bool = False ) -> CalculatedCostWithDimensions: """Calculate the costs of a single FM invocation (excluding any session-level costs) Args: response: The InvocationResponse to estimate costs for - save: Set `True` to also store the result in `response.cost`, in addition to returning - it. Defaults to `False` + save: Set `True` to also store the result in `response.annotations` (under `cost_` + prefixed keys), in addition to returning it. Because `annotations` is a persisted + field, saved costs survive `Result` save/load. Defaults to `False` """ dim_costs = CalculatedCostWithDimensions( **{ @@ -112,8 +101,7 @@ async def calculate_request_cost( } ) if save: - dim_costs.save_on_namespace(response, key_prefix="cost_") - + dim_costs.save_on_namespace(response.annotations, key_prefix="cost_") return dim_costs async def calculate_run_cost( @@ -138,7 +126,6 @@ async def calculate_run_cost( run_cost = CalculatedCostWithDimensions( **{name: await dim.calculate(result) for name, dim in self.run_dims.items()} ) - if recalculate_request_costs: resp_costs = [ await self.calculate_request_cost(r, save=save) @@ -150,17 +137,16 @@ async def calculate_run_cost( lambda c: c, # Skip responses where no cost data was found at all ( CalculatedCostWithDimensions.load_from_namespace( - r, key_prefix="cost_" + r.annotations, key_prefix="cost_" ) for r in result.responses ), - ), + ) ) - # Merge the total request-level costs into the run-level costs: - # (Unless requests is empty, because sum([]) = 0 and not a CalculatedCostWithDimensions) if len(resp_costs): + # Merge the total request-level costs into the run-level costs: + # (Unless requests is empty, because sum([])=0 and not a CalculatedCostWithDimensions) run_cost.merge(sum(resp_costs)) # type: ignore - if save: # Save the overall run cost and breakdown on the main result object: run_cost.save_on_namespace(result, key_prefix="cost_") @@ -175,17 +161,13 @@ async def calculate_run_cost( ) run_cost.save_on_namespace(stats, key_prefix="cost_") result._update_contributed_stats(stats) - return run_cost - async def before_invoke(self, payload: dict) -> None: - """This LLMeter Callback hook is a no-op for CostModel""" - pass - async def after_invoke(self, response: InvocationResponse) -> None: """LLMeter Callback.after_invoke hook - Calls calculate_request_cost() with `save=True` to save the cost on the InvocationResponse. + Calls calculate_request_cost() with `save=True` to save the per-request cost into + `response.annotations` (under `cost_` prefixed keys), so it persists with the response. """ await self.calculate_request_cost(response, save=True) @@ -202,32 +184,3 @@ async def after_run(self, result: Result) -> None: await self.calculate_run_cost( result, recalculate_request_costs=False, save=True ) - - def save_to_file(self, path: WritablePathLike) -> None: - """Save the cost model (including all dimensions) to a JSON file""" - path = ensure_path(path) - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w") as f: - f.write(self.to_json()) - - @classmethod - def from_dict(cls, raw: dict, alt_classes: dict = {}, **kwargs) -> "CostModel": - dim_classes = { - **importlib.import_module("llmeter.callbacks.cost.dimensions").__dict__, - **alt_classes, - } - raw_args = {**raw} - 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) - - @classmethod - def _load_from_file(cls, path: ReadablePathLike): - """Load the cost model (including all dimensions) from a JSON file""" - path = ensure_path(path) - with path.open("r") as f: - return cls.from_json(f.read()) diff --git a/llmeter/callbacks/cost/serde.py b/llmeter/callbacks/cost/serde.py deleted file mode 100644 index 72025c1..0000000 --- a/llmeter/callbacks/cost/serde.py +++ /dev/null @@ -1,217 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 -"""(De/re)serialization interfaces for saving Cost Model objects to file and loading them back""" - -# Python Built-Ins: -import json -import logging -from dataclasses import is_dataclass -from typing import Any, Protocol, TypeVar - -# External Dependencies: -from upath import UPath as Path -from upath.types import ReadablePathLike, WritablePathLike - -# Local Dependencies: -from ...json_utils import llmeter_default_serializer -from ...utils import ensure_path - -logger = logging.getLogger(__name__) - -TSerializable = TypeVar("TSerializable", bound="ISerializable") - - -class ISerializable(Protocol): - def to_dict(self) -> dict: - """Represent this object with a fully JSON-serializable dictionary""" - ... - - @classmethod - def from_dict(cls, **kwargs) -> TSerializable: - """Load an instance of this class from a JSON-able dictionary representation""" - ... - - -def is_dataclass_instance(obj): - """Check whether `obj` is an instance of any dataclass - - See: https://docs.python.org/3/library/dataclasses.html#dataclasses.is_dataclass - """ - return is_dataclass(obj) and not isinstance(obj, type) - - -def to_dict_recursive_generic(obj: object, **kwargs) -> dict: - """Convert a vaguely dataclass-like object (with maybe IJSONable fields) to a JSON-ready dict - - The output dict is augmented with `_type` storing the `__class__.__name__` of the provided - `obj`. - - Args: - obj: The object to convert - **kwargs: Optional extra parameters to insert in the output dictionary - """ - obj_classname = obj.__class__.__name__ - result: dict = {} if obj_classname == "dict" else {"_type": obj.__class__.__name__} - if hasattr(obj, "__dict__"): - # We *don't* use dataclass asdict() here because we want our custom behaviour instead of - # its recursion: - result.update(obj.__dict__) - elif ( - isinstance(obj, dict) - or hasattr(obj, "__keys__") - and hasattr(obj, "__getitem__") - ): - result.update(obj) - else: - result.update({k: getattr(obj, k) for k in dir(obj)}) - result.update(kwargs) - for k, v in result.items(): - if hasattr(v, "to_dict"): - result[k] = v.to_dict() - elif isinstance(v, dict): - result[k] = to_dict_recursive_generic(v) - elif isinstance(v, (list, tuple)): - result[k] = [to_dict_recursive_generic(item) for item in v] - return result - - -TFromDict = TypeVar("TFromDict") - - -def from_dict_with_class(raw: dict, cls: type[TFromDict], **kwargs) -> TFromDict: - """Initialize an instance of a class from a plain dict (with optional extra kwargs) - - If the input dictionary contains a `_type` key, and this doesn't match the provided - `cls.__name__`, a warning will be logged. - - Args: - raw: A plain Python dict, for example loaded from a JSON file - cls: The class to create an instance of - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - raw_args = {k: v for k, v in raw.items()} - raw_type = raw_args.pop("_type", None) - if raw_type is not None and raw_type != cls.__name__: - logger.warning( - "from_dict: _type '%s' doesn't match class '%s' being loaded. %s" - % (raw_type, cls.__name__, raw) - ) - return cls(**raw_args, **kwargs) - - -def from_dict_with_class_map( - raw: dict, class_map: dict[str, type[TFromDict]], **kwargs -) -> TFromDict: - """Initialize an instance of a class from a plain dict (with optional extra kwargs) - - Args: - raw: A plain Python dict which must contain a `_type` key - class_map: A mapping from `_type` string to class to create an instance of - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - if "_type" not in raw: - raise ValueError("from_dict_with_class_map: No _type in raw dict: %s" % raw) - if raw["_type"] not in class_map: - raise ValueError( - "Object _type '%s' not found in provided class_map %s" - % (raw["_type"], class_map) - ) - return from_dict_with_class(raw, class_map[raw["_type"]], **kwargs) - - -TJSONable = TypeVar("TJSONable", bound="JSONableBase") - - -class JSONableBase: - """A base class for speeding up implementation of JSON-serializable objects""" - - @classmethod - def from_dict( - cls: type[TJSONable], - raw: dict, - alt_classes: dict[str, TJSONable] = {}, - **kwargs: Any, - ) -> TJSONable: - """Initialize an instance of this class from a plain dict (with optional extra kwargs) - - Args: - raw: A plain Python dict, for example loaded from a JSON file - alt_classes: By default, this method will only use the class of the current object - (i.e. `cls`). If you want to support loading of subclasses, provide a mapping - from your raw dict's `_type` field to class, for example `{cls.__name__: cls}`. - **kwargs: Optional extra keyword arguments to pass to the constructor - """ - if alt_classes: - return from_dict_with_class_map( - raw=raw, - class_map={cls.__name__: cls, **alt_classes}, - **kwargs, - ) - else: - return from_dict_with_class(raw=raw, cls=cls, **kwargs) - - @classmethod - def from_file( - cls: type[TJSONable], input_path: ReadablePathLike, **kwargs - ) -> TJSONable: - """Initialize an instance of this class from a (local or Cloud) JSON file - - Args: - input_path: The path to the JSON data file. - **kwargs: Optional extra keyword arguments to pass to `from_dict()` - """ - input_path = ensure_path(input_path) - with input_path.open("r") as f: - return cls.from_json(f.read(), **kwargs) - - @classmethod - def from_json(cls: type[TJSONable], json_string: str, **kwargs: Any) -> TJSONable: - """Initialize an instance of this class from a JSON string (with optional extra kwargs) - - Args: - json_string: A string containing valid JSON data - **kwargs: Optional extra keyword arguments to pass to `from_dict()` - """ - return cls.from_dict(json.loads(json_string), **kwargs) - - def to_dict(self, **kwargs) -> dict: - """Save the state of the object to a JSON-dumpable dictionary (with optional extra kwargs) - - Implementers of this method should ensure that the returned dict is fully JSON-compatible: - Mapping any child fields from Python classes to dicts if necessary, avoiding any circular - references, etc. - """ - return to_dict_recursive_generic(self, **kwargs) - - def to_file( - self, - output_path: WritablePathLike, - indent: int | str | None = 4, - **kwargs: Any, - ) -> Path: - """Save the state of the object to a (local or Cloud) JSON file - - Args: - output_path: The path where the configuration file will be saved. - indent: Optional indentation passed through to `to_json()` and therefore - `json.dumps()` - **kwargs: Optional extra keyword arguments to pass to `to_json()` - - Returns: - output_path: Universal Path representation of the target file. - """ - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w") as f: - f.write(self.to_json(indent=indent, **kwargs)) - return output_path - - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: - """Serialize this object to JSON. - - Args: - default: Fallback serializer. Defaults to - :func:`~llmeter.json_utils.llmeter_default_serializer`. - **kwargs: Extra keyword arguments passed to :func:`json.dumps`. - """ - return json.dumps(self.to_dict(), default=default, **kwargs) diff --git a/llmeter/callbacks/mlflow.py b/llmeter/callbacks/mlflow.py index bac5ba1..c9fba98 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,17 +66,6 @@ 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__ - @classmethod - def _load_from_file(cls, path: ReadablePathLike): - raise NotImplementedError( - "TODO: MlflowCallback does not yet support loading from file" - ) - - def save_to_file(self, path: WritablePathLike) -> None: - raise NotImplementedError( - "TODO: MlflowCallback does not yet support saving to file" - ) - async def _log_llmeter_run(self, result: Result): """Log parameters and metrics from an LLMeter run to MLflow. diff --git a/llmeter/endpoints/anthropic_messages.py b/llmeter/endpoints/anthropic_messages.py index 4a4189d..f314dfb 100644 --- a/llmeter/endpoints/anthropic_messages.py +++ b/llmeter/endpoints/anthropic_messages.py @@ -66,17 +66,22 @@ after all thinking is complete. """ +# Python Built-Ins: +import inspect import logging import time from typing import Any, Generic, Iterable, TypeVar +# External Dependencies: import anthropic from anthropic.types import ( Message, MessageCreateParams, RawMessageStreamEvent, ) +import httpx # (Indirect dependency of anthropic) +# Local Dependencies: from .base import Endpoint, InvocationResponse logger = logging.getLogger(__name__) @@ -142,8 +147,64 @@ def __init__( kwargs["api_key"] = api_key if aws_region is not None: kwargs["aws_region"] = aws_region + if isinstance(kwargs.get("timeout"), dict): + kwargs["timeout"] = httpx.Timeout(**kwargs["timeout"]) self._client = client_cls(**kwargs) + def _get_llmeter_state(self) -> dict: + """Extract serializable state by introspecting the underlying client. + + Since some constructor params are set on the underlying Anthropic client but *not* + persisted directly on the Endpoint object, we need to extend the default serialization + behaviour to preserve these extra config fields so they survive a serialization round trip. + + Rather than hardcoding per-provider fields (since there are different ANTHROPIC_CLIENTS + classes), we inspect the actual client class's __init__ signature and capture any matching + instance attributes — automatically adapting to whichever provider is in use. Secrets + (params containing 'key', 'token', or 'credential') and non-serializable objects are + excluded. + """ + skip_params = frozenset( + { + "http_client", + "_strict_response_validation", + # default_headers/default_query properties return merged SDK headers including + # secrets; we read _custom_headers/_custom_query directly below instead. + "default_headers", + "default_query", + } + ) + skip_substrings = ("key", "token", "credential") + + state = super()._get_llmeter_state() + sig = inspect.signature(type(self._client).__init__) + for name, param in sig.parameters.items(): + if name == "self" or param.kind in ( + param.VAR_POSITIONAL, + param.VAR_KEYWORD, + ): + continue + if name in skip_params or any(s in name for s in skip_substrings): + continue + if not hasattr(self._client, name): + continue + val = getattr(self._client, name) + if isinstance(val, httpx.URL): + val = str(val) + elif isinstance(val, httpx.Timeout): + val = { + "connect": val.connect, + "read": val.read, + "write": val.write, + "pool": val.pool, + } + state[name] = val + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state + def _parse_payload(self, payload: MessageCreateParams | dict) -> str: """Extract user message text from an Anthropic Messages API payload. diff --git a/llmeter/endpoints/base.py b/llmeter/endpoints/base.py index 5b1dff1..1c4d80c 100644 --- a/llmeter/endpoints/base.py +++ b/llmeter/endpoints/base.py @@ -8,12 +8,14 @@ import copy import functools import importlib +import inspect import json import logging import time +import warnings from abc import ABC, abstractmethod from collections.abc import Callable -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field, fields from datetime import datetime, timezone from typing import Any, Generic, TypeVar from uuid import uuid4 @@ -21,7 +23,13 @@ from upath import UPath as Path from upath.types import ReadablePathLike, WritablePathLike -from ..json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from ..serialization import ( + Serializable, + bytes_decoder, + json_default, + load_object, + restore_dataclass_types, +) from ..utils import ensure_path logger = logging.getLogger(__name__) @@ -50,6 +58,12 @@ class InvocationResponse: time_per_output_token (float): The average time taken to generate each token in the response. error (str): Any error that occurred during invocation. request_time: The wall-clock time when the request was sent. + annotations (dict): Free-form extra data attached to this response, for example by + callbacks. This is the **preferred** place for a `Callback` to store additional + per-response fields (rather than setting arbitrary attributes on the response), because + `annotations` is a declared field and therefore round-trips through + `to_json`/`from_json` and disk persistence. On load, any *unrecognized* top-level keys + (e.g. from older files or other LLMeter versions) are collected into `annotations` too. """ response_text: str | None @@ -66,6 +80,7 @@ class InvocationResponse: error: str | None = None retries: int | None = None request_time: datetime | None = None + annotations: dict[str, Any] = field(default_factory=dict) @classmethod def from_json(cls, json_str: str) -> "InvocationResponse": @@ -75,13 +90,18 @@ def from_json(cls, json_str: str) -> "InvocationResponse": correctly restores types that the default JSON round-trip would leave as strings or marker objects: - * `request_time` is parsed from an ISO-8601 string back to a Python `datetime` - * `payload`s containing `bytes` (as `__llmeter_bytes__` markers) are correctly loaded back - as bytes. + * `datetime`-annotated fields are parsed from ISO-8601 strings back to Python `datetime` + * `bytes`-typed fields and `__llmeter_bytes__` markers in nested payloads are restored + + For legacy compatability, any top-level keys that are *not* recognized fields are currently + collected into [`annotations`][llmeter.endpoints.base.InvocationResponse] rather than being + dropped or raising an error. This supports older files where CostModel callbacks wrote + extra fields (such as `cost_*`) directly onto the response - but may be dropped in a future + version. Args: json_str: A JSON string representation of an InvocationResponse (produced by `to_json` - or similar). + or similar). Returns: InvocationResponse: The deserialized response. @@ -93,16 +113,28 @@ def from_json(cls, json_str: str) -> "InvocationResponse": restored = InvocationResponse.from_json(original.to_json()) ``` """ - data = json.loads(json_str, object_hook=llmeter_bytes_decoder) - rt = data.get("request_time") - if rt is not None and isinstance(rt, str): - data["request_time"] = datetime.fromisoformat(rt.replace("Z", "+00:00")) + data = json.loads(json_str, object_hook=bytes_decoder) + restore_dataclass_types(cls, data) + # Route any unrecognized top-level keys into `annotations` rather than failing. This keeps + # forward/backward compatibility: e.g. older files where callbacks wrote extra fields (like + # `cost_*`) directly onto the response, or fields written by a different LLMeter version. + known = {f.name for f in fields(cls)} + extras = {k: data.pop(k) for k in list(data) if k not in known} + if extras: + data["annotations"] = {**extras, **(data.get("annotations") or {})} + logger.debug( + "Loaded %d unrecognized InvocationResponse field(s) into `annotations`: %s", + len(extras), + ", ".join(sorted(extras)), + ) return cls(**data) - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: + def to_json( + self, default: Callable[[Any], Any] | None = json_default, **kwargs: Any + ) -> str: """Serialize this response to a JSON string. - Uses [`llmeter_default_serializer`][llmeter.json_utils.llmeter_default_serializer] by + Uses [`json_default`][llmeter.serialization.json_default] by default, which handles `bytes`, `datetime`, `PathLike`, and other common non-serializable types. @@ -151,7 +183,7 @@ def to_dict(self) -> dict: `datetime` comparisons and arithmetic. For JSON output, use [`to_json`][llmeter.endpoints.base.InvocationResponse.to_json], (which - delegates to [`llmeter_default_serializer`][llmeter.json_utils.llmeter_default_serializer] + delegates to [`json_default`][llmeter.serialization.json_default] by default, for non-JSON-serializable data types). Returns: @@ -163,7 +195,7 @@ def to_dict(self) -> dict: TRawResponse = TypeVar("TRawResponse", bound=Any) -class Endpoint(ABC, Generic[TRawResponse]): +class Endpoint(Serializable, ABC, Generic[TRawResponse]): """ An abstract base class for endpoint implementations. @@ -474,23 +506,26 @@ def __subclasshook__(cls, C: type) -> bool: return NotImplemented def save(self, output_path: WritablePathLike) -> Path: - """ - Save the endpoint configuration to a JSON file. + """Save the endpoint configuration to a JSON file. - This method serializes the endpoint's configuration (excluding private attributes) - to a JSON file at the specified path. + .. deprecated:: + Use :meth:`~llmeter.serialization.Serializable.save_to_file` instead, which + provides the same behavior with a consistent name across all serializable + LLMeter objects. This alias will be removed in a future major version. Args: output_path (str | UPath): The path where the configuration file will be saved. Returns: - None + Path: The path the file was written to. """ - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w") as f: - json.dump(self, f, indent=4, default=llmeter_default_serializer) - return output_path + warnings.warn( + "Endpoint.save() is deprecated and will be removed in a future version; " + "use Endpoint.save_to_file() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.save_to_file(output_path) def to_dict(self) -> dict: """ @@ -504,41 +539,46 @@ def to_dict(self) -> dict: return endpoint_conf @classmethod - def load_from_file(cls, input_path: ReadablePathLike) -> "Endpoint": - """ - Load an endpoint configuration from a JSON file. + def load_from_file(cls, path: ReadablePathLike) -> "Endpoint": + """Load an endpoint configuration from a JSON file. This class method reads a JSON file containing an endpoint configuration, determines the appropriate endpoint class, and instantiates it with the loaded configuration. Args: - input_path (str|UPath): The path to the JSON configuration file. + path (str | UPath): The path to the JSON configuration file. Returns: Endpoint: An instance of the appropriate endpoint class, initialized with the configuration from the file. """ - - input_path = ensure_path(input_path) - with input_path.open("r") as f: + path = ensure_path(path) + with path.open("r") as f: data = json.load(f) + if "__llmeter_class__" in data: + return load_object(data) endpoint_type = data.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") endpoint_class = getattr(endpoint_module, endpoint_type) - return endpoint_class(**data) + return endpoint_class(**_filter_legacy_ctor_kwargs(endpoint_class, data)) @classmethod def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore - """ - Load an endpoint configuration from a dictionary. + """Load an endpoint configuration from a dictionary. This class method reads a dictionary containing an endpoint configuration, determines the appropriate endpoint class, and instantiates it with the loaded configuration. + !!! warning "Deprecated" + This supports the legacy `{"endpoint_type": ...}` format. New code should + use [`load_object`][llmeter.serialization.load_object] with dicts produced by + [`dump_object`][llmeter.serialization.dump_object]. + Args: - endpoint_config (Dict): A dictionary containing the endpoint configuration. + endpoint_config (dict): A dictionary containing the endpoint configuration. + Must include at minimum an `endpoint_type` key. Returns: Endpoint: An instance of the appropriate endpoint class, initialized @@ -547,4 +587,37 @@ def load(cls, endpoint_config: dict) -> "Endpoint": # type: ignore endpoint_type = endpoint_config.pop("endpoint_type") endpoint_module = importlib.import_module("llmeter.endpoints") endpoint_class = getattr(endpoint_module, endpoint_type) - return endpoint_class(**endpoint_config) + return endpoint_class( + **_filter_legacy_ctor_kwargs(endpoint_class, endpoint_config) + ) + + +def _filter_legacy_ctor_kwargs(endpoint_class: type, config: dict) -> dict: + """Drop legacy config keys that the target endpoint constructor won't accept. + + Older LLMeter configs persisted derived/read-only attributes (notably `provider`, which + endpoints now set internally) alongside the real constructor arguments. Passing those through + to a modern `__init__` raises `TypeError`, so we filter the dict down to the parameters the + constructor actually declares. + + If the constructor accepts `**kwargs` the dict is passed through unchanged. + + Args: + endpoint_class: The endpoint class about to be instantiated. + config: The legacy configuration dict (already stripped of ``endpoint_type``). + + Returns: + A copy of `config` containing only keys the constructor accepts. + """ + params = inspect.signature(endpoint_class.__init__).parameters + if any(p.kind == p.VAR_KEYWORD for p in params.values()): + return config + accepted = {name for name in params if name != "self"} + dropped = set(config) - accepted + if dropped: + logger.debug( + "Ignoring legacy config field(s) not accepted by %s.__init__: %s", + endpoint_class.__name__, + ", ".join(sorted(dropped)), + ) + return {k: v for k, v in config.items() if k in accepted} diff --git a/llmeter/endpoints/bedrock.py b/llmeter/endpoints/bedrock.py index 9253c6f..52cf0fc 100644 --- a/llmeter/endpoints/bedrock.py +++ b/llmeter/endpoints/bedrock.py @@ -161,6 +161,8 @@ def __init__( ) self.region = region or boto3.session.Session().region_name + # Persist extra config on the class just so de/serialization catches it: + self._max_attempts = max_attempts logger.info(f"Using AWS region: {self.region}") self._bedrock_client = bedrock_boto3_client diff --git a/llmeter/endpoints/bedrock_invoke.py b/llmeter/endpoints/bedrock_invoke.py index a98312f..af3ea86 100644 --- a/llmeter/endpoints/bedrock_invoke.py +++ b/llmeter/endpoints/bedrock_invoke.py @@ -58,6 +58,8 @@ def __init__( self.generated_token_count_jmespath = generated_token_count_jmespath self.input_text_jmespath = input_text_jmespath self.input_token_count_jmespath = input_token_count_jmespath + # Persist extra config on the class just so de/serialization catches it: + self._max_attempts = max_attempts self.region = ( region @@ -348,7 +350,7 @@ def process_raw_response( """Parse the streaming response from Bedrock InvokeModelWithResponseStream API. Args: - client_response: The raw response from the Bedrock API. + raw_response: The raw response from the Bedrock API. start_t: The timestamp when the request was initiated. Returns: diff --git a/llmeter/endpoints/openai.py b/llmeter/endpoints/openai.py index 79bcf72..13fd05d 100644 --- a/llmeter/endpoints/openai.py +++ b/llmeter/endpoints/openai.py @@ -2,13 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 """LLMeter targets for testing OpenAI ChatCompletions-compatible endpoints (wherever they're hosted)""" +# Python Built-Ins: import base64 import logging import os import time +from collections.abc import Mapping from typing import Any, Literal, Generic, Iterable, TypeVar, cast +# External Dependencies: +import httpx # (Indirect dependency of OpenAI) from openai import OpenAI +from openai._constants import DEFAULT_MAX_RETRIES from openai.types.chat import ( ChatCompletion, ChatCompletionChunk, @@ -23,6 +28,7 @@ File as ChatCompletionFile, ) +# Local Dependencies: from ..prompt_utils import ( ContentItem, MediaContent, @@ -123,6 +129,14 @@ def __init__( endpoint_name: str = "openai", api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize OpenAI endpoint. @@ -130,12 +144,77 @@ def __init__( Args: model_id: ID of the OpenAI model to use endpoint_name: Name of the endpoint. Defaults to "openai". - api_key: OpenAI API key. Defaults to None. + api_key: OpenAI API key. Defaults to None (reads from OPENAI_API_KEY env var). provider: Provider name. Defaults to "openai". + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. If a + dict is passed (as is the case after serialization), it'll be interpreted as a set + of arguments to create an httpx.Timeout + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__(endpoint_name, model_id, provider=provider) - self._client = OpenAI(api_key=api_key, **kwargs) + client_kwargs: dict[str, Any] = { + "api_key": api_key, + "organization": organization, + "project": project, + "max_retries": max_retries, + **kwargs, + } + if base_url is not None: + client_kwargs["base_url"] = base_url + if websocket_base_url is not None: + client_kwargs["websocket_base_url"] = websocket_base_url + if timeout is not None: + if isinstance(timeout, dict): + timeout = httpx.Timeout(**timeout) + client_kwargs["timeout"] = timeout + if default_headers is not None: + client_kwargs["default_headers"] = default_headers + if default_query is not None: + client_kwargs["default_query"] = default_query + self._client = OpenAI(**client_kwargs) + + @property + def project(self) -> str | None: + """The OpenAI project ID configured on the client.""" + return self._client.project + + def _get_llmeter_state(self) -> dict: + """Extract serializable state from the endpoint and its underlying client. + + This override extends the default state (which pulls immediate properties of the Endpoint) + to include other params set on the OpenAI client but *not* persisted directly on the + Endpoint object. Any config we don't persist will get lost on a serialization round trip, + but anything we pull in should be restored because __init__ passes extra **kwargs to the + OpenAI client constructor. + """ + state = super()._get_llmeter_state() + state["organization"] = self._client.organization + state["base_url"] = str(self._client.base_url) + state["max_retries"] = self._client.max_retries + if self._client.websocket_base_url is not None: + state["websocket_base_url"] = str(self._client.websocket_base_url) + timeout = self._client.timeout + if isinstance(timeout, httpx.Timeout): + state["timeout"] = { + "connect": timeout.connect, + "read": timeout.read, + "write": timeout.write, + "pool": timeout.pool, + } + elif timeout is not None: + state["timeout"] = timeout + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state def _parse_payload(self, payload: CompletionCreateParams | dict) -> str: """Extract the user message text from a ChatCompletions request payload diff --git a/llmeter/endpoints/openai_response.py b/llmeter/endpoints/openai_response.py index 5b78072..a9cb03f 100644 --- a/llmeter/endpoints/openai_response.py +++ b/llmeter/endpoints/openai_response.py @@ -1,18 +1,23 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +# Python Built-Ins: import logging import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Generic, Iterable, TypeVar, cast +# External Dependencies: +import httpx # (Indirect dependency of OpenAI) from openai import OpenAI +from openai._constants import DEFAULT_MAX_RETRIES from openai.types.responses import Response, ResponseCreateParams, ResponseStreamEvent from openai.types.responses.response_create_params import ( ResponseCreateParamsNonStreaming, ResponseCreateParamsStreaming, ) +# Local Dependencies: from .base import Endpoint, InvocationResponse logger = logging.getLogger(__name__) @@ -31,6 +36,14 @@ def __init__( model_id: str, api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize Response API endpoint. @@ -40,17 +53,80 @@ def __init__( model_id: ID of the OpenAI model to use api_key: OpenAI API key (optional, uses OPENAI_API_KEY env var if not provided) provider: Provider name (default: "openai") + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. If a + dict is passed (as is the case after serialization), it'll be interpreted as a set + of arguments to create an httpx.Timeout. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__(endpoint_name, model_id, provider=provider) - self._client = OpenAI(api_key=api_key, **kwargs) + client_kwargs: dict[str, Any] = { + "api_key": api_key, + "organization": organization, + "project": project, + "max_retries": max_retries, + **kwargs, + } + if base_url is not None: + client_kwargs["base_url"] = base_url + if websocket_base_url is not None: + client_kwargs["websocket_base_url"] = websocket_base_url + if timeout is not None: + if isinstance(timeout, dict): + timeout = httpx.Timeout(**timeout) + client_kwargs["timeout"] = timeout + if default_headers is not None: + client_kwargs["default_headers"] = default_headers + if default_query is not None: + client_kwargs["default_query"] = default_query + self._client = OpenAI(**client_kwargs) + + @property + def project(self) -> str | None: + """The OpenAI project ID configured on the client.""" + return self._client.project + + def _get_llmeter_state(self) -> dict: + """Extract serializable state from the endpoint and its underlying client. + + This override extends the default state (which pulls immediate properties of the Endpoint) + to include other params set on the OpenAI client but *not* persisted directly on the + Endpoint object. + """ + state = super()._get_llmeter_state() + state["organization"] = self._client.organization + state["base_url"] = str(self._client.base_url) + state["max_retries"] = self._client.max_retries + if self._client.websocket_base_url is not None: + state["websocket_base_url"] = str(self._client.websocket_base_url) + timeout = self._client.timeout + if isinstance(timeout, httpx.Timeout): + state["timeout"] = { + "connect": timeout.connect, + "read": timeout.read, + "write": timeout.write, + "pool": timeout.pool, + } + elif timeout is not None: + state["timeout"] = timeout + if self._client._custom_headers: + state["default_headers"] = dict(self._client._custom_headers) + if self._client._custom_query: + state["default_query"] = dict(self._client._custom_query) + return state @staticmethod def create_payload( user_message: str | Sequence[str], max_output_tokens: int = 256, instructions: str | None = None, - **kwargs, + **kwargs: Any, ) -> ResponseCreateParams: """Create a payload for the Responses API request. @@ -129,6 +205,14 @@ def __init__( endpoint_name: str = "openai-response", api_key: str | None = None, provider: str = "openai", + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, **kwargs: Any, ): """Initialize Response API endpoint. @@ -138,10 +222,30 @@ def __init__( endpoint_name: Name of the endpoint (default: "openai-response") api_key: OpenAI API key (optional, uses OPENAI_API_KEY env var if not provided) provider: Provider name (default: "openai") + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__( - endpoint_name, model_id, api_key=api_key, provider=provider, **kwargs + endpoint_name, + model_id, + api_key=api_key, + provider=provider, + organization=organization, + project=project, + base_url=base_url, + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, + **kwargs, ) @OpenAIEndpointBase.llmeter_invoke @@ -216,7 +320,15 @@ def __init__( api_key: str | None = None, provider: str = "openai", ttft_visible_tokens_only: bool = True, - **kwargs, + organization: str | None = None, + project: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | httpx.Timeout | dict | None = None, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, + **kwargs: Any, ): """Initialize streaming Response API endpoint. @@ -227,13 +339,29 @@ def __init__( provider: Provider name (default: "openai") ttft_visible_tokens_only: When True (default), TTFT measures time to first visible text token. When False, TTFT includes reasoning token events. + organization: OpenAI organization ID. Defaults to None. + project: OpenAI project ID. Defaults to None. + base_url: Override the default base URL for the API. + websocket_base_url: Override the default base URL for websocket connections. + timeout: Request timeout in seconds, or an httpx.Timeout for granular control. + max_retries: Maximum number of retries for failed requests. + default_headers: Additional headers to send with every request. + default_query: Additional query parameters to send with every request. **kwargs: Additional arguments passed to OpenAI client """ super().__init__( - model_id=model_id, - endpoint_name=endpoint_name, + endpoint_name, + model_id, api_key=api_key, provider=provider, + organization=organization, + project=project, + base_url=base_url, + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, **kwargs, ) self.ttft_visible_tokens_only = ttft_visible_tokens_only @@ -270,10 +398,12 @@ def process_raw_response( `response.reasoning_text.delta`): when `ttft_visible_tokens_only` is ``False``, these set TTFT on the first reasoning token. """ - _REASONING_DELTA_TYPES = frozenset(( - "response.reasoning_summary_text.delta", - "response.reasoning_text.delta", - )) + _REASONING_DELTA_TYPES = frozenset( + ( + "response.reasoning_summary_text.delta", + "response.reasoning_text.delta", + ) + ) for event in raw_response: now = time.perf_counter() @@ -315,9 +445,7 @@ def process_raw_response( elif event.type == "response.failed": error_obj = getattr(event.response, "error", None) if error_obj is not None: - error_msg = ( - getattr(error_obj, "message", None) or str(error_obj) - ) + error_msg = getattr(error_obj, "message", None) or str(error_obj) error_code = getattr(error_obj, "code", None) if error_code: error_msg = f"{error_code}: {error_msg}" diff --git a/llmeter/endpoints/sagemaker.py b/llmeter/endpoints/sagemaker.py index 5a82935..458c90a 100644 --- a/llmeter/endpoints/sagemaker.py +++ b/llmeter/endpoints/sagemaker.py @@ -5,7 +5,7 @@ import json import logging import time -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar import warnings import boto3 @@ -141,8 +141,8 @@ def create_payload( input_text: str | list[ContentItem], max_tokens: int = 256, inference_parameters: dict = {}, - **kwargs, - ): + **kwargs: Any, + ) -> dict: """Create a payload for the SageMaker API request. Args: @@ -291,8 +291,8 @@ def create_payload( input_text: str | list[ContentItem], max_tokens: int = 256, inference_parameters: dict = {}, - **kwargs, - ): + **kwargs: Any, + ) -> dict: """Create a payload for the SageMaker streaming API request. Args: diff --git a/llmeter/json_utils.py b/llmeter/json_utils.py deleted file mode 100644 index 08e8bb7..0000000 --- a/llmeter/json_utils.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 -"""JSON encoding and decoding helpers used across LLMeter. - -Provides a ``default``-compatible serializer function and a matching decoder hook -for round-tripping binary content (``bytes``) through JSON via base64 marker -objects, while also handling ``datetime``, ``os.PathLike``, and objects that -implement ``to_dict()``. - -Example:: - - import json - from llmeter.json_utils import llmeter_default_serializer, llmeter_bytes_decoder - - payload = {"image": {"bytes": b"\\xff\\xd8\\xff\\xe0"}} - - # Serialize - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Deserialize (bytes are restored automatically) - decoded = json.loads(encoded, object_hook=llmeter_bytes_decoder) - assert decoded == payload -""" - -import base64 -import os -from datetime import date, datetime, time, timezone -from typing import Any - -from upath import UPath as Path - - -def llmeter_default_serializer(obj: Any) -> Any: - """Serialize a single non-JSON-serializable object. - - Intended for use as the ``default`` argument to :func:`json.dumps` or - :func:`json.dump`. - - Type handling (checked in order): - - * Objects with a ``to_dict()`` method — delegates to that method. - * ``bytes`` — wrapped in a ``{"__llmeter_bytes__": ""}`` marker so - that :func:`llmeter_bytes_decoder` can restore them on the way back. - * ``datetime`` — converted to a UTC ISO-8601 string with a ``Z`` suffix. - * ``date`` / ``time`` — converted via ``.isoformat()``. - * ``os.PathLike`` — converted to a POSIX path string. - * Anything else — ``str()`` fallback (returns ``None`` if that also fails). - - Args: - obj: The object that the default encoder could not handle. - - Returns: - A JSON-serializable representation of *obj*. - - Example:: - - >>> import json - >>> from llmeter.json_utils import llmeter_default_serializer - >>> json.dumps({"ts": datetime(2024, 1, 1)}, default=llmeter_default_serializer) - '{"ts": "2024-01-01T00:00:00"}' - """ - if hasattr(obj, "to_dict") and callable(obj.to_dict): - result = obj.to_dict() - if not isinstance(result, dict): - # This check guards against infinite recursion in case something tries to serialize a - # MagicMock object with this function (in which to_dict returns another mock) - raise TypeError( - f"{type(obj).__name__}.to_dict() returned {type(result).__name__}, expected dict" - ) - return result - if isinstance(obj, bytes): - return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")} - if isinstance(obj, datetime): - if obj.tzinfo is not None: - obj = obj.astimezone(timezone.utc) - return obj.isoformat(timespec="seconds").replace("+00:00", "Z") - if isinstance(obj, (date, time)): - return obj.isoformat() - if isinstance(obj, (os.PathLike, Path)): - return Path(obj).as_posix() - try: - return str(obj) - except Exception: - return None - - -def llmeter_bytes_decoder(dct: dict) -> dict | bytes: - """Decode ``__llmeter_bytes__`` marker objects back to ``bytes``. - - Intended for use as the ``object_hook`` argument to :func:`json.loads` or - :func:`json.load`. Marker objects produced by :func:`llmeter_default_serializer` - are detected and converted back to ``bytes``; all other dicts pass through - unchanged. - - Args: - dct: A dictionary produced by the JSON parser. - - Returns: - The original ``bytes`` if *dct* is a marker object, otherwise *dct* unchanged. - - Example:: - - >>> import json - >>> from llmeter.json_utils import llmeter_bytes_decoder - >>> json.loads('{"__llmeter_bytes__": "/9j/4A=="}', object_hook=llmeter_bytes_decoder) - b'\\xff\\xd8\\xff\\xe0' - """ - if "__llmeter_bytes__" in dct and len(dct) == 1: - return base64.b64decode(dct["__llmeter_bytes__"]) - return dct diff --git a/llmeter/prompt_utils.py b/llmeter/prompt_utils.py index 2408c23..49e7583 100644 --- a/llmeter/prompt_utils.py +++ b/llmeter/prompt_utils.py @@ -4,14 +4,15 @@ import json import logging import random +from collections.abc import Callable, Iterator from dataclasses import dataclass from itertools import product -from typing import Any, Callable, Iterator +from typing import Any from upath import UPath as Path from upath.types import ReadablePathLike, WritablePathLike -from .json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from .serialization import bytes_decoder, json_default from .tokenizers import DummyTokenizer, Tokenizer from .utils import DeferredError, ensure_path @@ -419,7 +420,7 @@ def load_payloads( This function reads JSON data from either a single file or multiple files in a directory. It supports both .json and .jsonl file formats. Binary content - (bytes objects) that were serialized using ``llmeter_default_serializer`` are automatically + (bytes objects) that were serialized using ``json_default`` are automatically restored during deserialization. Binary Content Handling: @@ -513,13 +514,11 @@ def _load_data_file(file: Path) -> Iterator[dict]: try: if not line.strip(): continue - yield json.loads( - line.strip(), object_hook=llmeter_bytes_decoder - ) + yield json.loads(line.strip(), object_hook=bytes_decoder) except json.JSONDecodeError as e: print(f"Error decoding JSON in {file}: {e}") else: # Assume it's a regular JSON file - yield json.load(f, object_hook=llmeter_bytes_decoder) + yield json.load(f, object_hook=bytes_decoder) except IOError as e: print(f"Error reading file {file}: {e}") except json.JSONDecodeError as e: @@ -638,5 +637,5 @@ def save_payloads( payloads = [payloads] with output_file_path.open(mode="w") as f: for payload in payloads: - f.write(json.dumps(payload, default=llmeter_default_serializer) + "\n") + f.write(json.dumps(payload, default=json_default) + "\n") return output_file_path diff --git a/llmeter/results.py b/llmeter/results.py index b40e9aa..0b01133 100644 --- a/llmeter/results.py +++ b/llmeter/results.py @@ -3,32 +3,22 @@ import json import logging -import types as _types -from dataclasses import asdict, dataclass, fields +from collections.abc import Callable, Sequence +from dataclasses import asdict, dataclass from datetime import datetime from numbers import Number -from typing import Any, Sequence +from typing import Any import jmespath from upath.types import ReadablePathLike, WritablePathLike from .endpoints import InvocationResponse -from .json_utils import llmeter_default_serializer +from .serialization import json_default, restore_dataclass_types from .utils import ensure_path, summary_stats_from_list logger = logging.getLogger(__name__) -def _get_type_args(tp) -> tuple: - """Return the members of a union type (e.g. ``datetime | None`` -> (datetime, NoneType)).""" - if isinstance(tp, _types.UnionType): - return tp.__args__ - origin = getattr(tp, "__origin__", None) - if origin is _types.UnionType: - return tp.__args__ - return (tp,) if isinstance(tp, type) else () - - @dataclass class Result: """Results of a test run.""" @@ -50,7 +40,7 @@ class Result: end_time: datetime | None = None def __str__(self): - return json.dumps(self.stats, indent=4, default=llmeter_default_serializer) + return json.dumps(self.stats, indent=4, default=json_default) def __post_init__(self): """Initialize the Result instance.""" @@ -58,24 +48,6 @@ def __post_init__(self): if not hasattr(self, "_preloaded_stats"): self._preloaded_stats = None - @classmethod - def _parse_datetime_fields(cls, d: dict) -> None: - """Convert any datetime fields on cls present in d from ISO-8601 strings to datetimes - - Introspects this (data)class to find all `datetime`-typed fields, and converts any matching - entries in `d` with ISO-8601 string values to datetimes instead. This is used for loading - JSON data (in which dates are stringified) into the Result class or stats. - """ - for f in fields(cls): - if datetime not in _get_type_args(f.type): - continue - val = d.get(f.name) - if val and isinstance(val, str): - try: - d[f.name] = datetime.fromisoformat(val.replace("Z", "+00:00")) - except ValueError: - pass - def _update_contributed_stats(self, stats: dict[str, Number]): """ Upsert externally-provided statistics for the `stats` property @@ -127,28 +99,35 @@ def save(self, output_path: WritablePathLike | None = None) -> None: summary_path = output_path / "summary.json" stats_path = output_path / "stats.json" - with summary_path.open("w") as f, stats_path.open("w") as s: - f.write(self.to_json(indent=4)) - s.write( - json.dumps(self.stats, indent=4, default=llmeter_default_serializer) + with summary_path.open("w") as f: + json.dump( + { + k: o + for k, o in asdict(self).items() + if k not in ["responses", "stats"] + }, + f, + default=json_default, + indent=4, ) + with stats_path.open("w") as f: + json.dump(self.stats, f, default=json_default, indent=4) responses_path = output_path / "responses.jsonl" if not responses_path.exists(): with responses_path.open("w") as f: for response in self.responses: - f.write( - json.dumps(asdict(response), default=llmeter_default_serializer) - + "\n" - ) + f.write(json.dumps(asdict(response), default=json_default) + "\n") - def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: + def to_json( + self, default: Callable[[Any], Any] | None = json_default, **kwargs: Any + ) -> str: """Return the results as a JSON string. Args: default: Fallback serializer. Defaults to - :func:`~llmeter.json_utils.llmeter_default_serializer`. - **kwargs: Extra keyword arguments passed to :func:`json.dumps`. + [`json_default`][llmeter.serialization.json_default]. + **kwargs: Extra keyword arguments passed to `json.dumps`. """ summary = { k: o for k, o in asdict(self).items() if k not in ["responses", "stats"] @@ -158,20 +137,18 @@ def to_json(self, default=llmeter_default_serializer, **kwargs) -> str: def to_dict(self, include_responses: bool = False) -> dict: """Return a dictionary representation of this result. - Returns a plain ``dict`` produced by :func:`dataclasses.asdict`, - preserving native Python types (``datetime``, ``UPath``, etc.). - This is suitable for programmatic access and internal data + Returns a plain `dict` produced by `dataclasses.asdict`, preserving native Python types + (`datetime`, `UPath`, etc.). This is suitable for programmatic access and internal data processing. - For JSON output, use :meth:`to_json` which delegates to - :func:`~llmeter.json_utils.llmeter_default_serializer` for - non-serializable types, or pass the dict through - ``json.dumps(result.to_dict(), default=llmeter_default_serializer)``. + For JSON output, use `to_json` which delegates to + [`json_default`][llmeter.serialization.json_default] for non-serializable types, or pass + the dict through `json.dumps(result.to_dict(), default=json_default)`. Args: - include_responses: If ``True``, include the full list of - :class:`~llmeter.endpoints.base.InvocationResponse` dicts - and the ``stats`` key. Defaults to ``False``. + include_responses: Set `True` to include the full list of + [`InvocationResponse`][llmeter.endpoints.base.InvocationResponse] dicts and the + `stats` key. Returns: dict: A dictionary of result fields with native Python types. @@ -269,7 +246,7 @@ def _load_summary(cls, result_path) -> dict[str, Any]: with summary_path.open("r") as f: summary = json.load(f) - cls._parse_datetime_fields(summary) + restore_dataclass_types(cls, summary) if "output_path" not in summary or summary["output_path"] is None: summary["output_path"] = str(result_path) @@ -316,9 +293,20 @@ def _recover_metadata(cls, result_path) -> dict[str, Any]: metadata[passthru_field] = config[passthru_field] endpoint = config.get("endpoint", {}) if isinstance(endpoint, dict): - metadata["model_id"] = endpoint.get("model_id") - metadata["endpoint_name"] = endpoint.get("endpoint_name") - metadata["provider"] = endpoint.get("provider") + # Modern format nests endpoint fields under `__llmeter_state__`; legacy + # format keeps them at the top level. + if "__llmeter_class__" in endpoint: + endpoint_fields = endpoint.get("__llmeter_state__", {}) + else: + logger.warning( + "Recovering endpoint metadata from a legacy data format. Support for " + "this will be removed in a future major version. Consider re-saving " + "the file to update." + ) + endpoint_fields = endpoint + metadata["model_id"] = endpoint_fields.get("model_id") + metadata["endpoint_name"] = endpoint_fields.get("endpoint_name") + metadata["provider"] = endpoint_fields.get("provider") except (json.JSONDecodeError, KeyError) as e: logger.warning(f"Could not parse run_config.json: {e}") @@ -395,7 +383,7 @@ def _resolve_stats(cls, result: "Result", stats_path) -> None: try: with stats_path.open("r") as s: saved_stats = json.loads(s.read()) - cls._parse_datetime_fields(saved_stats) + restore_dataclass_types(cls, saved_stats) except (json.JSONDecodeError, OSError) as e: logger.warning(f"Could not load stats.json: {e}") saved_stats = None diff --git a/llmeter/runner.py b/llmeter/runner.py index 88637c5..743d797 100644 --- a/llmeter/runner.py +++ b/llmeter/runner.py @@ -23,6 +23,7 @@ from upath.types import ReadablePathLike, WritablePathLike from .live_display import LiveStatsDisplay +from .serialization import dump_object, json_default, load_object from .utils import RunningStats, ensure_path, now_utc if TYPE_CHECKING: @@ -30,7 +31,6 @@ from .callbacks.base import Callback from .endpoints.base import Endpoint, InvocationResponse -from .json_utils import llmeter_default_serializer from .prompt_utils import load_payloads, save_payloads from .results import Result from .tokenizers import DummyTokenizer, Tokenizer @@ -145,7 +145,15 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba assert self.endpoint is not None, "Endpoint cannot be None" if isinstance(self.endpoint, dict): - self._endpoint: Endpoint = Endpoint.load(self.endpoint) + if "__llmeter_class__" in self.endpoint: + self._endpoint: Endpoint = load_object(self.endpoint) + else: + # Legacy load path for <=v0.1.12 data files: + logger.warning( + "Loading endpoint config from a legacy data format. Support for this will be " + "removed in a future major version. Consider re-saving the file to update." + ) + self._endpoint = Endpoint.load(self.endpoint) else: self._endpoint = self.endpoint @@ -155,7 +163,15 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba if self.tokenizer is None: self.tokenizer = DummyTokenizer() if isinstance(self.tokenizer, dict): - self._tokenizer: Tokenizer = Tokenizer.load(self.tokenizer) + if "__llmeter_class__" in self.tokenizer: + self._tokenizer: Tokenizer = load_object(self.tokenizer) + else: + # Legacy load path for <=v0.1.12 data files: + logger.warning( + "Loading tokenizer from a legacy LLMeter data format. Support for this will " + "be removed in a future major version. Consider re-saving the file to update." + ) + self._tokenizer = Tokenizer.load(self.tokenizer) else: self._tokenizer = self.tokenizer @@ -171,6 +187,9 @@ def save( file_name: File name to create under `output_path`. """ output_path = ensure_path(output_path or self.output_path) + if output_path is None: + logger.info("No output_path provided - skipping saving run config") + return output_path.mkdir(parents=True, exist_ok=True) run_config_path = output_path / file_name @@ -180,19 +199,17 @@ def save( payload_path = save_payloads(self.payload, output_path) config_copy.payload = payload_path + # Serialize callable objects using dump_object assert self.endpoint is not None, "Endpoint cannot be None" - if not isinstance(self.endpoint, dict): - config_copy.endpoint = self.endpoint.to_dict() + config_copy.endpoint = dump_object(self._endpoint) - if not isinstance(self.tokenizer, dict): - config_copy.tokenizer = Tokenizer.to_dict(self.tokenizer) + config_copy.tokenizer = dump_object(self._tokenizer) + + if self.callbacks: + config_copy.callbacks = [dump_object(cb) for cb in self.callbacks] with run_config_path.open("w") as f: - f.write( - json.dumps( - asdict(config_copy), default=llmeter_default_serializer, indent=4 - ) - ) + json.dump(asdict(config_copy), f, default=json_default, indent=4) @classmethod def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): @@ -205,8 +222,28 @@ def load(cls, load_path: ReadablePathLike, file_name: str = "run_config.json"): load_path = ensure_path(load_path) with (load_path / file_name).open() as f: config = json.load(f) - config["endpoint"] = Endpoint.load(config["endpoint"]) - config["tokenizer"] = Tokenizer.load(config["tokenizer"]) + + # Restore endpoint and tokenizer. Only unwrap the new-format (`__llmeter_class__`) envelope + # here; legacy dicts (`endpoint_type` / `tokenizer_module`) are passed through untouched so + # `__post_init__` can route them to `Endpoint.load` / `Tokenizer.load`. + ep = config.get("endpoint") + if isinstance(ep, dict) and "__llmeter_class__" in ep: + config["endpoint"] = load_object(ep) + + tok = config.get("tokenizer") + if isinstance(tok, dict) and "__llmeter_class__" in tok: + config["tokenizer"] = load_object(tok) + + # Restore callbacks (new-format envelopes only; anything else is left as-is) + cbs = config.get("callbacks") + if isinstance(cbs, list): + config["callbacks"] = [ + load_object(cb) + if isinstance(cb, dict) and "__llmeter_class__" in cb + else cb + for cb in cbs + ] + return cls(**config) diff --git a/llmeter/serialization.py b/llmeter/serialization.py new file mode 100644 index 0000000..a6b9b85 --- /dev/null +++ b/llmeter/serialization.py @@ -0,0 +1,420 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unified JSON-based serialization for LLMeter objects. + +Most of the objects we'd like to save to and load from file in LLMeter are configuration-like and +*almost* JSON-compatible, but with some extensions (like datetimes, binary image/etc payloads, +and callback objects). Rather than falling back to Pickle (which is Python-specific, +non-human-readable, and Python version fragile) - we implement in this module a JSON-based scheme +for saving and loading compatible LLMeter objects. + +!!! warning "A warning on security" + + [`load_object`][llmeter.serialization.load_object] imports and instantiates whatever class path + is in the `__llmeter_class__` field. Like unpickle, this has the potential to run arbitrary + code. Do not load configs from untrusted sources! + +### Key components + +- **[`Serializable`][llmeter.serialization.Serializable]**: Mixin to give any class an automatic + *state* protocol (`_get_llmeter_state`/`_set_llmeter_state`) by introspecting `__init__`. + Deliberately distinct from the pickle protocol so `pickle` / `copy` / `deepcopy` keep their + native behavior. +- **[`dump_object`][llmeter.serialization.dump_object]** and + **[`load_object`][llmeter.serialization.load_object]**: Full round-trip persistence for + `Serializable`-compatible objects, using a + `{"__llmeter_class__": ..., "__llmeter_state__": ...}` envelope. + + +### Implementation details to be aware of + +**State vs. identity:** We split an object's serialized form into two layers. The *state* +(`_get_llmeter_state`) contains the object's own field values and does *not* describe which class +those values belong to. The **envelope** built by +[`dump_object`][llmeter.serialization.dump_object] adds the identity (`__llmeter_class__`) around +that state. The "state" dictionaries handled by `_get_llmeter_state` / `_set_llmeter_state` are +not a fully self-describing representation of the object - only of its state. Keeping identity out +of the state dict avoids polluting the user's field namespace and lets nested polymorphic values +each carry their own envelope. +""" + +import base64 +import importlib +import inspect +import json +import logging +import os +import re +import types as _types +from dataclasses import asdict, fields, is_dataclass +from datetime import date, datetime, time, timezone +from typing import Any + +from upath import UPath as Path +from upath.types import ReadablePathLike, WritablePathLike + +from .utils import ensure_path + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Datetime helpers +# --------------------------------------------------------------------------- + + +def datetime_to_str(dt: datetime) -> str: + """Convert a datetime to a UTC ISO-8601 string with ``Z`` suffix. + + Timezone-aware datetimes are converted to UTC first. Naive datetimes are serialized as-is + (assumed UTC). + """ + if dt.tzinfo is not None: + dt = dt.astimezone(timezone.utc) + return dt.isoformat(timespec="seconds").replace("+00:00", "Z") + + +def str_to_datetime(s: str) -> datetime: + """Parse an ISO-8601 string (with optional ``Z`` suffix) to a datetime.""" + return datetime.fromisoformat(s.replace("Z", "+00:00")) + + +# --------------------------------------------------------------------------- +# JSON helpers +# --------------------------------------------------------------------------- + + +def json_default(obj: Any) -> Any: + """Serialize a single non-natively-JSON-serializable object. + + Intended for use as the `default` argument to `json.dump` or `json.dumps`. This **does not** + handle full recursive LLMeter serialization - see + [`dump_object`][llmeter.serialization.dump_object] instead. + + Type handling (checked in order): + + * `bytes` — wrapped in a `{"__llmeter_bytes__": ""}` marker so that + [`bytes_decoder`][llmeter.serialization.bytes_decoder] can restore them on the way back. + * `datetime` — converted to a UTC ISO-8601 string with a `Z` suffix. + * `date` / `time` — converted via `.isoformat()`. + * `os.PathLike` — converted to a POSIX path string. + * Anything else — `str()` fallback (returns `None` if that also fails). + + Args: + obj: The object that the default JSON encoder could not handle. + + Returns: + A JSON-serializable representation of *obj*. + """ + if isinstance(obj, bytes): + return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")} + if isinstance(obj, datetime): + return datetime_to_str(obj) + if isinstance(obj, (date, time)): + return obj.isoformat() + if isinstance(obj, (os.PathLike, Path)): + return Path(obj).as_posix() + try: + return str(obj) + except Exception: + return None + + +def bytes_decoder(dct: dict) -> dict | bytes: + """Decode `__llmeter_bytes__` marker objects back to Python `bytes`. + + Intended for use as the `object_hook` argument to `json.load` or `json.loads`. Marker objects + produced by [`json_default`][llmeter.serialization.json_default] are detected and converted + back to `bytes`; all other dicts pass through unchanged. + + Args: + dct: A dictionary produced by the JSON parser. + + Returns: + The original `bytes` if *dct* is an LLMeter bytes marker object, otherwise *dct* unchanged. + """ + if "__llmeter_bytes__" in dct and len(dct) == 1: + return base64.b64decode(dct["__llmeter_bytes__"]) + return dct + + +def _get_type_args(tp) -> tuple: + """Return the members of a union type (e.g. `datetime | None` -> (datetime, NoneType)).""" + if isinstance(tp, _types.UnionType): + return tp.__args__ + origin = getattr(tp, "__origin__", None) + if origin is _types.UnionType: + return tp.__args__ + return (tp,) if isinstance(tp, type) else () + + +def restore_dataclass_types(cls: type, data: dict) -> None: + """Restore typed fields in a dict destined for a dataclass constructor. + + Introspects `cls` (a dataclass) and converts JSON-native values back to their annotated Python + types. Currently handles: + + * `datetime` fields — parses ISO-8601 strings via `str_to_datetime`. + * `bytes` fields — decodes `__llmeter_bytes__` markers via base64. + + Only fields declared on `cls` are touched — nested user payloads (e.g. `input_payload`) are + left unchanged. Mutates *data* in place. + + Args: + cls: A dataclass type to introspect for field type annotations. + data: A dictionary of field values (e.g. from `json.load`) to coerce. + """ + for f in fields(cls): + val = data.get(f.name) + if val is None: + continue + type_args = _get_type_args(f.type) + match val: + case str() if datetime in type_args: + try: + data[f.name] = str_to_datetime(val) + except ValueError: + pass + case {"__llmeter_bytes__": b64} if bytes in type_args and len(val) == 1: + data[f.name] = base64.b64decode(b64) + + +# --------------------------------------------------------------------------- +# Serializable mixin +# --------------------------------------------------------------------------- + + +class Serializable: + """Mixin providing a state extraction protocol compatible with LLMeter serialization. + + Serialization in LLMeter uses a state extraction protocol somewhat similar to, but deliberately + separate from, the (`__getstate__` / `__setstate__`) interface used by `pickle`, `copy.copy`, + and `copy.deepcopy`. Subclasses remain natively picklable/copyable, but can also + be saved to and loaded from LLMeter's JSON-based format. + + This default implementation works with plain classes, `@dataclass`, and any class whose + `__init__` parameters correspond to instance attributes (`self.x` or `self._x`). Nested + `Serializable`-like objects are recursively persisted and loaded via + [`dump_object`][llmeter.serialization.dump_object] and + [`load_object`][llmeter.serialization.load_object]. + + If your class needs more custom logic to represent and restore its state, customize the + provided methods by overriding or implementing your own from scratch. + """ + + def _get_llmeter_state(self) -> dict: + """Extract a JSON-serializable state dict by introspecting `__init__` parameters. + + Returns the object's constructor arguments (looked up as `self.` or `self._`), + recursively serialized. Note that: + + 1. This is state only — it carries no class identity; `dump_object` wraps it with + `__llmeter_class__`. + 2. Any properties not exposed as `__init__` arguments will not be persisted. Override your + class' state get and set methods if you need different behaviour. + """ + sig = inspect.signature(self.__init__) + state = {} + for name, param in sig.parameters.items(): + if name in ("self", "args", "kwargs"): + continue + if param.kind in (param.VAR_POSITIONAL, param.VAR_KEYWORD): + continue + if hasattr(self, name): + state[name] = _serialize_value(getattr(self, name)) + elif hasattr(self, f"_{name}"): + state[name] = _serialize_value(getattr(self, f"_{name}")) + return state + + def _set_llmeter_state(self, state: dict) -> None: + """Restore this instance from a state dict produced by `_get_llmeter_state`. + + Note this rebuilds the object by *calling the constructor* with the state as keyword + arguments. [`load_object`][llmeter.serialization.load_object] first creates a bare instance + with `__new__`, then calls this method to populate it. + """ + deserialized = {k: _deserialize_value(v) for k, v in state.items()} + self.__init__(**deserialized) + + def save_to_file(self, path: WritablePathLike) -> Path: + """Save this object to a JSON file. + + Uses the `_get_llmeter_state` protocol. Override `_get_llmeter_state` (not this method) if + custom JSON-based serialization is needed. + + Args: + path: (Local or Cloud) path where the object will be saved. + + Returns: + The (validated/normalized) path the object was written to. + """ + path = ensure_path(path) + path.parent.mkdir(parents=True, exist_ok=True) + data = dump_object(self) + with path.open("w") as f: + json.dump(data, f, indent=4, default=json_default) + return path + + @classmethod + def load_from_file(cls, path: ReadablePathLike) -> "Serializable": + """Load an object from a JSON file. + + Detects the type from the `__llmeter_class__` field and reconstructs it. + + Args: + path: (Local or Cloud) path where the object was saved. + Returns: + The loaded instance. + """ + path = ensure_path(path) + with path.open("r") as f: + data = json.load(f) + return load_object(data) + + +# --------------------------------------------------------------------------- +# Object serialization API +# --------------------------------------------------------------------------- + + +def dump_object(obj: Any) -> dict: + """Serialize an object to a type-tagged dict for round-trip persistence. + + The returned envelope has the form + `{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}`. + + Serialization strategy (checked in order): + + 1. If the object implements the LLMeter state protocol (as in + [`Serializable._get_llmeter_state`][llmeter.serialization.Serializable._get_llmeter_state]) + this will call it to obtain the state dict. + 2. If the object is a dataclass, uses `dataclasses.asdict`. + 3. Otherwise, takes all public (non-underscore-prefixed) entries from `__dict__`. + + Args: + obj: The object to serialize. + + Returns: + dct: A JSON-serializable dict that [`load_object`][llmeter.serialization.load_object] can + reconstruct. + """ + class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" + if hasattr(obj, "_get_llmeter_state"): + state = obj._get_llmeter_state() + elif is_dataclass(obj) and not isinstance(obj, type): + state = asdict(obj) + elif hasattr(obj, "__dict__"): + state = {k: v for k, v in obj.__dict__.items() if not k.startswith("_")} + else: + state = {} + return {"__llmeter_class__": class_path, "__llmeter_state__": state} + + +def load_object(data: dict) -> Any: + """Restore an object from a type-tagged dict produced by `dump_object`. + + !!! warning + This method imports and instantiates class paths specified by the input data, which (like + pickle) can enable arbitrary running arbitrary code. Do not run it on data from unstrusted + sources! + + Imports the module identified by `__llmeter_class__`, instantiates the class (bypassing + `__init__` via `__new__`), and restores its state via + [`Serializable._set_llmeter_state`][llmeter.serialization.Serializable._set_llmeter_state]. + Objects that do not implement the protocol are restored by assigning the (deserialized) state + onto their `__dict__`. + + Args: + data: A dict with `__llmeter_class__` and `__llmeter_state__` keys, as + produced by [`dump_object`][llmeter.serialization.dump_object]. + + Returns: + The reconstructed object instance. + """ + class_path = data["__llmeter_class__"] + module_path, class_name = class_path.rsplit(".", 1) + module = importlib.import_module(module_path) + cls = getattr(module, class_name) + + state = data["__llmeter_state__"] + obj = cls.__new__(cls) + if hasattr(obj, "_set_llmeter_state"): + obj._set_llmeter_state(state) + else: + # Fallback for non-Serializable classes (state came from dump_object's dataclass / + # __dict__ branches). We assign attributes directly rather than calling __setstate__: + # `object` has no __setstate__ (so it would raise for a plain class), and if the class + # *does* define one it's the pickle hook, which may expect a different state shape than + # our JSON dict. Direct assignment (with per-value deserialization) is the safe generic + # restore, and mirrors pickle's own "restore __dict__ without calling __init__" default. + for key, value in state.items(): + setattr(obj, key, _deserialize_value(value)) + return obj + + +# --------------------------------------------------------------------------- +# Internal helpers for recursive serialization +# --------------------------------------------------------------------------- + +_SERIALIZERS: list[tuple[type | tuple[type, ...], Any]] = [ + (bytes, lambda v: {"__llmeter_bytes__": base64.b64encode(v).decode("utf-8")}), + (datetime, datetime_to_str), + (os.PathLike, lambda v: Path(v).as_posix()), +] + +_DATETIME_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$") + + +def _serialize_value(val: Any) -> Any: + """Recursively prepare a value for JSON persistence. + + Handles primitives, known types (bytes, datetime, PathLike), nested + [`Serializable`][llmeter.serialization.Serializable] objects (via + [`dump_object`][llmeter.serialization.dump_object]), dicts, and lists/tuples. + + Raises: + TypeError: for objects it cannot serialize. + """ + if val is None or isinstance(val, (str, int, float, bool)): + return val + for types, fn in _SERIALIZERS: + if isinstance(val, types): + return fn(val) + if hasattr(val, "_get_llmeter_state"): + return dump_object(val) + if isinstance(val, dict): + return {k: _serialize_value(v) for k, v in val.items()} + if isinstance(val, (list, tuple)): + return [_serialize_value(item) for item in val] + raise TypeError( + f"Cannot serialize {type(val).__name__!r} object: {val!r}. " + "Only JSON primitives, bytes, datetime, PathLike, dicts, lists/tuples, and " + "LLMeter-serializable objects are supported. To make a custom object " + "serializable (e.g. a custom cost dimension or callback), have its class " + "inherit from llmeter.serialization.Serializable." + ) + + +def _deserialize_value(val: Any) -> Any: + """Recursively restore a value from JSON persistence. + + Recognizes type-tagged dicts (`__llmeter_class__`), bytes markers (`__llmeter_bytes__`), + ISO-8601 datetime strings, and recursively processes nested dicts and lists. + """ + match val: + case None | bool() | int() | float(): + return val + case str() if _DATETIME_RE.fullmatch(val): + return str_to_datetime(val) + case str(): + return val + case {"__llmeter_class__": _, "__llmeter_state__": _}: + return load_object(val) + case {"__llmeter_bytes__": b64} if len(val) == 1: + return base64.b64decode(b64) + case dict(): + return {k: _deserialize_value(v) for k, v in val.items()} + case list() | tuple(): + return [_deserialize_value(item) for item in val] + case _: + return val diff --git a/llmeter/tokenizers.py b/llmeter/tokenizers.py index 0b2bbe9..89e259d 100644 --- a/llmeter/tokenizers.py +++ b/llmeter/tokenizers.py @@ -3,17 +3,12 @@ from __future__ import annotations -import json from abc import ABC, abstractmethod -from typing import Any -from upath import UPath -from upath.types import ReadablePathLike, WritablePathLike +from .serialization import Serializable -from .utils import ensure_path - -class Tokenizer(ABC): +class Tokenizer(Serializable, ABC): def __init__(self, *args, **kwargs): pass @@ -46,123 +41,52 @@ def __subclasscheck__(cls, subclass): return True @staticmethod - def load_from_file(tokenizer_path: ReadablePathLike | None) -> Tokenizer: - """ - Loads a tokenizer from a file. - - Args: - tokenizer_path (ReadablePathLike): The path to the serialized tokenizer file. - - Returns: - Tokenizer: The loaded tokenizer. - """ - if tokenizer_path is None: - return DummyTokenizer() - tokenizer_path = ensure_path(tokenizer_path) - with tokenizer_path.open("r") as f: - tokenizer_info = json.load(f) - - return _load_tokenizer_from_info(tokenizer_info) + def load(tokenizer_info: dict) -> "Tokenizer": + """Load a tokenizer from a dictionary. - @staticmethod - def load(tokenizer_info: dict) -> Tokenizer: - """ - Loads a tokenizer from a dictionary. + This supports configs saved before the unified ``dump_object``/``load_object`` + serialization was introduced. New code should use + :func:`~llmeter.serialization.load_object` instead. Args: - tokenizer_info (Dict): The tokenizer information to load. + tokenizer_info (dict): The tokenizer information to load. Must include at minimum + a ``tokenizer_module`` key. Returns: Tokenizer: The loaded tokenizer. """ return _load_tokenizer_from_info(tokenizer_info) - @staticmethod - def to_dict(tokenizer: Any) -> dict: - """ - Serializes a tokenizer to a dictionary. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - - Returns: - Dict: The serialized tokenizer. - """ - return _to_dict(tokenizer) - - -def _to_dict(tokenizer: Any) -> dict: - """ - Serializes a tokenizer to a dictionary. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - - Returns: - Dict: The serialized tokenizer. - """ - if tokenizer.__module__.split(".")[0] == "transformers": - return {"tokenizer_module": "transformers", "name": tokenizer.name_or_path} - - if tokenizer.__module__.split(".")[0] == "tiktoken": - return {"tokenizer_module": "tiktoken", "name": tokenizer.name} - - if tokenizer.__module__.split(".")[0] == "llmeter": - return {"tokenizer_module": "llmeter"} - - raise ValueError(f"Unknown tokenizer module: {tokenizer.__module__}") - - -def save_tokenizer(tokenizer: Any, output_path: WritablePathLike) -> UPath: - """ - Save a tokenizer information to a file. - - Args: - tokenizer (Tokenizer): The tokenizer to serialize. - output_path (WritablePathLike): The path to save the serialized tokenizer to. - - Returns: - UPath: The path to the serialized tokenizer file. - """ - tokenizer_info = _to_dict(tokenizer) - - output_path = ensure_path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - with output_path.open("w") as f: - json.dump(tokenizer_info, f) - - return output_path - def _load_tokenizer_from_info(tokenizer_info: dict) -> Tokenizer: - """ - Loads a tokenizer from a file. + """Load a tokenizer from a legacy info dictionary. Args: - tokenizer_info (Dict): The tokenizer information to load. + tokenizer_info (dict): The tokenizer information to load. Returns: Tokenizer: The loaded tokenizer. """ - if tokenizer_info["tokenizer_module"] == "transformers": + module = tokenizer_info["tokenizer_module"] + + if module == "transformers": from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(tokenizer_info["name"]) # type: ignore - if tokenizer_info["tokenizer_module"] == "tiktoken": + if module == "tiktoken": from tiktoken import get_encoding return get_encoding(tokenizer_info["name"]) # type: ignore - if tokenizer_info["tokenizer_module"] == "llmeter": + if module == "llmeter": return DummyTokenizer() - raise ValueError(f"Unknown tokenizer module: {tokenizer_info['tokenizer_module']}") + raise ValueError(f"Unknown tokenizer module: {module}") class DummyTokenizer(Tokenizer): - """ - A dummy tokenizer that splits the input text on whitespace and returns the tokens as is. + """A dummy tokenizer that splits the input text on whitespace and returns the tokens as is. This tokenizer will generally under-estimate token counts in English and latin languages (where words comprise more than one token on average), and will give very poor results for languages @@ -174,8 +98,8 @@ class DummyTokenizer(Tokenizer): def __init__(self, *args, **kwargs): pass - def encode(self, text: str): + def encode(self, text: str) -> list[str]: return [k for k in text.split()] - def decode(self, tokens: list[str]): + def decode(self, tokens: list[str]) -> str: return " ".join(k for k in tokens) diff --git a/mkdocs.yml b/mkdocs.yml index 27fb90b..a3028dc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -42,6 +42,7 @@ nav: - user_guide/callbacks/cache_buster.md - user_guide/callbacks/cost.md - user_guide/callbacks/mlflow.md + - Save (and Reload) Results: user_guide/serialization.md - API Reference: - reference/index.md - callbacks: @@ -55,7 +56,6 @@ nav: - reference/callbacks/cost/providers/index.md - sagemaker: reference/callbacks/cost/providers/sagemaker.md - results: reference/callbacks/cost/results.md - - serde: reference/callbacks/cost/serde.md - mlflow: reference/callbacks/mlflow.md - endpoints: - reference/endpoints/index.md @@ -68,12 +68,12 @@ nav: - openai_response: reference/endpoints/openai_response.md - sagemaker: reference/endpoints/sagemaker.md - experiments: reference/experiments.md - - json_utils: reference/json_utils.md - plotting: - reference/plotting/index.md - prompt_utils: reference/prompt_utils.md - results: reference/results.md - runner: reference/runner.md + - serialization: reference/serialization.md - tokenizers: reference/tokenizers.md - utils: reference/utils.md repo_url: https://github.com/awslabs/llmeter diff --git a/tests/unit/callbacks/cost/providers/test_sagemaker.py b/tests/unit/callbacks/cost/providers/test_sagemaker.py index 2bf7909..3325eec 100644 --- a/tests/unit/callbacks/cost/providers/test_sagemaker.py +++ b/tests/unit/callbacks/cost/providers/test_sagemaker.py @@ -11,6 +11,7 @@ from moto import mock_aws # Local Dependencies: +from llmeter.serialization import dump_object, load_object from llmeter.callbacks.cost.providers.sagemaker import ( SageMakerRTEndpointCompute, SageMakerRTEndpointStorage, @@ -43,16 +44,15 @@ async def test_real_time_endpoint_compute_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.to_dict() + dim_ser = dim._get_llmeter_state() assert dim_ser == { - "_type": "SageMakerRTEndpointCompute", "instance_count": 1, "instance_type": "ml.doesnotexist", "price_per_hour": 6, "region": "us-nope-0", } - dim = SageMakerRTEndpointCompute.from_dict(dim_ser) + dim = load_object(dump_object(dim)) result.total_test_time = 90 assert await dim.calculate(result) == 6 * 90 / 3600 @@ -237,15 +237,14 @@ async def test_real_time_endpoint_storage_explicit_price(): result.total_test_time = None assert await dim.calculate(result) is None - dim_ser = dim.to_dict() + dim_ser = dim._get_llmeter_state() assert dim_ser == { - "_type": "SageMakerRTEndpointStorage", "gbs_provisioned": 5, "price_per_gb_hour": 9, "region": "us-nope-0", } - dim = SageMakerRTEndpointStorage.from_dict(dim_ser) + dim = load_object(dump_object(dim)) result.total_test_time = 90 assert await dim.calculate(result) == 9 * 5 * 90 / 3600 diff --git a/tests/unit/callbacks/cost/test_dimensions.py b/tests/unit/callbacks/cost/test_dimensions.py index 68cbe1c..e60ce3a 100644 --- a/tests/unit/callbacks/cost/test_dimensions.py +++ b/tests/unit/callbacks/cost/test_dimensions.py @@ -30,11 +30,10 @@ class TestRunDim(RunCostDimensionBase): @pytest.mark.asyncio async def test_cost_per_input_token(): spec = { - "_type": "InputTokens", "price_per_million": 30, "granularity": 10, } - dim_valid = InputTokens.from_dict(spec) + dim_valid = InputTokens(**spec) success_response = InvocationResponse(response_text="hi", num_tokens_input=199999) assert await dim_valid.calculate(success_response) == 6 @@ -47,8 +46,7 @@ async def test_cost_per_input_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.to_dict() == { - "_type": "InputTokens", + assert dim_valid._get_llmeter_state() == { "price_per_million": 30, "granularity": 10, } @@ -57,11 +55,10 @@ async def test_cost_per_input_token(): @pytest.mark.asyncio async def test_cost_per_output_token(): spec = { - "_type": "OutputTokens", "price_per_million": 40, "granularity": 10, } - dim_valid = OutputTokens.from_dict(spec) + dim_valid = OutputTokens(**spec) success_response = InvocationResponse(response_text="hi", num_tokens_output=199999) assert await dim_valid.calculate(success_response) == 8 @@ -74,8 +71,7 @@ async def test_cost_per_output_token(): err_response = InvocationResponse.error_output() assert await dim_valid.calculate(err_response) is None - assert dim_valid.to_dict() == { - "_type": "OutputTokens", + assert dim_valid._get_llmeter_state() == { "price_per_million": 40, "granularity": 10, } @@ -84,11 +80,10 @@ async def test_cost_per_output_token(): @pytest.mark.asyncio async def test_cost_per_hour(): spec = { - "_type": "EndpointTime", "price_per_hour": 30, "granularity_secs": 60, } - dim_valid = EndpointTime.from_dict(spec) + dim_valid = EndpointTime(**spec) result = Result( [], @@ -105,8 +100,7 @@ async def test_cost_per_hour(): result.total_test_time = None assert await dim_valid.calculate(result) is None - assert dim_valid.to_dict() == { - "_type": "EndpointTime", + assert dim_valid._get_llmeter_state() == { "price_per_hour": 30, "granularity_secs": 60, } diff --git a/tests/unit/callbacks/cost/test_model.py b/tests/unit/callbacks/cost/test_model.py index b1ea1a8..fa31532 100644 --- a/tests/unit/callbacks/cost/test_model.py +++ b/tests/unit/callbacks/cost/test_model.py @@ -1,45 +1,111 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import json from dataclasses import dataclass from unittest.mock import AsyncMock, Mock, NonCallableMock import pytest +from llmeter.serialization import dump_object, load_object from llmeter.callbacks.cost.model import CostModel from llmeter.callbacks.cost.results import CalculatedCostWithDimensions +from llmeter.endpoints.base import InvocationResponse def test_cost_model_serialization(): - """Cost models can be serialized & de-serialized""" - spec = { - "_type": "CostModel", - "request_dims": { - "TokensIn": {"_type": "InputTokens", "price_per_million": 30}, - }, - "run_dims": { - "ComputeSeconds": {"_type": "EndpointTime", "price_per_hour": 50}, - }, - } - model = CostModel.from_dict(spec) - assert model.request_dims["TokensIn"].price_per_million == 30 - assert model.run_dims["ComputeSeconds"].price_per_hour == 50 - assert model.to_dict() == { - "_type": "CostModel", - "request_dims": { - "TokensIn": { - "_type": "InputTokens", - "price_per_million": 30, - "granularity": 1, - }, - }, - "run_dims": { - "ComputeSeconds": { - "_type": "EndpointTime", - "price_per_hour": 50, - "granularity_secs": 1, - }, + """Cost models can be serialized & de-serialized via dump_object/load_object""" + from llmeter.callbacks.cost.dimensions import EndpointTime, InputTokens + + model = CostModel( + request_dims={"TokensIn": InputTokens(price_per_million=30)}, + run_dims={"ComputeSeconds": EndpointTime(price_per_hour=50)}, + ) + + # Round-trip via dump_object / load_object + data = dump_object(model) + restored = load_object(data) + + assert restored.request_dims["TokensIn"].price_per_million == 30 + assert restored.run_dims["ComputeSeconds"].price_per_hour == 50 + + # get state produces a plain dict representation + d = model._get_llmeter_state() + assert "request_dims" in d + assert "run_dims" in d + + +def test_cost_model_save_load_file_roundtrip(tmp_path): + """A CostModel can be saved to and loaded from a file, re-initialising its dimensions + + Exercises the inherited Serializable.save_to_file / load_from_file API (file I/O plus + json_default), which the in-memory dump_object/load_object test does not cover. + """ + from llmeter.callbacks.cost.dimensions import ( + EndpointTime, + InputTokens, + OutputTokens, + ) + + model = CostModel( + request_dims={ + "TokensIn": InputTokens(price_per_million=30, granularity=10), + "TokensOut": OutputTokens(price_per_million=60), }, - } + run_dims={"ComputeSeconds": EndpointTime(price_per_hour=50)}, + ) + + path = tmp_path / "cost_model.json" + saved_path = model.save_to_file(path) + # save_to_file returns the (validated/normalized) path it wrote to + assert str(saved_path) == str(path) + assert path.is_file() + + # The file is valid JSON tagged with the CostModel class path + with open(path) as f: + raw = json.load(f) + assert raw["__llmeter_class__"] == "llmeter.callbacks.cost.model.CostModel" + + restored = CostModel.load_from_file(path) + + # The correct concrete type is reconstructed from the file + assert isinstance(restored, CostModel) + + # Each dimension is re-initialised as the correct type, under the same name, with values intact + assert set(restored.request_dims) == {"TokensIn", "TokensOut"} + assert set(restored.run_dims) == {"ComputeSeconds"} + + tokens_in = restored.request_dims["TokensIn"] + assert isinstance(tokens_in, InputTokens) + assert tokens_in.price_per_million == 30 + assert tokens_in.granularity == 10 + + tokens_out = restored.request_dims["TokensOut"] + assert isinstance(tokens_out, OutputTokens) + assert tokens_out.price_per_million == 60 + + compute = restored.run_dims["ComputeSeconds"] + assert isinstance(compute, EndpointTime) + assert compute.price_per_hour == 50 + + +def test_cost_model_load_from_file_dispatches_via_base_class(tmp_path): + """load_from_file resolves the concrete type from the file, even when called on Serializable + + The class is detected from the ``__llmeter_class__`` marker rather than the class the + classmethod is invoked on, so loading via the base mixin still yields a CostModel. + """ + from llmeter.callbacks.cost.dimensions import InputTokens + from llmeter.serialization import Serializable + + model = CostModel(request_dims={"TokensIn": InputTokens(price_per_million=15)}) + + path = tmp_path / "cost_model.json" + model.save_to_file(path) + + restored = Serializable.load_from_file(path) + assert isinstance(restored, CostModel) + assert restored.request_dims["TokensIn"].price_per_million == 15 + assert restored.run_dims == {} def test_cost_model_detects_duplicate_cost_dim_names(): @@ -79,7 +145,7 @@ class DummyCostDimension: @pytest.mark.asyncio async def test_cost_model_callback_saves_request_costs(): - """By default, CostModel callbacks save request cost calculations to InvocationResponse""" + """By default, CostModel callbacks save request cost calculations to response.annotations""" dummy_req_dim = Mock() dummy_req_dim.calculate = AsyncMock(return_value=42) @@ -88,17 +154,54 @@ async def test_cost_model_callback_saves_request_costs(): run_dims=[], ) - response_mock = NonCallableMock() - assert await model.after_invoke(response_mock) is None - assert response_mock.cost_total == 42 - assert response_mock.cost_Mock == 42 # Class name is the default dimension name + response = InvocationResponse(response_text="hi") + assert await model.after_invoke(response) is None + # Costs are stored in the (persisted) annotations dict, not as loose attributes + assert response.annotations["cost_total"] == 42 + assert ( + response.annotations["cost_Mock"] == 42 + ) # Class name is the default dimension name # Check calculate_* fn produces same result as callback: assert await model.calculate_request_cost( - response_mock + response ) == CalculatedCostWithDimensions.load_from_namespace( - response_mock, key_prefix="cost_" + response.annotations, key_prefix="cost_" + ) + + +@pytest.mark.asyncio +async def test_cost_model_request_costs_survive_response_roundtrip(): + """Per-response costs saved by CostModel persist through InvocationResponse to_json/from_json. + + This is the behavior that regressed when response serialization moved to ``asdict`` (which + drops loose attributes): costs stored in ``annotations`` must round-trip so a saved Result can + be reloaded with its per-request costs intact. + """ + from llmeter.callbacks.cost.dimensions import InputTokens, OutputTokens + + model = CostModel( + request_dims=[ + InputTokens(price_per_million=3.0), + OutputTokens(price_per_million=15.0), + ] + ) + response = InvocationResponse( + response_text="hello", num_tokens_input=1000, num_tokens_output=500 + ) + await model.after_invoke(response) + assert response.annotations["cost_total"] == pytest.approx(0.003 + 0.0075) + + # Round-trip through JSON (what gets written to responses.jsonl) + restored = InvocationResponse.from_json(response.to_json()) + assert restored.annotations == response.annotations + + # ...and the cost model can read the costs back off the restored response + reloaded = CalculatedCostWithDimensions.load_from_namespace( + restored.annotations, key_prefix="cost_" ) + assert reloaded["InputTokens"] == pytest.approx(0.003) + assert reloaded["OutputTokens"] == pytest.approx(0.0075) @pytest.mark.asyncio @@ -152,13 +255,17 @@ async def test_cost_model_combines_req_and_run_dims(): # Run the dummy test: run_mock = NonCallableMock() await model.before_run(run_mock) - response_mocks = [NonCallableMock(), NonCallableMock(), NonCallableMock()] - for r in response_mocks: + responses = [ + InvocationResponse(response_text="a"), + InvocationResponse(response_text="b"), + InvocationResponse(response_text="c"), + ] + for r in responses: await model.after_invoke(r) results_mock = NonCallableMock() update_contrib_stats_mock = Mock() results_mock._update_contributed_stats = update_contrib_stats_mock - results_mock.responses = response_mocks + results_mock.responses = responses results_mock.additional_metrics_for_aggregation = None await model.after_run(results_mock) diff --git a/tests/unit/callbacks/test_base.py b/tests/unit/callbacks/test_base.py index b3aa824..2b450bf 100644 --- a/tests/unit/callbacks/test_base.py +++ b/tests/unit/callbacks/test_base.py @@ -1,51 +1,46 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest +import json from llmeter.callbacks.base import Callback +from llmeter.callbacks.mlflow import MlflowCallback 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") - - assert ( - str(excinfo.value) - == "TODO: Callback.load_from_file is not yet implemented!" - ) - - 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") - - assert ( - str(excinfo.value) - == "TODO: Callback._load_from_file is not yet implemented!" - ) + def test_save_to_file_creates_valid_json(self, tmp_path): + """save_to_file creates a valid JSON file with __llmeter_class__ and __llmeter_state__.""" + cb = MlflowCallback(step=5, nested=True) + path = tmp_path / "callback.json" + cb.save_to_file(path) + + with open(path) as f: + data = json.load(f) + + assert "__llmeter_class__" in data + assert "__llmeter_state__" in data + assert data["__llmeter_class__"] == "llmeter.callbacks.mlflow.MlflowCallback" + assert data["__llmeter_state__"] == {"step": 5, "nested": True} + + def test_load_from_file_restores_callback(self, tmp_path): + """load_from_file correctly restores a callback instance.""" + cb = MlflowCallback(step=3, nested=False) + path = tmp_path / "callback.json" + cb.save_to_file(path) + + restored = Callback.load_from_file(path) + + assert isinstance(restored, MlflowCallback) + assert restored.step == 3 + assert restored.nested is False + + def test_getstate_setstate_roundtrip(self): + """get state / set state round-trips correctly.""" + cb = MlflowCallback(step=7, nested=True) + state = cb._get_llmeter_state() + + restored = MlflowCallback.__new__(MlflowCallback) + restored._set_llmeter_state(state) + + assert restored.step == 7 + assert restored.nested is True diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..871b09a --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,16 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared fixtures for unit tests.""" + +from pathlib import Path + +import pytest + +RESULT_SNAPSHOTS_DIR = Path(__file__).parent / "fixtures" / "result_snapshots" + + +@pytest.fixture +def snapshots_dir() -> Path: + """Base directory holding the static synthetic result snapshots.""" + return RESULT_SNAPSHOTS_DIR diff --git a/tests/unit/endpoints/test_anthropic_messages.py b/tests/unit/endpoints/test_anthropic_messages.py index 3273129..00b59a6 100644 --- a/tests/unit/endpoints/test_anthropic_messages.py +++ b/tests/unit/endpoints/test_anthropic_messages.py @@ -1,11 +1,16 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +# Python Built-Ins: +from pathlib import Path +import tempfile import time from unittest.mock import Mock, patch +# External Dependencies: import pytest +# Local Dependencies: from llmeter.endpoints.anthropic_messages import ( _ANTHROPIC_CLIENTS, AnthropicMessages, @@ -457,9 +462,7 @@ def test_invoke_success(self, mock_client): assert result.error is None def test_invoke_api_error(self, mock_client): - mock_client.return_value.messages.create.side_effect = Exception( - "Stream error" - ) + mock_client.return_value.messages.create.side_effect = Exception("Stream error") endpoint = AnthropicMessagesStream(model_id="test-model") result = endpoint.invoke( @@ -705,3 +708,172 @@ def test_custom_endpoint_name(self, mock_client): model_id="test-model", endpoint_name="my-custom-endpoint" ) assert endpoint.endpoint_name == "my-custom-endpoint" + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestAnthropicSerialization: + """Test that client configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state includes provider and model.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="test") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "claude-opus-4-7" + assert state["provider"] == "anthropic" + assert state["endpoint_name"] == "anthropic-messages" + assert state["max_retries"] == 2 + assert "base_url" in state + + def test_secrets_excluded(self): + """Test that secrets (keys, tokens) are not persisted.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="sk-secret") + state = dump_object(endpoint)["__llmeter_state__"] + for key in state: + assert "key" not in key, f"Secret-like key '{key}' found in state" + assert "token" not in key, f"Secret-like key '{key}' found in state" + assert "sk-secret" not in str(state.values()) + + def test_bedrock_mantle_serializes_region(self): + """Test Bedrock Mantle provider captures aws_region.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="anthropic.claude-opus-4-7", + provider="bedrock-mantle", + aws_region="us-west-2", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["aws_region"] == "us-west-2" + assert state["provider"] == "bedrock-mantle" + assert "aws_secret_key" not in state + assert "aws_access_key" not in state + assert "aws_session_token" not in state + + def test_custom_headers_serialized(self): + """Test that user-supplied headers round-trip but SDK headers don't.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + default_headers={"X-Custom": "val"}, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["default_headers"] == {"X-Custom": "val"} + + def test_no_headers_when_none_custom(self): + """Test that default_headers is absent when no custom headers set.""" + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages(model_id="claude-opus-4-7", api_key="test") + state = dump_object(endpoint)["__llmeter_state__"] + assert "default_headers" not in state + + def test_timeout_serialized_as_dict(self): + """Test httpx.Timeout is serialized as a dict.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + timeout=httpx.Timeout(30.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 30.0, + "write": 30.0, + "pool": 30.0, + } + + def test_round_trip_direct_provider(self): + """Test save/load round-trip for direct Anthropic provider.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + max_retries=5, + base_url="https://custom.anthropic.com", + default_headers={"X-Foo": "bar"}, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessages) + assert loaded.model_id == "claude-opus-4-7" + assert loaded.provider == "anthropic" + assert loaded._client.max_retries == 5 + assert "custom.anthropic.com" in str(loaded._client.base_url) + assert loaded._client._custom_headers == {"X-Foo": "bar"} + + def test_round_trip_bedrock_mantle(self): + """Test save/load round-trip for Bedrock Mantle provider.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="anthropic.claude-opus-4-7", + provider="bedrock-mantle", + aws_region="eu-west-1", + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessages) + assert loaded.provider == "bedrock-mantle" + assert loaded._client.aws_region == "eu-west-1" + + def test_round_trip_stream_endpoint(self): + """Test save/load round-trip for streaming endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessagesStream( + model_id="claude-opus-4-7", + api_key="test", + ttft_visible_tokens_only=False, + max_retries=4, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, AnthropicMessagesStream) + assert loaded.ttft_visible_tokens_only is False + assert loaded._client.max_retries == 4 + + def test_round_trip_with_granular_timeout(self): + """Test httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = AnthropicMessages( + model_id="claude-opus-4-7", + api_key="test", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) diff --git a/tests/unit/endpoints/test_bedrock.py b/tests/unit/endpoints/test_bedrock.py index 1dd8066..b3949a0 100644 --- a/tests/unit/endpoints/test_bedrock.py +++ b/tests/unit/endpoints/test_bedrock.py @@ -833,3 +833,106 @@ def dropping_stream(): assert response.error is not None assert "connection" in response.error.lower() assert response.input_payload is not None + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestBedrockSerialization: + """Test that Bedrock endpoint configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state when only required params are set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockConverse(model_id="test-model") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "test-model" + assert state["max_attempts"] == 3 + assert state["inference_config"] is None + assert "region" in state + + def test_serializes_all_fields(self): + """Test serialized state includes all optional fields when set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockConverse( + model_id="amazon.nova-pro-v1:0", + region="us-west-2", + max_attempts=7, + inference_config={"maxTokens": 512, "temperature": 0.7}, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "amazon.nova-pro-v1:0" + assert state["region"] == "us-west-2" + assert state["max_attempts"] == 7 + assert state["inference_config"] == {"maxTokens": 512, "temperature": 0.7} + + def test_round_trip_converse_all_fields(self): + """Test save/load round-trip preserves all fields for BedrockConverse.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverse( + model_id="amazon.nova-pro-v1:0", + region="us-west-2", + max_attempts=5, + inference_config={"maxTokens": 1024}, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverse) + assert loaded.model_id == "amazon.nova-pro-v1:0" + assert loaded.region == "us-west-2" + assert loaded._max_attempts == 5 + assert loaded._inference_config == {"maxTokens": 1024} + + def test_round_trip_converse_stream(self): + """Test save/load round-trip preserves all fields for BedrockConverseStream.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverseStream( + model_id="amazon.nova-lite-v1:0", + region="eu-west-1", + max_attempts=10, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverseStream) + assert loaded.model_id == "amazon.nova-lite-v1:0" + assert loaded.region == "eu-west-1" + assert loaded._max_attempts == 10 + + def test_round_trip_without_optional_fields(self): + """Test save/load works with only required params.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockConverse(model_id="test-model") + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockConverse) + assert loaded.model_id == "test-model" + assert loaded._max_attempts == 3 + assert loaded._inference_config is None diff --git a/tests/unit/endpoints/test_bedrock_invoke.py b/tests/unit/endpoints/test_bedrock_invoke.py index af101bc..2a76fbb 100644 --- a/tests/unit/endpoints/test_bedrock_invoke.py +++ b/tests/unit/endpoints/test_bedrock_invoke.py @@ -655,3 +655,104 @@ def dropping_body(): assert response.error is not None assert "connection" in response.error.lower() assert response.input_payload is not None + + +# --------------------------------------------------------------------------- +# Tests: Serialization +# --------------------------------------------------------------------------- + + +class TestBedrockInvokeSerialization: + """Test that BedrockInvoke endpoint configuration round-trips through serialization.""" + + def test_serializes_basic_fields(self): + """Test serialized state includes required fields with defaults.""" + from llmeter.serialization import dump_object + + endpoint = BedrockInvoke(model_id="test-model") + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "test-model" + assert state["generated_text_jmespath"] == "choices[0].message.content" + assert state["input_text_jmespath"] == "messages[].content[].text" + assert state["max_attempts"] == 3 + assert state["generated_token_count_jmespath"] == "usage.completion_tokens" + assert state["input_token_count_jmespath"] == "usage.prompt_tokens" + assert "region" in state + + def test_serializes_all_fields(self): + """Test serialized state includes all optional fields when set.""" + from llmeter.serialization import dump_object + + endpoint = BedrockInvoke( + model_id="amazon.titan-text-v1", + generated_text_jmespath="results[0].outputText", + input_text_jmespath="inputText", + generated_token_count_jmespath="results[0].tokenCount", + input_token_count_jmespath="inputTokenCount", + region="us-west-2", + max_attempts=7, + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["model_id"] == "amazon.titan-text-v1" + assert state["generated_text_jmespath"] == "results[0].outputText" + assert state["input_text_jmespath"] == "inputText" + assert state["generated_token_count_jmespath"] == "results[0].tokenCount" + assert state["input_token_count_jmespath"] == "inputTokenCount" + assert state["region"] == "us-west-2" + assert state["max_attempts"] == 7 + + def test_round_trip_invoke_all_fields(self): + """Test save/load round-trip preserves all fields for BedrockInvoke.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockInvoke( + model_id="amazon.titan-text-v1", + generated_text_jmespath="results[0].outputText", + input_text_jmespath="inputText", + generated_token_count_jmespath="results[0].tokenCount", + region="us-west-2", + max_attempts=5, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockInvoke) + assert loaded.model_id == "amazon.titan-text-v1" + assert loaded.generated_text_jmespath == "results[0].outputText" + assert loaded.input_text_jmespath == "inputText" + assert loaded.generated_token_count_jmespath == "results[0].tokenCount" + assert loaded.region == "us-west-2" + assert loaded._max_attempts == 5 + + def test_round_trip_invoke_stream(self): + """Test save/load round-trip preserves all fields for BedrockInvokeStream.""" + import tempfile + from pathlib import Path + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = BedrockInvokeStream( + model_id="amazon.titan-text-v1", + generated_text_jmespath="outputText", + input_text_jmespath="inputText", + region="eu-west-1", + max_attempts=8, + ) + path = Path(tmpdir) / "endpoint.json" + original.save_to_file(path) + + loaded = Endpoint.load_from_file(path) + + assert isinstance(loaded, BedrockInvokeStream) + assert loaded.model_id == "amazon.titan-text-v1" + assert loaded.generated_text_jmespath == "outputText" + assert loaded.input_text_jmespath == "inputText" + assert loaded.region == "eu-west-1" + assert loaded._max_attempts == 8 diff --git a/tests/unit/endpoints/test_endpoints.py b/tests/unit/endpoints/test_endpoints.py index bdaad99..2c54445 100644 --- a/tests/unit/endpoints/test_endpoints.py +++ b/tests/unit/endpoints/test_endpoints.py @@ -1,6 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +import json + import pytest import llmeter @@ -186,6 +188,57 @@ def test_endpoint_load_from_dict(): assert loaded_endpoint.provider == "test_provider" +def test_invocation_response_annotations_roundtrip(): + """The `annotations` dict round-trips through to_json/from_json.""" + resp = InvocationResponse( + response_text="hi", annotations={"cost_total": 0.01, "label": "x"} + ) + restored = InvocationResponse.from_json(resp.to_json()) + assert restored.annotations == {"cost_total": 0.01, "label": "x"} + + +def test_invocation_response_annotations_default_is_empty_dict(): + """annotations defaults to an independent empty dict per instance.""" + a = InvocationResponse(response_text="a") + b = InvocationResponse(response_text="b") + assert a.annotations == {} + a.annotations["x"] = 1 + assert b.annotations == {} # no shared mutable default + + +def test_from_json_collects_unknown_fields_into_annotations(): + """Unknown top-level keys (e.g. legacy callback-written fields) are captured into annotations.""" + raw = json.dumps( + { + "response_text": "hi", + "id": "r1", + "num_tokens_input": 10, + "cost_total": 0.02, # legacy callback-written field + "cost_InputTokens": 0.01, + "end_time": "2024-01-01T00:00:00Z", # since-removed field + } + ) + resp = InvocationResponse.from_json(raw) + + # Known fields are parsed normally + assert resp.response_text == "hi" + assert resp.id == "r1" + assert resp.num_tokens_input == 10 + # Unknown fields are preserved (not dropped, not raising) under annotations + assert resp.annotations == { + "cost_total": 0.02, + "cost_InputTokens": 0.01, + "end_time": "2024-01-01T00:00:00Z", + } + + +def test_from_json_merges_unknown_fields_with_existing_annotations(): + """Explicit annotations and stray unknown keys both survive load.""" + raw = json.dumps({"response_text": "hi", "annotations": {"kept": 1}, "stray": 2}) + resp = InvocationResponse.from_json(raw) + assert resp.annotations == {"kept": 1, "stray": 2} + + def test_invocation_response_to_dict(): response = InvocationResponse( id="test_id", diff --git a/tests/unit/endpoints/test_multimodal_properties.py b/tests/unit/endpoints/test_multimodal_properties.py index a5f7c3d..4b0fbbc 100644 --- a/tests/unit/endpoints/test_multimodal_properties.py +++ b/tests/unit/endpoints/test_multimodal_properties.py @@ -17,7 +17,7 @@ from hypothesis import strategies as st from llmeter.endpoints.bedrock import BedrockBase -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default from llmeter.prompt_utils import ( AudioContent, DocumentContent, @@ -602,11 +602,11 @@ def create_multimodal_payload(input_text, **kwargs): assert len(image_blocks) == 1 # Verify the payload can be serialized - json_str = json.dumps(payload, default=llmeter_default_serializer) + json_str = json.dumps(payload, default=json_default) assert len(json_str) > 0 # Verify it can be deserialized - deserialized = json.loads(json_str, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(json_str, object_hook=bytes_decoder) assert deserialized == payload diff --git a/tests/unit/endpoints/test_openai.py b/tests/unit/endpoints/test_openai.py index 5934c13..a6fdc8e 100644 --- a/tests/unit/endpoints/test_openai.py +++ b/tests/unit/endpoints/test_openai.py @@ -1,6 +1,8 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from pathlib import Path +import tempfile import time from unittest.mock import MagicMock, Mock, patch @@ -35,6 +37,8 @@ def test_initialization(self): assert endpoint.endpoint_name == "test_openai" assert endpoint.provider == "openai" assert endpoint._client is not None + assert endpoint._client.project is None + assert endpoint.project is None def test_initialization_with_custom_provider(self): """Test OpenAI endpoint initialization with custom provider.""" @@ -47,6 +51,18 @@ def test_initialization_with_custom_provider(self): assert endpoint.provider == "custom_openai" assert endpoint.model_id == "gpt-4" + def test_initialization_with_project_id(self): + """Test OpenAI endpoint initialization with custom project ID.""" + endpoint = OpenAICompletionEndpoint( + model_id="openai.glm-5", + api_key="test_key", + project="proj_DUMMY", + ) + + assert endpoint.model_id == "openai.glm-5" + assert endpoint._client.project == "proj_DUMMY" + assert endpoint.project == "proj_DUMMY" + def test_initialization_without_api_key(self): """Test OpenAI endpoint initialization without API key.""" endpoint = OpenAICompletionEndpoint( @@ -758,6 +774,193 @@ def test_stream_response_usage_none_value(self): assert response.num_tokens_output is None +class TestOpenAIEndpointSerialization: + """Test serialization behaves as expected for both endpoint types.""" + + def test_completion_endpoint_serializes_basic_fields(self): + """Test sync endpoint serialized state when optional parameters not set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "openai" + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + assert "base_url" in state + assert state["max_retries"] == 2 + # timeout is always serialized (as dict when it's an httpx.Timeout) + assert "timeout" in state + + def test_completion_endpoint_serializes_all_client_fields(self): + """Test sync endpoint serializes all client configuration fields.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint( + model_id="openai.gpt-oss-120b", + api_key="test_key", + endpoint_name="test-endpoint", + organization="org-abc", + project="proj_TEST", + base_url="https://custom.example.com/v1", + websocket_base_url="wss://custom.example.com/ws", + timeout=30.0, + max_retries=5, + default_headers={"X-Custom": "val"}, + default_query={"version": "2"}, + provider="test-provider", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "test-endpoint" + assert state["model_id"] == "openai.gpt-oss-120b" + assert state["organization"] == "org-abc" + assert state["project"] == "proj_TEST" + assert "custom.example.com" in state["base_url"] + assert state["websocket_base_url"] == "wss://custom.example.com/ws" + assert state["timeout"] == 30.0 + assert state["max_retries"] == 5 + assert state["default_headers"] == {"X-Custom": "val"} + assert state["default_query"] == {"version": "2"} + assert state["provider"] == "test-provider" + + def test_timeout_serialized_as_dict_when_granular(self): + """Test httpx.Timeout is serialized as a dict with connect/read/write/pool.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 10.0, + "write": 10.0, + "pool": 10.0, + } + + def test_completion_endpoint_round_trip(self): + """Test save/load round-trip preserves all fields for completion endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="proj-test", + organization="org-rt", + project="proj_roundtrip", + base_url="https://custom.test/v1", + websocket_base_url="wss://custom.test/ws", + timeout=45.0, + max_retries=3, + default_headers={"X-Foo": "bar"}, + default_query={"q": "1"}, + provider="test-provider", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.endpoint_name == "proj-test" + assert loaded._client.organization == "org-rt" + assert loaded.project == "proj_roundtrip" + assert "custom.test" in str(loaded._client.base_url) + assert loaded._client.websocket_base_url == "wss://custom.test/ws" + assert loaded._client.timeout == 45.0 + assert loaded._client.max_retries == 3 + assert loaded._client._custom_headers == {"X-Foo": "bar"} + assert loaded._client._custom_query == {"q": "1"} + assert loaded.provider == "test-provider" + + def test_round_trip_with_granular_timeout(self): + """Test that httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) + + def test_stream_endpoint_serializes_basic_fields(self): + """Test stream endpoint serialized state when optional parameters not set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAICompletionStreamEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["endpoint_name"] == "openai" + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + + def test_stream_endpoint_round_trip(self): + """Test save/load round-trip preserves all fields for streaming endpoint.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionStreamEndpoint( + model_id="zai.glm-5", + api_key="test_key", + endpoint_name="proj-test-stream", + organization="org-stream", + project="proj_roundtrip_stream", + base_url="https://stream.test/v1", + max_retries=4, + default_headers={"X-Stream": "yes"}, + provider="test-provider-stream", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionStreamEndpoint) + assert loaded.model_id == "zai.glm-5" + assert loaded.endpoint_name == "proj-test-stream" + assert loaded._client.organization == "org-stream" + assert loaded.project == "proj_roundtrip_stream" + assert "stream.test" in str(loaded._client.base_url) + assert loaded._client.max_retries == 4 + assert loaded._client._custom_headers == {"X-Stream": "yes"} + assert loaded.provider == "test-provider-stream" + + def test_round_trip_without_optional_fields(self): + """Test save/load round-trip works when only required params are set.""" + from llmeter.endpoints.base import Endpoint + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAICompletionEndpoint(model_id="gpt-4", api_key="test_key") + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAICompletionEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.project is None + assert loaded._client.organization is None + + class TestStreamMidStreamErrors: """Verify that errors during stream consumption are caught by the invoke wrapper.""" diff --git a/tests/unit/endpoints/test_openai_response_serialization.py b/tests/unit/endpoints/test_openai_response_serialization.py index fee2e94..20e0f1a 100644 --- a/tests/unit/endpoints/test_openai_response_serialization.py +++ b/tests/unit/endpoints/test_openai_response_serialization.py @@ -79,11 +79,12 @@ def test_save_creates_file_with_correct_content(self, mock_openai_class): with open(saved_path, "r") as f: data = json.load(f) - # Verify content matches endpoint configuration - assert data["model_id"] == "gpt-4-turbo" - assert data["endpoint_name"] == "saved-endpoint" - assert data["provider"] == "openai" - assert data["endpoint_type"] == "OpenAIResponseEndpoint" + # Verify envelope format with correct state + assert "__llmeter_class__" in data + state = data["__llmeter_state__"] + assert state["model_id"] == "gpt-4-turbo" + assert state["endpoint_name"] == "saved-endpoint" + assert state["provider"] == "openai" @patch("llmeter.endpoints.openai_response.OpenAI") def test_load_from_file_reconstructs_endpoint(self, mock_openai_class): @@ -223,6 +224,167 @@ def test_stream_endpoint_round_trip(self, mock_openai_class): assert loaded.provider == original.provider +class TestOpenAIResponseEndpointSerialization: + """Test that OpenAI client configuration round-trips through serialization.""" + + def test_response_endpoint_serializes_all_client_fields(self): + """Test all client config fields appear in serialized state.""" + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint( + model_id="gpt-4", + endpoint_name="test-response", + organization="org-abc", + project="proj_TEST", + base_url="https://custom.example.com/v1", + websocket_base_url="wss://custom.example.com/ws", + timeout=30.0, + max_retries=5, + default_headers={"X-Custom": "val"}, + default_query={"version": "2"}, + provider="test-provider", + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["model_id"] == "gpt-4" + assert state["organization"] == "org-abc" + assert state["project"] == "proj_TEST" + assert "custom.example.com" in state["base_url"] + assert state["websocket_base_url"] == "wss://custom.example.com/ws" + assert state["timeout"] == 30.0 + assert state["max_retries"] == 5 + assert state["default_headers"] == {"X-Custom": "val"} + assert state["default_query"] == {"version": "2"} + assert state["provider"] == "test-provider" + + def test_response_endpoint_serializes_defaults(self): + """Test serialized state is correct when only required params are set.""" + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint(model_id="gpt-4") + state = dump_object(endpoint)["__llmeter_state__"] + assert "api_key" not in state + assert state["model_id"] == "gpt-4" + assert state.get("project") is None + assert state.get("organization") is None + assert state["provider"] == "openai" + assert state["max_retries"] == 2 + assert "base_url" in state + + def test_timeout_serialized_as_dict_when_granular(self): + """Test httpx.Timeout is serialized as a dict with connect/read/write/pool.""" + import httpx + + from llmeter.serialization import dump_object + + endpoint = OpenAIResponseEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=5.0), + ) + state = dump_object(endpoint)["__llmeter_state__"] + assert state["timeout"] == { + "connect": 5.0, + "read": 10.0, + "write": 10.0, + "pool": 10.0, + } + + def test_response_endpoint_round_trip_with_all_fields(self): + """Test save/load round-trip preserves all client config fields.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="rt-test", + organization="org-rt", + project="proj_roundtrip", + base_url="https://custom.test/v1", + websocket_base_url="wss://custom.test/ws", + timeout=45.0, + max_retries=3, + default_headers={"X-Foo": "bar"}, + default_query={"q": "1"}, + provider="test-provider", + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded._client.organization == "org-rt" + assert loaded.project == "proj_roundtrip" + assert "custom.test" in str(loaded._client.base_url) + assert loaded._client.websocket_base_url == "wss://custom.test/ws" + assert loaded._client.timeout == 45.0 + assert loaded._client.max_retries == 3 + assert loaded._client._custom_headers == {"X-Foo": "bar"} + assert loaded._client._custom_query == {"q": "1"} + assert loaded.provider == "test-provider" + + def test_stream_endpoint_round_trip_with_all_fields(self): + """Test save/load round-trip preserves all client config for streaming endpoint.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseStreamEndpoint( + model_id="gpt-4", + api_key="test_key", + endpoint_name="stream-rt-test", + organization="org-stream", + project="proj_stream_rt", + base_url="https://stream.test/v1", + max_retries=4, + default_headers={"X-Stream": "yes"}, + provider="test-stream", + ttft_visible_tokens_only=False, + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseStreamEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded._client.organization == "org-stream" + assert loaded.project == "proj_stream_rt" + assert "stream.test" in str(loaded._client.base_url) + assert loaded._client.max_retries == 4 + assert loaded._client._custom_headers == {"X-Stream": "yes"} + assert loaded.provider == "test-stream" + assert loaded.ttft_visible_tokens_only is False + + def test_round_trip_with_granular_timeout(self): + """Test httpx.Timeout round-trips correctly through dict serialization.""" + import httpx + + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint( + model_id="gpt-4", + timeout=httpx.Timeout(10.0, connect=2.0), + ) + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded._client.timeout == httpx.Timeout(10.0, connect=2.0) + + def test_round_trip_without_optional_fields(self): + """Test save/load works when only required params are set.""" + with tempfile.TemporaryDirectory() as tmpdir: + original = OpenAIResponseEndpoint(model_id="gpt-4", api_key="test_key") + output_path = Path(tmpdir) / "endpoint.json" + original.save(output_path) + + loaded = Endpoint.load_from_file(output_path) + + assert isinstance(loaded, OpenAIResponseEndpoint) + assert loaded.model_id == "gpt-4" + assert loaded.project is None + assert loaded._client.organization is None + + class TestSerializationPropertyTests: """Property-based tests for endpoint serialization. diff --git a/tests/unit/fixtures/_generate_fixtures.py b/tests/unit/fixtures/_generate_fixtures.py new file mode 100644 index 0000000..30cee45 --- /dev/null +++ b/tests/unit/fixtures/_generate_fixtures.py @@ -0,0 +1,901 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-shot script to generate fixture data with mathematically consistent stats. + +Run once from the repo root: + uv run python tests/unit/fixtures/_generate_fixtures.py + +This writes all fixture files under tests/unit/fixtures/result_snapshots/. +After running, verify the tests pass, then this script can be kept as documentation +of how the fixtures were generated (or deleted if preferred). +""" + +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from llmeter.endpoints.base import InvocationResponse +from llmeter.results import Result +from llmeter.serialization import json_default, restore_dataclass_types + +FIXTURES_DIR = Path(__file__).parent / "result_snapshots" + + +def _write_json(path: Path, data: dict): + path.write_text(json.dumps(data, indent=4, default=json_default) + "\n") + + +def _write_jsonl(path: Path, lines: list[str]): + path.write_text("\n".join(lines) + "\n") + + +def _compute_and_write_stats( + responses: list[InvocationResponse], summary: dict, path: Path +): + """Build a Result from the summary metadata, compute stats, write stats.json. + + Passing the full summary metadata (parsed to native types) makes the resulting + stats.json mirror what a real ``Result.save()`` produces. + """ + meta = {k: v for k, v in summary.items() if k not in ("total_requests",)} + restore_dataclass_types(Result, meta) + result = Result( + responses=responses, total_requests=summary["total_requests"], **meta + ) + stats = Result._compute_stats(result) + _write_json(path, stats) + return stats + + +# ============================================================================= +# Scenario: base (modern format, OpenAI endpoint, CostModel + Mlflow callbacks) +# ============================================================================= + + +def generate_base(): + out = FIXTURES_DIR / "base" + + responses = [ + InvocationResponse( + id="synth-base-001", + response_text="Lorem ipsum dolor sit amet, consectetur adipiscing elit.", + input_prompt="Summarize the key concepts of distributed systems.", + input_payload={ + "messages": [ + { + "role": "user", + "content": "Summarize the key concepts of distributed systems.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.452, + time_to_last_token=1.203, + num_tokens_input=42, + num_tokens_output=87, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00863, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-002", + response_text="Sed ut perspiciatis unde omnis iste natus error sit voluptatem.", + input_prompt="Explain the CAP theorem briefly.", + input_payload={ + "messages": [ + {"role": "user", "content": "Explain the CAP theorem briefly."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.389, + time_to_last_token=0.951, + num_tokens_input=38, + num_tokens_output=64, + num_tokens_input_cached=12, + num_tokens_output_reasoning=None, + time_per_output_token=0.00879, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 1, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-003", + response_text="Nemo enim ipsam voluptatem quia voluptas sit aspernatur.", + input_prompt="What is eventual consistency?", + input_payload={ + "messages": [ + {"role": "user", "content": "What is eventual consistency?"} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.521, + time_to_last_token=1.445, + num_tokens_input=35, + num_tokens_output=102, + num_tokens_input_cached=None, + num_tokens_output_reasoning=8, + time_per_output_token=0.00906, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 2, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-004", + response_text="Ut enim ad minima veniam, quis nostrum exercitationem.", + input_prompt="Define partition tolerance.", + input_payload={ + "messages": [ + {"role": "user", "content": "Define partition tolerance."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.298, + time_to_last_token=0.812, + num_tokens_input=30, + num_tokens_output=55, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00934, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 3, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-base-005", + response_text="Quis autem vel eum iure reprehenderit qui in ea voluptate.", + input_prompt="Compare synchronous and asynchronous replication.", + input_payload={ + "messages": [ + { + "role": "user", + "content": "Compare synchronous and asynchronous replication.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + time_to_first_token=0.415, + time_to_last_token=1.102, + num_tokens_input=40, + num_tokens_output=78, + num_tokens_input_cached=None, + num_tokens_output_reasoning=None, + time_per_output_token=0.00881, + error=None, + retries=0, + request_time=datetime(2025, 1, 15, 10, 0, 4, tzinfo=timezone.utc), + annotations={}, + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": None, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json (modern __llmeter_class__ envelopes) + run_config = { + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.openai.OpenAICompletionStreamEndpoint", + "__llmeter_state__": { + "model_id": "synthetic.test-model-v1", + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "organization": None, + "base_url": "https://api.synthetic-provider.example.com/v1/", + "max_retries": 3, + "timeout": { + "connect": 5.0, + "read": 60.0, + "write": 5.0, + "pool": 5.0, + }, + "api_key": "sk-synthetic-test-key-not-real", + }, + }, + "output_path": "stale/path/that/should/be/overridden", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {}, + }, + "clients": 2, + "n_requests": 3, + "run_duration": None, + "payload": "tests/unit/fixtures/result_snapshots/base/payload.jsonl", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "timeout": 60, + "callbacks": [ + { + "__llmeter_class__": "llmeter.callbacks.cost.model.CostModel", + "__llmeter_state__": { + "request_dims": { + "InputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.InputTokens", + "__llmeter_state__": { + "price_per_million": 3.0, + "granularity": 1, + }, + }, + "OutputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.OutputTokens", + "__llmeter_state__": { + "price_per_million": 15.0, + "granularity": 1, + }, + }, + }, + "run_dims": {}, + }, + }, + { + "__llmeter_class__": "llmeter.callbacks.mlflow.MlflowCallback", + "__llmeter_state__": { + "step": None, + "nested": True, + }, + }, + ], + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + # payload.jsonl + payloads = [ + { + "messages": [ + { + "role": "user", + "content": "Summarize the key concepts of distributed systems.", + } + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + { + "messages": [ + {"role": "user", "content": "Explain the CAP theorem briefly."} + ], + "max_tokens": 256, + "model": "synthetic.test-model-v1", + "stream": True, + "stream_options": {"include_usage": True}, + }, + ] + _write_jsonl(out / "payload.jsonl", [json.dumps(p) for p in payloads]) + + print(f" base/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: legacy/v0_1_endpoint_type +# ============================================================================= + + +def generate_legacy_endpoint_type(): + out = FIXTURES_DIR / "legacy" / "v0_1_endpoint_type" + + responses = [ + InvocationResponse( + id="synth-legacy-001", + response_text="Consectetur adipiscing elit, sed do eiusmod tempor.", + input_prompt="What is a load balancer?", + time_to_first_token=0.612, + time_to_last_token=1.534, + num_tokens_input=28, + num_tokens_output=72, + time_per_output_token=0.01280, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 0, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-002", + response_text="Duis aute irure dolor in reprehenderit in voluptate velit.", + input_prompt="Explain horizontal scaling.", + time_to_first_token=0.548, + time_to_last_token=1.289, + num_tokens_input=24, + num_tokens_output=65, + time_per_output_token=0.01140, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 1, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-003", + response_text="Excepteur sint occaecat cupidatat non proident.", + input_prompt="What is auto-scaling?", + time_to_first_token=0.701, + time_to_last_token=1.678, + num_tokens_input=22, + num_tokens_output=58, + time_per_output_token=0.01684, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 2, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-004", + response_text="Sunt in culpa qui officia deserunt mollit anim id est laborum.", + input_prompt="Define microservices architecture.", + time_to_first_token=0.489, + time_to_last_token=1.102, + num_tokens_input=26, + num_tokens_output=71, + time_per_output_token=0.00863, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 3, tzinfo=timezone.utc), + ), + InvocationResponse( + id="synth-legacy-005", + response_text="Neque porro quisquam est qui dolorem ipsum.", + input_prompt="What is service mesh?", + time_to_first_token=0.555, + time_to_last_token=1.345, + num_tokens_input=20, + num_tokens_output=48, + time_per_output_token=0.01646, + error=None, + request_time=datetime(2025, 1, 10, 14, 0, 4, tzinfo=timezone.utc), + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": None, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl — legacy format with flat cost annotations as top-level keys + legacy_responses = [] + costs = [ + { + "cost_InputTokens": 0.000084, + "cost_OutputTokens": 0.00108, + "cost_total": 0.001164, + }, + { + "cost_InputTokens": 0.000072, + "cost_OutputTokens": 0.000975, + "cost_total": 0.001047, + }, + { + "cost_InputTokens": 0.000066, + "cost_OutputTokens": 0.00087, + "cost_total": 0.000936, + }, + { + "cost_InputTokens": 0.000078, + "cost_OutputTokens": 0.001065, + "cost_total": 0.001143, + }, + { + "cost_InputTokens": 0.00006, + "cost_OutputTokens": 0.00072, + "cost_total": 0.00078, + }, + ] + for resp, cost in zip(responses, costs): + d = json.loads(resp.to_json()) + # Remove modern fields that didn't exist in v0.1 + d.pop("num_tokens_input_cached", None) + d.pop("num_tokens_output_reasoning", None) + d.pop("retries", None) + d.pop("annotations", None) + # Add legacy flat cost annotations + d.update(cost) + legacy_responses.append(json.dumps(d)) + _write_jsonl(out / "responses.jsonl", legacy_responses) + + # run_config.json — legacy format + run_config = { + "endpoint": { + "endpoint_name": "synthetic-legacy-endpoint", + "model_id": "synthetic.legacy-model-v1", + "provider": "bedrock", + "region": "us-west-2", + "endpoint_type": "BedrockConverse", + }, + "output_path": "stale/legacy/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 3, + "n_requests": 2, + "payload": "stale/legacy/payload.jsonl", + "run_name": "synthetic-legacy-test", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" legacy/v0_1_endpoint_type/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: legacy/v0_1_str_callbacks +# ============================================================================= + + +def generate_legacy_str_callbacks(): + out = FIXTURES_DIR / "legacy" / "v0_1_str_callbacks" + + responses = [ + InvocationResponse( + id=f"synth-strcb-{i:03d}", + response_text=f"Synthetic response number {i} for string callback test.", + input_prompt=f"Synthetic prompt {i}.", + time_to_first_token=0.4 + 0.1 * i, + time_to_last_token=1.0 + 0.2 * i, + num_tokens_input=30 + 5 * i, + num_tokens_output=50 + 10 * i, + time_per_output_token=0.01, + error=None, + request_time=datetime(2025, 1, 12, 9, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": None, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z", + } + _write_json(out / "summary.json", summary) + + # stats.json + _compute_and_write_stats(responses, summary, out / "stats.json") + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json — legacy with string repr callbacks + run_config = { + "endpoint": { + "endpoint_name": "synthetic-strcb-endpoint", + "model_id": "synthetic.strcb-model-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream", + }, + "output_path": "stale/strcb/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 1, + "n_requests": 5, + "payload": "stale/strcb/payload.jsonl", + "run_name": "synthetic-strcb-test", + "run_description": None, + "timeout": 60, + "callbacks": [ + "<__main__.SyntheticCallback object at 0x1234567890>", + "", + ], + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" legacy/v0_1_str_callbacks/ done ({len(responses)} responses)") + + +# ============================================================================= +# Scenario: interrupted_run (no summary.json) +# ============================================================================= + + +def generate_interrupted_run(): + out = FIXTURES_DIR / "interrupted_run" + + responses = [ + InvocationResponse( + id=f"synth-intr-{i:03d}", + response_text=f"Interrupted run response {i}: tempor incididunt ut labore.", + input_prompt=f"Interrupted test prompt {i}.", + time_to_first_token=0.3 + 0.05 * i, + time_to_last_token=0.8 + 0.15 * i, + num_tokens_input=25 + 3 * i, + num_tokens_output=40 + 8 * i, + time_per_output_token=0.009, + error=None, + retries=0, + request_time=datetime(2025, 1, 20, 16, 30, i * 2, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # NO summary.json — this is the point of this scenario + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # run_config.json — modern __llmeter_class__ format (what a *current* interrupted run leaves) + run_config = { + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverseStream", + "__llmeter_state__": { + # Mirrors real dump_object() output: fields the constructor accepts, with + # derived values like `provider` deliberately absent (recomputed on load). + "model_id": "synthetic.interrupted-model-v1", + "endpoint_name": "synthetic-interrupted-endpoint", + "region": "eu-west-1", + "inference_config": None, + "max_attempts": 3, + "ttft_visible_tokens_only": True, + }, + }, + "output_path": "stale/interrupted/path", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {}, + }, + "clients": 4, + "n_requests": 20, + "payload": "stale/interrupted/payload.jsonl", + "run_name": "synthetic-interrupted-test", + "run_description": "This run was interrupted before completion", + "timeout": 120, + "callbacks": None, + "low_memory": False, + "progress_bar_stats": None, + } + _write_json(out / "run_config.json", run_config) + + print(f" interrupted_run/ done ({len(responses)} responses, no summary.json)") + + +# ============================================================================= +# Scenario: legacy/v0_1_interrupted (legacy-format config, no summary.json) +# ============================================================================= + + +def generate_legacy_interrupted_run(): + """A legacy-format interrupted run: no summary.json, top-level endpoint fields. + + Exercises the legacy branch of ``Result._recover_metadata`` - endpoint fields + (including the derived ``provider``) live at the top level of the endpoint dict, + and loading should recover them while warning that the format is deprecated. + """ + out = FIXTURES_DIR / "legacy" / "v0_1_interrupted" + + responses = [ + InvocationResponse( + id=f"synth-legacy-intr-{i:03d}", + response_text=f"Legacy interrupted response {i}.", + input_prompt=f"Legacy interrupted prompt {i}.", + time_to_first_token=0.35 + 0.05 * i, + time_to_last_token=0.9 + 0.1 * i, + num_tokens_input=20 + 4 * i, + num_tokens_output=45 + 6 * i, + time_per_output_token=0.01, + error=None, + request_time=datetime(2025, 1, 5, 12, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + # NO summary.json — recovery path + + # responses.jsonl (legacy: no retries/annotations/cached/reasoning fields) + legacy_responses = [] + for resp in responses: + d = json.loads(resp.to_json()) + d.pop("num_tokens_input_cached", None) + d.pop("num_tokens_output_reasoning", None) + d.pop("retries", None) + d.pop("annotations", None) + legacy_responses.append(json.dumps(d)) + _write_jsonl(out / "responses.jsonl", legacy_responses) + + # run_config.json — legacy format with top-level endpoint fields (incl. derived provider) + run_config = { + "endpoint": { + "endpoint_name": "synthetic-legacy-intr-endpoint", + "model_id": "synthetic.legacy-interrupted-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverse", + }, + "output_path": "stale/legacy/interrupted/path", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 2, + "n_requests": 4, + "payload": "stale/legacy/interrupted/payload.jsonl", + "run_name": "synthetic-legacy-interrupted-test", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + _write_json(out / "run_config.json", run_config) + + print( + f" legacy/v0_1_interrupted/ done ({len(responses)} responses, no summary.json)" + ) + + +# ============================================================================= +# Scenario: errors_and_annotations +# ============================================================================= + + +def generate_errors_and_annotations(): + out = FIXTURES_DIR / "errors_and_annotations" + + responses = [ + InvocationResponse( + id="synth-err-001", + response_text="Successful response with custom annotations.", + input_prompt="Test prompt with annotations.", + time_to_first_token=0.350, + time_to_last_token=0.920, + num_tokens_input=32, + num_tokens_output=56, + time_per_output_token=0.01018, + error=None, + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 0, tzinfo=timezone.utc), + annotations={"custom_metric": 42.5, "experiment_tag": "baseline"}, + ), + InvocationResponse( + id="synth-err-002", + response_text="", + input_prompt="This request was throttled.", + time_to_first_token=None, + time_to_last_token=2.105, + num_tokens_input=25, + num_tokens_output=0, + time_per_output_token=None, + error="ThrottlingException: Rate exceeded", + retries=2, + request_time=datetime(2025, 1, 18, 8, 0, 1, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-003", + response_text="Another successful response after the error.", + input_prompt="Normal prompt after throttling.", + time_to_first_token=0.410, + time_to_last_token=1.050, + num_tokens_input=30, + num_tokens_output=62, + time_per_output_token=0.01032, + error=None, + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 3, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-004", + response_text="", + input_prompt="This request timed out.", + time_to_first_token=None, + time_to_last_token=None, + num_tokens_input=22, + num_tokens_output=0, + time_per_output_token=None, + error="ReadTimeoutError: timed out after 60s", + retries=0, + request_time=datetime(2025, 1, 18, 8, 0, 5, tzinfo=timezone.utc), + annotations={}, + ), + InvocationResponse( + id="synth-err-005", + response_text="Response with cached input tokens.", + input_prompt="Prompt that leverages prompt caching.", + time_to_first_token=0.180, + time_to_last_token=0.650, + num_tokens_input=45, + num_tokens_output=38, + num_tokens_input_cached=30, + num_tokens_output_reasoning=None, + time_per_output_token=0.01237, + error=None, + retries=1, + request_time=datetime(2025, 1, 18, 8, 0, 7, tzinfo=timezone.utc), + annotations={"cache_hit_ratio": 0.667}, + ), + ] + + # summary.json + summary = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z", + } + _write_json(out / "summary.json", summary) + + # responses.jsonl + _write_jsonl(out / "responses.jsonl", [r.to_json() for r in responses]) + + # stats.json — computed from the responses (including errors) + _compute_and_write_stats(responses, summary, out / "stats.json") + + print(f" errors_and_annotations/ done ({len(responses)} responses, 2 errors)") + + +# ============================================================================= +# Scenario: load_test (multi-concurrency) +# ============================================================================= + + +def generate_load_test(): + out = FIXTURES_DIR / "load_test" + + # 00001-clients + responses_1 = [ + InvocationResponse( + id=f"synth-lt1-{i:03d}", + response_text=f"Load test single-client response {i}.", + input_prompt=f"Load test prompt {i}.", + time_to_first_token=0.200 + 0.05 * i, + time_to_last_token=0.600 + 0.1 * i, + num_tokens_input=20 + 2 * i, + num_tokens_output=35 + 5 * i, + time_per_output_token=0.008, + error=None, + retries=0, + request_time=datetime(2025, 1, 22, 12, 0, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + out_1 = out / "00001-clients" + summary_1 = { + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": None, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z", + } + _write_json(out_1 / "summary.json", summary_1) + _compute_and_write_stats(responses_1, summary_1, out_1 / "stats.json") + _write_jsonl(out_1 / "responses.jsonl", [r.to_json() for r in responses_1]) + + # 00003-clients + responses_3 = [ + InvocationResponse( + id=f"synth-lt3-{i:03d}", + response_text=f"Load test triple-client response {i}.", + input_prompt=f"Load test prompt {i}.", + time_to_first_token=0.350 + 0.08 * i, + time_to_last_token=0.900 + 0.15 * i, + num_tokens_input=20 + 2 * i, + num_tokens_output=35 + 5 * i, + time_per_output_token=0.012, + error=None, + retries=0, + request_time=datetime(2025, 1, 22, 12, 5, i, tzinfo=timezone.utc), + ) + for i in range(1, 6) + ] + + out_3 = out / "00003-clients" + summary_3 = { + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": None, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z", + } + _write_json(out_3 / "summary.json", summary_3) + _compute_and_write_stats(responses_3, summary_3, out_3 / "stats.json") + _write_jsonl(out_3 / "responses.jsonl", [r.to_json() for r in responses_3]) + + print( + f" load_test/ done (2 subdirs, {len(responses_1)}+{len(responses_3)} responses)" + ) + + +# ============================================================================= + +if __name__ == "__main__": + print("Generating fixture data...") + generate_base() + generate_legacy_endpoint_type() + generate_legacy_str_callbacks() + generate_interrupted_run() + generate_legacy_interrupted_run() + generate_errors_and_annotations() + generate_load_test() + print("Done!") diff --git a/tests/unit/fixtures/result_snapshots/base/payload.jsonl b/tests/unit/fixtures/result_snapshots/base/payload.jsonl new file mode 100644 index 0000000..6168fbe --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/payload.jsonl @@ -0,0 +1,2 @@ +{"messages": [{"role": "user", "content": "Summarize the key concepts of distributed systems."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}} +{"messages": [{"role": "user", "content": "Explain the CAP theorem briefly."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}} diff --git a/tests/unit/fixtures/result_snapshots/base/responses.jsonl b/tests/unit/fixtures/result_snapshots/base/responses.jsonl new file mode 100644 index 0000000..85d79db --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit.", "input_payload": {"messages": [{"role": "user", "content": "Summarize the key concepts of distributed systems."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-001", "input_prompt": "Summarize the key concepts of distributed systems.", "time_to_first_token": 0.452, "time_to_last_token": 1.203, "num_tokens_input": 42, "num_tokens_output": 87, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00863, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:00Z", "annotations": {}} +{"response_text": "Sed ut perspiciatis unde omnis iste natus error sit voluptatem.", "input_payload": {"messages": [{"role": "user", "content": "Explain the CAP theorem briefly."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-002", "input_prompt": "Explain the CAP theorem briefly.", "time_to_first_token": 0.389, "time_to_last_token": 0.951, "num_tokens_input": 38, "num_tokens_output": 64, "num_tokens_input_cached": 12, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00879, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:01Z", "annotations": {}} +{"response_text": "Nemo enim ipsam voluptatem quia voluptas sit aspernatur.", "input_payload": {"messages": [{"role": "user", "content": "What is eventual consistency?"}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-003", "input_prompt": "What is eventual consistency?", "time_to_first_token": 0.521, "time_to_last_token": 1.445, "num_tokens_input": 35, "num_tokens_output": 102, "num_tokens_input_cached": null, "num_tokens_output_reasoning": 8, "time_per_output_token": 0.00906, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:02Z", "annotations": {}} +{"response_text": "Ut enim ad minima veniam, quis nostrum exercitationem.", "input_payload": {"messages": [{"role": "user", "content": "Define partition tolerance."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-004", "input_prompt": "Define partition tolerance.", "time_to_first_token": 0.298, "time_to_last_token": 0.812, "num_tokens_input": 30, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00934, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:03Z", "annotations": {}} +{"response_text": "Quis autem vel eum iure reprehenderit qui in ea voluptate.", "input_payload": {"messages": [{"role": "user", "content": "Compare synchronous and asynchronous replication."}], "max_tokens": 256, "model": "synthetic.test-model-v1", "stream": true, "stream_options": {"include_usage": true}}, "id": "synth-base-005", "input_prompt": "Compare synchronous and asynchronous replication.", "time_to_first_token": 0.415, "time_to_last_token": 1.102, "num_tokens_input": 40, "num_tokens_output": 78, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.00881, "error": null, "retries": 0, "request_time": "2025-01-15T10:00:04Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/base/run_config.json b/tests/unit/fixtures/result_snapshots/base/run_config.json new file mode 100644 index 0000000..4056dfd --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/run_config.json @@ -0,0 +1,65 @@ +{ + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.openai.OpenAICompletionStreamEndpoint", + "__llmeter_state__": { + "model_id": "synthetic.test-model-v1", + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "organization": null, + "base_url": "https://api.synthetic-provider.example.com/v1/", + "max_retries": 3, + "timeout": { + "connect": 5.0, + "read": 60.0, + "write": 5.0, + "pool": 5.0 + }, + "api_key": "sk-synthetic-test-key-not-real" + } + }, + "output_path": "stale/path/that/should/be/overridden", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {} + }, + "clients": 2, + "n_requests": 3, + "run_duration": null, + "payload": "tests/unit/fixtures/result_snapshots/base/payload.jsonl", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "timeout": 60, + "callbacks": [ + { + "__llmeter_class__": "llmeter.callbacks.cost.model.CostModel", + "__llmeter_state__": { + "request_dims": { + "InputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.InputTokens", + "__llmeter_state__": { + "price_per_million": 3.0, + "granularity": 1 + } + }, + "OutputTokens": { + "__llmeter_class__": "llmeter.callbacks.cost.dimensions.OutputTokens", + "__llmeter_state__": { + "price_per_million": 15.0, + "granularity": 1 + } + } + }, + "run_dims": {} + } + }, + { + "__llmeter_class__": "llmeter.callbacks.mlflow.MlflowCallback", + "__llmeter_state__": { + "step": null, + "nested": true + } + } + ], + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/base/stats.json b/tests/unit/fixtures/result_snapshots/base/stats.json new file mode 100644 index 0000000..01744b4 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/stats.json @@ -0,0 +1,45 @@ +{ + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": null, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 66.35700066357, + "total_input_tokens": 185, + "total_output_tokens": 386, + "total_cached_input_tokens": 12, + "total_reasoning_output_tokens": 8, + "average_input_tokens_per_minute": 2455.2090245520903, + "average_output_tokens_per_minute": 5122.760451227605, + "time_to_last_token-average": 1.1026, + "time_to_last_token-p50": 1.102, + "time_to_last_token-p90": 1.5417999999999998, + "time_to_last_token-p99": 1.67248, + "time_to_first_token-average": 0.415, + "time_to_first_token-p50": 0.415, + "time_to_first_token-p90": 0.5486000000000001, + "time_to_first_token-p99": 0.5858599999999999, + "num_tokens_output-average": 77.2, + "num_tokens_output-p50": 78, + "num_tokens_output-p90": 108.0, + "num_tokens_output-p99": 116.1, + "num_tokens_input-average": 37, + "num_tokens_input-p50": 38, + "num_tokens_input-p90": 42.8, + "num_tokens_input-p99": 43.88, + "num_tokens_input_cached-average": 12, + "num_tokens_input_cached-p50": 12, + "num_tokens_input_cached-p90": 12, + "num_tokens_input_cached-p99": 12 +} diff --git a/tests/unit/fixtures/result_snapshots/base/summary.json b/tests/unit/fixtures/result_snapshots/base/summary.json new file mode 100644 index 0000000..8ecc92e --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/base/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 2, + "n_requests": 3, + "total_test_time": 4.521, + "model_id": "synthetic.test-model-v1", + "output_path": null, + "endpoint_name": "synthetic-openai-stream", + "provider": "openai", + "run_name": "synthetic-base-test", + "run_description": "Synthetic fixture for deserialization regression tests", + "start_time": "2025-01-15T10:00:00Z", + "first_request_time": "2025-01-15T10:00:00Z", + "last_request_time": "2025-01-15T10:00:04Z", + "end_time": "2025-01-15T10:00:04Z" +} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl b/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl new file mode 100644 index 0000000..0b31f9c --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Successful response with custom annotations.", "input_payload": null, "id": "synth-err-001", "input_prompt": "Test prompt with annotations.", "time_to_first_token": 0.35, "time_to_last_token": 0.92, "num_tokens_input": 32, "num_tokens_output": 56, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01018, "error": null, "retries": 0, "request_time": "2025-01-18T08:00:00Z", "annotations": {"custom_metric": 42.5, "experiment_tag": "baseline"}} +{"response_text": "", "input_payload": null, "id": "synth-err-002", "input_prompt": "This request was throttled.", "time_to_first_token": null, "time_to_last_token": 2.105, "num_tokens_input": 25, "num_tokens_output": 0, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": null, "error": "ThrottlingException: Rate exceeded", "retries": 2, "request_time": "2025-01-18T08:00:01Z", "annotations": {}} +{"response_text": "Another successful response after the error.", "input_payload": null, "id": "synth-err-003", "input_prompt": "Normal prompt after throttling.", "time_to_first_token": 0.41, "time_to_last_token": 1.05, "num_tokens_input": 30, "num_tokens_output": 62, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01032, "error": null, "retries": 0, "request_time": "2025-01-18T08:00:03Z", "annotations": {}} +{"response_text": "", "input_payload": null, "id": "synth-err-004", "input_prompt": "This request timed out.", "time_to_first_token": null, "time_to_last_token": null, "num_tokens_input": 22, "num_tokens_output": 0, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": null, "error": "ReadTimeoutError: timed out after 60s", "retries": 0, "request_time": "2025-01-18T08:00:05Z", "annotations": {}} +{"response_text": "Response with cached input tokens.", "input_payload": null, "id": "synth-err-005", "input_prompt": "Prompt that leverages prompt caching.", "time_to_first_token": 0.18, "time_to_last_token": 0.65, "num_tokens_input": 45, "num_tokens_output": 38, "num_tokens_input_cached": 30, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01237, "error": null, "retries": 1, "request_time": "2025-01-18T08:00:07Z", "annotations": {"cache_hit_ratio": 0.667}} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json b/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json new file mode 100644 index 0000000..8786a50 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/stats.json @@ -0,0 +1,45 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z", + "failed_requests": 2, + "failed_requests_rate": 0.4, + "requests_per_minute": 38.46153846153847, + "total_input_tokens": 154, + "total_output_tokens": 156, + "total_cached_input_tokens": 30, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1184.6153846153848, + "average_output_tokens_per_minute": 1200.0, + "time_to_last_token-average": 1.18125, + "time_to_last_token-p50": 0.9850000000000001, + "time_to_last_token-p90": 2.6325, + "time_to_last_token-p99": 3.10725, + "time_to_first_token-average": 0.3133333333333333, + "time_to_first_token-p50": 0.35, + "time_to_first_token-p90": 0.446, + "time_to_first_token-p99": 0.46760000000000007, + "num_tokens_output-average": 31.2, + "num_tokens_output-p50": 38, + "num_tokens_output-p90": 64.4, + "num_tokens_output-p99": 67.64, + "num_tokens_input-average": 30.8, + "num_tokens_input-p50": 30, + "num_tokens_input-p90": 50.2, + "num_tokens_input-p99": 57.22, + "num_tokens_input_cached-average": 30, + "num_tokens_input_cached-p50": 30, + "num_tokens_input_cached-p90": 30, + "num_tokens_input_cached-p99": 30 +} diff --git a/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json b/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json new file mode 100644 index 0000000..3836355 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/errors_and_annotations/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 7.8, + "model_id": "synthetic.errors-model-v1", + "output_path": "stale/errors/path", + "endpoint_name": "synthetic-errors-endpoint", + "provider": "bedrock", + "run_name": "synthetic-errors-test", + "run_description": "Fixture with errors, retries, and annotations", + "start_time": "2025-01-18T08:00:00Z", + "first_request_time": "2025-01-18T08:00:00Z", + "last_request_time": "2025-01-18T08:00:07Z", + "end_time": "2025-01-18T08:00:07Z" +} diff --git a/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl b/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl new file mode 100644 index 0000000..bd2988a --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/interrupted_run/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Interrupted run response 1: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-001", "input_prompt": "Interrupted test prompt 1.", "time_to_first_token": 0.35, "time_to_last_token": 0.9500000000000001, "num_tokens_input": 28, "num_tokens_output": 48, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:02Z", "annotations": {}} +{"response_text": "Interrupted run response 2: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-002", "input_prompt": "Interrupted test prompt 2.", "time_to_first_token": 0.4, "time_to_last_token": 1.1, "num_tokens_input": 31, "num_tokens_output": 56, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:04Z", "annotations": {}} +{"response_text": "Interrupted run response 3: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-003", "input_prompt": "Interrupted test prompt 3.", "time_to_first_token": 0.45, "time_to_last_token": 1.25, "num_tokens_input": 34, "num_tokens_output": 64, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:06Z", "annotations": {}} +{"response_text": "Interrupted run response 4: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-004", "input_prompt": "Interrupted test prompt 4.", "time_to_first_token": 0.5, "time_to_last_token": 1.4, "num_tokens_input": 37, "num_tokens_output": 72, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:08Z", "annotations": {}} +{"response_text": "Interrupted run response 5: tempor incididunt ut labore.", "input_payload": null, "id": "synth-intr-005", "input_prompt": "Interrupted test prompt 5.", "time_to_first_token": 0.55, "time_to_last_token": 1.55, "num_tokens_input": 40, "num_tokens_output": 80, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.009, "error": null, "retries": 0, "request_time": "2025-01-20T16:30:10Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json b/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json new file mode 100644 index 0000000..1d441a0 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/interrupted_run/run_config.json @@ -0,0 +1,27 @@ +{ + "endpoint": { + "__llmeter_class__": "llmeter.endpoints.bedrock.BedrockConverseStream", + "__llmeter_state__": { + "model_id": "synthetic.interrupted-model-v1", + "endpoint_name": "synthetic-interrupted-endpoint", + "region": "eu-west-1", + "inference_config": null, + "max_attempts": 3, + "ttft_visible_tokens_only": true + } + }, + "output_path": "stale/interrupted/path", + "tokenizer": { + "__llmeter_class__": "llmeter.tokenizers.DummyTokenizer", + "__llmeter_state__": {} + }, + "clients": 4, + "n_requests": 20, + "payload": "stale/interrupted/payload.jsonl", + "run_name": "synthetic-interrupted-test", + "run_description": "This run was interrupted before completion", + "timeout": 120, + "callbacks": null, + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl new file mode 100644 index 0000000..8d8dddf --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Consectetur adipiscing elit, sed do eiusmod tempor.", "input_payload": null, "id": "synth-legacy-001", "input_prompt": "What is a load balancer?", "time_to_first_token": 0.612, "time_to_last_token": 1.534, "num_tokens_input": 28, "num_tokens_output": 72, "time_per_output_token": 0.0128, "error": null, "request_time": "2025-01-10T14:00:00Z", "cost_InputTokens": 8.4e-05, "cost_OutputTokens": 0.00108, "cost_total": 0.001164} +{"response_text": "Duis aute irure dolor in reprehenderit in voluptate velit.", "input_payload": null, "id": "synth-legacy-002", "input_prompt": "Explain horizontal scaling.", "time_to_first_token": 0.548, "time_to_last_token": 1.289, "num_tokens_input": 24, "num_tokens_output": 65, "time_per_output_token": 0.0114, "error": null, "request_time": "2025-01-10T14:00:01Z", "cost_InputTokens": 7.2e-05, "cost_OutputTokens": 0.000975, "cost_total": 0.001047} +{"response_text": "Excepteur sint occaecat cupidatat non proident.", "input_payload": null, "id": "synth-legacy-003", "input_prompt": "What is auto-scaling?", "time_to_first_token": 0.701, "time_to_last_token": 1.678, "num_tokens_input": 22, "num_tokens_output": 58, "time_per_output_token": 0.01684, "error": null, "request_time": "2025-01-10T14:00:02Z", "cost_InputTokens": 6.6e-05, "cost_OutputTokens": 0.00087, "cost_total": 0.000936} +{"response_text": "Sunt in culpa qui officia deserunt mollit anim id est laborum.", "input_payload": null, "id": "synth-legacy-004", "input_prompt": "Define microservices architecture.", "time_to_first_token": 0.489, "time_to_last_token": 1.102, "num_tokens_input": 26, "num_tokens_output": 71, "time_per_output_token": 0.00863, "error": null, "request_time": "2025-01-10T14:00:03Z", "cost_InputTokens": 7.8e-05, "cost_OutputTokens": 0.001065, "cost_total": 0.001143} +{"response_text": "Neque porro quisquam est qui dolorem ipsum.", "input_payload": null, "id": "synth-legacy-005", "input_prompt": "What is service mesh?", "time_to_first_token": 0.555, "time_to_last_token": 1.345, "num_tokens_input": 20, "num_tokens_output": 48, "time_per_output_token": 0.01646, "error": null, "request_time": "2025-01-10T14:00:04Z", "cost_InputTokens": 6e-05, "cost_OutputTokens": 0.00072, "cost_total": 0.00078} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json new file mode 100644 index 0000000..9d8ee73 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/run_config.json @@ -0,0 +1,20 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-legacy-endpoint", + "model_id": "synthetic.legacy-model-v1", + "provider": "bedrock", + "region": "us-west-2", + "endpoint_type": "BedrockConverse" + }, + "output_path": "stale/legacy/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 3, + "n_requests": 2, + "payload": "stale/legacy/payload.jsonl", + "run_name": "synthetic-legacy-test", + "run_description": null, + "timeout": 60, + "callbacks": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json new file mode 100644 index 0000000..50f3609 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": null, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 58.80047040376323, + "total_input_tokens": 120, + "total_output_tokens": 314, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1411.2112896903175, + "average_output_tokens_per_minute": 3692.6695413563307, + "time_to_last_token-average": 1.3896, + "time_to_last_token-p50": 1.345, + "time_to_last_token-p90": 1.7355999999999998, + "time_to_last_token-p99": 1.8133599999999999, + "time_to_first_token-average": 0.581, + "time_to_first_token-p50": 0.555, + "time_to_first_token-p90": 0.7365999999999999, + "time_to_first_token-p99": 0.7846600000000001, + "num_tokens_output-average": 62.8, + "num_tokens_output-p50": 65, + "num_tokens_output-p90": 72.4, + "num_tokens_output-p99": 72.94, + "num_tokens_input-average": 24, + "num_tokens_input-p50": 24, + "num_tokens_input-p90": 28.8, + "num_tokens_input-p99": 29.88 +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json new file mode 100644 index 0000000..d46b45d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_endpoint_type/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 5.102, + "model_id": "synthetic.legacy-model-v1", + "output_path": "stale/legacy/path", + "endpoint_name": "synthetic-legacy-endpoint", + "provider": "bedrock", + "run_name": "synthetic-legacy-test", + "run_description": null, + "start_time": "2025-01-10T14:00:00Z", + "first_request_time": "2025-01-10T14:00:00Z", + "last_request_time": "2025-01-10T14:00:04Z", + "end_time": "2025-01-10T14:00:05Z" +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl new file mode 100644 index 0000000..24246b0 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Legacy interrupted response 1.", "input_payload": null, "id": "synth-legacy-intr-001", "input_prompt": "Legacy interrupted prompt 1.", "time_to_first_token": 0.39999999999999997, "time_to_last_token": 1.0, "num_tokens_input": 24, "num_tokens_output": 51, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:01Z"} +{"response_text": "Legacy interrupted response 2.", "input_payload": null, "id": "synth-legacy-intr-002", "input_prompt": "Legacy interrupted prompt 2.", "time_to_first_token": 0.44999999999999996, "time_to_last_token": 1.1, "num_tokens_input": 28, "num_tokens_output": 57, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:02Z"} +{"response_text": "Legacy interrupted response 3.", "input_payload": null, "id": "synth-legacy-intr-003", "input_prompt": "Legacy interrupted prompt 3.", "time_to_first_token": 0.5, "time_to_last_token": 1.2000000000000002, "num_tokens_input": 32, "num_tokens_output": 63, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:03Z"} +{"response_text": "Legacy interrupted response 4.", "input_payload": null, "id": "synth-legacy-intr-004", "input_prompt": "Legacy interrupted prompt 4.", "time_to_first_token": 0.55, "time_to_last_token": 1.3, "num_tokens_input": 36, "num_tokens_output": 69, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:04Z"} +{"response_text": "Legacy interrupted response 5.", "input_payload": null, "id": "synth-legacy-intr-005", "input_prompt": "Legacy interrupted prompt 5.", "time_to_first_token": 0.6, "time_to_last_token": 1.4, "num_tokens_input": 40, "num_tokens_output": 75, "time_per_output_token": 0.01, "error": null, "request_time": "2025-01-05T12:00:05Z"} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json new file mode 100644 index 0000000..357eb32 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_interrupted/run_config.json @@ -0,0 +1,20 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-legacy-intr-endpoint", + "model_id": "synthetic.legacy-interrupted-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverse" + }, + "output_path": "stale/legacy/interrupted/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 2, + "n_requests": 4, + "payload": "stale/legacy/interrupted/payload.jsonl", + "run_name": "synthetic-legacy-interrupted-test", + "run_description": null, + "timeout": 60, + "callbacks": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl new file mode 100644 index 0000000..51fcf3f --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Synthetic response number 1 for string callback test.", "input_payload": null, "id": "synth-strcb-001", "input_prompt": "Synthetic prompt 1.", "time_to_first_token": 0.5, "time_to_last_token": 1.2, "num_tokens_input": 35, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:01Z", "annotations": {}} +{"response_text": "Synthetic response number 2 for string callback test.", "input_payload": null, "id": "synth-strcb-002", "input_prompt": "Synthetic prompt 2.", "time_to_first_token": 0.6000000000000001, "time_to_last_token": 1.4, "num_tokens_input": 40, "num_tokens_output": 70, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:02Z", "annotations": {}} +{"response_text": "Synthetic response number 3 for string callback test.", "input_payload": null, "id": "synth-strcb-003", "input_prompt": "Synthetic prompt 3.", "time_to_first_token": 0.7000000000000001, "time_to_last_token": 1.6, "num_tokens_input": 45, "num_tokens_output": 80, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:03Z", "annotations": {}} +{"response_text": "Synthetic response number 4 for string callback test.", "input_payload": null, "id": "synth-strcb-004", "input_prompt": "Synthetic prompt 4.", "time_to_first_token": 0.8, "time_to_last_token": 1.8, "num_tokens_input": 50, "num_tokens_output": 90, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:04Z", "annotations": {}} +{"response_text": "Synthetic response number 5 for string callback test.", "input_payload": null, "id": "synth-strcb-005", "input_prompt": "Synthetic prompt 5.", "time_to_first_token": 0.9, "time_to_last_token": 2.0, "num_tokens_input": 55, "num_tokens_output": 100, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.01, "error": null, "retries": null, "request_time": "2025-01-12T09:00:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json new file mode 100644 index 0000000..15aca5e --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/run_config.json @@ -0,0 +1,25 @@ +{ + "endpoint": { + "endpoint_name": "synthetic-strcb-endpoint", + "model_id": "synthetic.strcb-model-v1", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream" + }, + "output_path": "stale/strcb/path", + "tokenizer": { + "tokenizer_module": "llmeter" + }, + "clients": 1, + "n_requests": 5, + "payload": "stale/strcb/payload.jsonl", + "run_name": "synthetic-strcb-test", + "run_description": null, + "timeout": 60, + "callbacks": [ + "<__main__.SyntheticCallback object at 0x1234567890>", + "" + ], + "low_memory": false, + "progress_bar_stats": null +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json new file mode 100644 index 0000000..94eaf21 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": null, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 46.15384615384615, + "total_input_tokens": 225, + "total_output_tokens": 400, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 2076.9230769230767, + "average_output_tokens_per_minute": 3692.3076923076924, + "time_to_last_token-average": 1.6, + "time_to_last_token-p50": 1.6, + "time_to_last_token-p90": 2.08, + "time_to_last_token-p99": 2.1879999999999997, + "time_to_first_token-average": 0.7000000000000001, + "time_to_first_token-p50": 0.7000000000000001, + "time_to_first_token-p90": 0.9399999999999998, + "time_to_first_token-p99": 0.9939999999999999, + "num_tokens_output-average": 80, + "num_tokens_output-p50": 80, + "num_tokens_output-p90": 104.0, + "num_tokens_output-p99": 109.4, + "num_tokens_input-average": 45, + "num_tokens_input-p50": 45, + "num_tokens_input-p90": 57.0, + "num_tokens_input-p99": 59.7 +} diff --git a/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json new file mode 100644 index 0000000..7a0e968 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/legacy/v0_1_str_callbacks/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 6.5, + "model_id": "synthetic.strcb-model-v1", + "output_path": "stale/strcb/path", + "endpoint_name": "synthetic-strcb-endpoint", + "provider": "bedrock", + "run_name": "synthetic-strcb-test", + "run_description": null, + "start_time": "2025-01-12T09:00:01Z", + "first_request_time": "2025-01-12T09:00:01Z", + "last_request_time": "2025-01-12T09:00:05Z", + "end_time": "2025-01-12T09:00:06Z" +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl new file mode 100644 index 0000000..174b71d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Load test single-client response 1.", "input_payload": null, "id": "synth-lt1-001", "input_prompt": "Load test prompt 1.", "time_to_first_token": 0.25, "time_to_last_token": 0.7, "num_tokens_input": 22, "num_tokens_output": 40, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:01Z", "annotations": {}} +{"response_text": "Load test single-client response 2.", "input_payload": null, "id": "synth-lt1-002", "input_prompt": "Load test prompt 2.", "time_to_first_token": 0.30000000000000004, "time_to_last_token": 0.8, "num_tokens_input": 24, "num_tokens_output": 45, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:02Z", "annotations": {}} +{"response_text": "Load test single-client response 3.", "input_payload": null, "id": "synth-lt1-003", "input_prompt": "Load test prompt 3.", "time_to_first_token": 0.35000000000000003, "time_to_last_token": 0.9, "num_tokens_input": 26, "num_tokens_output": 50, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:03Z", "annotations": {}} +{"response_text": "Load test single-client response 4.", "input_payload": null, "id": "synth-lt1-004", "input_prompt": "Load test prompt 4.", "time_to_first_token": 0.4, "time_to_last_token": 1.0, "num_tokens_input": 28, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:04Z", "annotations": {}} +{"response_text": "Load test single-client response 5.", "input_payload": null, "id": "synth-lt1-005", "input_prompt": "Load test prompt 5.", "time_to_first_token": 0.45, "time_to_last_token": 1.1, "num_tokens_input": 30, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.008, "error": null, "retries": 0, "request_time": "2025-01-22T12:00:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json new file mode 100644 index 0000000..77fe968 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": null, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 54.54545454545455, + "total_input_tokens": 130, + "total_output_tokens": 250, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 1418.1818181818182, + "average_output_tokens_per_minute": 2727.272727272727, + "time_to_last_token-average": 0.9, + "time_to_last_token-p50": 0.9, + "time_to_last_token-p90": 1.1400000000000001, + "time_to_last_token-p99": 1.194, + "time_to_first_token-average": 0.35000000000000003, + "time_to_first_token-p50": 0.35000000000000003, + "time_to_first_token-p90": 0.4699999999999999, + "time_to_first_token-p99": 0.49699999999999994, + "num_tokens_output-average": 50, + "num_tokens_output-p50": 50, + "num_tokens_output-p90": 62.0, + "num_tokens_output-p99": 64.7, + "num_tokens_input-average": 26, + "num_tokens_input-p50": 26, + "num_tokens_input-p90": 30.8, + "num_tokens_input-p99": 31.88 +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json new file mode 100644 index 0000000..030ec9d --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00001-clients/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 1, + "n_requests": 5, + "total_test_time": 5.5, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00001-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00001-clients", + "run_description": null, + "start_time": "2025-01-22T12:00:01Z", + "first_request_time": "2025-01-22T12:00:01Z", + "last_request_time": "2025-01-22T12:00:05Z", + "end_time": "2025-01-22T12:00:05Z" +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl new file mode 100644 index 0000000..2041ea1 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/responses.jsonl @@ -0,0 +1,5 @@ +{"response_text": "Load test triple-client response 1.", "input_payload": null, "id": "synth-lt3-001", "input_prompt": "Load test prompt 1.", "time_to_first_token": 0.43, "time_to_last_token": 1.05, "num_tokens_input": 22, "num_tokens_output": 40, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:01Z", "annotations": {}} +{"response_text": "Load test triple-client response 2.", "input_payload": null, "id": "synth-lt3-002", "input_prompt": "Load test prompt 2.", "time_to_first_token": 0.51, "time_to_last_token": 1.2, "num_tokens_input": 24, "num_tokens_output": 45, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:02Z", "annotations": {}} +{"response_text": "Load test triple-client response 3.", "input_payload": null, "id": "synth-lt3-003", "input_prompt": "Load test prompt 3.", "time_to_first_token": 0.59, "time_to_last_token": 1.35, "num_tokens_input": 26, "num_tokens_output": 50, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:03Z", "annotations": {}} +{"response_text": "Load test triple-client response 4.", "input_payload": null, "id": "synth-lt3-004", "input_prompt": "Load test prompt 4.", "time_to_first_token": 0.6699999999999999, "time_to_last_token": 1.5, "num_tokens_input": 28, "num_tokens_output": 55, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:04Z", "annotations": {}} +{"response_text": "Load test triple-client response 5.", "input_payload": null, "id": "synth-lt3-005", "input_prompt": "Load test prompt 5.", "time_to_first_token": 0.75, "time_to_last_token": 1.65, "num_tokens_input": 30, "num_tokens_output": 60, "num_tokens_input_cached": null, "num_tokens_output_reasoning": null, "time_per_output_token": 0.012, "error": null, "retries": 0, "request_time": "2025-01-22T12:05:05Z", "annotations": {}} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json new file mode 100644 index 0000000..ad2a2c5 --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/stats.json @@ -0,0 +1,41 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": null, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z", + "failed_requests": 0, + "failed_requests_rate": 0.0, + "requests_per_minute": 93.75, + "total_input_tokens": 130, + "total_output_tokens": 250, + "total_cached_input_tokens": 0, + "total_reasoning_output_tokens": 0, + "average_input_tokens_per_minute": 2437.5, + "average_output_tokens_per_minute": 4687.5, + "time_to_last_token-average": 1.35, + "time_to_last_token-p50": 1.35, + "time_to_last_token-p90": 1.7099999999999997, + "time_to_last_token-p99": 1.7909999999999997, + "time_to_first_token-average": 0.59, + "time_to_first_token-p50": 0.59, + "time_to_first_token-p90": 0.782, + "time_to_first_token-p99": 0.8252000000000002, + "num_tokens_output-average": 50, + "num_tokens_output-p50": 50, + "num_tokens_output-p90": 62.0, + "num_tokens_output-p99": 64.7, + "num_tokens_input-average": 26, + "num_tokens_input-p50": 26, + "num_tokens_input-p90": 30.8, + "num_tokens_input-p99": 31.88 +} diff --git a/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json new file mode 100644 index 0000000..06f98af --- /dev/null +++ b/tests/unit/fixtures/result_snapshots/load_test/00003-clients/summary.json @@ -0,0 +1,16 @@ +{ + "total_requests": 5, + "clients": 3, + "n_requests": 2, + "total_test_time": 3.2, + "model_id": "synthetic.loadtest-model-v1", + "output_path": "stale/loadtest/00003-clients", + "endpoint_name": "synthetic-loadtest-endpoint", + "provider": "bedrock", + "run_name": "00003-clients", + "run_description": null, + "start_time": "2025-01-22T12:05:01Z", + "first_request_time": "2025-01-22T12:05:01Z", + "last_request_time": "2025-01-22T12:05:05Z", + "end_time": "2025-01-22T12:05:05Z" +} diff --git a/tests/unit/test_interrupted_run.py b/tests/unit/test_interrupted_run.py index a57fedf..b813d38 100644 --- a/tests/unit/test_interrupted_run.py +++ b/tests/unit/test_interrupted_run.py @@ -29,9 +29,7 @@ def sample_responses(): time_to_last_token=0.3 * (i + 1), num_tokens_input=10 * (i + 1), num_tokens_output=20 * (i + 1), - request_time=datetime( - 2025, 6, 1, 10, 0, i * 2, tzinfo=timezone.utc - ), + request_time=datetime(2025, 6, 1, 10, 0, i * 2, tzinfo=timezone.utc), ) for i in range(5) ] @@ -44,7 +42,7 @@ def interrupted_run_dir(tmp_path, sample_responses): Contains only responses.jsonl and run_config.json, as would be left behind when a run is interrupted before Result.save() completes. """ - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default run_dir = UPath(tmp_path / "interrupted_run") run_dir.mkdir(parents=True) @@ -78,7 +76,7 @@ def interrupted_run_dir(tmp_path, sample_responses): } config_path = run_dir / "run_config.json" with config_path.open("w") as f: - json.dump(config, f, default=llmeter_default_serializer) + json.dump(config, f, default=json_default) return run_dir @@ -123,9 +121,7 @@ def test_load_derives_timestamps_from_responses(self, interrupted_run_dir): assert result.last_request_time == datetime( 2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc ) - assert result.end_time == datetime( - 2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc - ) + assert result.end_time == datetime(2025, 6, 1, 10, 0, 8, tzinfo=timezone.utc) def test_load_sets_total_requests_from_responses(self, interrupted_run_dir): """total_requests should be set to the number of recovered responses.""" @@ -214,7 +210,7 @@ def test_load_nonexistent_directory_raises(self, tmp_path): def test_load_recovers_interrupted_run_with_stats(self, tmp_path, sample_responses): """Simulates the case where both stats.json and responses.jsonl exist (e.g. the interrupt handler managed to write stats before exiting).""" - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default run_dir = UPath(tmp_path / "partial_with_stats") run_dir.mkdir(parents=True) @@ -232,7 +228,7 @@ def test_load_recovers_interrupted_run_with_stats(self, tmp_path, sample_respons } stats_path = run_dir / "stats.json" with stats_path.open("w") as f: - json.dump(stats, f, default=llmeter_default_serializer) + json.dump(stats, f, default=json_default) result = Result.load(run_dir) @@ -686,9 +682,7 @@ def slow_invoke(payload): import time time.sleep(0.05) - return InvocationResponse( - id="x", input_prompt="test", response_text="resp" - ) + return InvocationResponse(id="x", input_prompt="test", response_text="resp") mock_endpoint.invoke = slow_invoke diff --git a/tests/unit/test_invocation_response.py b/tests/unit/test_invocation_response.py index 67ef4d9..212e909 100644 --- a/tests/unit/test_invocation_response.py +++ b/tests/unit/test_invocation_response.py @@ -7,7 +7,7 @@ from datetime import datetime, timedelta, timezone from llmeter.endpoints.base import InvocationResponse -from llmeter.json_utils import llmeter_default_serializer +from llmeter.serialization import json_default class TestToJson: @@ -72,6 +72,20 @@ def test_save_request_time_with_offset(self): parsed = json.loads(response.to_json()) assert parsed["request_time"] == "2026-01-01T04:00:00Z" + def test_to_json_passes_kwargs_to_json_dumps(self): + """to_json() forwards kwargs (e.g. indent) to json.dumps.""" + response = InvocationResponse( + response_text="Test", + input_payload={"data": b"\x01\x02\x03"}, + id="test-kwargs", + ) + json_str = response.to_json(indent=2) + # Indentation implies newlines and leading spaces in the output. + assert "\n" in json_str + assert " " in json_str + parsed = json.loads(json_str) + assert parsed["id"] == "test-kwargs" + class TestFromJson: def test_round_trip_minimal(self): @@ -134,7 +148,7 @@ def test_round_trip_all_fields(self): def test_load_request_time_with_offset(self): json_str = json.dumps( {"response_text": "hi", "request_time": "2026-01-01T14:00:00+02:00"}, - default=llmeter_default_serializer, + default=json_default, ) restored = InvocationResponse.from_json(json_str) assert isinstance(restored.request_time, datetime) diff --git a/tests/unit/test_legacy_data_load.py b/tests/unit/test_legacy_data_load.py new file mode 100644 index 0000000..7d80ce9 --- /dev/null +++ b/tests/unit/test_legacy_data_load.py @@ -0,0 +1,245 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests for loading *legacy-format* LLMeter data written by older releases. + +This covers artifacts produced before the current ``dump_object`` serialization, across the three +data types that can be reloaded: + +* **Run configs** (``run_config.json``) and **endpoint configs** - identified the endpoint class + with an ``endpoint_type`` key (rather than the current ``__llmeter_class__`` envelope) and the + tokenizer with a ``tokenizer_module`` key, and stored derived / read-only attributes such as + ``provider`` alongside real constructor arguments. Modern endpoint constructors compute + ``provider`` internally and reject it as a keyword argument. +* **Responses** (``responses.jsonl``, via ``Result.load``) - carried extra per-response fields that + callbacks wrote directly onto the response (e.g. ``cost_InputTokens`` / ``cost_total``), plus a + since-removed ``end_time`` field. Current ``InvocationResponse`` doesn't declare these, so a + strict load would raise; they are now collected into ``annotations`` instead. + +All of these previously broke loading of real, local data. These tests pin the backward-compatible +behavior so it can't silently regress on a minor-version release. If you intend to drop +legacy-format support, that is a breaking change and these tests should be updated deliberately. +""" + +import json + +import pytest + +from llmeter.endpoints.base import Endpoint, _filter_legacy_ctor_kwargs +from llmeter.endpoints.bedrock import BedrockConverse, BedrockConverseStream +from llmeter.results import Result +from llmeter.runner import _RunConfig +from llmeter.tokenizers import DummyTokenizer + + +def _legacy_endpoint_config() -> dict: + """A legacy endpoint config dict, matching the on-disk ``examples/outputs`` format. + + Note ``provider`` is present here but is *not* a parameter of + ``BedrockConverseStream.__init__`` (it is derived internally). + """ + return { + "endpoint_name": "amazon bedrock", + "model_id": "apac.amazon.nova-pro-v1:0", + "provider": "bedrock", + "region": "us-east-1", + "endpoint_type": "BedrockConverseStream", + } + + +def _legacy_run_config() -> dict: + """A full legacy ``run_config.json`` payload, mirroring real ``examples/outputs`` data.""" + return { + "endpoint": _legacy_endpoint_config(), + "output_path": "outputs/apac.amazon.nova-pro-v1:0/20251210-0115", + "tokenizer": {"tokenizer_module": "llmeter"}, + "clients": 20, + "n_requests": 10, + "payload": "outputs/apac.amazon.nova-pro-v1:0/20251210-0115/payload.jsonl", + "run_name": "20251210-0115", + "run_description": None, + "timeout": 60, + "callbacks": None, + } + + +class TestLegacyEndpointConfigLoad: + """Endpoint.load / load_from_file must accept legacy ``endpoint_type`` configs.""" + + def test_load_drops_derived_provider_field(self): + """A legacy dict carrying a derived ``provider`` must load, not raise TypeError.""" + endpoint = Endpoint.load(_legacy_endpoint_config()) + + assert isinstance(endpoint, BedrockConverseStream) + assert endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert endpoint.region == "us-east-1" + assert endpoint.endpoint_name == "amazon bedrock" + # provider is derived by the constructor, and should round-trip to the same value + assert endpoint.provider == "bedrock" + + def test_load_from_file_drops_derived_provider_field(self, tmp_path): + """The file-based legacy loader must behave the same as the dict loader.""" + path = tmp_path / "endpoint.json" + path.write_text(json.dumps(_legacy_endpoint_config())) + + endpoint = Endpoint.load_from_file(path) + + assert isinstance(endpoint, BedrockConverseStream) + assert endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert endpoint.region == "us-east-1" + assert endpoint.provider == "bedrock" + + def test_load_missing_endpoint_type_still_raises(self): + """Legacy loading still requires ``endpoint_type`` to identify the class.""" + config = _legacy_endpoint_config() + del config["endpoint_type"] + with pytest.raises(KeyError): + Endpoint.load(config) + + +class TestLegacyRunConfigLoad: + """_RunConfig.load must reconstruct endpoint + tokenizer from a legacy run_config.json.""" + + def test_load_reconstructs_legacy_run_config(self, tmp_path): + (tmp_path / "run_config.json").write_text(json.dumps(_legacy_run_config())) + + cfg = _RunConfig.load(tmp_path) + + # Endpoint: legacy endpoint_type -> concrete class, provider derived + assert isinstance(cfg._endpoint, BedrockConverseStream) + assert cfg._endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert cfg._endpoint.region == "us-east-1" + assert cfg._endpoint.provider == "bedrock" + + # Tokenizer: legacy tokenizer_module -> DummyTokenizer + assert isinstance(cfg._tokenizer, DummyTokenizer) + + # Scalar fields survive unchanged + assert cfg.clients == 20 + assert cfg.n_requests == 10 + assert cfg.timeout == 60 + assert cfg.run_name == "20251210-0115" + + def test_load_legacy_run_config_custom_filename(self, tmp_path): + """The legacy loader honors a custom config file name.""" + (tmp_path / "my_config.json").write_text(json.dumps(_legacy_run_config())) + + cfg = _RunConfig.load(tmp_path, file_name="my_config.json") + + assert isinstance(cfg._endpoint, BedrockConverseStream) + assert isinstance(cfg._tokenizer, DummyTokenizer) + + +class TestLegacyResponsesLoad: + """Result.load must tolerate older responses.jsonl carrying extra per-response keys. + + Older LLMeter versions let callbacks (e.g. the cost callback) write extra fields such as + ``cost_InputTokens`` / ``cost_total`` directly onto each response, and serialized them into + ``responses.jsonl``. A couple of runs also carry a since-removed ``end_time`` response field. + Current ``InvocationResponse`` doesn't declare these, so a strict ``cls(**data)`` load would + raise ``TypeError``. These are collected into ``annotations`` instead, so the data is preserved + and the file still loads. + """ + + def test_load_result_with_legacy_response_extra_keys(self, tmp_path): + summary = { + "total_requests": 2, + "clients": 1, + "n_requests": 2, + "model_id": "some-model", + } + (tmp_path / "summary.json").write_text(json.dumps(summary)) + + lines = [ + json.dumps( + { + "response_text": "a", + "id": "r1", + "num_tokens_input": 10, + "num_tokens_output": 5, + "cost_InputTokens": 0.01, + "cost_OutputTokens": 0.02, + "cost_total": 0.03, + } + ), + json.dumps( + { + "response_text": "b", + "id": "r2", + "num_tokens_input": 20, + "num_tokens_output": 8, + "cost_total": 0.05, + "end_time": "2024-01-01T00:00:00Z", # since-removed field + } + ), + ] + (tmp_path / "responses.jsonl").write_text("\n".join(lines) + "\n") + + # Previously raised TypeError: unexpected keyword argument 'cost_InputTokens' + result = Result.load(tmp_path) + + assert len(result.responses) == 2 + r1, r2 = result.responses + + # Core fields are still parsed normally + assert r1.num_tokens_input == 10 + assert r2.id == "r2" + + # Legacy extra keys are preserved under annotations rather than dropped + assert r1.annotations == { + "cost_InputTokens": 0.01, + "cost_OutputTokens": 0.02, + "cost_total": 0.03, + } + assert r2.annotations["cost_total"] == 0.05 + assert r2.annotations["end_time"] == "2024-01-01T00:00:00Z" + + def test_load_result_without_summary_recovers_legacy_responses(self, tmp_path): + """Metadata recovery (no summary.json) also streams responses and must tolerate extras.""" + lines = [ + json.dumps( + { + "response_text": "a", + "id": "r1", + "request_time": "2024-01-01T00:00:00Z", + "cost_total": 0.03, + } + ), + ] + (tmp_path / "responses.jsonl").write_text("\n".join(lines) + "\n") + + result = Result.load(tmp_path) # no summary.json -> _recover_metadata path + + assert len(result.responses) == 1 + assert result.responses[0].annotations["cost_total"] == 0.03 + + +class TestFilterLegacyCtorKwargs: + """Unit tests for the constructor-kwarg filtering helper.""" + + def test_drops_keys_not_accepted_by_constructor(self): + config = { + "model_id": "m", + "region": "us-east-1", + "provider": "bedrock", # not a BedrockConverseStream.__init__ param + } + filtered = _filter_legacy_ctor_kwargs(BedrockConverseStream, config) + + assert "provider" not in filtered + assert filtered == {"model_id": "m", "region": "us-east-1"} + # Original dict is not mutated + assert "provider" in config + + def test_keeps_keys_the_constructor_accepts(self): + config = {"model_id": "m", "endpoint_name": "n", "region": "us-east-1"} + filtered = _filter_legacy_ctor_kwargs(BedrockConverse, config) + assert filtered == config + + def test_passthrough_when_constructor_accepts_var_keyword(self): + class AcceptsKwargs: + def __init__(self, model_id, **kwargs): + self.model_id = model_id + + config = {"model_id": "m", "anything": 1, "provider": "x"} + # **kwargs means every key is acceptable, so nothing is dropped + assert _filter_legacy_ctor_kwargs(AcceptsKwargs, config) == config diff --git a/tests/unit/test_prompt_utils.py b/tests/unit/test_prompt_utils.py index ac4ca98..72f55cb 100644 --- a/tests/unit/test_prompt_utils.py +++ b/tests/unit/test_prompt_utils.py @@ -11,7 +11,7 @@ from hypothesis import given, settings from hypothesis import strategies as st -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default from llmeter.prompt_utils import ( CreatePromptCollection, load_payloads, @@ -22,9 +22,9 @@ class TestLLMeterDefaultSerializer: - """Unit tests for llmeter_default_serializer function. + """Unit tests for json_default function. - These tests verify specific examples and edge cases for the llmeter_default_serializer + These tests verify specific examples and edge cases for the json_default function, complementing the property-based tests. Requirements: 1.1, 1.2, 1.3, 1.6 @@ -38,7 +38,7 @@ def test_simple_bytes_object_serialization(self): payload = {"data": b"hello world"} # Serialize using the encoder - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -78,7 +78,7 @@ def test_nested_bytes_in_dict_structure(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -107,7 +107,7 @@ def test_empty_bytes_object(self): payload = {"empty": b""} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -134,7 +134,7 @@ def test_large_binary_data_1mb(self): payload = {"large_image": large_data} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -161,7 +161,7 @@ def test_multiple_bytes_objects_in_payload(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -195,7 +195,7 @@ def test_bytes_in_list(self): payload = {"images": [b"image1", b"image2", b"image3"]} # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -223,7 +223,7 @@ def test_mixed_types_with_bytes(self): } # Serialize - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON parsed = json.loads(serialized) @@ -242,9 +242,9 @@ def test_mixed_types_with_bytes(self): class TestLLMeterBytesDecoder: - """Unit tests for llmeter_bytes_decoder function. + """Unit tests for bytes_decoder function. - These tests verify specific examples and edge cases for the llmeter_bytes_decoder + These tests verify specific examples and edge cases for the bytes_decoder function, complementing the property-based tests. Requirements: 2.1, 2.2, 2.3, 2.4, 6.2 @@ -260,7 +260,7 @@ def test_marker_object_decoding(self): marker = {"__llmeter_bytes__": "aGVsbG8gd29ybGQ="} # "hello world" in base64 # Decode the marker - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify it returns bytes assert isinstance(result, bytes) @@ -276,7 +276,7 @@ def test_non_marker_dict_passthrough(self): regular_dict = {"key": "value", "number": 42, "nested": {"data": "test"}} # Decode should return unchanged - result = llmeter_bytes_decoder(regular_dict) + result = bytes_decoder(regular_dict) # Verify it's the same dict assert result == regular_dict @@ -294,7 +294,7 @@ def test_invalid_base64_error_handling(self): # Should raise binascii.Error when trying to decode with pytest.raises(binascii.Error): - llmeter_bytes_decoder(invalid_marker) + bytes_decoder(invalid_marker) def test_multi_key_dict_with_marker_key_not_decoded(self): """Test that multi-key dict containing marker key is not decoded. @@ -311,7 +311,7 @@ def test_multi_key_dict_with_marker_key_not_decoded(self): } # Should return unchanged (not decode) - result = llmeter_bytes_decoder(multi_key_dict) + result = bytes_decoder(multi_key_dict) # Verify it's returned as-is assert result == multi_key_dict @@ -329,7 +329,7 @@ def test_empty_bytes_decoding(self): empty_marker = {"__llmeter_bytes__": ""} # Decode - result = llmeter_bytes_decoder(empty_marker) + result = bytes_decoder(empty_marker) # Verify it returns empty bytes assert isinstance(result, bytes) @@ -352,7 +352,7 @@ def test_large_binary_data_decoding(self): marker = {"__llmeter_bytes__": base64_encoded} # Decode - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify it matches original data assert isinstance(result, bytes) @@ -390,7 +390,7 @@ def test_nested_structure_with_marker(self): ) # Load with decoder - result = json.loads(json_str, object_hook=llmeter_bytes_decoder) + result = json.loads(json_str, object_hook=bytes_decoder) # Verify structure is preserved assert result["modelId"] == "test-model" @@ -411,7 +411,7 @@ def test_dict_without_marker_key(self): normal_dict = {"data": "value", "count": 123} # Should return unchanged - result = llmeter_bytes_decoder(normal_dict) + result = bytes_decoder(normal_dict) assert result == normal_dict assert isinstance(result, dict) @@ -431,7 +431,7 @@ def test_marker_with_special_characters(self): marker = {"__llmeter_bytes__": base64_encoded} # Decode - result = llmeter_bytes_decoder(marker) + result = bytes_decoder(marker) # Verify assert isinstance(result, bytes) @@ -1174,10 +1174,10 @@ def test_save_payloads_file_contains_valid_json_with_markers(self): else: assert "__llmeter_bytes__" in parsed["nested"]["data"] - def test_save_payloads_handles_to_dict_objects(self): - """Test that payloads with to_dict() objects are serialized correctly. + def test_save_payloads_uses_str_fallback_for_unknown_objects(self): + """Test that payloads with unknown objects fall back to str(). - Objects implementing to_dict() are handled by llmeter_default_serializer automatically. + json_default uses str() for objects without a dedicated handler. Validates: Requirements 5.1, 5.3 """ @@ -1185,8 +1185,8 @@ def test_save_payloads_handles_to_dict_objects(self): output_path = Path(tmpdir) class CustomObj: - def to_dict(self): - return {"custom": "value"} + def __str__(self): + return "custom_repr" payload = {"test": "data", "obj": CustomObj()} @@ -1197,7 +1197,7 @@ def to_dict(self): parsed = json.loads(line) assert parsed["test"] == "data" - assert parsed["obj"] == {"custom": "value"} + assert parsed["obj"] == "custom_repr" def test_save_payloads_backward_compatibility_no_bytes(self): """Test backward compatibility when payload has no bytes. @@ -1364,7 +1364,7 @@ def __init__(self, value): } # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) assert "test-model" in result def test_invalid_json_error_handling(self): @@ -1400,7 +1400,7 @@ def test_invalid_base64_error_handling(self): # Should raise binascii.Error when trying to decode with pytest.raises(binascii.Error): - llmeter_bytes_decoder(invalid_marker) + bytes_decoder(invalid_marker) def test_invalid_base64_in_load_payloads(self): """Test that invalid base64 in saved payload is handled during load. @@ -1428,14 +1428,14 @@ def test_invalid_base64_in_load_payloads(self): list(load_payloads(jsonl_file)) def test_save_payloads_encoder_error_propagation(self): - """Test that encoder errors from llmeter_default_serializer are propagated correctly. + """Test that encoder errors from json_default are propagated correctly. Validates: Requirements 6.4 """ with tempfile.TemporaryDirectory() as tmpdir: output_path = Path(tmpdir) - # An object whose str() raises — llmeter_default_serializer returns None for these, + # An object whose str() raises — json_default returns None for these, # which is valid JSON, so no error is raised. class FailingStr: def __str__(self): @@ -1490,7 +1490,7 @@ def __init__(self): } # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) assert "test-model" in result def test_bytes_serialization_with_encoding_error(self): @@ -1506,7 +1506,7 @@ def test_bytes_serialization_with_encoding_error(self): payload = {"data": all_bytes} # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1516,7 +1516,7 @@ def test_bytes_serialization_with_encoding_error(self): # Verify round-trip works - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) assert deserialized["data"] == all_bytes def test_empty_payload_serialization(self): @@ -1528,7 +1528,7 @@ def test_empty_payload_serialization(self): empty_payload = {} # Should serialize without errors - serialized = json.dumps(empty_payload, default=llmeter_default_serializer) + serialized = json.dumps(empty_payload, default=json_default) assert serialized == "{}" # Empty list @@ -1556,7 +1556,7 @@ def test_none_value_serialization(self): } # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1577,7 +1577,7 @@ def test_unicode_in_payload_with_bytes(self): } # Should serialize without errors - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Should be valid JSON parsed = json.loads(serialized) @@ -1629,7 +1629,7 @@ def test_1mb_image_serialization_performance(self): # Measure serialization time start_time = time.perf_counter() - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) end_time = time.perf_counter() # Calculate elapsed time in milliseconds @@ -1664,12 +1664,12 @@ def test_1mb_image_deserialization_performance(self): } # First serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Measure deserialization time start_time = time.perf_counter() - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) end_time = time.perf_counter() # Calculate elapsed time in milliseconds @@ -1704,7 +1704,7 @@ def test_serialization_no_unnecessary_copies(self): initial_size = sys.getsizeof(binary_data) # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Parse to verify structure parsed = json.loads(serialized) @@ -1740,11 +1740,11 @@ def test_deserialization_no_unnecessary_copies(self): payload = {"data": binary_data} # Serialize first - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Verify the deserialized bytes match original assert deserialized["data"] == binary_data @@ -1781,13 +1781,13 @@ def test_round_trip_performance_with_multiple_images(self): # Measure serialization time start_time = time.perf_counter() - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) serialize_time = (time.perf_counter() - start_time) * 1000 # Measure deserialization time start_time = time.perf_counter() - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) deserialize_time = (time.perf_counter() - start_time) * 1000 # With 3 images, we allow proportionally more time (but still reasonable) @@ -1820,7 +1820,7 @@ def test_serialization_performance_scales_linearly(self): small_payload = {"data": small_data} start_time = time.perf_counter() - json.dumps(small_payload, default=llmeter_default_serializer) + json.dumps(small_payload, default=json_default) small_time = time.perf_counter() - start_time # Test with 512KB (2x size) @@ -1828,7 +1828,7 @@ def test_serialization_performance_scales_linearly(self): large_payload = {"data": large_data} start_time = time.perf_counter() - json.dumps(large_payload, default=llmeter_default_serializer) + json.dumps(large_payload, default=json_default) large_time = time.perf_counter() - start_time # Large should take roughly 2x the time (allow 3x for variance) @@ -1851,19 +1851,19 @@ def test_deserialization_performance_scales_linearly(self): # Test with 256KB small_data = os.urandom(256 * 1024) small_payload = {"data": small_data} - small_serialized = json.dumps(small_payload, default=llmeter_default_serializer) + small_serialized = json.dumps(small_payload, default=json_default) start_time = time.perf_counter() - json.loads(small_serialized, object_hook=llmeter_bytes_decoder) + json.loads(small_serialized, object_hook=bytes_decoder) small_time = time.perf_counter() - start_time # Test with 512KB (2x size) large_data = os.urandom(512 * 1024) large_payload = {"data": large_data} - large_serialized = json.dumps(large_payload, default=llmeter_default_serializer) + large_serialized = json.dumps(large_payload, default=json_default) start_time = time.perf_counter() - json.loads(large_serialized, object_hook=llmeter_bytes_decoder) + json.loads(large_serialized, object_hook=bytes_decoder) large_time = time.perf_counter() - start_time # Large should take roughly 2x the time (allow 3x for variance) diff --git a/tests/unit/test_property_save_load.py b/tests/unit/test_property_save_load.py index 990333c..804f334 100644 --- a/tests/unit/test_property_save_load.py +++ b/tests/unit/test_property_save_load.py @@ -18,7 +18,7 @@ from llmeter.prompt_utils import load_payloads, save_payloads from llmeter.results import Result from llmeter.runner import _RunConfig -from llmeter.tokenizers import DummyTokenizer, Tokenizer, save_tokenizer +from llmeter.tokenizers import DummyTokenizer # Custom strategies @@ -88,65 +88,36 @@ def valid_result(draw): # Tokenizer save/load property tests class TestTokenizerSaveLoadProperties: - """Property-based tests for tokenizer serialization.""" + """Property-based tests for tokenizer serialization via dump_object/load_object.""" @given(st.text(min_size=0, max_size=500)) @settings(deadline=None) - def test_save_load_tokenizer_roundtrip(self, text): - """Saving and loading tokenizer should preserve behavior.""" - with tempfile.TemporaryDirectory() as tmpdir: - original = DummyTokenizer() - output_path = Path(tmpdir) / "tokenizer.json" - - # Save - saved_path = save_tokenizer(original, output_path) - assert saved_path.exists() - - # Load - loaded = Tokenizer.load_from_file(saved_path) - - # Should encode the same way - original_tokens = original.encode(text) - loaded_tokens = loaded.encode(text) - assert original_tokens == loaded_tokens - - def test_save_tokenizer_creates_valid_json(self): - """Saved tokenizer should be valid JSON.""" - with tempfile.TemporaryDirectory() as tmpdir: - tokenizer = DummyTokenizer() - output_path = Path(tmpdir) / "tokenizer.json" + def test_dump_load_tokenizer_roundtrip(self, text): + """dump_object/load_object should preserve tokenizer behavior.""" + from llmeter.serialization import dump_object, load_object - saved_path = save_tokenizer(tokenizer, output_path) + original = DummyTokenizer() - # Should be valid JSON - with open(saved_path, "r") as f: - data = json.load(f) + # Serialize and restore + data = dump_object(original) + loaded = load_object(data) - assert isinstance(data, dict) - assert "tokenizer_module" in data + # Should encode the same way + original_tokens = original.encode(text) + loaded_tokens = loaded.encode(text) + assert original_tokens == loaded_tokens - @given( - st.text( - min_size=1, - max_size=100, - alphabet=st.characters( - min_codepoint=32, max_codepoint=126, blacklist_characters='\\/:*?"<>|' - ), - ) - ) - @settings(deadline=None) - def test_save_tokenizer_with_various_filenames(self, filename): - """Tokenizer should save with various valid filenames.""" - with tempfile.TemporaryDirectory() as tmpdir: - tokenizer = DummyTokenizer() - output_path = Path(tmpdir) / f"{filename}.json" + def test_dump_object_produces_valid_structure(self): + """dump_object should produce a dict with __llmeter_class__ and __llmeter_state__.""" + from llmeter.serialization import dump_object - saved_path = save_tokenizer(tokenizer, output_path) - assert saved_path.exists() + tokenizer = DummyTokenizer() + data = dump_object(tokenizer) - # Should be loadable - loaded = Tokenizer.load_from_file(saved_path) - assert isinstance(loaded, DummyTokenizer) + assert isinstance(data, dict) + assert "__llmeter_class__" in data + assert "__llmeter_state__" in data + assert "DummyTokenizer" in data["__llmeter_class__"] # Payload save/load property tests @@ -302,6 +273,40 @@ def test_run_config_creates_parent_directories(self, config): assert config_file.exists() assert config_file.parent.exists() + def test_run_config_save_load_roundtrip(self): + """save/load should round-trip a real endpoint via the ``__llmeter_class__`` envelope. + + The property tests above only exercise the *save* side (with a mock endpoint), so this + is the coverage for ``_RunConfig.load`` reconstructing a real endpoint from the current + serialization format. It also pins that format (no legacy ``endpoint_type`` key), so a + future backward-compat change can't silently break current-format loading. Legacy-format + loading is covered separately in ``test_legacy_data_load.py``. + """ + from llmeter.endpoints.bedrock import BedrockConverseStream + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = Path(tmpdir) + config = _RunConfig( + endpoint=BedrockConverseStream( + model_id="apac.amazon.nova-pro-v1:0", region="us-east-1" + ), + clients=3, + n_requests=7, + ) + config.save(output_path) + + # New saves use the type-tagged envelope, not the legacy endpoint_type key + saved = json.loads((output_path / "run_config.json").read_text()) + assert "__llmeter_class__" in saved["endpoint"] + assert "endpoint_type" not in saved["endpoint"] + + loaded = _RunConfig.load(output_path) + assert isinstance(loaded._endpoint, BedrockConverseStream) + assert loaded._endpoint.model_id == "apac.amazon.nova-pro-v1:0" + assert loaded._endpoint.region == "us-east-1" + assert loaded.clients == 3 + assert loaded.n_requests == 7 + # Result save/load property tests class TestResultSaveLoadProperties: @@ -393,10 +398,11 @@ def test_endpoint_save_load_roundtrip(self, model_id): saved_path = endpoint.save(output_path) assert saved_path.exists() - # Verify file contains model_id + # Verify file uses envelope format with model_id in state with open(saved_path, "r") as f: data = json.load(f) - assert data["model_id"] == model_id + assert "__llmeter_class__" in data + assert data["__llmeter_state__"]["model_id"] == model_id @given(st.text(min_size=1, max_size=50)) @settings(deadline=None) diff --git a/tests/unit/test_property_simple.py b/tests/unit/test_property_simple.py index bce566f..184434f 100644 --- a/tests/unit/test_property_simple.py +++ b/tests/unit/test_property_simple.py @@ -14,7 +14,7 @@ from llmeter.endpoints.base import InvocationResponse from llmeter.endpoints.openai import OpenAICompletionEndpoint from llmeter.prompt_utils import load_payloads, save_payloads -from llmeter.tokenizers import DummyTokenizer, _to_dict +from llmeter.tokenizers import DummyTokenizer # Custom strategies @@ -69,13 +69,15 @@ def test_dummy_tokenizer_decode_returns_string(self, tokens): @given(st.text(min_size=0, max_size=1000)) def test_tokenizer_serialization(self, text): - """Tokenizer should serialize correctly.""" + """Tokenizer should serialize via dump_object correctly.""" + from llmeter.serialization import dump_object + tokenizer = DummyTokenizer() - tokenizer_dict = _to_dict(tokenizer) + data = dump_object(tokenizer) - assert isinstance(tokenizer_dict, dict) - assert "tokenizer_module" in tokenizer_dict - assert tokenizer_dict["tokenizer_module"] == "llmeter" + assert isinstance(data, dict) + assert "__llmeter_class__" in data + assert "DummyTokenizer" in data["__llmeter_class__"] # OpenAI endpoint property tests diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py index 467f51a..5c39c56 100644 --- a/tests/unit/test_results.py +++ b/tests/unit/test_results.py @@ -221,10 +221,10 @@ def test_stats_property_empty_result(): def test_stats_json_serializable_with_datetimes(): - """stats dict should be JSON-serializable via llmeter_default_serializer.""" + """stats dict should be JSON-serializable via json_default.""" from datetime import datetime, timezone - from llmeter.json_utils import llmeter_default_serializer + from llmeter.serialization import json_default result = Result( responses=sample_responses_successful[:2], @@ -237,8 +237,8 @@ def test_stats_json_serializable_with_datetimes(): ) stats = result.stats # stats contains raw datetime objects; serialization is handled at the - # boundary (e.g. save()) via llmeter_default_serializer. - json_str = json.dumps(stats, default=llmeter_default_serializer) + # boundary (e.g. save()) via json_default. + json_str = json.dumps(stats, default=json_default) parsed = json.loads(json_str) assert parsed["start_time"] == "2025-01-01T12:00:00Z" assert parsed["end_time"] == "2025-01-01T12:00:01Z" @@ -407,11 +407,13 @@ def test_save_method_existing_responses(sample_result: Result, temp_dir: UPath): assert responses[-1]["id"] == "extra_response" -# ── _parse_datetime_fields introspection ─────────────────────────────────────── +# ── restore_dataclass_types introspection ────────────────────────────────────── -def test_parse_datetime_fields_converts_iso_strings(): - """_parse_datetime_fields should convert all datetime-typed field keys.""" +def test_restore_dataclass_types_converts_iso_strings(): + """restore_dataclass_types should convert all datetime-typed field keys.""" + from llmeter.serialization import restore_dataclass_types + d = { "start_time": "2025-06-01T10:00:00Z", "end_time": "2025-06-01T10:05:00+00:00", @@ -419,7 +421,7 @@ def test_parse_datetime_fields_converts_iso_strings(): "last_request_time": None, "total_requests": 5, # non-datetime field, should be left alone } - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert isinstance(d["start_time"], datetime) assert isinstance(d["end_time"], datetime) @@ -428,20 +430,24 @@ def test_parse_datetime_fields_converts_iso_strings(): assert d["total_requests"] == 5 -def test_parse_datetime_fields_skips_already_parsed(): +def test_restore_dataclass_types_skips_already_parsed(): """Already-datetime values should pass through unchanged.""" from datetime import timezone + from llmeter.serialization import restore_dataclass_types + dt = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc) d = {"start_time": dt} - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert d["start_time"] is dt -def test_parse_datetime_fields_handles_invalid_string(): +def test_restore_dataclass_types_handles_invalid_string(): """Invalid date strings should be left as-is (not raise).""" + from llmeter.serialization import restore_dataclass_types + d = {"start_time": "not-a-date"} - Result._parse_datetime_fields(d) + restore_dataclass_types(Result, d) assert d["start_time"] == "not-a-date" @@ -497,279 +503,6 @@ def test_load_missing_responses_jsonl_with_summary(sample_result: Result, temp_d assert "time_to_first_token-p50" in loaded.stats -# Tests for llmeter_default_serializer -# Validates Requirements: 7.1, 7.2 - - -def test_llmeter_encoder_handles_bytes(): - """Test that llmeter_default_serializer handles bytes objects via base64 marker. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Test bytes object encoding - payload = {"image": {"bytes": b"\xff\xd8\xff\xe0"}} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify marker object is present - assert "__llmeter_bytes__" in encoded - - # Verify it can be decoded back - decoded = json.loads(encoded, object_hook=lambda d: d) - assert decoded["image"]["bytes"]["__llmeter_bytes__"] == "/9j/4A==" - - -def test_llmeter_encoder_str_fallback(): - """Test that llmeter_default_serializer falls back to str() for custom objects. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Create a custom object with __str__ method - class CustomObject: - def __str__(self): - return "custom_string_representation" - - payload = {"custom": CustomObject()} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify str() fallback was used - decoded = json.loads(encoded) - assert decoded["custom"] == "custom_string_representation" - - -def test_llmeter_encoder_none_on_str_failure(): - """Test that llmeter_default_serializer returns None when str() conversion fails. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - # Create a custom object that raises exception in __str__ - class FailingObject: - def __str__(self): - raise RuntimeError("Cannot convert to string") - - payload = {"failing": FailingObject()} - encoded = json.dumps(payload, default=llmeter_default_serializer) - - # Verify None was returned - decoded = json.loads(encoded) - assert decoded["failing"] is None - - -def test_llmeter_encoder_mixed_types(): - """Test that llmeter_default_serializer handles mixed types correctly. - - Validates: Requirements 7.1, 7.2 - """ - from llmeter.json_utils import llmeter_default_serializer - - class CustomObject: - def __str__(self): - return "custom" - - # Mix of bytes, custom objects, and standard types - payload = { - "bytes_field": b"\x00\x01\x02", - "custom_field": CustomObject(), - "string_field": "normal string", - "int_field": 42, - "nested": {"bytes": b"\xff\xfe", "custom": CustomObject()}, - } - - encoded = json.dumps(payload, default=llmeter_default_serializer) - decoded = json.loads(encoded) - - # Verify bytes were encoded with marker - assert "__llmeter_bytes__" in encoded - - # Verify custom objects were converted to strings - assert decoded["custom_field"] == "custom" - assert decoded["nested"]["custom"] == "custom" - - # Verify standard types remain unchanged - assert decoded["string_field"] == "normal string" - assert decoded["int_field"] == 42 - - -# Tests for InvocationResponse.to_json -# Validates Requirements: 7.1, 7.2, 7.3 - - -def test_invocation_response_to_json_with_binary_content(): - """Test InvocationResponse.to_json with binary content in input_payload. - - Validates: Requirements 7.1, 7.2 - """ - # Create InvocationResponse with binary content in input_payload - response = InvocationResponse( - response_text="This is an image of a cat", - input_payload={ - "modelId": "anthropic.claude-3-haiku-20240307-v1:0", - "messages": [ - { - "role": "user", - "content": [ - {"text": "What is in this image?"}, - { - "image": { - "format": "jpeg", - "source": {"bytes": b"\xff\xd8\xff\xe0\x00\x10JFIF"}, - } - }, - ], - } - ], - }, - id="test-123", - time_to_first_token=0.5, - time_to_last_token=1.2, - num_tokens_input=15, - num_tokens_output=8, - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify it's valid JSON - parsed = json.loads(json_str) - - # Verify marker object is present for bytes - assert "__llmeter_bytes__" in json_str - - # Verify structure is preserved - assert parsed["response_text"] == "This is an image of a cat" - assert parsed["id"] == "test-123" - assert ( - parsed["input_payload"]["modelId"] == "anthropic.claude-3-haiku-20240307-v1:0" - ) - - # Verify bytes were encoded with marker - bytes_marker = parsed["input_payload"]["messages"][0]["content"][1]["image"][ - "source" - ]["bytes"] - assert "__llmeter_bytes__" in bytes_marker - assert bytes_marker["__llmeter_bytes__"] == "/9j/4AAQSkZJRg==" - - -def test_invocation_response_to_json_valid_json_output(): - """Test that InvocationResponse.to_json produces valid JSON. - - Validates: Requirements 7.1, 7.2 - """ - response = InvocationResponse( - response_text="Test response", - input_payload={ - "image_data": b"\x89PNG\r\n\x1a\n", - "text": "Analyze this image", - }, - id="test-456", - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify it's valid JSON (should not raise exception) - parsed = json.loads(json_str) - - # Verify all fields are present - assert "response_text" in parsed - assert "input_payload" in parsed - assert "id" in parsed - - # Verify bytes were properly encoded - assert isinstance(parsed["input_payload"]["image_data"], dict) - assert "__llmeter_bytes__" in parsed["input_payload"]["image_data"] - - -def test_invocation_response_to_json_round_trip(): - """Test round-trip serialization/deserialization with InvocationResponse. - - Validates: Requirements 7.1, 7.2, 7.3 - """ - from llmeter.json_utils import llmeter_bytes_decoder - - # Create original response with binary content - original_payload = { - "video": {"format": "mp4", "source": {"bytes": b"\x00\x00\x00\x18ftypmp42"}}, - "prompt": "Describe this video", - } - - response = InvocationResponse( - response_text="A video of a sunset", - input_payload=original_payload, - id="test-789", - time_to_first_token=1.0, - time_to_last_token=2.5, - ) - - # Serialize to JSON - json_str = response.to_json() - - # Deserialize back - parsed = json.loads(json_str, object_hook=llmeter_bytes_decoder) - - # Verify input_payload was restored correctly - assert parsed["input_payload"] == original_payload - assert isinstance(parsed["input_payload"]["video"]["source"]["bytes"], bytes) - assert ( - parsed["input_payload"]["video"]["source"]["bytes"] - == b"\x00\x00\x00\x18ftypmp42" - ) - - # Verify other fields - assert parsed["response_text"] == "A video of a sunset" - assert parsed["id"] == "test-789" - assert parsed["time_to_first_token"] == 1.0 - assert parsed["time_to_last_token"] == 2.5 - - -def test_invocation_response_to_json_no_binary_content(): - """Test InvocationResponse.to_json with no binary content (backward compatibility). - - Validates: Requirements 7.1, 7.2 - """ - response = InvocationResponse( - response_text="Simple text response", - input_payload={"modelId": "test-model", "prompt": "Hello, world!"}, - id="test-no-binary", - ) - - # Serialize to JSON - json_str = response.to_json() - - # Verify no marker objects are present - assert "__llmeter_bytes__" not in json_str - - # Verify it's valid JSON - parsed = json.loads(json_str) - assert parsed["input_payload"]["prompt"] == "Hello, world!" - - -def test_invocation_response_to_json_with_kwargs(): - """Test that InvocationResponse.to_json passes through kwargs to json.dumps. - - Validates: Requirements 7.4 - """ - response = InvocationResponse( - response_text="Test", input_payload={"data": b"\x01\x02\x03"}, id="test-kwargs" - ) - - # Test with indent parameter - json_str = response.to_json(indent=2) - - # Verify indentation is present - assert "\n" in json_str - assert " " in json_str - - # Verify it's still valid JSON - parsed = json.loads(json_str) - assert parsed["id"] == "test-kwargs" - - # ── Contributed stats round-trip ───────────────────────────────────────────── diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py index 616988e..66f4532 100644 --- a/tests/unit/test_runner.py +++ b/tests/unit/test_runner.py @@ -13,7 +13,7 @@ import llmeter.endpoints from llmeter.endpoints.base import Endpoint, InvocationResponse from llmeter.runner import Runner, _Run, _RunConfig -from llmeter.tokenizers import Tokenizer +from llmeter.tokenizers import DummyTokenizer @pytest.fixture @@ -33,11 +33,7 @@ def mock_endpoint(): @pytest.fixture def mock_tokenizer(): - with patch( - "llmeter.tokenizers.Tokenizer.to_dict", - return_value={"tokenizer_module": "mock_tokenizer"}, - ): - yield MagicMock(spec=Tokenizer) + yield DummyTokenizer() @pytest.fixture @@ -91,11 +87,9 @@ def test_runner_initialization(runner: Runner): def test_count_tokens_no_wait(runner: Runner): - # Test the tokenizer encode method directly since _count_tokens_no_wait doesn't exist - runner._tokenizer.encode.return_value = [1, 2, 3] - result = len(runner._tokenizer.encode("test text")) - assert result == 3 - runner._tokenizer.encode.assert_called_once_with("test text") + # Test the tokenizer encode method directly + result = len(runner._tokenizer.encode("test text here")) + assert result == 3 # DummyTokenizer splits on whitespace @pytest.mark.asyncio @@ -473,7 +467,11 @@ async def test_count_tokens_from_q_timeout(run: _Run): def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): - llmeter.endpoints.mock_endpoint = mock_endpoint # type: ignore + """Runner config can be saved and loaded with real endpoints.""" + from llmeter.endpoints import BedrockConverse + + # Use a real (lightweight) endpoint instead of a mock for save/load round-trip + endpoint = BedrockConverse(model_id="test-model", region="us-east-1") config = Runner( payload={"prompt": "test"}, @@ -482,7 +480,7 @@ def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): output_path=Path(tmp_path), run_name="test_run", run_description="Test run description", - endpoint=mock_endpoint, + endpoint=endpoint, ) config.save(output_path=tmp_path) @@ -498,6 +496,8 @@ def test_run_config_save_load(tmp_path: Path, mock_endpoint: Endpoint): assert loaded_config.output_path == config.output_path assert loaded_config.run_name == config.run_name assert loaded_config.run_description == config.run_description + assert loaded_config._endpoint.model_id == "test-model" + assert loaded_config._endpoint.region == "us-east-1" @pytest.mark.parametrize( diff --git a/tests/unit/test_runner_backlog_ui.py b/tests/unit/test_runner_backlog_ui.py index c684793..fd73c9a 100644 --- a/tests/unit/test_runner_backlog_ui.py +++ b/tests/unit/test_runner_backlog_ui.py @@ -37,11 +37,10 @@ def mock_endpoint(): @pytest.fixture def mock_tokenizer(): - with patch( - "llmeter.tokenizers.Tokenizer.to_dict", - return_value={"tokenizer_module": "mock_tokenizer"}, - ): - yield MagicMock(spec=Tokenizer) + # Tokenizer now serializes via the Serializable mixin (_get_llmeter_state), not a + # to_dict() method. These tests run without an output_path, so the tokenizer + # is never actually serialized — a bare spec'd mock is all that's needed. + yield MagicMock(spec=Tokenizer) @pytest.fixture @@ -382,9 +381,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() diff --git a/tests/unit/test_serialization_properties.py b/tests/unit/test_serialization_properties.py index 9452710..cd72a8e 100644 --- a/tests/unit/test_serialization_properties.py +++ b/tests/unit/test_serialization_properties.py @@ -1,12 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Property-based tests for llmeter_default_serializer and llmeter_bytes_decoder. +"""Property-based tests for json_default and bytes_decoder. This module contains property-based tests using Hypothesis to verify that -llmeter_default_serializer correctly handles all supported types (bytes, datetime, date, -time, PathLike, to_dict() objects, and str() fallback) and that -llmeter_bytes_decoder restores binary content from marker objects. +json_default correctly handles all supported types (bytes, datetime, date, +time, PathLike, and str() fallback) and that bytes_decoder restores binary +content from marker objects. Feature: json-serialization-optimization """ @@ -20,7 +20,7 @@ from hypothesis import strategies as st from hypothesis.strategies import composite -from llmeter.json_utils import llmeter_bytes_decoder, llmeter_default_serializer +from llmeter.serialization import bytes_decoder, json_default # Test infrastructure is set up and ready for property test implementation # This file will contain property-based tests for: @@ -111,7 +111,7 @@ def test_property_1_serialization_produces_valid_json_with_marker_objects( base64-encoded string value. """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Verify it's valid JSON by parsing it parsed = json.loads(serialized) @@ -171,7 +171,7 @@ def test_property_3_deserialization_restores_bytes_from_markers( marker_object = {"__llmeter_bytes__": base64_str} # Decode the marker object - decoded_bytes = llmeter_bytes_decoder(marker_object) + decoded_bytes = bytes_decoder(marker_object) # Verify the decoded bytes match the original assert isinstance(decoded_bytes, bytes), "Decoder should return bytes" @@ -195,7 +195,7 @@ def apply_decoder_recursively(obj): """Apply decoder to all dicts in the structure.""" if isinstance(obj, dict): # Apply decoder to this dict - decoded = llmeter_bytes_decoder(obj) + decoded = bytes_decoder(obj) # If it's still a dict (not converted to bytes), recurse if isinstance(decoded, dict): return {k: apply_decoder_recursively(v) for k, v in decoded.items()} @@ -245,7 +245,7 @@ def test_property_2_serialization_preserves_non_binary_structure(self, payload): # Feature: json-serialization-optimization, Property 2: Serialization preserves non-binary structure """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Parse the JSON (without decoding markers back to bytes) parsed = json.loads(serialized) @@ -326,10 +326,10 @@ def test_property_5_round_trip_serialization_preserves_data_integrity( """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize the payload - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Verify round-trip equality assert deserialized == payload, ( @@ -385,10 +385,10 @@ def test_property_6_round_trip_preserves_dictionary_key_ordering(self, payload): """ # Serialize the payload - serialized = json.dumps(payload, default=llmeter_default_serializer) + serialized = json.dumps(payload, default=json_default) # Deserialize the payload - deserialized = json.loads(serialized, object_hook=llmeter_bytes_decoder) + deserialized = json.loads(serialized, object_hook=bytes_decoder) # Helper function to verify key ordering def verify_key_ordering(original, restored, path=""): @@ -444,16 +444,14 @@ def test_property_7_non_binary_payloads_are_backward_compatible(self, payload): # Feature: json-serialization-optimization, Property 7: Non-binary payloads are backward compatible """ # Serialize with the new encoder - serialized_with_encoder = json.dumps( - payload, default=llmeter_default_serializer - ) + serialized_with_encoder = json.dumps(payload, default=json_default) # Serialize with standard json.dumps serialized_standard = json.dumps(payload) # Verify they produce identical output assert serialized_with_encoder == serialized_standard, ( - "Serialization with llmeter_default_serializer should produce identical output " + "Serialization with json_default should produce identical output " "to standard json.dumps for payloads without bytes objects" ) @@ -535,7 +533,7 @@ def test_property_8_serialization_handles_unknown_types(self, data): payload = payload_creator(unserializable_obj) # The unified encoder falls back to str() for unknown types - result = json.dumps(payload, default=llmeter_default_serializer) + result = json.dumps(payload, default=json_default) # Result should be valid JSON parsed = json.loads(result) assert isinstance(parsed, dict) @@ -578,7 +576,7 @@ def test_property_9_deserialization_errors_are_descriptive(self, data): # Attempt to deserialize and verify it raises JSONDecodeError try: - json.loads(invalid_json, object_hook=llmeter_bytes_decoder) + json.loads(invalid_json, object_hook=bytes_decoder) # If we get here, deserialization succeeded when it shouldn't have raise AssertionError( f"Expected JSONDecodeError for invalid JSON: {invalid_json}" @@ -749,27 +747,6 @@ def verify_bytes_serialized(original, serialized_obj, path=""): ) -@composite -def to_dict_object_strategy(draw): - """Generate an object with a to_dict() method returning a JSON-safe dict.""" - inner = draw( - st.dictionaries( - keys=st.text(min_size=1, max_size=10), - values=st.one_of(st.integers(), st.text(max_size=20), st.booleans()), - max_size=5, - ) - ) - - class _Obj: - def __init__(self, d): - self._d = d - - def to_dict(self): - return self._d - - return _Obj(inner), inner - - class TestDatetimeSerializationProperties: """Property-based tests for datetime/date/time encoding.""" @@ -782,7 +759,7 @@ def test_datetime_produces_iso_string_with_z_suffix(self, dt): Naive datetimes are serialized as-is with no timezone indicator. Microseconds are truncated (the encoder uses timespec="seconds"). """ - result = json.loads(json.dumps({"v": dt}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": dt}, default=json_default)) assert isinstance(result["v"], str) if dt.tzinfo is not None: assert result["v"].endswith("Z") @@ -797,7 +774,7 @@ def test_datetime_produces_iso_string_with_z_suffix(self, dt): @settings(max_examples=100) def test_date_produces_iso_string(self, d): """Date values are serialized via isoformat().""" - result = json.loads(json.dumps({"v": d}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": d}, default=json_default)) assert result["v"] == d.isoformat() assert date.fromisoformat(result["v"]) == d @@ -805,7 +782,7 @@ def test_date_produces_iso_string(self, d): @settings(max_examples=100) def test_time_produces_iso_string(self, t): """Time values are serialized via isoformat().""" - result = json.loads(json.dumps({"v": t}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": t}, default=json_default)) assert result["v"] == t.isoformat() assert time.fromisoformat(result["v"]) == t @@ -817,29 +794,34 @@ class TestPathSerializationProperties: @settings(max_examples=100) def test_pathlike_produces_posix_string(self, p): """PathLike objects are serialized to POSIX path strings.""" - result = json.loads(json.dumps({"v": p}, default=llmeter_default_serializer)) + result = json.loads(json.dumps({"v": p}, default=json_default)) assert isinstance(result["v"], str) assert result["v"] == p.as_posix() -class TestToDictSerializationProperties: - """Property-based tests for to_dict() delegation.""" +class TestJsonDefaultFallback: + """Example-based tests for the str()/None fallback branch of json_default.""" - @given(to_dict_object_strategy()) - @settings(max_examples=100) - def test_to_dict_delegation_produces_expected_dict(self, obj_and_expected): - """Objects with to_dict() are serialized by calling that method.""" - obj, expected = obj_and_expected - result = json.loads(json.dumps({"v": obj}, default=llmeter_default_serializer)) - assert result["v"] == expected + def test_str_fallback_for_custom_object(self): + """Unknown objects fall back to their str() representation.""" - @given(to_dict_object_strategy(), st.binary(min_size=1, max_size=100)) - @settings(max_examples=100) - def test_to_dict_and_bytes_coexist(self, obj_and_expected, raw_bytes): - """Payloads mixing to_dict() objects and bytes round-trip correctly.""" - obj, expected = obj_and_expected - payload = {"obj": obj, "data": raw_bytes} - serialized = json.dumps(payload, default=llmeter_default_serializer) - restored = json.loads(serialized, object_hook=llmeter_bytes_decoder) - assert restored["obj"] == expected - assert restored["data"] == raw_bytes + class CustomObject: + def __str__(self): + return "custom_string_representation" + + decoded = json.loads( + json.dumps({"custom": CustomObject()}, default=json_default) + ) + assert decoded["custom"] == "custom_string_representation" + + def test_none_when_str_conversion_fails(self): + """Objects whose __str__ raises are serialized as None.""" + + class FailingObject: + def __str__(self): + raise RuntimeError("Cannot convert to string") + + decoded = json.loads( + json.dumps({"failing": FailingObject()}, default=json_default) + ) + assert decoded["failing"] is None diff --git a/tests/unit/test_serialize_value_roundtrip.py b/tests/unit/test_serialize_value_roundtrip.py new file mode 100644 index 0000000..68265cd --- /dev/null +++ b/tests/unit/test_serialize_value_roundtrip.py @@ -0,0 +1,195 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests that _serialize_value/_deserialize_value correctly round-trip datetime and bytes.""" + +import base64 +from datetime import datetime, timezone + +import pytest + +from llmeter.serialization import ( + Serializable, + _deserialize_value, + _serialize_value, +) + + +class TestBytesRoundTrip: + """_serialize_value must produce the __llmeter_bytes__ marker for bytes, + and _deserialize_value must restore it.""" + + def test_bytes_serializes_to_marker(self): + val = b"\x00\x01\x02\xff" + result = _serialize_value(val) + assert isinstance(result, dict) + assert "__llmeter_bytes__" in result + assert len(result) == 1 + assert base64.b64decode(result["__llmeter_bytes__"]) == val + + def test_bytes_roundtrip(self): + val = b"hello world" + assert _deserialize_value(_serialize_value(val)) == val + + def test_empty_bytes_roundtrip(self): + val = b"" + assert _deserialize_value(_serialize_value(val)) == val + + def test_bytes_nested_in_dict(self): + val = {"key": b"\xde\xad\xbe\xef", "other": "text"} + serialized = _serialize_value(val) + assert serialized["key"] == { + "__llmeter_bytes__": base64.b64encode(b"\xde\xad\xbe\xef").decode() + } + assert serialized["other"] == "text" + assert _deserialize_value(serialized) == val + + def test_bytes_nested_in_list(self): + val = [b"first", "middle", b"last"] + serialized = _serialize_value(val) + restored = _deserialize_value(serialized) + assert restored == val + + def test_bytes_deeply_nested(self): + val = {"a": {"b": [{"c": b"\x01\x02\x03"}]}} + assert _deserialize_value(_serialize_value(val)) == val + + +class TestDatetimeRoundTrip: + """_serialize_value produces ISO-8601 Z-suffixed strings for datetime, + and _deserialize_value recognizes the format and restores datetime objects.""" + + def test_datetime_serializes_to_iso_string(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + result = _serialize_value(dt) + assert isinstance(result, str) + assert result == "2024-06-15T10:30:00Z" + + def test_datetime_utc_roundtrip(self): + dt = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + restored = _deserialize_value(_serialize_value(dt)) + assert restored == dt + assert isinstance(restored, datetime) + + def test_datetime_naive_roundtrip(self): + dt = datetime(2024, 3, 15, 8, 45, 30) + serialized = _serialize_value(dt) + # Naive datetimes produce no Z suffix, so they stay as plain strings + assert serialized == "2024-03-15T08:45:30" + + def test_datetime_is_not_silently_lost(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + result = _serialize_value(dt) + # Must be the canonical format, not repr/str() of the datetime object + assert "datetime" not in result + assert result.endswith("Z") + + def test_datetime_nested_in_dict(self): + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + val = {"timestamp": dt, "label": "test"} + serialized = _serialize_value(val) + assert serialized["timestamp"] == "2024-06-15T10:30:00Z" + restored = _deserialize_value(serialized) + assert restored["timestamp"] == dt + assert restored["label"] == "test" + + def test_datetime_nested_in_list(self): + dt1 = datetime(2024, 1, 1, tzinfo=timezone.utc) + dt2 = datetime(2024, 12, 31, tzinfo=timezone.utc) + val = [dt1, "gap", dt2] + restored = _deserialize_value(_serialize_value(val)) + assert restored[0] == dt1 + assert restored[1] == "gap" + assert restored[2] == dt2 + + def test_plain_string_not_confused_for_datetime(self): + val = "hello world" + assert _deserialize_value(val) == "hello world" + + def test_partial_iso_string_stays_as_string(self): + val = "2024-06-15" + assert _deserialize_value(val) == "2024-06-15" + + +class TestPathSerialization: + """PathLike objects serialize to POSIX path strings.""" + + def test_path_serializes_to_string(self): + from pathlib import PurePosixPath + + p = PurePosixPath("/tmp/foo/bar") + result = _serialize_value(p) + assert isinstance(result, str) + assert result == "/tmp/foo/bar" + + +class TestSerializableWithDatetimeAndBytes: + """End-to-end: a Serializable class with datetime/bytes fields round-trips + via _get_llmeter_state/_set_llmeter_state.""" + + def test_serializable_with_datetime_field(self): + class MyObj(Serializable): + def __init__(self, created_at: datetime, name: str = "test"): + self.created_at = created_at + self.name = name + + dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + obj = MyObj(created_at=dt, name="hello") + state = obj._get_llmeter_state() + + restored = MyObj.__new__(MyObj) + restored._set_llmeter_state(state) + + assert restored.created_at == dt + assert isinstance(restored.created_at, datetime) + assert restored.name == "hello" + + def test_serializable_with_bytes_field(self): + class MyObj(Serializable): + def __init__(self, payload: bytes, label: str = "x"): + self.payload = payload + self.label = label + + obj = MyObj(payload=b"\xde\xad\xbe\xef", label="binary") + state = obj._get_llmeter_state() + + restored = MyObj.__new__(MyObj) + restored._set_llmeter_state(state) + + assert restored.payload == b"\xde\xad\xbe\xef" + assert isinstance(restored.payload, bytes) + assert restored.label == "binary" + + def test_serializable_with_mixed_fields(self): + class MyObj(Serializable): + def __init__(self, ts: datetime, data: bytes, info: dict): + self.ts = ts + self.data = data + self.info = info + + dt = datetime(2024, 1, 1, tzinfo=timezone.utc) + obj = MyObj(ts=dt, data=b"raw", info={"nested_bytes": b"\x01", "nested_ts": dt}) + state = obj._get_llmeter_state() + + restored = MyObj.__new__(MyObj) + restored._set_llmeter_state(state) + + assert restored.ts == dt + assert restored.data == b"raw" + assert restored.info["nested_bytes"] == b"\x01" + assert restored.info["nested_ts"] == dt + + +class TestUnserializableTypeRaises: + """_serialize_value must raise TypeError for unsupported types, + never silently convert to str.""" + + def test_set_raises(self): + with pytest.raises(TypeError, match="Cannot serialize"): + _serialize_value({1, 2, 3}) + + def test_custom_object_without_getstate_raises(self): + class Opaque: + pass + + with pytest.raises(TypeError, match="Cannot serialize"): + _serialize_value(Opaque()) diff --git a/tests/unit/test_snapshot_load.py b/tests/unit/test_snapshot_load.py new file mode 100644 index 0000000..1cd55c8 --- /dev/null +++ b/tests/unit/test_snapshot_load.py @@ -0,0 +1,354 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests that load *static, synthetic* result snapshots from disk. + +Unlike the other load tests (which build data in-memory at test time), these load committed +fixture files under ``tests/unit/fixtures/result_snapshots/``. They pin the on-disk de/serialization +contract across the whole loading pipeline - ``Result.load``, ``_RunConfig.load``, +``LoadTestResult.load`` and the endpoint / tokenizer / callback restoration underneath - so format +drift can't silently break loading of previously-saved runs. + +The fixtures are deliberately synthetic (fake model IDs like ``synthetic.test-model-v1``, +lorem-ipsum responses, timestamps in the past) so nobody mistakes them for real benchmark data. +They cover a base "happy path" modern-format snapshot plus variations: legacy formats, an +interrupted run, error/annotation edge cases, and a multi-concurrency load test. + +To regenerate the fixtures (e.g. after an intentional format change), run: + uv run python tests/unit/fixtures/_generate_fixtures.py +""" + +from datetime import datetime, timezone + +import httpx +import pytest + +from llmeter.callbacks.cost.dimensions import InputTokens, OutputTokens +from llmeter.callbacks.cost.model import CostModel +from llmeter.callbacks.mlflow import MlflowCallback +from llmeter.endpoints.bedrock import BedrockConverse, BedrockConverseStream +from llmeter.endpoints.openai import OpenAICompletionStreamEndpoint +from llmeter.experiments import LoadTestResult +from llmeter.results import Result +from llmeter.runner import _RunConfig +from llmeter.tokenizers import DummyTokenizer + + +class TestBaseLoad: + """The modern-format baseline: OpenAI endpoint + CostModel + MlflowCallback.""" + + @pytest.fixture + def base_dir(self, snapshots_dir): + return snapshots_dir / "base" + + def test_result_loads(self, base_dir): + result = Result.load(base_dir) + assert len(result.responses) == 5 + assert result.total_requests == 5 + assert result.clients == 2 + assert result.n_requests == 3 + assert result.model_id == "synthetic.test-model-v1" + assert result.run_name == "synthetic-base-test" + + def test_null_output_path_filled_from_load_path(self, base_dir): + """A null output_path in summary.json must be filled with the real load path. + + This is what makes a relocated/copied result directory reloadable - the on-disk + path won't match wherever the data now lives. + """ + result = Result.load(base_dir) + assert result.output_path is not None + assert str(base_dir) in str(result.output_path) + + def test_timestamps_parsed_as_datetime(self, base_dir): + result = Result.load(base_dir) + assert result.start_time == datetime(2025, 1, 15, 10, 0, 0, tzinfo=timezone.utc) + assert result.end_time == datetime(2025, 1, 15, 10, 0, 4, tzinfo=timezone.utc) + assert isinstance(result.responses[0].request_time, datetime) + + def test_stats_match_frozen_values(self, base_dir): + """Stats loaded from stats.json (no responses) should match the frozen values.""" + result = Result.load(base_dir, load_responses=False) + stats = result.stats + assert stats["total_requests"] == 5 + assert stats["failed_requests"] == 0 + assert stats["total_input_tokens"] == 185 + assert stats["total_output_tokens"] == 386 + assert stats["time_to_first_token-p50"] == pytest.approx(0.415) + assert stats["requests_per_minute"] == pytest.approx(66.35700066357) + + def test_response_optional_token_fields(self, base_dir): + result = Result.load(base_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-base-002"].num_tokens_input_cached == 12 + assert by_id["synth-base-003"].num_tokens_output_reasoning == 8 + + # --- run_config.json / endpoint / tokenizer / callback restoration --- + + @pytest.fixture + def base_config(self, base_dir): + return _RunConfig.load(base_dir) + + def test_endpoint_restored(self, base_config): + ep = base_config._endpoint + assert isinstance(ep, OpenAICompletionStreamEndpoint) + assert ep.model_id == "synthetic.test-model-v1" + assert ep.endpoint_name == "synthetic-openai-stream" + assert ep.provider == "openai" + + def test_endpoint_client_config_restored(self, base_config): + """The OpenAI client's config surface (base_url, retries, timeout) must round-trip.""" + client = base_config._endpoint._client + assert str(client.base_url) == "https://api.synthetic-provider.example.com/v1/" + assert client.max_retries == 3 + assert isinstance(client.timeout, httpx.Timeout) + assert client.timeout.connect == 5.0 + assert client.timeout.read == 60.0 + assert client.timeout.write == 5.0 + assert client.timeout.pool == 5.0 + + def test_tokenizer_restored(self, base_config): + assert isinstance(base_config._tokenizer, DummyTokenizer) + + def test_cost_model_callback_restored(self, base_config): + cost_model = base_config.callbacks[0] + assert isinstance(cost_model, CostModel) + + input_dim = cost_model.request_dims["InputTokens"] + output_dim = cost_model.request_dims["OutputTokens"] + assert isinstance(input_dim, InputTokens) + assert isinstance(output_dim, OutputTokens) + assert input_dim.price_per_million == 3.0 + assert output_dim.price_per_million == 15.0 + assert cost_model.run_dims == {} + + def test_mlflow_callback_restored(self, base_config): + mlflow_cb = base_config.callbacks[1] + assert isinstance(mlflow_cb, MlflowCallback) + assert mlflow_cb.step is None + assert mlflow_cb.nested is True + + +class TestLegacyV01EndpointTypeLoad: + """Legacy v0.1 format: ``endpoint_type`` dispatch + flat cost annotations on responses.""" + + @pytest.fixture + def legacy_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_endpoint_type" + + def test_result_loads(self, legacy_dir): + result = Result.load(legacy_dir) + assert len(result.responses) == 5 + assert result.model_id == "synthetic.legacy-model-v1" + assert result.run_name == "synthetic-legacy-test" + + def test_legacy_cost_keys_collected_into_annotations(self, legacy_dir): + """Flat ``cost_*`` keys that current InvocationResponse doesn't declare go to annotations.""" + result = Result.load(legacy_dir) + first = result.responses[0] + assert first.annotations["cost_InputTokens"] == pytest.approx(0.000084) + assert first.annotations["cost_OutputTokens"] == pytest.approx(0.00108) + assert first.annotations["cost_total"] == pytest.approx(0.001164) + + def test_endpoint_restored_via_endpoint_type(self, legacy_dir): + cfg = _RunConfig.load(legacy_dir) + ep = cfg._endpoint + assert isinstance(ep, BedrockConverse) + assert ep.model_id == "synthetic.legacy-model-v1" + assert ep.region == "us-west-2" + # provider is derived by the constructor, not read from the (legacy) config dict + assert ep.provider == "bedrock" + + def test_tokenizer_restored_via_tokenizer_module(self, legacy_dir): + cfg = _RunConfig.load(legacy_dir) + assert isinstance(cfg._tokenizer, DummyTokenizer) + + +class TestLegacyV01StrCallbacksLoad: + """Legacy v0.1 data where callbacks were serialized as Python repr strings.""" + + @pytest.fixture + def strcb_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_str_callbacks" + + def test_result_loads(self, strcb_dir): + result = Result.load(strcb_dir) + assert len(result.responses) == 5 + assert result.model_id == "synthetic.strcb-model-v1" + + def test_config_loads_without_raising_on_str_callbacks(self, strcb_dir): + """Unrestorable string callbacks must not break run_config loading.""" + cfg = _RunConfig.load(strcb_dir) + assert isinstance(cfg._endpoint, BedrockConverseStream) + # String callbacks are left untouched (not restored to objects, but not fatal either) + assert cfg.callbacks is not None + assert len(cfg.callbacks) == 2 + assert all(isinstance(cb, str) for cb in cfg.callbacks) + + +class TestLegacyV01InterruptedRunLoad: + """A legacy-format interrupted run: no summary.json + top-level endpoint fields. + + Exercises the legacy branch of ``Result._recover_metadata``: older data kept endpoint + fields (including the derived ``provider``) at the top level rather than nested in a + ``__llmeter_state__`` envelope. Recovery must handle it - consistent with the legacy + handling in ``_RunConfig.load`` / ``Endpoint.load`` - and warn that the format is + deprecated. + """ + + @pytest.fixture + def legacy_interrupted_dir(self, snapshots_dir): + return snapshots_dir / "legacy" / "v0_1_interrupted" + + def test_no_summary_file_present(self, legacy_interrupted_dir): + assert not (legacy_interrupted_dir / "summary.json").exists() + + def test_recovers_top_level_endpoint_fields(self, legacy_interrupted_dir): + result = Result.load(legacy_interrupted_dir) + assert result.model_id == "synthetic.legacy-interrupted-v1" + assert result.endpoint_name == "synthetic-legacy-intr-endpoint" + # Legacy configs persist the derived provider at the top level, so it recovers + assert result.provider == "bedrock" + assert result.clients == 2 + assert result.n_requests == 4 + + def test_recovery_warns_legacy_format(self, legacy_interrupted_dir, caplog): + import logging + + with caplog.at_level(logging.WARNING, logger="llmeter.results"): + Result.load(legacy_interrupted_dir) + assert any("legacy data format" in r.message for r in caplog.records) + + +class TestInterruptedRunLoad: + """A snapshot with no summary.json - exercises metadata recovery.""" + + @pytest.fixture + def interrupted_dir(self, snapshots_dir): + return snapshots_dir / "interrupted_run" + + def test_no_summary_file_present(self, interrupted_dir): + assert not (interrupted_dir / "summary.json").exists() + + def test_recovers_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + assert len(result.responses) == 5 + assert result.total_requests == 5 + + def test_recovers_metadata_from_run_config(self, interrupted_dir): + result = Result.load(interrupted_dir) + assert result.clients == 4 + assert result.n_requests == 20 + assert result.run_name == "synthetic-interrupted-test" + + def test_recovers_endpoint_fields_from_modern_config(self, interrupted_dir): + """Endpoint fields must be read from the modern ``__llmeter_state__`` envelope. + + A *current* interrupted run writes run_config.json via ``dump_object`` (nested + state), so recovery must look inside ``__llmeter_state__`` and not just the top + level. ``provider`` is derived by the endpoint constructor and not persisted, so + it stays ``None`` on this lightweight recovery path (no endpoint instantiation). + """ + result = Result.load(interrupted_dir) + assert result.model_id == "synthetic.interrupted-model-v1" + assert result.endpoint_name == "synthetic-interrupted-endpoint" + assert result.provider is None + + def test_derives_timestamps_from_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + # request_time runs 16:30:02 .. 16:30:10 (i*2 for i in 1..5) + assert result.start_time == datetime( + 2025, 1, 20, 16, 30, 2, tzinfo=timezone.utc + ) + assert result.end_time == datetime(2025, 1, 20, 16, 30, 10, tzinfo=timezone.utc) + + def test_modern_config_recovery_does_not_warn_legacy(self, interrupted_dir, caplog): + """The modern-format recovery path must not emit the legacy-format warning.""" + import logging + + with caplog.at_level(logging.WARNING, logger="llmeter.results"): + Result.load(interrupted_dir) + assert not any("legacy data format" in r.message for r in caplog.records) + + def test_computes_stats_from_responses(self, interrupted_dir): + result = Result.load(interrupted_dir) + stats = result.stats + assert stats["failed_requests"] == 0 + assert "time_to_first_token-p50" in stats + + +class TestErrorsAndAnnotationsLoad: + """Responses carrying errors, retries, cached tokens, and custom annotations.""" + + @pytest.fixture + def errors_dir(self, snapshots_dir): + return snapshots_dir / "errors_and_annotations" + + def test_result_loads(self, errors_dir): + result = Result.load(errors_dir) + assert len(result.responses) == 5 + + def test_errors_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-002"].error == "ThrottlingException: Rate exceeded" + assert by_id["synth-err-004"].error == "ReadTimeoutError: timed out after 60s" + assert by_id["synth-err-001"].error is None + + def test_failed_requests_counted(self, errors_dir): + result = Result.load(errors_dir) + assert result.stats["failed_requests"] == 2 + + def test_retries_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-002"].retries == 2 + assert by_id["synth-err-005"].retries == 1 + assert by_id["synth-err-001"].retries == 0 + + def test_custom_annotations_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-001"].annotations == { + "custom_metric": 42.5, + "experiment_tag": "baseline", + } + assert by_id["synth-err-005"].annotations["cache_hit_ratio"] == pytest.approx( + 0.667 + ) + + def test_cached_input_tokens_preserved(self, errors_dir): + result = Result.load(errors_dir) + by_id = {r.id: r for r in result.responses} + assert by_id["synth-err-005"].num_tokens_input_cached == 30 + + +class TestLoadTestResultLoad: + """A multi-concurrency snapshot loaded via LoadTestResult.load.""" + + @pytest.fixture + def load_test_dir(self, snapshots_dir): + return snapshots_dir / "load_test" + + def test_loads_all_concurrency_levels(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + assert set(lt.results.keys()) == {1, 3} + + def test_sub_results_have_correct_data(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + assert len(lt.results[1].responses) == 5 + assert len(lt.results[3].responses) == 5 + assert lt.results[1].clients == 1 + assert lt.results[3].clients == 3 + + def test_sub_result_stats_populated(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir) + for clients in (1, 3): + stats = lt.results[clients].stats + assert stats["total_requests"] == 5 + assert "time_to_first_token-p50" in stats + + def test_load_without_responses(self, load_test_dir): + lt = LoadTestResult.load(load_test_dir, load_responses=False) + assert lt.results[1].responses == [] + assert lt.results[1].stats["total_requests"] == 5 diff --git a/tests/unit/test_tokenizers.py b/tests/unit/test_tokenizers.py index 0aceedb..6df5918 100644 --- a/tests/unit/test_tokenizers.py +++ b/tests/unit/test_tokenizers.py @@ -1,16 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -import json - import pytest from llmeter.tokenizers import ( DummyTokenizer, Tokenizer, _load_tokenizer_from_info, - _to_dict, - save_tokenizer, ) # Check for optional dependencies @@ -65,77 +61,16 @@ def test_dummy_tokenizer(): assert tokenizer.decode(tokens) == text -# Test Tokenizer.load_from_file -def test_load_from_file(tmp_path): - file_path = tmp_path / "tokenizer.json" - tokenizer_info = {"tokenizer_module": "llmeter"} - with open(file_path, "w") as f: - json.dump(tokenizer_info, f) - - tokenizer = Tokenizer.load_from_file(file_path) - assert isinstance(tokenizer, DummyTokenizer) - - -def test_load_from_file_none(): - tokenizer = Tokenizer.load_from_file(None) - assert isinstance(tokenizer, DummyTokenizer) - - -# Test Tokenizer.load +# Test Tokenizer.load (legacy format) def test_load(): tokenizer_info = {"tokenizer_module": "llmeter"} tokenizer = Tokenizer.load(tokenizer_info) assert isinstance(tokenizer, DummyTokenizer) -# Test Tokenizer.to_dict -def test_to_dict(): - dummy_tokenizer = DummyTokenizer() - tokenizer_dict = Tokenizer.to_dict(dummy_tokenizer) - assert tokenizer_dict == {"tokenizer_module": "llmeter"} - - -# Test _to_dict function -def test_to_dict_transformers(monkeypatch): - # monkeypatch.setattr("llmeter.tokenizers.AutoTokenizer", MockTransformersTokenizer) - tokenizer = MockTransformersTokenizer() - tokenizer_dict = _to_dict(tokenizer) - assert tokenizer_dict == { - "tokenizer_module": "transformers", - "name": "mock-transformer", - } - - -def test_to_dict_tiktoken(monkeypatch): - tokenizer = MockTiktokenTokenizer() - tokenizer_dict = _to_dict(tokenizer) - assert tokenizer_dict == {"tokenizer_module": "tiktoken", "name": "mock-tiktoken"} - - -def test_to_dict_unknown(): - class UnknownTokenizer: - pass - - with pytest.raises(ValueError, match="Unknown tokenizer module"): - _to_dict(UnknownTokenizer()) - - -# Test save_tokenizer function -def test_save_tokenizer(tmp_path): - dummy_tokenizer = DummyTokenizer() - output_path = tmp_path / "tokenizer.json" - saved_path = save_tokenizer(dummy_tokenizer, output_path) - assert saved_path == output_path - assert output_path.exists() - with open(output_path, "r") as f: - saved_info = json.load(f) - assert saved_info == {"tokenizer_module": "llmeter"} - - # Test _load_tokenizer_from_info function @pytest.mark.skipif(not TRANSFORMERS_AVAILABLE, reason="transformers is not installed") def test_load_tokenizer_from_info_transformers(monkeypatch): - # Mock AutoTokenizer.from_pretrained to return our mock def mock_from_pretrained(name): return MockTransformersTokenizer() @@ -150,7 +85,6 @@ def mock_from_pretrained(name): @pytest.mark.skipif(not TIKTOKEN_AVAILABLE, reason="tiktoken is not installed") def test_load_tokenizer_from_info_tiktoken(monkeypatch): - # Mock get_encoding to return our mock def mock_get_encoding(name): return MockTiktokenTokenizer() @@ -173,3 +107,21 @@ def test_load_tokenizer_from_info_unknown(): tokenizer_info = {"tokenizer_module": "unknown"} with pytest.raises(ValueError, match="Unknown tokenizer module"): _load_tokenizer_from_info(tokenizer_info) + + +# Test dump_object/load_object round-trip (new serialization) +def test_serialization_roundtrip(): + from llmeter.serialization import dump_object, load_object + + original = DummyTokenizer() + data = dump_object(original) + + assert "__llmeter_class__" in data + assert "llmeter.tokenizers.DummyTokenizer" in data["__llmeter_class__"] + + restored = load_object(data) + assert isinstance(restored, DummyTokenizer) + + # Behavior preserved + text = "hello world test" + assert original.encode(text) == restored.encode(text)