open-swe/agent/utils/deferred_model.py

from __future__ import annotations

import logging
from typing import Any

from langchain_core.language_models import BaseChatModel

from .model import make_model

logger = logging.getLogger(__name__)


class DeferredErrorModel(BaseChatModel):
    """Model placeholder that raises a stored setup error on first invocation."""

    error_message: str
    model_id: str | None = None

    @property
    def _llm_type(self) -> str:
        return "deferred-error"

    def _get_ls_params(self, stop: Any = None, **kwargs: Any) -> dict[str, Any]:
        params = super()._get_ls_params(stop=stop, **kwargs)
        if self.model_id:
            params["ls_model_name"] = self.model_id
            if ":" in self.model_id:
                params["ls_provider"] = self.model_id.split(":", 1)[0]
        return params

    def bind_tools(self, tools: Any, **kwargs: Any) -> DeferredErrorModel:
        return self

    def _generate(self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any):
        raise ValueError(self.error_message)


def make_deferred_error_model(
    error: BaseException, *, model_id: str | None = None
) -> BaseChatModel:
    return DeferredErrorModel(error_message=f"{type(error).__name__}: {error}", model_id=model_id)


def make_model_or_defer(model_id: str, **kwargs: Any) -> BaseChatModel:
    """Call ``make_model`` and wrap any setup error in a ``DeferredErrorModel``.

    This lets graph factories pass a model placeholder into the agent graph on
    startup without crashing the process.  The error is raised later inside the
    agent run when the model is first invoked, so the caller can recover
    gracefully (e.g. via fallback middleware or by surfacing the error to the
    user).
    """
    try:
        return make_model(model_id, **kwargs)
    except Exception as e:  # noqa: BLE001
        logger.warning("Deferring model setup failure for %s", model_id, exc_info=True)
        return make_deferred_error_model(e, model_id=model_id)