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._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 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,
                    OpenAIInstrumentation,
                    OpenAIAgentsInstrumentation,
                    HTTPXInstrumentation,
                ),
            ):
                self._instrumentation.append(item)
            else:
                self._instrumentors.append(item)
        return self

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

        automatic = AutoInstrumentation(exclude=tuple(exclude), strict=strict)
        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,
            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,
        concurrency_limit: int | None = 1,
    ) -> ExperimentResult:
        from autobench.runtime.pipeline import run_benchmark_spec

        return await run_benchmark_spec(
            self.to_spec(),
            experiment_id=experiment_id,
            concurrency_limit=concurrency_limit,
            instrumentors=self._instrumentors,
        )

    def run(
        self,
        *,
        experiment_id: str | None = None,
        concurrency_limit: int | None = 1,
    ) -> ExperimentResult:
        return run_sync(
            self.run_async(
                experiment_id=experiment_id,
                concurrency_limit=concurrency_limit,
            )
        )

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 CaseGeneratorInput(BaseModel):
    seed_cases: tuple[Case, ...] = ()
    prompt: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

class GeneratedCaseBatch(BaseModel):
    generator_asset_version: str | None = None
    model_provider: str | None = None
    model_name: str | None = None
    cases: tuple[Case, ...] = ()

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 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,
            )
        ]

class TokenCostDeriverSpec(BaseModel):
    kind: Literal["token_cost"] = "token_cost"
    output: DerivedMetricOutput = Field(
        default_factory=lambda: DerivedMetricOutput(
            name="cost",
            semantic_type=Semantic.MONEY_COST,
            unit="usd",
            direction=Direction.MINIMIZE,
            role=ObservationRole.CONSTRAINT,
        )
    )
    inputs: TokenCostInputs = Field(default_factory=TokenCostInputs)
    pricing: str = Field(min_length=1)

class TokenCostInputs(BaseModel):
    input_tokens: SemanticType = Semantic.LLM_TOKENS_INPUT
    output_tokens: SemanticType = Semantic.LLM_TOKENS_OUTPUT
    provider: SemanticType = Semantic.LLM_PROVIDER
    model: SemanticType = Semantic.LLM_MODEL_NAME

class CompositeExtractor:
    def __init__(self, *extractors: TraceExtractor) -> None:
        self.extractors = extractors or (
            SignalExtractor(),
            SpanExtractor(),
            UsageExtractor(),
        )
        self.name = "abp.default" if not extractors else "abp.composite"
        self.version = "+".join(
            f"{extractor.name}@{extractor.version}" for extractor in self.extractors
        )

    def extract(
        self,
        trace: Trace,
        *,
        registry: SemanticRegistry,
        context: ExtractionContext,
    ) -> ExtractionResult:
        observations: dict[str, Observation] = {}
        diagnostics: list[Diagnostic] = []
        references: dict[tuple[ReferenceKind, str, str | None], EvidenceRef] = {}
        for extractor in self.extractors:
            result = extractor.extract(trace, registry=registry, context=context)
            observations.update(
                (observation.id, observation) for observation in result.observations
            )
            diagnostics.extend(result.diagnostics)
            for reference in result.references:
                references[(reference.kind, reference.id, reference.version)] = reference
        unique_diagnostics = {
            (
                diagnostic.code,
                diagnostic.signal_id,
                diagnostic.span_id,
                diagnostic.sequence,
            ): diagnostic
            for diagnostic in diagnostics
        }
        return ExtractionResult(
            observations=tuple(observations.values()),
            diagnostics=tuple(unique_diagnostics.values()),
            references=tuple(references.values()),
        )

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

    run_id: str
    benchmark_id: str
    experiment_id: str
    case_id: str
    variant_id: str

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

    extractor: str
    version: str
    observation_ids: tuple[str, ...] = ()
    diagnostics: tuple[Diagnostic, ...] = ()
    references: tuple[EvidenceRef, ...] = ()

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

    observations: tuple[Observation, ...] = ()
    diagnostics: tuple[Diagnostic, ...] = ()
    references: tuple[EvidenceRef, ...] = ()

