open-swe/agent/middleware/sanitize_thinking_blocks.py

"""Middleware that removes malformed Anthropic thinking blocks before model calls."""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import Any

from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
from langchain_anthropic import ChatAnthropic
from langchain_aws import ChatBedrockConverse
from langchain_core.messages import AIMessage


def _is_chat_anthropic(model: object) -> bool:
    seen: set[int] = set()
    current = model
    for _ in range(10):
        if isinstance(current, (ChatAnthropic, ChatBedrockConverse)):
            return True
        current_id = id(current)
        if current_id in seen:
            return False
        seen.add(current_id)
        bound = getattr(current, "bound", None)
        if bound is None or bound is current:
            return False
        current = bound
    return False


def _is_empty_thinking_block(block: object) -> bool:
    """True for an empty Anthropic ``thinking`` block or an empty Bedrock
    ``reasoning_content`` block (Bedrock Converse emits the latter shape)."""
    if not isinstance(block, dict):
        return False
    block_type = block.get("type")
    if block_type == "thinking":
        return not block.get("thinking")
    if block_type == "reasoning_content":
        payload = block.get("reasoning_content")
        if isinstance(payload, dict):
            return not (
                payload.get("text")
                or payload.get("signature")
                or payload.get("redactedContent")
                or payload.get("redacted_content")
            )
        return not payload
    return False


def _sanitize_messages(messages: list[Any]) -> None:
    for message in messages:
        if not isinstance(message, AIMessage) or not isinstance(message.content, list):
            continue
        content = [block for block in message.content if not _is_empty_thinking_block(block)]
        if len(content) != len(message.content):
            message.content = content


class SanitizeThinkingBlocksMiddleware(AgentMiddleware):
    """Drop empty Anthropic thinking blocks before provider validation."""

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelCallResult:
        if _is_chat_anthropic(request.model):
            _sanitize_messages(request.messages)
        return handler(request)

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> Any:
        if _is_chat_anthropic(request.model):
            _sanitize_messages(request.messages)
        return await handler(request)