mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import tempfile
|
|
from collections.abc import Mapping
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Final, Literal, Protocol, TypeVar, cast
|
|
|
|
from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError
|
|
|
|
from . import JSON_OBJECT_ADAPTER
|
|
from .cassette import deserialize_cassette, serialize_cassette
|
|
from .recording import RecordedInteraction
|
|
|
|
FIXTURE_SCHEMA_VERSION: Final = 1
|
|
|
|
|
|
class FixtureInput(Protocol):
|
|
def canonical_input(self) -> dict[str, object]: ...
|
|
|
|
|
|
CaseT = TypeVar("CaseT", bound=BaseModel)
|
|
|
|
|
|
class FixtureEnvelope(BaseModel):
|
|
model_config = ConfigDict(frozen=True, extra="forbid")
|
|
|
|
schema_version: int
|
|
recorded_at: AwareDatetime
|
|
case: dict[str, object]
|
|
|
|
|
|
def canonical_json(value: Mapping[str, object]) -> str:
|
|
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
|
|
|
|
|
def fixture_cache_key(case_input: FixtureInput) -> dict[str, object]:
|
|
return case_input.canonical_input()
|
|
|
|
|
|
def _fixture_digest(case_input: FixtureInput) -> str:
|
|
input_json: Final = canonical_json(fixture_cache_key(case_input))
|
|
return hashlib.sha256(input_json.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def fixture_path(directory: Path, case_input: FixtureInput) -> Path:
|
|
return directory / f"{_fixture_digest(case_input)}.yaml"
|
|
|
|
|
|
def load_fixture(directory: Path, case_input: FixtureInput, case_type: type[CaseT]) -> CaseT | None:
|
|
path: Final = fixture_path(directory, case_input)
|
|
if path.is_file():
|
|
return read_fixture(path, case_type)
|
|
legacy_path: Final = path.with_suffix(".json")
|
|
if not legacy_path.is_file():
|
|
return None
|
|
return read_fixture(legacy_path, case_type)
|
|
|
|
|
|
def save_fixture(
|
|
directory: Path,
|
|
case_input: FixtureInput,
|
|
case: BaseModel,
|
|
interactions: tuple[RecordedInteraction, ...],
|
|
*,
|
|
recorded_at: datetime | None = None,
|
|
request_source: Literal["recorded", "python_replay"] = "recorded",
|
|
) -> Path:
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
path: Final = fixture_path(directory, case_input)
|
|
serialized: Final = serialize_cassette(
|
|
cast(dict[str, object], case.model_dump(mode="json", exclude_unset=True)),
|
|
interactions,
|
|
recorded_at or datetime.now(timezone.utc),
|
|
request_source,
|
|
)
|
|
with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=directory, delete=False) as temporary:
|
|
temporary_path: Final = Path(temporary.name)
|
|
try:
|
|
temporary.write(serialized)
|
|
temporary.close()
|
|
temporary_path.replace(path)
|
|
finally:
|
|
temporary_path.unlink(missing_ok=True)
|
|
return path
|
|
|
|
|
|
def read_fixture(path: Path, case_type: type[CaseT]) -> CaseT:
|
|
contents: Final = path.read_text(encoding="utf-8")
|
|
if path.suffix == ".json":
|
|
return _load_fixture(JSON_OBJECT_ADAPTER.validate_json(contents), path, case_type)
|
|
try:
|
|
cassette: Final = deserialize_cassette(contents)
|
|
return case_type.model_validate(cassette.case_data())
|
|
except ValueError as error:
|
|
raise ValueError(f"invalid parity cassette {path}") from error
|
|
|
|
|
|
def _load_fixture(raw_fixture: dict[str, object], path: Path, case_type: type[CaseT]) -> CaseT:
|
|
schema_version: Final = raw_fixture.get("schema_version")
|
|
if schema_version != FIXTURE_SCHEMA_VERSION:
|
|
raise ValueError(
|
|
f"fixture {path} has schema_version {schema_version!r}, expected {FIXTURE_SCHEMA_VERSION}; "
|
|
"delete it and regenerate the fixture bundle"
|
|
)
|
|
try:
|
|
envelope: Final = FixtureEnvelope.model_validate(raw_fixture)
|
|
return case_type.model_validate(envelope.case)
|
|
except ValidationError as error:
|
|
raise ValueError(f"invalid parity fixture {path} ({len(error.errors())} validation errors)") from error
|
|
|
|
|
|
def recorded_fixtures(directory: Path, case_type: type[CaseT]) -> tuple[CaseT, ...]:
|
|
if not directory.is_dir():
|
|
return ()
|
|
paths: Final = tuple(sorted((*directory.rglob("*.yaml"), *directory.rglob("*.json"))))
|
|
return tuple(read_fixture(path, case_type) for path in paths)
|
|
|
|
|
|
def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path:
|
|
if configured is not None:
|
|
return configured.expanduser()
|
|
if env_value == "":
|
|
raise ValueError("fixture directory environment variable is set but empty")
|
|
return Path(env_value).expanduser() if env_value is not None else default.expanduser()
|
|
|
|
|
|
def fixture_id(case_input: FixtureInput, prefix: str) -> str:
|
|
return f"{prefix}-{_fixture_digest(case_input)[:8]}"
|