Providers¶

async def complete(
    self,
    messages: list[ChatMessage],
    model: str,
    *,
    tools: list[ToolDefinition] | None = None,
    config: CompletionConfig | None = None,
) -> CompletionResponse:
    """Execute a non-streaming chat completion.

    Args:
        messages: Conversation history.
        model: Model identifier to use.
        tools: Available tools for function calling.
        config: Optional completion parameters.

    Returns:
        The full completion response.
    """
    ...

async def stream(
    self,
    messages: list[ChatMessage],
    model: str,
    *,
    tools: list[ToolDefinition] | None = None,
    config: CompletionConfig | None = None,
) -> AsyncIterator[StreamChunk]:
    """Execute a streaming chat completion.

    Args:
        messages: Conversation history.
        model: Model identifier to use.
        tools: Available tools for function calling.
        config: Optional completion parameters.

    Returns:
        Async iterator of stream chunks.
    """
    ...

async def get_model_capabilities(self, model: str) -> ModelCapabilities:
    """Return capability metadata for the given model.

    Args:
        model: Model identifier.

    Returns:
        Static capability and cost information.
    """
    ...

async def batch_get_capabilities(
    self,
    models: tuple[str, ...],
) -> Mapping[str, ModelCapabilities | None]:
    """Return capability metadata for many models in one call.

    Failures degrade per-model: models whose lookup fails surface as
    ``None`` entries so callers preserve graceful per-model fallback.
    The returned mapping keys are exactly the input ``models`` tuple.

    Args:
        models: Tuple of model identifiers to look up.

    Returns:
        Mapping from model id to capabilities (or ``None`` on failure).
    """
    ...

async def complete(
    self,
    messages: list[ChatMessage],
    model: str,
    *,
    tools: list[ToolDefinition] | None = None,
    config: CompletionConfig | None = None,
) -> CompletionResponse:
    """Validate inputs, delegate to ``_do_complete``.

    Applies rate limiting and retry automatically when configured.

    Args:
        messages: Conversation history.
        model: Model identifier to use.
        tools: Available tools for function calling.
        config: Optional completion parameters.

    Returns:
        The completion response.

    Raises:
        InvalidRequestError: If messages are empty or model is blank.
        RetryExhaustedError: If all retries are exhausted.
    """
    self._validate_messages(messages)
    self._validate_model(model)
    logger.debug(
        PROVIDER_CALL_START,
        model=model,
        message_count=len(messages),
    )

    async def _attempt() -> CompletionResponse:
        return await self._rate_limited_call(
            self._do_complete,
            messages,
            model,
            tools=tools,
            config=config,
        )

    from .resilience.retry import RetryResult  # noqa: PLC0415, TC001

    # Per-call child span under whatever parent span the caller
    # owns (typically ``agent.execution`` from AgentEngine).
    # ``record_exception=False`` and ``set_status_on_exception=False``
    # opt out of the auto-instrumentation that would otherwise
    # stamp the unscrubbed ``str(exc)`` into the span; we set
    # ``exception.message`` via ``safe_error_description`` instead
    # so attacker-controlled provider error strings are scrubbed
    # before reaching the OTLP exporter.
    provider_label = self._provider_label()
    span_attributes: dict[str, Any] = {
        "provider.name": provider_label,
        "provider.model": model,
        "provider.message_count": len(messages),
        "provider.tool_count": len(tools) if tools else 0,
    }
    with _tracer.start_as_current_span(
        "provider.complete",
        attributes=span_attributes,
        record_exception=False,
        set_status_on_exception=False,
    ) as span:
        t_start = self._clock.monotonic()
        retry_info: RetryResult[CompletionResponse] | None = None
        try:
            if self._retry_handler is not None:
                retry_info = await self._retry_handler.execute(_attempt)
                result = retry_info.value
            else:
                result = await _attempt()
        except MemoryError, RecursionError:
            raise
        except Exception as exc:
            latency_ms = (
                self._clock.monotonic() - t_start
            ) * _MILLISECONDS_PER_SECOND
            # ``logger.exception`` (what TRY400 suggests) would
            # attach a traceback whose serialized frame-locals can
            # leak provider credentials (API keys in headers,
            # connection URLs with user:pass). Use ``logger.error``
            # with the structured ``error_type`` + scrubbed
            # ``error`` fields instead.
            logger.error(
                PROVIDER_CALL_ERROR,
                model=model,
                latency_ms=latency_ms,
                error_type=type(exc).__name__,
                error=safe_error_description(exc),
            )
            span.set_attribute("exception.type", type(exc).__name__)
            span.set_attribute(
                "exception.message",
                safe_error_description(exc),
            )
            span.set_attribute("provider.latency_ms", latency_ms)
            # ``set_status_on_exception=False`` opts out of the
            # auto-instrumentation that would have stamped an
            # un-scrubbed ``str(exc)`` into the span status, so the
            # ERROR status must be set manually here. The scrubbed
            # error description is exposed via ``exception.message``
            # above; ``Status.description`` is intentionally left
            # unset so the OTLP exporter never carries the raw
            # provider string.
            span.set_status(Status(StatusCode.ERROR))
            record_provider_error(
                provider=provider_label,
                model=model,
                error_class=classify_provider_error(exc),
            )
            raise
        latency_ms = (self._clock.monotonic() - t_start) * _MILLISECONDS_PER_SECOND
        span.set_attribute("provider.latency_ms", latency_ms)
        if retry_info is not None:
            span.set_attribute(
                "provider.retry_count",
                max(0, retry_info.attempt_count - 1),
            )

    metadata: dict[str, object] = {"_synthorg_latency_ms": latency_ms}
    if retry_info is not None:
        metadata["_synthorg_retry_count"] = max(
            0,
            retry_info.attempt_count - 1,
        )
        if retry_info.retry_reason is not None:
            metadata["_synthorg_retry_reason"] = retry_info.retry_reason

    merged_metadata = dict(result.provider_metadata or {})
    merged_metadata.update(metadata)
    result = result.model_copy(update={"provider_metadata": merged_metadata})
    logger.debug(
        PROVIDER_CALL_SUCCESS,
        model=model,
    )

    # Cost recording chokepoint: when a ``cost_recording_scope`` is
    # open in the current asyncio task, emit a CostRecord. Sites
    # without a scope (probes, tests, and the engine path which
    # records via ``record_execution_costs`` post-execution) see no
    # change. Recording errors are logged and swallowed inside the
    # helper -- never surface to the caller.
    ctx = current_cost_context()
    if ctx is not None:
        await emit_cost_record_from_context(
            ctx,
            result,
            model=model,
            provider=self._provider_label(),
        )
    return result

