Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Standardize object (de)-serialization in LLMeter #7

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Standardize object (de)-serialization in LLMeter #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 22 additions & 83 deletions llmeter/endpoints/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,20 @@
# SPDX-License-Identifier: Apache-2.0

import importlib
import json
import os
from abc import ABC, abstractmethod
from dataclasses import asdict, dataclass
from dataclasses import dataclass
from typing import Dict, TypeVar
from uuid import uuid4

from upath import UPath as Path
from llmeter.serde import JSONableBase

Self = TypeVar(
"Self", bound="Endpoint"
) # for python >= 3.11 can be replaced with direct import of `Self`


@dataclass
class InvocationResponse:
class InvocationResponse(JSONableBase):
"""
A class representing a invocation result.

Expand All @@ -43,9 +41,6 @@ class InvocationResponse:
time_per_output_token: float | None = None
error: str | None = None

def to_json(self, **kwargs) -> str:
return json.dumps(self.__dict__, **kwargs)

@staticmethod
def error_output(
input_prompt: str | None = None, error=None, id: str | None = None
Expand All @@ -64,11 +59,8 @@ def __repr__(self):
def __str__(self):
return self.to_json(indent=4, default=str)

def to_dict(self):
return asdict(self)


class Endpoint(ABC):
class Endpoint(JSONableBase, ABC):
"""
An abstract base class for endpoint implementations.

Expand Down Expand Up @@ -154,79 +146,26 @@ def __subclasshook__(cls, C):
return True
return NotImplemented

def save(self, output_path: os.PathLike) -> os.PathLike:
"""
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.

Args:
output_path (str | UPath): The path where the configuration file will be saved.

Returns:
None
"""
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w") as f:
endpoint_conf = self.to_dict()
json.dump(endpoint_conf, f, indent=4, default=str)
return output_path

def to_dict(self) -> Dict:
"""
Convert the endpoint configuration to a dictionary.

Returns:
Dict: A dictionary representation of the endpoint configuration.
"""
endpoint_conf = {k: v for k, v in vars(self).items() if not k.startswith("_")}
endpoint_conf["endpoint_type"] = self.__class__.__name__
return endpoint_conf

@classmethod
def load_from_file(cls, input_path: os.PathLike) -> Self:
"""
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.

Returns:
Endpoint: An instance of the appropriate endpoint class, initialized
with the configuration from the file.
"""

input_path = Path(input_path)
with input_path.open("r") as f:
data = json.load(f)
endpoint_type = data.pop("endpoint_type")
endpoint_module = importlib.import_module("llmeter.endpoints")
endpoint_class = getattr(endpoint_module, endpoint_type)
return endpoint_class(**data)

@classmethod
def load(cls, endpoint_config: Dict) -> Self: # type: ignore
"""
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.
def from_dict(
cls: Self, raw: Dict, alt_classes: Dict[str, Self] = {}, **kwargs
) -> Self:
"""Load any built-in Endpoint type (or custom ones) from a plain JSON dictionary

Args:
data (Dict): A dictionary containing the endpoint configuration.
raw: A plain Endpoint config dictionary, as created with `to_dict()`, `to_json`, etc.
alt_classes (Dict[str, type[Endpoint]]): A dictionary mapping additional custom type
names (beyond those in `llmeter.endpoints`, which are included automatically), to
corresponding classes for loading custom endpoint types.
**kwargs: Optional extra keyword arguments to pass to the constructor

Returns:
Endpoint: An instance of the appropriate endpoint class, initialized
with the configuration from the dictionary.
"""
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)
endpoint: An instance of the appropriate endpoint class, initialized with the
configuration from the file.
"""
builtin_endpoint_types = importlib.import_module("llmeter.endpoints")
class_map = {
**builtin_endpoint_types,
**alt_classes,
}
return super().from_dict(raw, alt_classes=class_map, **kwargs)
142 changes: 111 additions & 31 deletions llmeter/results.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,22 @@
from math import isnan
import os
from statistics import StatisticsError, mean, median, quantiles
from typing import Dict, Sequence
from typing import Any, Callable, Dict, Sequence, TypeVar

import jmespath
from upath import UPath as Path

from .endpoints import InvocationResponse
from .serde import from_dict_with_class_map, JSONableBase

logger = logging.getLogger(__name__)


TResult = TypeVar("TResult", bound="Result")


@dataclass
class Result:
class Result(JSONableBase):
"""Results of an experiment run."""

