diff --git a/agent/middleware/__init__.py b/agent/middleware/__init__.py index 1298af4e..b402bd16 100644 --- a/agent/middleware/__init__.py +++ b/agent/middleware/__init__.py @@ -1,6 +1,7 @@ from .check_message_queue import check_message_queue_before_model from .ensure_no_empty_msg import ensure_no_empty_msg from .exclude_tools import ExcludeToolsMiddleware +from .model_fallback import ModelFallbackMiddleware from .notify_step_limit import notify_step_limit_reached from .refresh_slack_status import SlackAssistantStatusMiddleware from .sandbox_circuit_breaker import SandboxCircuitBreakerMiddleware @@ -9,6 +10,7 @@ from .tool_error_handler import ToolErrorMiddleware __all__ = [ "ExcludeToolsMiddleware", + "ModelFallbackMiddleware", "SanitizeToolInputsMiddleware", "ToolErrorMiddleware", "SandboxCircuitBreakerMiddleware", diff --git a/agent/middleware/model_fallback.py b/agent/middleware/model_fallback.py new file mode 100644 index 00000000..79bae7a0 --- /dev/null +++ b/agent/middleware/model_fallback.py @@ -0,0 +1,93 @@ +"""Middleware that falls back to a secondary model when the primary fails transiently. + +Wraps the model call. When the primary model raises a transient provider error +(5xx, 429, connection/timeout), the same request is retried once against the +configured fallback model. The fallback is bound to tools by the agent factory +on the second call, so swapping ``request.model`` is sufficient. + +Bidirectional: if the primary is Anthropic the fallback is typically OpenAI, +and vice versa. The middleware itself is provider-agnostic — it inspects the +exception type/status code to decide whether to fall over. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable +from typing import Any + +import anthropic +import openai +from langchain.agents.middleware import AgentMiddleware +from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse +from langchain_core.language_models import BaseChatModel + +logger = logging.getLogger(__name__) + +_RETRYABLE_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504, 529} + +_TRANSIENT_EXCEPTIONS: tuple[type[BaseException], ...] = ( + anthropic.APIConnectionError, + anthropic.APITimeoutError, + anthropic.RateLimitError, + anthropic.InternalServerError, + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + openai.InternalServerError, +) + + +def _should_fallback(exc: BaseException) -> bool: + if isinstance(exc, _TRANSIENT_EXCEPTIONS): + return True + # Catches OverloadedError (529) and other 5xx/429 surfaced as APIStatusError. + if isinstance(exc, (anthropic.APIStatusError, openai.APIStatusError)): + status = getattr(exc, "status_code", None) + if isinstance(status, int) and status in _RETRYABLE_STATUS_CODES: + return True + return False + + +class ModelFallbackMiddleware(AgentMiddleware): + """Retry the model call against a fallback provider on transient errors.""" + + def __init__(self, fallback_model: BaseChatModel) -> None: + super().__init__() + self._fallback_model = fallback_model + + def wrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelCallResult: + try: + return handler(request) + except Exception as exc: + if not _should_fallback(exc): + raise + logger.warning( + "Primary model failed (%s); falling back to %s", + type(exc).__name__, + getattr(self._fallback_model, "model_name", None) + or getattr(self._fallback_model, "model", "fallback"), + ) + return handler(request.override(model=self._fallback_model)) + + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> Any: + try: + return await handler(request) + except Exception as exc: + if not _should_fallback(exc): + raise + logger.warning( + "Primary model failed (%s); falling back to %s", + type(exc).__name__, + getattr(self._fallback_model, "model_name", None) + or getattr(self._fallback_model, "model", "fallback"), + ) + return await handler(request.override(model=self._fallback_model)) diff --git a/agent/server.py b/agent/server.py index 892c9038..06d66f97 100644 --- a/agent/server.py +++ b/agent/server.py @@ -6,6 +6,7 @@ import logging import os import warnings +from typing import Any logger = logging.getLogger(__name__) @@ -29,6 +30,7 @@ from langsmith.sandbox import SandboxClientError from .integrations.langsmith import _configure_github_proxy from .middleware import ( + ModelFallbackMiddleware, SandboxCircuitBreakerMiddleware, SanitizeToolInputsMiddleware, SlackAssistantStatusMiddleware, @@ -56,7 +58,7 @@ from .tools import ( from .utils.auth import resolve_github_token from .utils.authorship import resolve_triggering_user_identity from .utils.github_app import get_github_app_installation_token -from .utils.model import ModelKwargs, OpenAIReasoning, make_model +from .utils.model import ModelKwargs, OpenAIReasoning, fallback_model_id_for, make_model from .utils.sandbox import create_sandbox from .utils.sandbox_paths import aresolve_sandbox_work_dir @@ -362,6 +364,17 @@ async def get_agent(config: RunnableConfig) -> Pregel: if model_id == DEFAULT_LLM_MODEL_ID: model_kwargs["reasoning"] = DEFAULT_LLM_REASONING + fallback_model_id = os.environ.get("LLM_FALLBACK_MODEL_ID") or fallback_model_id_for(model_id) + fallback_middleware: list[Any] = [] + if fallback_model_id and fallback_model_id != model_id: + fallback_kwargs: ModelKwargs = {"max_tokens": DEFAULT_LLM_MAX_TOKENS} + if fallback_model_id.startswith("openai:"): + fallback_kwargs["reasoning"] = DEFAULT_LLM_REASONING + fallback_middleware.append( + ModelFallbackMiddleware(make_model(fallback_model_id, **fallback_kwargs)) + ) + logger.info("Configured model fallback %s -> %s", model_id, fallback_model_id) + logger.info("Returning agent with sandbox for thread %s", thread_id) return create_deep_agent( model=make_model(model_id, **model_kwargs), @@ -396,5 +409,6 @@ async def get_agent(config: RunnableConfig) -> Pregel: ensure_no_empty_msg, notify_step_limit_reached, SandboxCircuitBreakerMiddleware(), + *fallback_middleware, ], ).with_config(config) diff --git a/agent/utils/model.py b/agent/utils/model.py index d5f70dec..6bc84dff 100644 --- a/agent/utils/model.py +++ b/agent/utils/model.py @@ -4,6 +4,10 @@ from langchain.chat_models import init_chat_model OPENAI_RESPONSES_WS_BASE_URL = "wss://api.openai.com/v1" +# Anthropic SDK default is 2; a 529 burst can outlive that. Bump to give the +# primary provider a fair chance before the fallback middleware kicks in. +DEFAULT_MAX_RETRIES = 6 + OpenAIReasoningEffort = Literal["none", "low", "medium", "high", "xhigh"] @@ -16,13 +20,29 @@ class ModelKwargs(TypedDict, total=False): max_tokens: int | None reasoning: OpenAIReasoning | None temperature: float | None + max_retries: int | None def make_model(model_id: str, **kwargs: Unpack[ModelKwargs]): model_kwargs: dict[str, object] = kwargs.copy() + model_kwargs.setdefault("max_retries", DEFAULT_MAX_RETRIES) if model_id.startswith("openai:"): model_kwargs["base_url"] = OPENAI_RESPONSES_WS_BASE_URL model_kwargs["use_responses_api"] = True return init_chat_model(model=model_id, **model_kwargs) + + +def fallback_model_id_for(primary_model_id: str) -> str | None: + """Return the cross-provider fallback model id for a given primary, if any. + + Anthropic primaries fall back to OpenAI and vice versa. Returns ``None`` + when the provider has no configured cross-provider fallback (e.g. local + or self-hosted providers we don't want to silently route off-host). + """ + if primary_model_id.startswith("anthropic:"): + return "openai:gpt-5.5" + if primary_model_id.startswith("openai:"): + return "anthropic:claude-opus-4-5" + return None diff --git a/tests/test_model_fallback_middleware.py b/tests/test_model_fallback_middleware.py new file mode 100644 index 00000000..c39bab64 --- /dev/null +++ b/tests/test_model_fallback_middleware.py @@ -0,0 +1,134 @@ +"""Tests for ModelFallbackMiddleware.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import anthropic +import httpx +import openai +import pytest +from langchain_core.messages import AIMessage + +from agent.middleware.model_fallback import ( + ModelFallbackMiddleware, + _should_fallback, +) + + +def _anthropic_overloaded() -> anthropic.APIStatusError: + request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + response = httpx.Response( + 529, + request=request, + json={"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}, + ) + body = response.json() + return anthropic.APIStatusError("Overloaded", response=response, body=body) + + +def _openai_5xx() -> openai.APIStatusError: + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + response = httpx.Response(503, request=request, json={"error": {"message": "unavailable"}}) + return openai.APIStatusError("unavailable", response=response, body=response.json()) + + +def _make_request() -> MagicMock: + request = MagicMock() + request.override = MagicMock(return_value=MagicMock(name="overridden_request")) + return request + + +class TestShouldFallback: + def test_anthropic_529_overload_falls_back(self) -> None: + assert _should_fallback(_anthropic_overloaded()) is True + + def test_openai_503_falls_back(self) -> None: + assert _should_fallback(_openai_5xx()) is True + + def test_anthropic_rate_limit_falls_back(self) -> None: + request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + response = httpx.Response(429, request=request, json={"error": {}}) + exc = anthropic.RateLimitError("rate", response=response, body={}) + assert _should_fallback(exc) is True + + def test_anthropic_400_does_not_fall_back(self) -> None: + request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + response = httpx.Response(400, request=request, json={"error": {}}) + exc = anthropic.BadRequestError("bad", response=response, body={}) + assert _should_fallback(exc) is False + + def test_value_error_does_not_fall_back(self) -> None: + assert _should_fallback(ValueError("nope")) is False + + +class TestModelFallbackMiddleware: + @pytest.mark.asyncio + async def test_async_falls_over_on_overloaded(self) -> None: + fallback_model = MagicMock(name="fallback_model") + middleware = ModelFallbackMiddleware(fallback_model) + + calls: list[object] = [] + good_response = MagicMock(result=[AIMessage(content="ok from fallback")]) + + async def handler(req: object) -> object: + calls.append(req) + if len(calls) == 1: + raise _anthropic_overloaded() + return good_response + + request = _make_request() + result = await middleware.awrap_model_call(request, handler) + + assert result is good_response + assert len(calls) == 2 + request.override.assert_called_once_with(model=fallback_model) + assert calls[1] is request.override.return_value + + @pytest.mark.asyncio + async def test_async_propagates_non_transient_error(self) -> None: + middleware = ModelFallbackMiddleware(MagicMock()) + calls: list[object] = [] + + async def handler(req: object) -> object: + calls.append(req) + raise ValueError("not transient") + + with pytest.raises(ValueError, match="not transient"): + await middleware.awrap_model_call(_make_request(), handler) + + assert len(calls) == 1 + + @pytest.mark.asyncio + async def test_async_does_not_double_fall_back(self) -> None: + """If the fallback also fails transiently, the error propagates.""" + middleware = ModelFallbackMiddleware(MagicMock()) + calls: list[object] = [] + + async def handler(req: object) -> object: + calls.append(req) + raise _openai_5xx() + + with pytest.raises(openai.APIStatusError): + await middleware.awrap_model_call(_make_request(), handler) + + assert len(calls) == 2 + + def test_sync_falls_over_on_overloaded(self) -> None: + fallback_model = MagicMock(name="fallback_model") + middleware = ModelFallbackMiddleware(fallback_model) + calls: list[object] = [] + good_response = MagicMock(result=[AIMessage(content="ok")]) + + def handler(req: object) -> object: + calls.append(req) + if len(calls) == 1: + raise _anthropic_overloaded() + return good_response + + request = _make_request() + result = middleware.wrap_model_call(request, handler) + + assert result is good_response + assert len(calls) == 2 + request.override.assert_called_once_with(model=fallback_model)