open-swe/agent/middleware/sanitize_thinking_blocks.py
Adam Moussa 1c93e5c62b
feat: switch model providers to AWS Bedrock (Claude) and Fireworks (non-Claude)
Migrate off direct provider APIs: AWS Bedrock for Anthropic/Claude via the
cross-region inference profile us.anthropic.claude-opus-4-8, Fireworks AI for
all non-Claude models. Drop OpenAI (gpt-5.5) and Google (gemini-3.5-flash)
entirely. DEFAULT_MODEL_ID is now Bedrock Claude; all Fireworks models stay
freely selectable for the agent and reviewer graphs and via team/profile
defaults.

- pyproject: add langchain-aws (ChatBedrockConverse + boto3)
- options.py: Bedrock Claude entry + default; remove openai/google entries
- model.py: bedrock_converse provider_model_kwargs (effort -> thinking budget),
  region pin in make_model, bedrock<->fireworks fallback pairing, AWS_REGION/
  FIREWORKS_API_KEY local-dev validation
- server.py: provider-aware fallback kwargs build
- sanitize_thinking_blocks: also sanitize ChatBedrockConverse thinking blocks
- model_fallback: treat transient botocore ClientError codes as fallback-worthy
- eval_jobs: repoint hardcoded eval model id to Bedrock Claude
- tests: repoint dropped model ids; drop obsolete google test module
2026-06-29 14:57:14 -04:00

68 lines
2.2 KiB
Python

"""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 _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 (
isinstance(block, dict)
and block.get("type") == "thinking"
and not block.get("thinking")
)
]
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)