spec_loader.py

"""Load BTX LLM-span spec YAML files.

Handles the custom YAML tags used in the spec:
  !fn <name-or-lambda>  — named predicate or arbitrary lambda (eval'd in Python)
  !starts_with <prefix> — string prefix check
  !or [...]             — at-least-one-of validator
  !gen <generator-name> — value generated by the test runner (e.g. test_runner_client)
"""

from __future__ import annotations

import dataclasses
import os
import uuid
from pathlib import Path
from typing import Any

import yaml


# ---------------------------------------------------------------------------
# Matcher / generator types
# ---------------------------------------------------------------------------


@dataclasses.dataclass
class FnMatcher:
    """A named or lambda-expression validator."""

    expr: str  # e.g. "is_non_negative_number" or "lambda value: value > 0"


@dataclasses.dataclass
class StartsWithMatcher:
    prefix: str


@dataclasses.dataclass
class OrMatcher:
    alternatives: list[Any]


@dataclasses.dataclass
class GenValue:
    """A value generated by the test runner at execution time.

    The generator name determines what value is produced:
      test_runner_client — a string identifying this SDK/client (e.g. "python-openai")
      vcr_nonce          — a random string that changes every run (busts caches)
    """

    generator: str  # e.g. "test_runner_client", "vcr_nonce"


# ---------------------------------------------------------------------------
# YAML custom constructors
# ---------------------------------------------------------------------------


def _fn_constructor(loader: yaml.SafeLoader, node: yaml.Node) -> FnMatcher:
    expr = loader.construct_scalar(node)  # type: ignore[arg-type]
    return FnMatcher(expr=expr)


def _starts_with_constructor(loader: yaml.SafeLoader, node: yaml.Node) -> StartsWithMatcher:
    prefix = loader.construct_scalar(node)  # type: ignore[arg-type]
    return StartsWithMatcher(prefix=prefix)


def _or_constructor(loader: yaml.SafeLoader, node: yaml.Node) -> OrMatcher:
    alternatives = loader.construct_sequence(node, deep=True)
    return OrMatcher(alternatives=alternatives)


def _gen_constructor(loader: yaml.SafeLoader, node: yaml.Node) -> GenValue:
    generator = loader.construct_scalar(node)  # type: ignore[arg-type]
    return GenValue(generator=generator)


def _make_loader() -> type:
    """Return a SafeLoader subclass with BTX custom tags registered."""

    class BtxLoader(yaml.SafeLoader):
        pass

    BtxLoader.add_constructor("!fn", _fn_constructor)
    BtxLoader.add_constructor("!starts_with", _starts_with_constructor)
    BtxLoader.add_constructor("!or", _or_constructor)
    BtxLoader.add_constructor("!gen", _gen_constructor)
    return BtxLoader


# ---------------------------------------------------------------------------
# Generator resolution
# ---------------------------------------------------------------------------

# Stable client identifier for this SDK implementation.
_CLIENT_ID = "python-btx"

# Per-process nonce — constant within a run so cassette body matching is stable,
# but differs across runs so cache-busting specs actually bust caches.
_VCR_NONCE = str(uuid.uuid4())[:8]

_GENERATORS: dict[str, str] = {
    "test_runner_client": _CLIENT_ID,
    "vcr_nonce": _VCR_NONCE,
}


def _resolve_gen(value: GenValue) -> str:
    if value.generator in _GENERATORS:
        return _GENERATORS[value.generator]
    raise ValueError(f"Unknown !gen generator: {value.generator!r}")


def _resolve_variables(variables: dict[str, Any]) -> dict[str, str]:
    """Resolve all !gen values in the variables map to concrete strings."""
    resolved: dict[str, str] = {}
    for key, val in variables.items():
        if isinstance(val, GenValue):
            resolved[key] = _resolve_gen(val)
        else:
            resolved[key] = str(val)
    return resolved


def _substitute_templates(obj: Any, variables: dict[str, str]) -> Any:
    """Recursively substitute {{var}} placeholders in strings."""
    if isinstance(obj, str):
        for key, value in variables.items():
            obj = obj.replace(f"{{{{{key}}}}}", value)
        return obj
    if isinstance(obj, dict):
        return {k: _substitute_templates(v, variables) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_substitute_templates(item, variables) for item in obj]
    return obj


# ---------------------------------------------------------------------------
# Spec dataclass
# ---------------------------------------------------------------------------


@dataclasses.dataclass
class LlmSpanSpec:
    name: str
    type: str
    provider: str
    endpoint: str
    requests: list[dict[str, Any]]
    expected_brainstore_spans: list[dict[str, Any]]
    headers: dict[str, str]
    source_path: Path

    @property
    def display_name(self) -> str:
        """pytest ID: <provider>/<name>"""
        return f"{self.provider}/{self.name}"

    @classmethod
    def from_dict(cls, data: dict[str, Any], source_path: Path) -> "LlmSpanSpec":
        # Resolve variables and substitute templates in requests
        raw_variables = data.get("variables", {})
        variables = _resolve_variables(raw_variables)
        requests = _substitute_templates(data.get("requests", []), variables)

        return cls(
            name=data["name"],
            type=data["type"],
            provider=data["provider"],
            endpoint=data["endpoint"],
            requests=requests,
            expected_brainstore_spans=data.get("expected_brainstore_spans", []),
            headers=data.get("headers", {}),
            source_path=source_path,
        )


# ---------------------------------------------------------------------------
# Loader
# ---------------------------------------------------------------------------

_BTX_DIR = Path(__file__).parent


def _spec_root(override: str | None = None) -> Path:
    if override:
        return Path(override)
    env = os.environ.get("BTX_SPEC_ROOT")
    if env:
        return Path(env)
    return _BTX_DIR / "spec" / "test" / "llm_span"


def load_specs(
    spec_root: str | Path | None = None,
    providers: list[str] | None = None,
) -> list[LlmSpanSpec]:
    """Load all YAML spec files under *spec_root*.

    Args:
        spec_root: Path to the ``test/llm_span`` directory.
        providers: Optional allow-list of provider names (e.g. ``["openai"]``).

    Returns:
        Sorted list of :class:`LlmSpanSpec` instances.
    """
    root = Path(spec_root) if spec_root is not None else _spec_root()

    if not root.exists():
        raise FileNotFoundError(
            f"BTX spec root not found: {root}\n"
            "Run the spec-fetch fixture or set BTX_SPEC_ROOT to the llm_span directory."
        )

    loader_cls = _make_loader()
    specs: list[LlmSpanSpec] = []

    for yaml_path in sorted(root.rglob("*.yaml")):
        provider_dir = yaml_path.parent.name
        if providers is not None and provider_dir not in providers:
            continue

        with open(yaml_path) as f:
            data = yaml.load(f, Loader=loader_cls)  # noqa: S506 — intentional, custom loader

        spec = LlmSpanSpec.from_dict(data, source_path=yaml_path)
        specs.append(spec)

    return specs