API Reference¶

class BenchContext(BaseModel):
    values: dict[str, Any] = Field(default_factory=dict)

    def set(self, key: str, value: Any) -> None:
        self.values[key] = value

    def get(self, key: str) -> Any:
        return self.values[key]

class Benchmark:
    def __init__(self, benchmark_id: str) -> None:
        self._benchmark = BenchmarkInfo(id=benchmark_id)
        self._capture: CapturePolicy | None = None
        self._execution = ExecutionSpec()
        self._dataset = DatasetSpec()
        self._task: TaskSpec | None = None
        self._variants: list[Variant] = []
        self._scoring: list[ScoringSpec] = []
        self._derive: list[DeriverSpec] = []
        self._instrumentation: list[InstrumentationConfig] = []
        self._instrumentors: list[Instrumentor] = []

    def description(self, value: str) -> Benchmark:
        self._benchmark = self._benchmark.model_copy(update={"description": value})
        return self

    def capture(self, policy: CapturePolicy | Mapping[str, Any]) -> Benchmark:
        self._capture = (
            policy if isinstance(policy, CapturePolicy) else CapturePolicy.model_validate(policy)
        )
        return self

    def correlation(
        self,
        value: ExecutionCorrelation | Mapping[str, Any],
    ) -> Benchmark:
        correlation = (
            value
            if isinstance(value, ExecutionCorrelation)
            else ExecutionCorrelation.model_validate(value)
        )
        self._execution = ExecutionSpec(correlation=correlation)
        return self

    def dataset(
        self,
        cases: list[Case | dict[str, Any]] | None = None,
        *,
        source: str | Path | None = None,
        dataset_id: str | None = None,
        version: str | None = None,
        metadata: dict[str, Any] | None = None,
        case_defaults: CaseDefaults | dict[str, Any] | None = None,
    ) -> Benchmark:
        defaults = (
            case_defaults
            if isinstance(case_defaults, CaseDefaults)
            else CaseDefaults.model_validate(case_defaults or {})
        )
        self._dataset = DatasetSpec(
            id=dataset_id,
            source=str(source) if source is not None else None,
            version=version,
            metadata=metadata or {},
            cases=[
                case if isinstance(case, Case) else Case.model_validate(case)
                for case in cases or []
            ],
            case_defaults=defaults,
        )
        return self

    def variants(self, variants: list[Variant | dict[str, Any]]) -> Benchmark:
        self._variants = [_normalize_variant(variant) for variant in variants]
        return self

    def task(self, target: str | TaskSpec, *, kind: str = "python") -> Benchmark:
        self._task = target if isinstance(target, TaskSpec) else TaskSpec(kind=kind, target=target)
        return self

    def scoring(self, scoring: list[ScoringSpec]) -> Benchmark:
        self._scoring = list(scoring)
        return self

    def derive(self, derive: list[DeriverSpec]) -> Benchmark:
        self._derive = list(derive)
        return self

    def instrument(
        self,
        *instrumentation: InstrumentationConfig | Instrumentor,
    ) -> Benchmark:
        """Add serializable settings or a custom runtime instrumentor."""

        for item in instrumentation:
            if isinstance(
                item,
                (
                    AutoInstrumentation,
                    PydanticAIInstrumentation,
                    PydanticGEPAInstrumentation,
                    OpenAIInstrumentation,
                    OpenAIAgentsInstrumentation,
                    HTTPXInstrumentation,
                ),
            ):
                self._instrumentation.append(item)
            else:
                self._instrumentors.append(item)
        return self

    def instrument_all(
        self,
        *,
        exclude: Collection[InstrumentorName] = (),
        strict: bool = False,
        assets: AssetDiscoverySettings | Mapping[str, Any] | None = None,
    ) -> Benchmark:
        """Enable every compatible built-in instrumentor available at runtime."""

        automatic = AutoInstrumentation(
            exclude=tuple(exclude),
            strict=strict,
            assets=(None if assets is None else AssetDiscoverySettings.model_validate(assets)),
        )
        self._instrumentation = [
            automatic,
            *(
                config
                for config in self._instrumentation
                if not isinstance(config, AutoInstrumentation)
            ),
        ]
        return self

    def to_spec(self) -> BenchmarkSpec:
        return BenchmarkSpec(
            benchmark=self._benchmark,
            capture=self._capture,
            execution=self._execution,
            dataset=self._dataset,
            task=self._task,
            variants=self._variants,
            scoring=self._scoring,
            derive=self._derive,
            instrumentation=self._instrumentation,
        )

    async def run_async(
        self,
        *,
        experiment_id: str | None = None,
        correlation: ExecutionCorrelation | None = None,
        concurrency_limit: int | None = 1,
        progress_handlers: Sequence[ProgressHandler] = (),
        progress_error_policy: ProgressErrorPolicy = ProgressErrorPolicy.STRICT,
        progress_error_handler: ProgressErrorHandler | None = None,
    ) -> ExperimentResult:
        from autobench.runtime.pipeline import run_benchmark_spec

        return await run_benchmark_spec(
            self.to_spec(),
            experiment_id=experiment_id,
            correlation=correlation,
            concurrency_limit=concurrency_limit,
            instrumentors=self._instrumentors,
            progress_handlers=progress_handlers,
            progress_error_policy=progress_error_policy,
            progress_error_handler=progress_error_handler,
        )

    def run(
        self,
        *,
        experiment_id: str | None = None,
        correlation: ExecutionCorrelation | None = None,
        concurrency_limit: int | None = 1,
        progress_handlers: Sequence[ProgressHandler] = (),
        progress_error_policy: ProgressErrorPolicy = ProgressErrorPolicy.STRICT,
        progress_error_handler: ProgressErrorHandler | None = None,
    ) -> ExperimentResult:
        return run_sync(
            self.run_async(
                experiment_id=experiment_id,
                correlation=correlation,
                concurrency_limit=concurrency_limit,
                progress_handlers=progress_handlers,
                progress_error_policy=progress_error_policy,
                progress_error_handler=progress_error_handler,
            )
        )