async def stream(
    self,
    messages: list[ChatMessage],
    model: str,
    *,
    tools: list[ToolDefinition] | None = None,
    config: CompletionConfig | None = None,
) -> AsyncIterator[StreamChunk]:
    """Validate inputs, delegate to ``_do_stream``.

    Only the initial connection setup is retried; mid-stream errors
    are not retried.

    .. note::

        Unlike :meth:`complete`, ``stream`` does **not** fire the
        cost-recording chokepoint.  Streaming responses surface
        usage as a terminal ``StreamEventType.USAGE`` chunk, so the
        recording logic would have to consume the iterator to
        extract token counts -- conflating cost recording with the
        stream-consumption contract.  Until streaming becomes a
        mainstream LLM call path in this codebase, callers using
        ``stream()`` are responsible for emitting their own
        ``CostRecord`` from the final usage chunk.  No call site in
        the current diff uses ``stream()`` for paid LLM work.

    Args:
        messages: Conversation history.
        model: Model identifier to use.
        tools: Available tools for function calling.
        config: Optional completion parameters.

    Returns:
        Async iterator of stream chunks.

    Raises:
        InvalidRequestError: If messages are empty or model is blank.
        RetryExhaustedError: If all retries are exhausted.
    """
    self._validate_messages(messages)
    self._validate_model(model)
    logger.debug(
        PROVIDER_STREAM_START,
        model=model,
        message_count=len(messages),
    )

    async def _attempt() -> AsyncIterator[StreamChunk]:
        return await self._rate_limited_call(
            self._do_stream,
            messages,
            model,
            tools=tools,
            config=config,
        )

    try:
        return await self._resilient_execute(_attempt)
    except MemoryError, RecursionError:
        raise
    except Exception as exc:
        # See the ``complete`` sibling handler; ``logger.error``
        # + scrubbed fields instead of ``logger.exception``
        # prevents traceback frame-locals from leaking provider
        # credentials.
        logger.error(
            PROVIDER_CALL_ERROR,
            model=model,
            error_type=type(exc).__name__,
            error=safe_error_description(exc),
        )
        record_provider_error(
            provider=self._provider_label(),
            model=model,
            error_class=classify_provider_error(exc),
        )
        raise

