DocsGPT/docsgpt/llm/base.py

import logging
from abc import ABC, abstractmethod
from typing import ClassVar, Dict, Optional, Tuple

import httpx
import httpx2
import openai

from docsgpt.cache import gen_cache, stream_cache

from docsgpt.core.settings import settings
from docsgpt.usage import gen_token_usage, stream_token_usage

logger = logging.getLogger(__name__)


# Errors safe to retry the primary once on: the request either never
# reached the server (connect), or the peer closed the connection before
# a proper end-of-stream marker (read/protocol). None of these leave a
# side effect on the upstream, so a repeat call is cheap and idempotent.
# Excludes API-level errors (4xx / 5xx status) which are not transport
# retries — RateLimitError needs backoff, BadRequestError won't get any
# better on retry, and the existing fallback handles both.
# Both HTTP stacks are listed because the providers are split across them and
# the two libraries' exception classes are unrelated types: openai, anthropic
# and the MCP client run on httpx2, while google-genai, elevenlabs,
# qdrant-client and daytona are still on httpx. Naming only one silently stops
# matching for half the providers — the retry just never fires.
_TRANSPORT_ERROR_NAMES = (
    "RemoteProtocolError",
    "ReadError",
    "ReadTimeout",
    "WriteError",
    "WriteTimeout",
    "ConnectError",
    "ConnectTimeout",
    "PoolTimeout",
)
_STREAM_RETRYABLE_TRANSPORT_ERRORS = tuple(
    getattr(module, name)
    for module in (httpx2, httpx)
    for name in _TRANSPORT_ERROR_NAMES
) + (openai.APIConnectionError,)


def optional_int(value) -> Optional[int]:
    """Coerce a provider-reported count to int, keeping "unreported" as None.

    Cache bins persist as NULL when a provider says nothing and as 0 when it
    reports zero, so the two must stay distinguishable all the way from the
    usage object: providers report ``cached_tokens: 0`` on every uncached
    request, and folding that into "unknown" would file the ordinary case as
    no-data.

    Args:
        value: The raw attribute off a provider usage object.

    Returns:
        The integer value, or None when absent or not a number.
    """
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