class Component(Protocol):  # pragma: no cover
    id: str
    kind: str

class Stage(Protocol):  # pragma: no cover
    id: str
    consumes: set[str]
    produces: set[str]

    async def run(self, ctx: BenchContext) -> None: ...

class Case(BaseModel):
    id: str = Field(min_length=1)
    input: Any = None
    expected: Any = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    tags: list[str] = Field(default_factory=list)
    attachments: list[ArtifactRef] = Field(default_factory=list)

class CaseDefaults(BaseModel):
    input: Any = None
    expected: Any = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    tags: list[str] = Field(default_factory=list)
    attachments: list[ArtifactRef] = Field(default_factory=list)

class DatasetSpec(BaseModel):
    id: str | None = None
    source: str | None = None
    version: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    cases: list[Case] = Field(default_factory=list)
    case_defaults: CaseDefaults = Field(default_factory=CaseDefaults)

class CaseGenerator(Protocol):
    def __call__(
        self,
        request: CaseGeneratorInput,
        /,
    ) -> GeneratedCaseBatch | Awaitable[GeneratedCaseBatch]: ...

class CaseGeneratorInput(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    seed_cases: tuple[Case, ...] = ()
    prompt: str | None = None
    prompt_asset_version: str | None = None
    seed: int | str | None = None
    settings: dict[str, SerializedValue] = Field(default_factory=dict)
    metadata: dict[str, SerializedValue] = Field(default_factory=dict)

class GeneratedCaseBatch(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    generator_asset_version: str | None = None
    model_provider: str | None = None
    model_name: str | None = None
    determinism: GenerationDeterminism = GenerationDeterminism.UNKNOWN
    usage: GenerationUsage = Field(default_factory=GenerationUsage)
    cost: GenerationCost | None = None
    reviews: tuple[GeneratedCaseReview, ...] = ()
    complete: bool = True
    incomplete_reason: str | None = Field(default=None, min_length=1)
    cases: tuple[Case, ...] = ()

    @model_validator(mode="after")
    def validate_batch(self) -> GeneratedCaseBatch:
        case_ids = tuple(case.id for case in self.cases)
        if len(case_ids) != len(set(case_ids)):
            raise ValueError("generated case ids must be unique")
        review_ids = tuple(review.case_id for review in self.reviews)
        if len(review_ids) != len(set(review_ids)):
            raise ValueError("generated case reviews must be unique")
        unknown = sorted(set(review_ids).difference(case_ids))
        if unknown:
            raise ValueError(
                f"generated case reviews reference unknown cases: {', '.join(unknown)}"
            )
        if self.complete and self.incomplete_reason is not None:
            raise ValueError("complete generation batches cannot have an incomplete reason")
        if not self.complete and self.incomplete_reason is None:
            raise ValueError("incomplete generation batches require a reason")
        return self

class GeneratedCaseRecord(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    case: Case
    review_status: ReviewStatus
    rejection_reason: str | None = None
    content_hash: str = Field(pattern=r"^[0-9a-f]{64}$")

    @model_validator(mode="after")
    def validate_review(self) -> GeneratedCaseRecord:
        if self.review_status is ReviewStatus.REJECTED and self.rejection_reason is None:
            raise ValueError("rejected generated case records require a rejection reason")
        if self.review_status is not ReviewStatus.REJECTED and self.rejection_reason is not None:
            raise ValueError("only rejected generated case records may have a rejection reason")
        if generated_case_content_hash(self.case) != self.content_hash:
            raise ValueError("generated case record content hash does not match its case")
        return self

class GeneratedCaseReview(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    case_id: str = Field(min_length=1)
    status: ReviewStatus = ReviewStatus.CANDIDATE
    rejection_reason: str | None = Field(default=None, min_length=1)

    @model_validator(mode="after")
    def validate_rejection(self) -> GeneratedCaseReview:
        if self.status is ReviewStatus.REJECTED and self.rejection_reason is None:
            raise ValueError("rejected generated cases require a rejection reason")
        if self.status is not ReviewStatus.REJECTED and self.rejection_reason is not None:
            raise ValueError("only rejected generated cases may have a rejection reason")
        return self

class GenerationCost(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid", allow_inf_nan=False)

    amount: float = Field(ge=0)
    currency: str = Field(default="usd", min_length=1)

class GenerationDeterminism(StrEnum):
    GUARANTEED = "guaranteed"
    NOT_GUARANTEED = "not_guaranteed"
    UNKNOWN = "unknown"

class GenerationResult(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    generator_id: str = Field(min_length=1)
    started_at: datetime
    completed_at: datetime
    request: CaseGeneratorInput
    request_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    batch: GeneratedCaseBatch
    generated_cases: tuple[GeneratedCaseRecord, ...]
    dataset: DatasetSpec | None = None
    dataset_hash: str | None = Field(default=None, pattern=r"^[0-9a-f]{64}$")

    @model_validator(mode="after")
    def validate_completion(self) -> GenerationResult:
        if self.completed_at < self.started_at:
            raise ValueError("generation completion cannot precede its start")
        if generation_request_hash(self.request) != self.request_hash:
            raise ValueError("generation request hash does not match its request")

        batch_case_ids = tuple(case.id for case in self.batch.cases)
        record_case_ids = tuple(record.case.id for record in self.generated_cases)
        if record_case_ids != batch_case_ids:
            raise ValueError("generated case records must match batch cases in order")
        if any(
            record.case != batch_case
            for record, batch_case in zip(self.generated_cases, self.batch.cases, strict=True)
        ):
            raise ValueError("generated case records must contain the batch case payloads")

        if self.batch.complete and (self.dataset is None or self.dataset_hash is None):
            raise ValueError("complete generation requires a frozen dataset")
        if not self.batch.complete and (self.dataset is not None or self.dataset_hash is not None):
            raise ValueError("incomplete generation cannot publish a benchmark dataset")
        if self.dataset is not None:
            included_cases = [
                record.case
                for record in self.generated_cases
                if record.review_status is not ReviewStatus.REJECTED
            ]
            if self.dataset.cases != included_cases:
                raise ValueError(
                    "generated dataset must contain every non-rejected case in generation order"
                )
            if dataset_content_hash(self.dataset) != self.dataset_hash:
                raise ValueError("generated dataset hash does not match its dataset")
        return self

class GenerationUsage(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    input_tokens: int = Field(default=0, ge=0)
    output_tokens: int = Field(default=0, ge=0)
    cached_input_tokens: int = Field(default=0, ge=0)
    requests: int = Field(default=0, ge=0)
    metadata: dict[str, SerializedValue] = Field(default_factory=dict)

class GenerationWriteResult(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    complete: bool
    dataset_path: Path | None = None
    manifest_path: Path

class ProductionSample(BaseModel):
    id: str
    input: Any = None
    output: Any = None
    expected: Any = None
    trace: TraceEnvelope | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    timestamp: datetime | None = None
    privacy_tags: tuple[str, ...] = ()
    reason: SampleReason = SampleReason.RANDOM
    review_status: ReviewStatus = ReviewStatus.CANDIDATE

class ReviewStatus(StrEnum):
    CANDIDATE = "candidate"
    ACCEPTED = "accepted"
    REJECTED = "rejected"

class SampleReason(StrEnum):
    RANDOM = "random"
    FAILURE_ONLY = "failure_only"
    LOW_CONFIDENCE = "low_confidence"
    HIGH_COST = "high_cost"
    HIGH_LATENCY = "high_latency"

class SamplingPolicy(BaseModel):
    reasons: tuple[SampleReason, ...] = (SampleReason.RANDOM,)
    max_samples: int | None = None

class FactorValue(BaseModel):
    name: str = Field(min_length=1)
    value: Any
    semantic_type: SemanticType | None = None
    optimize: bool = False

class Variant(BaseModel):
    id: str = Field(min_length=1)
    label: str | None = None
    factors: list[FactorValue] = Field(default_factory=list)

class AutobenchError(Exception):
    """Base exception for Autobench."""

class ErrorRecord(BaseModel):
    error_type: str
    message: str
    traceback: str | None = None
    span_id: str | None = None

    @classmethod
    def from_exception(
        cls,
        exc: BaseException,
        *,
        span_id: str | None = None,
        include_traceback: bool = True,
    ) -> ErrorRecord:
        rendered_traceback: str | None = None
        if include_traceback:
            rendered_traceback = "".join(
                traceback_module.format_exception(type(exc), exc, exc.__traceback__)
            )
        return cls(
            error_type=type(exc).__name__,
            message=str(exc),
            traceback=rendered_traceback,
            span_id=span_id,
        )

class GenerationError(AutobenchError):
    """Raised when generated dataset preparation cannot produce valid evidence."""

class SpecLoadError(AutobenchError):
    """Raised when a YAML spec cannot be loaded."""

    def __init__(
        self,
        message: str,
        *,
        path: Path | None = None,
        line: int | None = None,
        column: int | None = None,
    ) -> None:
        super().__init__(message)
        self.path = path
        self.line = line
        self.column = column

class SpecValidationError(AutobenchError):
    """Raised when a loaded YAML spec does not match the Autobench model."""

class TaskResolutionError(AutobenchError):
    """Raised when a Python task target cannot be resolved."""

class ActionMatchResult(BaseModel):
    expected: ExpectedAction
    matched_span_id: str | None = None
    target_matched: bool = False
    input_matched: bool = False
    output_matched: bool = False

    @property
    def matched(self) -> bool:
        return self.target_matched and self.input_matched

class ExpectedAction(BaseModel):
    id: str
    kind: str = "tool"
    target: str
    input: Any = None
    output: Any = None
    order: int | None = None
    required: bool = True
    tolerance: dict[str, Any] = Field(default_factory=dict)

class ComparisonVerdictSpec(BaseModel):
    output: DerivedMetricOutput
    threshold: RelativeThreshold = Field(default_factory=lambda: RelativeThreshold(pct=0.0))

class PairedBaselineDeriverSpec(BaseModel):
    kind: Literal["paired_baseline"] = "paired_baseline"
    baseline_variant: str = Field(min_length=1)
    match_on: tuple[RunMatchKey, ...] = Field(default_factory=lambda: (RunMatchKey(),))
    metric: SemanticType
    output: DerivedMetricOutput
    formula: PairedBaselineFormula = "baseline_over_candidate"
    threshold: RelativeThreshold | None = None
    verdict: ComparisonVerdictSpec | None = None
    include_baseline: bool = False
    missing: PostDerivationMissingPolicy = "diagnostic"
    zero_division: PostDerivationMissingPolicy = "diagnostic"
    diagnostics_name: str = "paired_baseline_unavailable"

class RelativeThreshold(BaseModel):
    kind: Literal["relative_noise"] = "relative_noise"
    pct: float = Field(ge=0.0)

class RunMatchKey(BaseModel):
    kind: RunMatchKeyKind = "case_id"
    name: str = ""

    @model_validator(mode="after")
    def _validate_name(self) -> RunMatchKey:
        if self.kind == "factor" and not self.name:
            raise ValueError("factor match keys require name")
        if self.kind == "case_id" and self.name:
            raise ValueError("case_id match keys cannot declare name")
        return self

class DerivedMetricOutput(BaseModel):
    name: str = Field(min_length=1)
    semantic_type: SemanticType
    unit: str | None = None
    direction: Direction | None = None
    role: ObservationRole | None = None

class TokenCostDeriver:
    def __init__(self, spec: TokenCostDeriverSpec) -> None:
        self.spec = spec

    def derive(
        self,
        *,
        ctx: RunContext,
        observations: list[Observation],
        registry: SemanticRegistry,
    ) -> list[Observation]:
        query = ObservationQuery(observations=observations, registry=registry)
        metric_kinds = (ObservationKind.METRIC, ObservationKind.FACTOR)

        input_tokens = query.first_related(
            self.spec.inputs.input_tokens,
            kind=metric_kinds,
        )
        output_tokens = query.first_related(
            self.spec.inputs.output_tokens,
            kind=metric_kinds,
        )
        provider = query.first_related(
            self.spec.inputs.provider,
            kind=metric_kinds,
        )
        model = query.first_related(
            self.spec.inputs.model,
            kind=metric_kinds,
        )

        if input_tokens is None or output_tokens is None or provider is None or model is None:
            missing = [
                name
                for name, value in (
                    ("input_tokens", input_tokens),
                    ("output_tokens", output_tokens),
                    ("provider", provider),
                    ("model", model),
                )
                if value is None
            ]
            return [
                _diagnostic_observation(
                    ctx=ctx,
                    name="token_cost_missing_inputs",
                    message="Missing inputs required for token cost derivation.",
                    tags={"missing": missing},
                )
            ]

        pricing = load_pricing_table(Path(self.spec.pricing))
        resolved_pricing = pricing.resolve_model_pricing(
            provider=str(provider.value),
            model=str(model.value),
        )
        if resolved_pricing is None:
            return [
                _diagnostic_observation(
                    ctx=ctx,
                    name="token_cost_unknown_pricing",
                    message="No pricing entry found for model/provider.",
                    tags={
                        "provider": str(provider.value),
                        "model": str(model.value),
                    },
                )
            ]
        resolved_model_id, model_pricing = resolved_pricing
        input_rate = model_pricing.input_rate_for_tokens(float(input_tokens.value))
        output_rate = model_pricing.output_rate_for_tokens(float(output_tokens.value))
        if input_rate is None or output_rate is None:
            return [
                _diagnostic_observation(
                    ctx=ctx,
                    name="token_cost_missing_rates",
                    message="Pricing entry did not define input/output rates.",
                    tags={
                        "provider": str(provider.value),
                        "model": str(model.value),
                        "model_id": resolved_model_id,
                    },
                )
            ]

        cost = (float(input_tokens.value) / 1_000_000.0) * input_rate + (
            float(output_tokens.value) / 1_000_000.0
        ) * output_rate
        return [
            Observation(
                id=ctx._next_observation_id(),
                name=self.spec.output.name,
                kind=ObservationKind.METRIC,
                semantic_type=self.spec.output.semantic_type,
                value=cost,
                unit=self.spec.output.unit,
                direction=self.spec.output.direction,
                role=self.spec.output.role,
                source=ObservationSource.DERIVED,
                tags={
                    "provider": str(provider.value),
                    "model": str(model.value),
                    "model_id": resolved_model_id,
                    "pricing_path": self.spec.pricing,
                },
                case_id=ctx.case.id,
                variant_id=ctx.variant.id,
            )
        ]