pallas/pallas/mantle_shims.py

"""AWS Bedrock Mantle compatibility shims for fast-agent.

Mantle is AWS's Anthropic-Messages-API-compatible gateway, hosted at
``https://bedrock-mantle.{region}.api.aws/anthropic``. Fast-agent talks to it
via its built-in ``anthropic`` provider, but two layers of reshaping are needed
before the wire traffic is valid:

1. **Model-name prefix.** Mantle requires the full ``anthropic.<name>`` wire
   id (e.g. ``anthropic.claude-opus-4-7``). Fast-agent's model-spec parser
   treats the ``anthropic.`` prefix as the provider hint and strips it off
   the wire name. We re-register the prefixed forms via
   ``ModelDatabase._PROVIDER_WIRE_MODEL_NAMES`` so the right id goes out.

2. **``caller: null`` leakage on replayed ``tool_use`` blocks.** Anthropic
   SDK 0.100.x ``BetaToolUseBlock`` carries an optional ``caller`` field;
   the matching ``BetaToolUseBlockParam`` TypedDict declares it required.
   Fast-agent's multipart converter re-serialises assistant history with
   ``exclude_none=False``, producing ``{"type": "tool_use", ..., "caller": null}``.
   ``api.anthropic.com`` silently accepts that; Mantle rejects it as
   ``tool_use.caller: Input should be a valid dictionary or object``,
   breaking the tool-use loop on the second turn. We strip ``caller`` from
   any ``tool_use`` dict emitted by the two static methods that feed
   replayed history back into the wire.

   Upstream SDK tracker: https://github.com/anthropics/anthropic-sdk-python/issues/1454

Both shims are idempotent and may be installed at process startup before any
fast-agent ``FastAgent`` instance is constructed.
"""
from __future__ import annotations

import logging
from collections.abc import Mapping, Sequence
from typing import Any, cast

logger = logging.getLogger(__name__)


# ── Model ids known to be served on Mantle (keep in sync with AWS docs). ──
# The key is fast-agent's internal model_name (provider prefix stripped),
# the value is the wire id Mantle expects.
MANTLE_WIRE_NAMES: dict[str, str] = {
    "claude-haiku-4-5": "anthropic.claude-haiku-4-5",
    "claude-opus-4-7":  "anthropic.claude-opus-4-7",
    "claude-opus-4-8":  "anthropic.claude-opus-4-8",
    "claude-fable-5":   "anthropic.claude-fable-5",
}


def is_mantle_base_url(base_url: str | None) -> bool:
    """Return True if the given anthropic base_url points at Mantle."""
    if not base_url:
        return False
    return "bedrock-mantle" in base_url


# ── Shim 1: model-name prefix ────────────────────────────────────────────────

def install_wire_name_prefix() -> None:
    """Register the prefixed wire ids for known Mantle-hosted Claude models."""
    from fast_agent.llm.model_database import ModelDatabase
    from fast_agent.llm.provider_types import Provider

    for fa_name, wire_name in MANTLE_WIRE_NAMES.items():
        key = (Provider.ANTHROPIC, ModelDatabase.normalize_model_name(fa_name))
        ModelDatabase._PROVIDER_WIRE_MODEL_NAMES[key] = wire_name  # noqa: SLF001

    logger.info(
        "Mantle wire-name shim installed for models: %s",
        ", ".join(sorted(MANTLE_WIRE_NAMES.keys())),
    )


# ── Shim 2: strip `caller` from replayed tool_use blocks ─────────────────────

def _strip_tool_use_caller(blocks: list[Any]) -> list[Any]:
    """Remove the stray ``caller`` field Anthropic SDK 0.100.x leaks into
    replayed ``tool_use`` blocks. Idempotent; only touches dicts whose
    ``type == "tool_use"``.
    """
    for block in blocks:
        if isinstance(block, dict) and block.get("type") == "tool_use":
            block.pop("caller", None)
    return blocks


_tool_use_patch_installed = False


def install_tool_use_caller_strip() -> None:
    """Monkeypatch ``AnthropicConverter`` to drop ``caller`` from replayed
    ``tool_use`` blocks. Safe to call more than once; subsequent calls are
    no-ops.
    """
    global _tool_use_patch_installed
    if _tool_use_patch_installed:
        return

    from fast_agent.llm.provider.anthropic.multipart_converter_anthropic import (
        AnthropicConverter,
    )

    original_deserialize = AnthropicConverter._deserialize_assistant_raw_blocks  # noqa: SLF001

    def patched_deserialize(
        channels: Mapping[str, Sequence[Any]],
    ) -> list[Any]:
        result = original_deserialize(channels)
        return cast("list[Any]", _strip_tool_use_caller(list(result)))

    AnthropicConverter._deserialize_assistant_raw_blocks = staticmethod(  # noqa: SLF001
        patched_deserialize
    )

    original_append = AnthropicConverter._append_server_tool_channel_blocks  # noqa: SLF001

    def patched_append(
        channels: Mapping[str, Sequence[Any]] | None,
        destination: list[Any],
    ) -> None:
        original_append(channels, destination)
        _strip_tool_use_caller(destination)

    AnthropicConverter._append_server_tool_channel_blocks = staticmethod(  # noqa: SLF001
        patched_append
    )

    _tool_use_patch_installed = True
    logger.info("Mantle tool_use.caller strip shim installed")


# ── Orchestrator ─────────────────────────────────────────────────────────────

def install_all() -> None:
    """Install all Mantle shims. Call once at process startup."""
    install_wire_name_prefix()
    install_tool_use_caller_strip()


def maybe_install(anthropic_base_url: str | None) -> bool:
    """Install shims only when ``anthropic_base_url`` is a Mantle endpoint.
    Returns True if the shims were installed.
    """
    if not is_mantle_base_url(anthropic_base_url):
        return False
    install_all()
    return True