From 259842243d6438f393768e684882211ca0bcc4e7 Mon Sep 17 00:00:00 2001 From: Jan Piotrowski Date: Wed, 25 Mar 2026 16:41:06 +0100 Subject: [PATCH] feat: add context retry middleware when API returns 4xx errors cause by context exceeding limits (#92) --- .gitignore | 1 + EvoScientist/EvoScientist.py | 29 +++-- EvoScientist/middleware/__init__.py | 2 + EvoScientist/middleware/context_overflow.py | 84 +++++++++++++++ tests/test_context_overflow_middleware.py | 112 ++++++++++++++++++++ 5 files changed, 220 insertions(+), 8 deletions(-) create mode 100644 EvoScientist/middleware/context_overflow.py create mode 100644 tests/test_context_overflow_middleware.py diff --git a/.gitignore b/.gitignore index 9d0e085..1a8ce31 100644 --- a/.gitignore +++ b/.gitignore @@ -37,6 +37,7 @@ workspace/ skills/ memory/ media/ +conversation_history/ .deno_cache/ *.ipynb *CLAUDE.md diff --git a/EvoScientist/EvoScientist.py b/EvoScientist/EvoScientist.py index ea5d6c4..9274edb 100644 --- a/EvoScientist/EvoScientist.py +++ b/EvoScientist/EvoScientist.py @@ -23,6 +23,8 @@ import os from datetime import datetime from pathlib import Path +from langchain.agents.middleware import AgentMiddleware + from . import paths as _paths_mod from .config import apply_config_to_env, get_effective_config from .paths import set_active_workspace, set_workspace_root @@ -136,10 +138,12 @@ def _inject_subagent_middleware(subs: list[dict]) -> None: ToolNode handler which produces terse messages without tracebacks or retry guidance — reducing the subagent's ability to self-recover. """ - from .middleware import ToolErrorHandlerMiddleware + from .middleware import ContextOverflowMapperMiddleware, ToolErrorHandlerMiddleware for sa in subs: - sa.setdefault("middleware", []).append(ToolErrorHandlerMiddleware()) + sa.setdefault("middleware", []).append( + ToolErrorHandlerMiddleware(), ContextOverflowMapperMiddleware() + ) def _build_prompt_refs() -> dict: @@ -270,14 +274,20 @@ def _get_default_backend(): def _get_default_middleware(): """Build the default middleware list.""" - from .middleware import ToolErrorHandlerMiddleware, create_memory_middleware + from .middleware import ( + ContextOverflowMapperMiddleware, + ToolErrorHandlerMiddleware, + create_memory_middleware, + ) cfg = _ensure_config() memory_dir = str(_paths_mod.MEMORY_DIR) mw = [ + ContextOverflowMapperMiddleware(), ToolErrorHandlerMiddleware(), create_memory_middleware(memory_dir, extraction_model=_ensure_chat_model()), ] + if cfg.enable_ask_user and not cfg.auto_approve: from .middleware.ask_user import AskUserMiddleware @@ -341,14 +351,16 @@ def create_cli_agent(workspace_dir: str | None = None, checkpointer=None, config from . import paths as _paths from .backends import CustomSandboxBackend, MergedReadOnlyBackend - from .middleware import ToolErrorHandlerMiddleware, create_memory_middleware + from .middleware import ( + ContextOverflowMapperMiddleware, + ToolErrorHandlerMiddleware, + create_memory_middleware, + ) cfg = _ensure_config(config) if checkpointer is None: - from langgraph.checkpoint.memory import ( - InMemorySaver, # type: ignore[import-untyped] - ) + from langgraph.checkpoint.memory import InMemorySaver checkpointer = InMemorySaver() @@ -392,7 +404,8 @@ def create_cli_agent(workspace_dir: str | None = None, checkpointer=None, config }, ) - mw = [ + mw: list[AgentMiddleware] = [ + ContextOverflowMapperMiddleware(), ToolErrorHandlerMiddleware(), create_memory_middleware(_mem_dir, extraction_model=_ensure_chat_model()), ] diff --git a/EvoScientist/middleware/__init__.py b/EvoScientist/middleware/__init__.py index ddbcc38..e728b2c 100644 --- a/EvoScientist/middleware/__init__.py +++ b/EvoScientist/middleware/__init__.py @@ -11,6 +11,7 @@ from .ask_user import ( Choice, Question, ) +from .context_overflow import ContextOverflowMapperMiddleware from .memory import ( EvoMemoryMiddleware, EvoMemoryState, @@ -24,6 +25,7 @@ __all__ = [ "AskUserRequest", "AskUserWidgetResult", "Choice", + "ContextOverflowMapperMiddleware", "EvoMemoryMiddleware", "EvoMemoryState", "ExtractedMemory", diff --git a/EvoScientist/middleware/context_overflow.py b/EvoScientist/middleware/context_overflow.py new file mode 100644 index 0000000..1976c15 --- /dev/null +++ b/EvoScientist/middleware/context_overflow.py @@ -0,0 +1,84 @@ +"""Middleware that maps provider-specific context limit errors to a standard ContextOverflowError. + +When the LLM returns a 400 error indicating the context length has been +exceeded, this middleware detects it and raises a standard +``langchain_core.exceptions.ContextOverflowError``. This allows the +underlying framework (e.g. deepagents' SummarizationMiddleware) to catch +the error, trigger summarization, and retry the request automatically. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable + +from langchain.agents.middleware.types import ( + AgentMiddleware, + ModelRequest, + ModelResponse, +) + +logger = logging.getLogger(__name__) + + +class ContextOverflowMapperMiddleware(AgentMiddleware): + """Map provider-specific context limit errors to ContextOverflowError.""" + + name = "context_overflow_mapper" + + def wrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], ModelResponse], + ) -> ModelResponse: + try: + return handler(request) + except Exception as exc: + if self._is_context_limit_error(exc): + from langchain_core.exceptions import ContextOverflowError + + logger.warning( + "Context limit exceeded. Mapping to ContextOverflowError..." + ) + + raise ContextOverflowError(str(exc)) from exc + raise + + async def awrap_model_call( + self, + request: ModelRequest, + handler: Callable[[ModelRequest], Awaitable[ModelResponse]], + ) -> ModelResponse: + try: + return await handler(request) + except Exception as exc: + if self._is_context_limit_error(exc): + from langchain_core.exceptions import ContextOverflowError + + logger.warning( + "Context limit exceeded. Mapping to ContextOverflowError..." + ) + + raise ContextOverflowError(str(exc)) from exc + raise + + def _is_context_limit_error(self, exc: Exception) -> bool: + """Detect if an exception is a context length/limit error. + + It triggers when there's an error 400 raised and one of specified patterns exists in the error message. + """ + err_msg = str(exc).lower() + + patterns = [ + "context_length_exceeded", + "context length exceeded", + "too many tokens", + "maximum context length", + "output too large", + "context_window_exceeded", + "string_too_long", + "max_tokens_exceeded", + ] + is_400 = "400" in err_msg or "bad request" in err_msg + + return any(p in err_msg for p in patterns) and is_400 diff --git a/tests/test_context_overflow_middleware.py b/tests/test_context_overflow_middleware.py new file mode 100644 index 0000000..51976da --- /dev/null +++ b/tests/test_context_overflow_middleware.py @@ -0,0 +1,112 @@ +"""Tests for ContextOverflowMapperMiddleware.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain.agents.middleware.types import ModelRequest +from langchain_core.exceptions import ContextOverflowError +from langchain_core.messages import HumanMessage + +from EvoScientist.middleware.context_overflow import ContextOverflowMapperMiddleware + + +def test_is_context_limit_error_openai(): + mw = ContextOverflowMapperMiddleware() + exc = Exception( + "Error code: 400 - {'error': {'message': 'This model's maximum context length is 8192 tokens. However, your messages resulted in 10000 tokens.', 'type': 'invalid_request_error', 'param': 'messages', 'code': 'context_length_exceeded'}}" + ) + assert mw._is_context_limit_error(exc) is True + + +def test_is_context_limit_error_anthropic(): + mw = ContextOverflowMapperMiddleware() + exc = Exception("HTTP 400 Bad Request: Output too large") + assert mw._is_context_limit_error(exc) is True + + +def test_is_not_context_limit_error_without_400(): + mw = ContextOverflowMapperMiddleware() + exc = Exception("context_length_exceeded, but no status code") + assert mw._is_context_limit_error(exc) is False + + +def test_is_not_context_limit_error_with_400_but_no_pattern(): + mw = ContextOverflowMapperMiddleware() + exc = Exception("HTTP 400 Bad Request: Some other API error") + assert mw._is_context_limit_error(exc) is False + + +def test_is_not_context_limit_error_other_status(): + mw = ContextOverflowMapperMiddleware() + exc = Exception("HTTP 401 Unauthorized") + assert mw._is_context_limit_error(exc) is False + + +def test_wrap_model_call_raises_context_overflow(): + # Setup mocks + msgs = [HumanMessage(content=f"msg {i}") for i in range(10)] + request = ModelRequest( + messages=msgs, + model=MagicMock(), + state={}, + runtime=MagicMock(), + system_message=MagicMock(), + ) + + # Mock handler that fails with context error + handler = MagicMock() + handler.side_effect = Exception("400 Bad Request: context_length_exceeded") + + mw = ContextOverflowMapperMiddleware() + + with pytest.raises(ContextOverflowError) as excinfo: + mw.wrap_model_call(request, handler) + + assert "context_length_exceeded" in str(excinfo.value) + assert handler.call_count == 1 + + +@pytest.mark.anyio +async def test_awrap_model_call_raises_context_overflow(): + # Setup mocks + msgs = [HumanMessage(content=f"msg {i}") for i in range(10)] + request = ModelRequest( + messages=msgs, + model=MagicMock(), + state={}, + runtime=MagicMock(), + system_message=MagicMock(), + ) + + # Mock handler that fails with context error + handler = AsyncMock() + handler.side_effect = Exception("400 Bad Request: context_length_exceeded") + + mw = ContextOverflowMapperMiddleware() + + with pytest.raises(ContextOverflowError) as excinfo: + await mw.awrap_model_call(request, handler) + + assert "context_length_exceeded" in str(excinfo.value) + assert handler.call_count == 1 + + +@pytest.mark.anyio +async def test_awrap_model_call_passes_through_other_errors(): + request = ModelRequest( + messages=[], + model=MagicMock(), + state={}, + runtime=MagicMock(), + system_message=MagicMock(), + ) + handler = AsyncMock() + handler.side_effect = RuntimeError("Something else") + + mw = ContextOverflowMapperMiddleware() + + with pytest.raises(RuntimeError) as excinfo: + await mw.awrap_model_call(request, handler) + + assert "Something else" in str(excinfo.value) + assert not isinstance(excinfo.value, ContextOverflowError)