async def get_model_capabilities(self, model: str) -> ModelCapabilities:
    """Validate model identifier, delegate to ``_do_get_model_capabilities``.

    Capability lookups go through the same retry handler and rate
    limiter as ``complete()`` / ``stream()`` so the contract "all
    provider calls go through BaseCompletionProvider" stays honest
    for any future driver whose ``_do_get_model_capabilities``
    does network I/O.  Same budget as completions: capability
    lookups consume a rate-limiter slot and are retried on
    retryable errors.

    Args:
        model: Model identifier.

    Returns:
        Static capability and cost information.

    Raises:
        InvalidRequestError: If model is blank.
        RetryExhaustedError: If all retries are exhausted.
    """
    self._validate_model(model)

    async def _attempt() -> ModelCapabilities:
        return await self._rate_limited_call(
            self._do_get_model_capabilities,
            model,
        )

    try:
        return await self._resilient_execute(_attempt)
    except MemoryError, RecursionError:
        raise
    except Exception as exc:
        # ``logger.exception`` would attach a traceback whose
        # frame-locals can leak provider credentials; use
        # ``logger.error`` with the structured ``error_type`` +
        # scrubbed ``error`` fields, mirroring ``complete()`` /
        # ``stream()``. ``record_provider_error`` keeps the
        # provider-error metric in sync with the other call paths
        # so dashboards do not under-count capability failures.
        logger.error(
            PROVIDER_CALL_ERROR,
            model=model,
            phase="get_model_capabilities",
            error_type=type(exc).__name__,
            error=safe_error_description(exc),
        )
        record_provider_error(
            provider=self._provider_label(),
            model=model,
            error_class=classify_provider_error(exc),
        )
        raise

async def batch_get_capabilities(
    self,
    models: tuple[str, ...],
) -> Mapping[str, ModelCapabilities | None]:
    """Fan out capability lookups across many models in parallel.

    The default implementation runs ``get_model_capabilities`` per
    model concurrently via :class:`asyncio.TaskGroup`.

    Per-model classification errors (model-not-found, validation,
    non-retryable provider errors) degrade to ``None`` entries so a
    single bad model id does not poison the whole batch.

    ``RetryExhaustedError`` propagates: retry exhaustion is a
    signal that the provider is unhealthy, not a per-model
    classification issue. Surfacing it lets the caller decide
    whether to fail the whole list-models request or retry the
    batch later. ``MemoryError`` and ``RecursionError`` also
    propagate unchanged.

    Subclasses that expose a cheaper bulk source (e.g. a static
    preset catalog) should override this to avoid the per-model
    round trip.
    """
    if not models:
        return {}

    async def _one(m: str) -> tuple[str, ModelCapabilities | None]:
        try:
            return m, await self.get_model_capabilities(m)
        except MemoryError, RecursionError, RetryExhaustedError:
            raise
        except Exception as exc:
            logger.warning(
                PROVIDER_BATCH_CAPABILITIES_PARTIAL,
                model=m,
                error_type=type(exc).__name__,
                error=safe_error_description(exc),
            )
            return m, None

    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(_one(m)) for m in models]
    return dict(t.result() for t in tasks)

@staticmethod
def compute_cost(
    input_tokens: int,
    output_tokens: int,
    *,
    cost_per_1k_input: float,
    cost_per_1k_output: float,
) -> TokenUsage:
    """Build a ``TokenUsage`` from raw token counts and per-1k rates.

    Args:
        input_tokens: Number of input tokens (must be >= 0).
        output_tokens: Number of output tokens (must be >= 0).
        cost_per_1k_input: Cost per 1,000 input tokens in the configured
            currency (finite and >= 0).
        cost_per_1k_output: Cost per 1,000 output tokens in the configured
            currency (finite and >= 0).

    Returns:
        Populated ``TokenUsage`` with computed cost.

    Raises:
        InvalidRequestError: If any parameter is negative or
            non-finite.
    """
    if input_tokens < 0:
        msg = "input_tokens must be non-negative"
        raise InvalidRequestError(
            msg,
            context={"input_tokens": input_tokens},
        )
    if output_tokens < 0:
        msg = "output_tokens must be non-negative"
        raise InvalidRequestError(
            msg,
            context={"output_tokens": output_tokens},
        )
    if cost_per_1k_input < 0 or not math.isfinite(cost_per_1k_input):
        msg = "cost_per_1k_input must be a finite non-negative number"
        raise InvalidRequestError(
            msg,
            context={"cost_per_1k_input": cost_per_1k_input},
        )
    if cost_per_1k_output < 0 or not math.isfinite(cost_per_1k_output):
        msg = "cost_per_1k_output must be a finite non-negative number"
        raise InvalidRequestError(
            msg,
            context={"cost_per_1k_output": cost_per_1k_output},
        )
    cost = (input_tokens / 1000) * cost_per_1k_input + (
        output_tokens / 1000
    ) * cost_per_1k_output
    return TokenUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        cost=round(cost, BUDGET_ROUNDING_PRECISION),
    )

def __str__(self) -> str:
    """Format error with optional context metadata.

    Sensitive keys (api_key, token, etc.) are redacted to prevent
    accidental secret leakage in logs and tracebacks.
    """
    if self.context:
        ctx = ", ".join(
            f"{k}='***'" if _is_sensitive_key(k) else f"{k}={v!r}"
            for k, v in self.context.items()
        )
        return f"{self.message} ({ctx})"
    return self.message