class SignalExtractor:
    name = "abp.signals"
    version = "2"

    def extract(
        self,
        trace: Trace,
        *,
        registry: SemanticRegistry,
        context: ExtractionContext,
    ) -> ExtractionResult:
        observations: list[Observation] = []
        references: dict[tuple[ReferenceKind, str, str | None], EvidenceRef] = {}
        for signal in trace.signals:
            if isinstance(signal, Measurement):
                attributes = signal.attributes
                kind_value = attributes.get("kind", ObservationKind.METRIC.value)
                kind = (
                    ObservationKind(kind_value)
                    if isinstance(kind_value, str)
                    and kind_value in {kind.value for kind in ObservationKind}
                    else ObservationKind.METRIC
                )
                source_value = attributes.get("source", ObservationSource.IMPORTED.value)
                source = (
                    ObservationSource(source_value)
                    if isinstance(source_value, str)
                    and source_value in {source.value for source in ObservationSource}
                    else ObservationSource.IMPORTED
                )
                tags_value = attributes.get("tags", {})
                tags = dict(tags_value) if isinstance(tags_value, dict) else {}
                tags.setdefault(MEASUREMENT_SCOPE_TAG, signal.measurement_scope.value)
                tags.setdefault(ABSTRACTION_LAYER_TAG, signal.layer.value)
                tags.setdefault(INSTRUMENTOR_TAG, signal.scope.instrumentor_name)
                logical_operation_id = _logical_operation_id(attributes)
                if logical_operation_id is not None:
                    tags.setdefault(LOGICAL_OPERATION_TAG, logical_operation_id)
                observations.append(
                    Observation(
                        id=_observation_id(signal.signal_id, attributes),
                        name=signal.name,
                        kind=kind,
                        semantic_type=registry.normalize(signal.semantic_type),
                        value=signal.value,
                        unit=signal.unit,
                        direction=signal.direction,
                        role=signal.role,
                        span_id=_span_id(signal.span_id, attributes),
                        source=source,
                        tags=tags,
                        case_id=context.case_id,
                        variant_id=context.variant_id,
                    )
                )
            elif isinstance(signal, Event):
                attributes = signal.attributes
                kind_value = attributes.get("kind", ObservationKind.EVENT.value)
                kind = (
                    ObservationKind(kind_value)
                    if isinstance(kind_value, str)
                    and kind_value in {kind.value for kind in ObservationKind}
                    else ObservationKind.EVENT
                )
                source_value = attributes.get("source", ObservationSource.IMPORTED.value)
                source = (
                    ObservationSource(source_value)
                    if isinstance(source_value, str)
                    and source_value in {source.value for source in ObservationSource}
                    else ObservationSource.IMPORTED
                )
                tags_value = attributes.get("tags", {})
                tags = dict(tags_value) if isinstance(tags_value, dict) else {}
                tags.setdefault(ABSTRACTION_LAYER_TAG, signal.scope.layer.value)
                tags.setdefault(INSTRUMENTOR_TAG, signal.scope.instrumentor_name)
                role = (
                    ObservationRole.DIAGNOSTIC
                    if registry.normalize(signal.semantic_type) == Semantic.DIAGNOSTIC_EVENT
                    else None
                )
                value = (
                    signal.reference.model_dump(mode="json")
                    if signal.reference is not None
                    else signal.body
                )
                observations.append(
                    Observation(
                        id=_observation_id(signal.signal_id, attributes),
                        name=signal.name,
                        kind=kind,
                        semantic_type=registry.normalize(signal.semantic_type),
                        value=value,
                        role=role,
                        span_id=_span_id(signal.span_id, attributes),
                        source=source,
                        tags=tags,
                        case_id=context.case_id,
                        variant_id=context.variant_id,
                    )
                )
                if signal.reference is not None:
                    reference = signal.reference
                    references[(reference.kind, reference.id, reference.version)] = reference

        for reference_signal in trace.references:
            reference = reference_signal.reference
            references[(reference.kind, reference.id, reference.version)] = reference
        for span in trace.spans:
            for reference_signal in span.references:
                reference = reference_signal.reference
                references[(reference.kind, reference.id, reference.version)] = reference
            for reference in span.errors:
                references[(reference.kind, reference.id, reference.version)] = reference

        return ExtractionResult(
            observations=tuple(observations),
            diagnostics=trace.diagnostics,
            references=tuple(references.values()),
        )