diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 93b7abcef60..b53a7520ee2 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -1356,6 +1356,74 @@ class AsyncAnthropicAuxiliaryClient: self._real_client = sync_wrapper._real_client +class _BedrockCompletionsAdapter: + """Translates ``chat.completions.create(**kwargs)`` into Bedrock Converse.""" + + def __init__(self, region: str, model: str): + self._region = region + self._model = model + + def create(self, **kwargs) -> Any: + from agent.bedrock_adapter import call_converse + + messages = kwargs.get("messages", []) + model = kwargs.get("model", self._model) + max_tokens = kwargs.get("max_tokens") or kwargs.get("max_completion_tokens") + return call_converse( + region=self._region, + model=model, + messages=messages, + tools=kwargs.get("tools"), + max_tokens=int(max_tokens) if max_tokens else 4096, + temperature=kwargs.get("temperature"), + top_p=kwargs.get("top_p"), + stop_sequences=kwargs.get("stop"), + ) + + +class _BedrockChatShim: + def __init__(self, adapter: "_BedrockCompletionsAdapter"): + self.completions = adapter + + +class BedrockAuxiliaryClient: + """OpenAI-client-compatible wrapper over AWS Bedrock Converse API.""" + + def __init__(self, region: str, model: str): + self._region = region + self._model = model + adapter = _BedrockCompletionsAdapter(region, model) + self.chat = _BedrockChatShim(adapter) + self.api_key = "aws-sdk" + self.base_url = f"https://bedrock-runtime.{region}.amazonaws.com" + + def close(self): + pass + + +class _AsyncBedrockCompletionsAdapter: + def __init__(self, sync_adapter: _BedrockCompletionsAdapter): + self._sync = sync_adapter + + async def create(self, **kwargs) -> Any: + import asyncio + return await asyncio.to_thread(self._sync.create, **kwargs) + + +class _AsyncBedrockChatShim: + def __init__(self, adapter: _AsyncBedrockCompletionsAdapter): + self.completions = adapter + + +class AsyncBedrockAuxiliaryClient: + def __init__(self, sync_wrapper: "BedrockAuxiliaryClient"): + sync_adapter = sync_wrapper.chat.completions + async_adapter = _AsyncBedrockCompletionsAdapter(sync_adapter) + self.chat = _AsyncBedrockChatShim(async_adapter) + self.api_key = sync_wrapper.api_key + self.base_url = sync_wrapper.base_url + + def _endpoint_speaks_anthropic_messages(base_url: str) -> bool: """True if the endpoint at ``base_url`` speaks the Anthropic Messages protocol instead of OpenAI chat.completions. @@ -1411,6 +1479,8 @@ def _maybe_wrap_anthropic( # Already wrapped — don't double-wrap. if _safe_isinstance(client_obj, AnthropicAuxiliaryClient): return client_obj + if _safe_isinstance(client_obj, BedrockAuxiliaryClient): + return client_obj # Other specialized adapters we should never re-dispatch. if _safe_isinstance(client_obj, CodexAuxiliaryClient): return client_obj @@ -4204,6 +4274,8 @@ def _to_async_client(sync_client, model: str, is_vision: bool = False): return AsyncCodexAuxiliaryClient(sync_client), model if isinstance(sync_client, AnthropicAuxiliaryClient): return AsyncAnthropicAuxiliaryClient(sync_client), model + if isinstance(sync_client, BedrockAuxiliaryClient): + return AsyncBedrockAuxiliaryClient(sync_client), model try: from agent.gemini_native_adapter import GeminiNativeClient, AsyncGeminiNativeClient @@ -4943,10 +5015,14 @@ def resolve_provider_client( else (client, final_model)) elif pconfig.auth_type == "aws_sdk": - # AWS SDK providers (Bedrock) — use the Anthropic Bedrock client via - # boto3's credential chain (IAM roles, SSO, env vars, instance metadata). + # AWS SDK providers (Bedrock) — Claude models use the Anthropic Bedrock + # SDK (prompt caching, thinking); non-Claude models use Converse API. try: - from agent.bedrock_adapter import has_aws_credentials, resolve_bedrock_region + from agent.bedrock_adapter import ( + has_aws_credentials, + is_anthropic_bedrock_model, + resolve_bedrock_region, + ) from agent.anthropic_adapter import build_anthropic_bedrock_client except ImportError: logger.warning("resolve_provider_client: bedrock requested but " @@ -4961,17 +5037,26 @@ def resolve_provider_client( region = resolve_bedrock_region() default_model = "anthropic.claude-haiku-4-5-20251001-v1:0" final_model = _normalize_resolved_model(model or default_model, provider) - try: - real_client = build_anthropic_bedrock_client(region) - except ImportError as exc: - logger.warning("resolve_provider_client: cannot create Bedrock " - "client: %s", exc) - return None, None - client = AnthropicAuxiliaryClient( - real_client, final_model, api_key="aws-sdk", - base_url=f"https://bedrock-runtime.{region}.amazonaws.com", - ) - logger.debug("resolve_provider_client: bedrock (%s, %s)", final_model, region) + base_url = f"https://bedrock-runtime.{region}.amazonaws.com" + + if is_anthropic_bedrock_model(final_model): + try: + real_client = build_anthropic_bedrock_client(region) + except ImportError as exc: + logger.warning("resolve_provider_client: cannot create Bedrock " + "client: %s", exc) + return None, None + client = AnthropicAuxiliaryClient( + real_client, final_model, api_key="aws-sdk", + base_url=base_url, + ) + logger.debug("resolve_provider_client: bedrock anthropic (%s, %s)", + final_model, region) + else: + client = BedrockAuxiliaryClient(region, final_model) + logger.debug("resolve_provider_client: bedrock converse (%s, %s)", + final_model, region) + return (_to_async_client(client, final_model, is_vision=is_vision) if async_mode else (client, final_model))