responses: list[InvocationResponse]
Expand Down Expand Up @@ -65,40 +69,108 @@ def save(self, output_path: os.PathLike | str | None = None):
The method uses the Universal Path (UPath) library for file operations,
which provides a unified interface for working with different file systems.
"""

try:
output_path = Path(self.output_path or output_path)
except TypeError:
raise ValueError("No output path provided")

output_path.mkdir(parents=True, exist_ok=True)

summary_path = output_path / "summary.json"
self.to_file(
output_path / "summary.json",
include_responses=False,
) # Already creates output_path folders if needed
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))
with stats_path.open("w") as s:
s.write(json.dumps(self.stats, indent=4, default=str))

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)) + "\n")

def to_json(self, **kwargs):
"""Return the results as a JSON string."""
summary = {
k: o for k, o in asdict(self).items() if k not in ["responses", "stats"]
}
return json.dumps(summary, default=str, **kwargs)
with responses_path.open("w") as f:
for response in self.responses:
f.write(json.dumps(asdict(response)) + "\n")

def to_dict(self, include_responses: bool = False):
"""Return the results as a dictionary."""
if include_responses:
return asdict(self)
return {
k: o for k, o in asdict(self).items() if k not in ["responses", "stats"]
}
@classmethod
def from_dict(
cls, raw: dict, alt_classes: dict[str, TResult] = {}, **kwargs
) -> TResult:
"""Load a run Result 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
"""
data = {**raw}
if "responses" in data:
data["responses"] = [
# Just in case users wanted to override InvocationResponse itself...
from_dict_with_class_map(
resp,
alt_classes={
InvocationResponse.__name__: InvocationResponse,
**alt_classes,
},
)
for resp in data["responses"]
]
else:
data["responses"] = []
data.pop("stats", None) # Calculated property should be omitted
return super().from_dict(data, alt_classes, **kwargs)

def to_dict(self, include_responses: bool = False, **kwargs) -> dict:
"""Save the results to a JSON-dumpable dictionary (with optional extra kwargs)

Args:
include_responses: Set `True` to include the `responses` and `stats` in the output.
By default, these will be omitted.
**kwargs: Additional fields to save in the output dictionary.
"""
result = super().to_dict(**kwargs)
if not include_responses:
result.pop("responses", None)
result.pop("stats", None)
return result

def to_file(
self,
output_path: os.PathLike,
include_responses: bool = False,
indent: int | str | None = 4,
default: Callable[[Any], Any] | None = {},
**kwargs,
) -> Path:
"""Save the Run Result to a (local or Cloud) JSON file

Args:
output_path: The path where the file will be saved.
include_responses: Set `True` to include the `responses` and `stats` in the output.
By default, these will be omitted.
indent: Optional indentation passed through to `to_json()` and therefore `json.dumps()`
default: Optional function to convert non-JSON-serializable objects to strings, passed
through to `to_json()` and therefore to `json.dumps()`
**kwargs: Optional extra keyword arguments to pass to `to_json()`

Returns:
output_path: Universal Path representation of the target file.
"""
return super().to_file(
output_path,
include_responses=include_responses,
indent=indent,
default=default,
**kwargs,
)

def to_json(self, include_responses: bool = False, **kwargs) -> str:
"""Serialize the results to JSON, with optional kwargs passed through to `json.dumps()`

Args:
include_responses: Set `True` to include the `responses` and `stats` in the output.
By default, these will be omitted.
**kwargs: Optional arguments to pass to `json.dumps()`.
"""
return json.dumps(self.to_dict(include_responses=include_responses), **kwargs)

@classmethod
def load(cls, result_path: os.PathLike | str):
Expand Down Expand Up @@ -126,12 +198,20 @@ def load(cls, result_path: os.PathLike | str):

"""
result_path = Path(result_path)
responses_path = result_path / "responses.jsonl"
summary_path = result_path / "summary.json"
with open(responses_path, "r") as f, summary_path.open("r") as g:
responses = [InvocationResponse(**json.loads(line)) for line in f if line]
summary = json.load(g)
return cls(responses=responses, **summary)
with summary_path.open("r") as f:
raw = json.load(f)

responses_path = result_path / "responses.jsonl"
try:
with responses_path.open("r") as f:
raw["responses"] = [
InvocationResponse.from_json(line) for line in f if line
]
except FileNotFoundError:
logger.info("Result.load: No responses data found at %s", responses_path)

return cls.from_dict(raw)

@cached_property
def stats(self) -> Dict:
Expand Down
Loading