class BaseLLM(ABC):
    # Stamped onto the ``llm_stream_start`` event so dashboards can group
    # calls by vendor. Subclasses override.
    provider_name: ClassVar[str] = "unknown"

    # Name of the gen kwarg this provider takes structured output on
    # ("response_format" for OpenAI-wire classes, "response_schema" for
    # Google); None = the provider has no structured-output kwarg.
    structured_output_kwarg: ClassVar[Optional[str]] = None

    # (json_schema, strict) last passed to ``prepare_structured_output_format``;
    # lets a cross-provider fallback re-prepare the schema in the backup's own
    # format instead of forwarding a kwarg the backup cannot read.
    _structured_output_source: Optional[Tuple[Dict, bool]] = None

    # Gen kwargs that carry structured output, across providers. Kept here
    # (not on the provider subclasses) because the adapter has to recognize
    # every provider's kwarg, not just its own.
    _STRUCTURED_OUTPUT_KWARGS: ClassVar[Tuple[str, ...]] = (
        "response_format",
        "response_schema",
    )

    def __init__(
        self,
        decoded_token=None,
        agent_id=None,
        model_id=None,
        base_url=None,
        backup_models=None,
        model_user_id=None,
        capabilities=None,
    ):
        self.decoded_token = decoded_token
        self.agent_id = str(agent_id) if agent_id else None
        self.model_id = model_id
        self.base_url = base_url
        self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
        self._backup_models = backup_models or []
        self._fallback_llm = None
        # Registry-resolved per-model capability overrides (BYOM caps,
        # operator YAML). None falls back to provider-class defaults.
        self.capabilities = capabilities
        # BYOM-resolution scope captured at LLM creation time so backup
        # / fallback lookups hit the same per-user layer as the primary.
        self.model_user_id = model_user_id
        # Provider whose model actually produced the most recent response.
        # Equals ``provider_name`` until a cross-provider fallback swaps the
        # responding model mid-call (see ``_stream_with_fallback``); the
        # handler layer reads this to parse chunks with the right provider's
        # handler instead of the primary's.
        self._responding_provider = self.provider_name

    @property
    def fallback_llm(self):
        """Lazy-loaded fallback LLM: tries per-agent backup models first,
        then the global FALLBACK_* settings."""
        if self._fallback_llm is not None:
            return self._fallback_llm

        from docsgpt.llm.llm_creator import LLMCreator
        from docsgpt.core.model_utils import (
            get_provider_from_model_id,
            get_api_key_for_provider,
        )

        # model_user_id (BYOM scope) takes precedence over the caller's
        # sub so shared-agent backups resolve under the owner's layer.
        caller_sub = (
            self.decoded_token.get("sub")
            if isinstance(self.decoded_token, dict)
            else None
        )
        backup_user_id = self.model_user_id or caller_sub
        for backup_model_id in self._backup_models:
            try:
                provider = get_provider_from_model_id(
                    backup_model_id, user_id=backup_user_id
                )
                if not provider:
                    logger.warning(
                        f"Could not resolve provider for backup model: {backup_model_id}"
                    )
                    continue
                api_key = get_api_key_for_provider(provider)
                self._fallback_llm = LLMCreator.create_llm(
                    provider,
                    api_key=api_key,
                    user_api_key=getattr(self, "user_api_key", None),
                    decoded_token=self.decoded_token,
                    model_id=backup_model_id,
                    agent_id=self.agent_id,
                    model_user_id=self.model_user_id,
                )
                # Tag the fallback LLM so its rows land as
                # ``source='fallback'`` in cost-attribution dashboards.
                # Propagate the parent's ``_request_id`` so a user
                # request that ran fallback is still grouped under one id.
                self._fallback_llm._token_usage_source = "fallback"
                self._fallback_llm._request_id = getattr(
                    self, "_request_id", None,
                )
                logger.info(
                    f"Fallback LLM initialized from agent backup model: "
                    f"{provider}/{backup_model_id}"
                )
                return self._fallback_llm
            except Exception as e:
                logger.warning(
                    f"Failed to initialize backup model {backup_model_id}: {str(e)}"
                )
                continue

        # Fall back to global FALLBACK_* settings. Forward
        # ``model_user_id`` here too: deployments can configure
        # ``FALLBACK_LLM_NAME`` to a BYOM UUID, and that UUID is owned
        # by the same user the primary model was resolved under.
        if settings.FALLBACK_LLM_PROVIDER:
            try:
                self._fallback_llm = LLMCreator.create_llm(
                    settings.FALLBACK_LLM_PROVIDER,
                    api_key=settings.FALLBACK_LLM_API_KEY or settings.API_KEY,
                    user_api_key=getattr(self, "user_api_key", None),
                    decoded_token=self.decoded_token,
                    model_id=settings.FALLBACK_LLM_NAME,
                    agent_id=self.agent_id,
                    model_user_id=self.model_user_id,
                )
                # Same rationale as the agent-backup branch.
                self._fallback_llm._token_usage_source = "fallback"
                self._fallback_llm._request_id = getattr(
                    self, "_request_id", None,
                )
                logger.info(
                    f"Fallback LLM initialized from global settings: "
                    f"{settings.FALLBACK_LLM_PROVIDER}/{settings.FALLBACK_LLM_NAME}"
                )
            except Exception as e:
                logger.error(
                    f"Failed to initialize fallback LLM: {str(e)}", exc_info=True
                )

        return self._fallback_llm

    @staticmethod
    def _remove_null_values(args_dict):
        if not isinstance(args_dict, dict):
            return args_dict
        return {k: v for k, v in args_dict.items() if v is not None}

    @staticmethod
    def _fallback_payload_fits(fallback, kwargs) -> bool:
        """Whether the failed request's payload can fit the fallback model.

        A primary rejected for size (context-length 400, capacity cap) hands
        the *same* oversized payload to the fallback, which then rejects it
        too — one guaranteed-failed provider call plus one estimated-prompt
        ``token_usage`` row for nothing. Skip the attempt when the estimated
        prompt already exceeds the fallback's context window. Estimation
        errors never block the attempt.
        """
        messages = kwargs.get("messages")
        if not messages:
            return True
        try:
            from docsgpt.core.model_utils import get_token_limit
            from docsgpt.usage import _count_prompt_tokens

            estimated = _count_prompt_tokens(messages, tools=kwargs.get("tools"))
            limit = get_token_limit(
                fallback.model_id,
                user_id=getattr(fallback, "model_user_id", None),
            )
            # 10% slack: the tiktoken estimate over-counts vs provider
            # tokenizers (and ``get_token_limit`` returns a conservative
            # default for unregistered models) — only skip when the payload
            # is decisively over, never on a borderline estimate.
            if estimated > int(limit * 1.1):
                logger.warning(
                    f"Skipping fallback to {fallback.model_id}: estimated "
                    f"prompt (~{estimated} tokens) cannot fit its context "
                    f"window ({limit} tokens)."
                )
                return False
        except Exception:
            logger.debug("Fallback payload size estimation failed", exc_info=True)
        return True

    @staticmethod
    def _fallback_attachment_texts(attachments):
        """Extracted attachment texts, in upload order, for file-part swaps."""
        texts = []
        for attachment in attachments or []:
            if not isinstance(attachment, dict):
                continue
            if attachment.get("content"):
                texts.append(attachment["content"])
        return texts

    def _prepare_fallback_messages(self, fallback, messages, attachments=None):
        """Rebuild primary-prepared messages so the fallback can accept them.

        ``prepare_messages_with_attachments`` ran against the *primary*
        model, so by the time a fallback engages the array can carry content
        parts the backup provider cannot take: ``file`` parts hold Files-API
        ids only the primary's endpoint+credential can resolve, and
        ``image_url`` parts 4xx on non-vision models. Handing them over
        unchanged makes the fallback die exactly like the primary did
        ("Fallback LLM also failed"). Swap what the fallback can't accept
        for text — the attachment's extracted content when available — and
        collapse all-text parts arrays to plain string content, the one
        shape every chat endpoint accepts.
        """
        if not messages:
            return messages
        try:
            supported = list(fallback.get_supported_attachment_types() or [])
        except Exception:
            supported = []
        keeps_images = any(str(t).startswith("image/") for t in supported)
        # A Files-API id resolves only against the endpoint + credential
        # that minted it (the invariant ``_scoped_file_id`` enforces), so a
        # pdf-capable fallback keeps file parts only on the same endpoint.
        keeps_files = "application/pdf" in supported
        if keeps_files:
            try:
                self_scope = getattr(self, "_endpoint_scope", None)
                fallback_scope = getattr(fallback, "_endpoint_scope", None)
                keeps_files = (
                    callable(self_scope)
                    and callable(fallback_scope)
                    and self_scope() == fallback_scope()
                )
            except Exception:
                keeps_files = False
        file_texts = self._fallback_attachment_texts(attachments)
        prepared = []
        for message in messages:
            content = message.get("content") if isinstance(message, dict) else None
            if not isinstance(content, list):
                prepared.append(message)
                continue
            parts = []
            all_text = True
            for part in content:
                part_type = part.get("type") if isinstance(part, dict) else None
                if part_type == "file" and not keeps_files:
                    if file_texts:
                        parts.append(
                            {
                                "type": "text",
                                "text": f"File content:\n\n{file_texts.pop(0)}",
                            }
                        )
                    else:
                        filename = (part.get("file") or {}).get(
                            "filename"
                        ) or "attachment"
                        parts.append(
                            {
                                "type": "text",
                                "text": f"[File '{filename}' could not be included]",
                            }
                        )
                elif part_type == "image_url" and not keeps_images:
                    parts.append(
                        {
                            "type": "text",
                            "text": "[Image attachment omitted: the responding "
                            "model does not support image input]",
                        }
                    )
                else:
                    parts.append(part)
                    if part_type != "text":
                        all_text = False
            if all_text:
                text = "\n\n".join(
                    p.get("text", "") for p in parts if isinstance(p, dict)
                )
                prepared.append({**message, "content": text})
            else:
                prepared.append({**message, "content": parts})
        return prepared

    @staticmethod
    def _fallback_enforces_structured_output(fallback) -> bool:
        """Whether the fallback's capabilities still allow schema enforcement."""
        supports = getattr(fallback, "_supports_structured_output", None)
        if supports is None:
            return True
        if not callable(supports):
            return bool(supports)
        try:
            return bool(supports())
        except Exception:
            logger.debug(
                "Structured-output capability check failed for %s; assuming supported",
                type(fallback).__name__,
                exc_info=True,
            )
            return True

    @staticmethod
    def _recover_structured_output_source(present: Dict) -> Optional[Tuple[Dict, bool]]:
        """Recover a raw (schema, strict) pair from already-prepared kwargs.

        Callers that hand-build an OpenAI ``response_format`` never went through
        ``prepare_structured_output_format``, so nothing was recorded; the raw
        schema is still readable out of the envelope. A Google
        ``response_schema`` is a lossy, type-mapped conversion — not reversible —
        so it yields nothing.
        """
        response_format = present.get("response_format")
        if not isinstance(response_format, dict):
            return None
        json_schema = response_format.get("json_schema")
        if not isinstance(json_schema, dict):
            return None
        schema = json_schema.get("schema")
        if not isinstance(schema, dict) or not schema:
            return None
        return schema, bool(json_schema.get("strict", True))

    def _adapt_structured_output_kwargs(self, fallback, kwargs: Dict) -> Dict:
        """Re-express the primary's structured-output kwargs for ``fallback``.

        Structured output is provider-specific: OpenAI-wire classes take
        ``response_format``, Google takes ``response_schema``. Forwarding the
        primary's kwarg verbatim to a different-family backup either loses
        enforcement silently (Google swallows ``response_format`` in
        ``**kwargs``) or raises ``TypeError`` inside the OpenAI SDK
        (``response_schema`` is not a Chat-Completions param), which turns the
        fallback into no fallback at all.

        Args:
            fallback: Backup LLM the request is about to be re-sent to.
            kwargs: The primary's generation kwargs.

        Returns:
            A copy of ``kwargs`` carrying the backup's own structured-output
            kwarg, or neither when the backup cannot enforce a schema.
        """
        adapted = dict(kwargs)
        present = {
            name: adapted.pop(name)
            for name in self._STRUCTURED_OUTPUT_KWARGS
            if name in adapted
        }
        if not present:
            return adapted

        target = getattr(fallback, "structured_output_kwarg", None)
        if not target or not self._fallback_enforces_structured_output(fallback):
            logger.warning(
                f"Fallback {fallback.model_id} cannot enforce structured "
                f"output; continuing unstructured"
            )
            return adapted

        response_format = present.get("response_format")
        if (
            isinstance(response_format, dict)
            and response_format.get("type") == "json_object"
        ):
            # json_object mode is an OpenAI-wire shape with no wired equivalent
            # elsewhere; keep it only within the same family.
            if target == "response_format":
                adapted[target] = response_format
            else:
                logger.warning(
                    f"Fallback {fallback.model_id} has no json_object mode; "
                    f"continuing unstructured"
                )
            return adapted

        if target in present:
            # Same wire family (OpenAI -> openai_compatible, Google -> Google):
            # the prepared value is already in the backup's format.
            adapted[target] = present[target]
            return adapted

        source = self._structured_output_source
        if not source:
            source = self._recover_structured_output_source(present)
        if not source:
            logger.warning(
                f"Fallback {fallback.model_id} needs {target} but the source "
                f"schema is unavailable; continuing unstructured"
            )
            return adapted

        schema, strict = source
        try:
            prepared = fallback.prepare_structured_output_format(schema, strict=strict)
        except Exception:
            logger.warning(
                f"Failed to prepare structured output for fallback "
                f"{fallback.model_id}; continuing unstructured",
                exc_info=True,
            )
            prepared = None
        if prepared:
            adapted[target] = prepared
        else:
            logger.warning(
                f"Fallback {fallback.model_id} cannot enforce structured "
                f"output; continuing unstructured"
            )
        return adapted

    def _execute_with_fallback(
        self, method_name: str, decorators: list, *args, **kwargs
    ):
        """
        Execute method with fallback support.

        Any error raised by the primary model triggers a single attempt on
        the fallback model, when one is configured. There is no error
        classification: 5xx/transient failures are obviously recoverable on a
        different model, and for client-side (4xx) errors — including rate
        limits (429) and provider-specific payload rejections — the one extra
        attempt is cheap insurance, since a second provider often accepts what
        the first refused. ``GeneratorExit``/cancellation are ``BaseException``
        subclasses and so bypass this handler (no fallback on client
        disconnect), which is intentional.

        Args:
            method_name: Name of the raw method ('_raw_gen' or '_raw_gen_stream')
            decorators: List of decorators to apply
            *args: Positional arguments
            **kwargs: Keyword arguments
        """

        def decorated_method():
            method = getattr(self, method_name)
            for decorator in decorators:
                method = decorator(method)
            return method(self, *args, **kwargs)

        is_stream = "stream" in method_name

        if is_stream:
            return self._stream_with_fallback(
                decorated_method, method_name, decorators, *args, **kwargs
            )

        self._responding_provider = self.provider_name
        try:
            return decorated_method()
        except Exception as e:
            if not self.fallback_llm:
                logger.error(f"Primary LLM failed and no fallback configured: {str(e)}")
                raise
            fallback = self.fallback_llm
            if not self._fallback_payload_fits(fallback, kwargs):
                raise
            self._responding_provider = fallback.provider_name
            logger.warning(
                f"Primary LLM failed. Falling back to "
                f"{fallback.model_id}. Error: {str(e)}"
            )

            # Mirror the streaming path: emit the fallback's own start event so
            # dashboards attribute the response to the backup provider, not the
            # failed primary.
            fallback._emit_gen_start_log(
                fallback.model_id,
                kwargs.get("messages"),
                kwargs.get("tools"),
                bool(
                    kwargs.get("_usage_attachments")
                    or kwargs.get("attachments")
                ),
            )
            # Apply decorators to fallback's raw method directly — calling
            # fallback.gen() would re-enter the orchestrator and recurse via
            # fallback.fallback_llm.
            fallback_method = getattr(fallback, method_name)
            for decorator in decorators:
                fallback_method = decorator(fallback_method)
            fallback_kwargs = {**kwargs, "model": fallback.model_id}
            fallback_kwargs = self._adapt_structured_output_kwargs(
                fallback, fallback_kwargs
            )
            if fallback_kwargs.get("messages"):
                fallback_kwargs["messages"] = self._prepare_fallback_messages(
                    fallback,
                    fallback_kwargs["messages"],
                    kwargs.get("_usage_attachments") or kwargs.get("attachments"),
                )
            try:
                return fallback_method(fallback, *args, **fallback_kwargs)
            except Exception as e2:
                logger.error(f"Fallback LLM also failed; giving up: {str(e2)}")
                raise

    def _stream_with_fallback(
        self, decorated_method, method_name, decorators, *args, **kwargs
    ):
        """
        Wrapper generator that catches mid-stream errors and falls back.

        Unlike non-streaming calls where exceptions are raised immediately,
        streaming generators raise exceptions during iteration. This wrapper
        ensures that if the primary LLM fails at any point during streaming
        (creation or mid-stream), we fall back to the backup model.

        Transport errors that fire before any chunk was yielded (Azure
        Front Door reset, backend deploy hiccup, TLS blip) get one same-
        primary retry before we engage the fallback — the request never
        made it to a state the peer produced output for, so a repeat is
        idempotent. Once we've yielded anything, retrying would duplicate
        delivered content; those errors go straight to the fallback path
        as before.
        """
        self._responding_provider = self.provider_name
        chunks_yielded = 0
        try:
            for chunk in decorated_method():
                chunks_yielded += 1
                yield chunk
            return
        except Exception as e:
            if getattr(self, "_stream_reached_finish", False):
                # The primary already delivered a finish signal — only
                # trailing frames (usage chunk, [DONE]) failed. Restreaming
                # from the fallback would duplicate the entire answer the
                # user already received (and re-run tool calls). Re-raise;
                # the streaming handler treats post-finish failures as
                # non-fatal.
                logger.warning(
                    f"Primary LLM failed after delivering its finish signal; "
                    f"not engaging fallback. Error: {str(e)}"
                )
                raise
            # Same-primary retry once for transport errors before any chunk
            # was yielded. Covers the observed Azure Responses-API pattern
            # where the peer resets the SSE stream within tens of seconds
            # with a RemoteProtocolError, well before a legitimate fallback
            # scenario. A yielded chunk = downstream already saw content,
            # so replaying would duplicate it — skip retry in that case.
            if (
                chunks_yielded == 0
                and isinstance(e, _STREAM_RETRYABLE_TRANSPORT_ERRORS)
            ):
                logger.warning(
                    f"Primary LLM transport error before any output; "
                    f"retrying once. Error: {str(e)}"
                )
                # Emit a fresh stream-start so dashboards get one start/
                # finish pair per attempt (stream_token_usage fires a
                # finish per decorated_method invocation).
                self._emit_stream_start_log(
                    kwargs.get("model") or getattr(self, "model_id", None),
                    kwargs.get("messages"),
                    kwargs.get("tools"),
                    bool(
                        kwargs.get("_usage_attachments")
                        or kwargs.get("attachments")
                    ),
                )
                try:
                    for chunk in decorated_method():
                        chunks_yielded += 1
                        yield chunk
                    return
                except Exception as retry_e:
                    # Retry delivered a full stream but died on a trailing
                    # frame (usage chunk, [DONE]) — same guard the outer
                    # except uses. Without this, the fallback would run
                    # and the user would receive the whole answer twice
                    # (and any tool calls would be executed twice).
                    if getattr(self, "_stream_reached_finish", False):
                        logger.warning(
                            f"Primary LLM retry delivered its finish "
                            f"signal, then failed on a trailing frame; "
                            f"not engaging fallback. Error: {str(retry_e)}"
                        )
                        raise
                    logger.warning(
                        f"Primary LLM retry also failed; engaging fallback. "
                        f"Error: {str(retry_e)}"
                    )
                    e = retry_e
            if not self.fallback_llm:
                logger.error(
                    f"Primary LLM failed and no fallback configured: {str(e)}"
                )
                raise
            fallback = self.fallback_llm
            if not self._fallback_payload_fits(fallback, kwargs):
                raise
            self._responding_provider = fallback.provider_name
            logger.warning(
                f"Primary LLM failed mid-stream. Falling back to "
                f"{fallback.model_id}. Error: {str(e)}"
            )
            # Apply decorators to fallback's raw stream method directly —
            # calling fallback.gen_stream() would re-enter the orchestrator
            # and recurse via fallback.fallback_llm. Emit the stream-start
            # event manually so dashboards still see the fallback's
            # provider/model when the response actually comes from it.
            fallback._emit_stream_start_log(
                fallback.model_id,
                kwargs.get("messages"),
                kwargs.get("tools"),
                bool(
                    kwargs.get("_usage_attachments")
                    or kwargs.get("attachments")
                ),
            )
            fallback_method = getattr(fallback, method_name)
            for decorator in decorators:
                fallback_method = decorator(fallback_method)
            fallback_kwargs = {**kwargs, "model": fallback.model_id}
            fallback_kwargs = self._adapt_structured_output_kwargs(
                fallback, fallback_kwargs
            )
            if fallback_kwargs.get("messages"):
                fallback_kwargs["messages"] = self._prepare_fallback_messages(
                    fallback,
                    fallback_kwargs["messages"],
                    kwargs.get("_usage_attachments") or kwargs.get("attachments"),
                )
            try:
                yield from fallback_method(fallback, *args, **fallback_kwargs)
            except Exception as e2:
                logger.error(
                    f"Fallback LLM also failed mid-stream; giving up: {str(e2)}"
                )
                raise

    def gen(self, model, messages, stream=False, tools=None, *args, **kwargs):
        # Mirror gen_stream: emit the start event before the decorators run so
        # ``_usage_attachments`` is still in kwargs (the gen decorators pop it).
        has_attachments = bool(
            kwargs.get("_usage_attachments") or kwargs.get("attachments")
        )
        self._emit_gen_start_log(model, messages, tools, has_attachments)
        decorators = [gen_token_usage, gen_cache]
        return self._execute_with_fallback(
            "_raw_gen",
            decorators,
            model=model,
            messages=messages,
            stream=stream,
            tools=tools,
            *args,
            **kwargs,
        )

    def _emit_gen_start_log(self, model, messages, tools, has_attachments):
        # Non-streaming counterpart to ``_emit_stream_start_log``. Emitted by
        # ``gen()`` before the call — and again for the fallback provider in
        # ``_execute_with_fallback`` — so non-streaming invocations are
        # observable from the first log line, not just streaming ones. A
        # distinct event name keeps non-stream calls out of stream dashboards.
        logging.info(
            "llm_gen_start",
            extra={
                "model": model,
                "provider": self.provider_name,
                "message_count": len(messages) if messages is not None else 0,
                "has_attachments": bool(has_attachments),
                "has_tools": bool(tools),
            },
        )

    def _emit_gen_finished_log(
        self,
        model,
        *,
        prompt_tokens,
        completion_tokens,
        latency_ms,
        cached_tokens=None,
        cache_write_tokens=None,
        error=None,
    ):
        # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired
        # with ``llm_gen_start`` so cost dashboards can join start/finish for
        # non-streaming calls just as they do for streams. Token counts come
        # from ``gen_token_usage``: provider-exact when the vendor reported
        # usage (OpenAI-family chat + Responses), tiktoken estimates
        # otherwise; ``status`` is ``"error"`` when the call raised. A
        # distinct event name keeps non-stream calls out of stream
        # dashboards.
        extra = {
            "model": model,
            "provider": self.provider_name,
            "prompt_tokens": int(prompt_tokens),
            "completion_tokens": int(completion_tokens),
            "latency_ms": int(latency_ms),
            "status": "error" if error is not None else "ok",
        }
        if cached_tokens is not None:
            extra["cached_tokens"] = int(cached_tokens)
        if cache_write_tokens is not None:
            extra["cache_write_tokens"] = int(cache_write_tokens)
        if error is not None:
            extra["error_class"] = type(error).__name__
        logging.info("llm_gen_finished", extra=extra)

    def _emit_stream_start_log(self, model, messages, tools, has_attachments):
        # Stamped with ``self.provider_name`` so dashboards can group calls
        # by vendor; the fallback path emits its own copy on the fallback
        # instance so the actual responding provider is recorded.
        logging.info(
            "llm_stream_start",
            extra={
                "model": model,
                "provider": self.provider_name,
                "message_count": len(messages) if messages is not None else 0,
                "has_attachments": bool(has_attachments),
                "has_tools": bool(tools),
            },
        )

    def _emit_stream_finished_log(
        self,
        model,
        *,
        prompt_tokens,
        completion_tokens,
        latency_ms,
        cached_tokens=None,
        cache_write_tokens=None,
        error=None,
    ):
        # Paired with ``llm_stream_start`` so cost dashboards can sum tokens
        # by user/agent/provider. Token counts come from
        # ``stream_token_usage``: provider-exact when the vendor reported
        # usage (OpenAI-family via ``stream_options.include_usage`` and the
        # Responses API), tiktoken estimates for providers that don't.
        extra = {
            "model": model,
            "provider": self.provider_name,
            "prompt_tokens": int(prompt_tokens),
            "completion_tokens": int(completion_tokens),
            "latency_ms": int(latency_ms),
            "status": "error" if error is not None else "ok",
        }
        if cached_tokens is not None:
            extra["cached_tokens"] = int(cached_tokens)
        if cache_write_tokens is not None:
            extra["cache_write_tokens"] = int(cache_write_tokens)
        if error is not None:
            extra["error_class"] = type(error).__name__
        logging.info("llm_stream_finished", extra=extra)

    def gen_stream(self, model, messages, stream=True, tools=None, *args, **kwargs):
        # Attachments arrive as ``_usage_attachments`` from ``Agent._llm_gen``;
        # the ``stream_token_usage`` decorator pops that key, but the log
        # fires before the decorator runs so it's still in ``kwargs`` here.
        has_attachments = bool(
            kwargs.get("_usage_attachments") or kwargs.get("attachments")
        )
        self._emit_stream_start_log(model, messages, tools, has_attachments)
        decorators = [stream_cache, stream_token_usage]
        return self._execute_with_fallback(
            "_raw_gen_stream",
            decorators,
            model=model,
            messages=messages,
            stream=stream,
            tools=tools,
            *args,
            **kwargs,
        )

    @abstractmethod
    def _raw_gen(self, model, messages, stream, tools, *args, **kwargs):
        pass

    @abstractmethod
    def _raw_gen_stream(self, model, messages, stream, *args, **kwargs):
        pass

    def supports_tools(self):
        return hasattr(self, "_supports_tools") and callable(
            getattr(self, "_supports_tools")
        )

    def _supports_tools(self):
        raise NotImplementedError("Subclass must implement _supports_tools method")

    def supports_structured_output(self):
        """Check if the LLM supports structured output/JSON schema enforcement"""
        return hasattr(self, "_supports_structured_output") and callable(
            getattr(self, "_supports_structured_output")
        )

    def _supports_structured_output(self):
        return False

    def prepare_structured_output_format(self, json_schema, strict=True):
        """Prepare structured output format specific to the LLM provider.

        Overrides must record ``self._structured_output_source`` so a
        cross-provider fallback can re-prepare the schema in its own format.
        """
        _ = (json_schema, strict)
        return None

    def get_supported_attachment_types(self):
        """
        Return a list of MIME types supported by this LLM for file uploads.

        Returns:
            list: List of supported MIME types
        """
        return []