guarded-llm¶

def call(
    self,
    prompt: str,
    *,
    system: str | None = None,
    messages: list[dict] | None = None,
    **kwargs: Any,
) -> Any:
    """Run an LLM call and return the validated instance.

    Args:
        prompt: user prompt string (ignored if `messages=` is passed).
        system: optional system prompt prepended as the first message.
        messages: optional fully-formed messages list (overrides `prompt`).
        **kwargs: forwarded to the provider (e.g. `api_key=`, `temperature=`).

    Returns:
        The validated instance (Pydantic model instance, dict, or legacy
        dataclass instance — whatever the schema returns).

    Raises:
        BudgetExceededError: if cumulative cost exceeds the Budget cap.
        RetryExhausted: if all attempts fail validation.
        LLMCallError: if the provider itself fails (network, auth, etc.)
            and all retries are exhausted.
    """
    from .providers import get_provider

    if messages is None:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

    prov = get_provider(self.provider)
    merged_kwargs = {**self.provider_kwargs, **kwargs}

    stats = GuardedCallStats()
    last_err: str | None = None

    for attempt in range(self.retry.max_attempts):
        if attempt > 0:
            sleep_s = self.retry.sleep_seconds(attempt)
            if sleep_s > 0:
                time.sleep(sleep_s)

        attempt_messages = list(messages)
        if last_err is not None and attempt_messages:
            hint = (
                f"\n\nPrevious output failed validation: {last_err}\n"
                "Output valid JSON only — no prose, no markdown fences."
            )
            last = dict(attempt_messages[-1])
            if last.get("role") == "user":
                last["content"] = last.get("content", "") + hint
                attempt_messages[-1] = last
            else:
                attempt_messages.append({"role": "user", "content": hint.strip()})

        stats.attempts += 1
        try:
            call_out = prov.call(
                messages=attempt_messages,
                model=self.model,
                max_tokens=self.max_tokens,
                schema=self.schema,
                **merged_kwargs,
            )
        except Exception as exc:
            err = f"attempt {attempt + 1}: provider {self.provider!r} raised: {exc}"
            stats.errors.append(err)
            last_err = err
            continue

        raw = (
            call_out.get("text", "") if isinstance(call_out, dict) else str(call_out)
        )
        cost = (
            float(call_out.get("cost_usd", 0.0))
            if isinstance(call_out, dict)
            else 0.0
        )
        stats.raw_outputs.append(raw)
        stats.cost_usd += cost

        # Charge the budget BEFORE checking schema, so a runaway loop
        # can't keep burning money on a bad prompt.
        if self.budget is not None:
            # Budget.consume() raises BudgetExceeded if we'd exceed any cap.
            # Reset stats first so caller can inspect partial spend.
            self.last_stats = stats
            self.budget.consume(cost)

        cleaned = state_machine_fix(raw)
        ok, err, parsed = validate_json(cleaned, self.schema)
        if ok:
            self.last_stats = stats
            return parsed
        err_msg = f"attempt {attempt + 1}: {err}"
        stats.errors.append(err_msg)
        last_err = err

    self.last_stats = stats
    raise RetryExhausted(
        f"all {self.retry.max_attempts} attempts failed validation",
        attempts=stats.errors,
        last_raw=stats.raw_outputs[-1] if stats.raw_outputs else None,
    )

def call_as_result(
    self,
    prompt: str,
    *,
    system: str | None = None,
    messages: list[dict] | None = None,
    **kwargs: Any,
) -> GuardrailResult:
    """Like `.call()` but returns a `GuardrailResult` (never raises on
    validation failure — errors are accumulated in the result)."""
    try:
        parsed = self.call(prompt, system=system, messages=messages, **kwargs)
    except RetryExhausted:
        return GuardrailResult(
            parsed=None,
            errors=list(self.last_stats.errors),
            attempts=self.last_stats.attempts,
            cost_usd=self.last_stats.cost_usd,
            raw_outputs=list(self.last_stats.raw_outputs),
        )
    except BudgetExceededError:
        # Re-raise — budget is a hard stop, not a soft validation failure.
        raise
    return GuardrailResult(
        parsed=parsed,
        errors=list(self.last_stats.errors),
        attempts=self.last_stats.attempts,
        cost_usd=self.last_stats.cost_usd,
        raw_outputs=list(self.last_stats.raw_outputs),
    )

def consume(self, usd: float) -> None:
    """Record a charge. Raise `BudgetExceeded` if it would exceed any cap.

    The charge is NOT recorded when an exception is raised — so a caller
    can catch `BudgetExceeded` and the Budget state stays consistent.
    """
    if not isinstance(usd, (int, float)) or isinstance(usd, bool):
        raise TypeError("usd must be a number")
    if usd < 0:
        raise ValueError("usd must be >= 0")
    if usd > self.usd_per_call:
        raise BudgetExceeded(
            f"single call cost ${usd:.4f} > per-call cap ${self.usd_per_call:.4f}",
            spent_usd=self.spent_usd,
            cap_usd=self.usd_per_call,
        )
    if self.spent_usd + usd > self.usd_total:
        raise BudgetExceeded(
            f"cumulative cost ${self.spent_usd + usd:.4f} "
            f"> total cap ${self.usd_total:.4f}",
            spent_usd=self.spent_usd + usd,
            cap_usd=self.usd_total,
        )
    self.spent_usd += usd

def reset(self) -> None:
    """Zero out the spent-so-far counter (kept caps unchanged)."""
    self.spent_usd = 0.0

def sleep_seconds(self, attempt: int, rng: random.Random | None = None) -> float:
    """Compute backoff for the given attempt number (0-indexed).

    Linear: attempt 0 → 0 sec (don't sleep before first call),
            attempt N → N * backoff_seconds (* jitter, if enabled).
    """
    if attempt <= 0:
        return 0.0
    base = attempt * self.backoff_seconds
    if not self.jitter:
        return base
    r = rng if rng is not None else random
    return base * r.uniform(0.5, 1.5)

def validate(self, d: Any) -> tuple[bool, str | None, Any]:
    """Validate `d` against the Pydantic model.

    Returns (ok, error_message_or_none, model_instance_or_none).
    """
    try:
        inst = self._model.model_validate(d)
    except ValidationError as e:
        # Pretty-print the first error path so it's useful in retry hints
        errs = e.errors()
        if errs:
            first = errs[0]
            path = ".".join(str(p) for p in first.get("loc", [])) or "<root>"
            msg = f"{path}: {first.get('msg', 'validation error')}"
        else:
            msg = str(e)
        return False, msg, None
    except Exception as e:
        return False, f"pydantic error: {e}", None
    return True, None, inst

def validate(self, d: Any) -> tuple[bool, str | None, Any]:
    """Validate `d` against this schema.

    Returns (ok, error_message_or_none, validated_instance_or_none).
    """
    errors = sorted(self._validator.iter_errors(d), key=lambda e: list(e.absolute_path))
    if not errors:
        return True, None, d
    first = errors[0]
    path = ".".join(str(p) for p in first.absolute_path) or "<root>"
    msg = f"{path}: {first.message}"
    return False, msg, None

@abstractmethod
def call(
    self,
    messages: list[dict],
    model: str,
    max_tokens: int,
    schema: Any = None,
    **kwargs: Any,
) -> dict:
    """Send `messages` to the LLM and return {"text": str, "cost_usd": float}."""
    ...