150 lines
6.0 KiB
Python
150 lines
6.0 KiB
Python
"""Tests for pallas.mantle_shims.
|
|
|
|
These tests exercise the module in isolation: they do not hit the network
|
|
and do not require fast-agent to be configured. They do, however, import
|
|
fast-agent so that the monkeypatch targets exist — so fast-agent-mcp must
|
|
be installed in the environment running pytest.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from pallas import mantle_shims
|
|
|
|
|
|
# ── is_mantle_base_url ───────────────────────────────────────────────────────
|
|
|
|
@pytest.mark.parametrize(
|
|
"url,expected",
|
|
[
|
|
("https://bedrock-mantle.us-east-1.api.aws/anthropic", True),
|
|
("https://bedrock-mantle.ca-central-1.api.aws/anthropic", True),
|
|
("https://api.anthropic.com", False),
|
|
("https://example.com/bedrock", False),
|
|
("", False),
|
|
(None, False),
|
|
],
|
|
)
|
|
def test_is_mantle_base_url(url: str | None, expected: bool) -> None:
|
|
assert mantle_shims.is_mantle_base_url(url) is expected
|
|
|
|
|
|
# ── install_wire_name_prefix ─────────────────────────────────────────────────
|
|
|
|
def test_install_wire_name_prefix_registers_prefixed_ids() -> None:
|
|
from fast_agent.llm.model_database import ModelDatabase
|
|
from fast_agent.llm.provider_types import Provider
|
|
|
|
mantle_shims.install_wire_name_prefix()
|
|
|
|
for fa_name, wire_name in mantle_shims.MANTLE_WIRE_NAMES.items():
|
|
key = (Provider.ANTHROPIC, ModelDatabase.normalize_model_name(fa_name))
|
|
assert ModelDatabase._PROVIDER_WIRE_MODEL_NAMES.get(key) == wire_name
|
|
|
|
|
|
def test_install_wire_name_prefix_is_idempotent() -> None:
|
|
mantle_shims.install_wire_name_prefix()
|
|
mantle_shims.install_wire_name_prefix() # must not raise
|
|
# Second call leaves the same mapping in place.
|
|
from fast_agent.llm.model_database import ModelDatabase
|
|
from fast_agent.llm.provider_types import Provider
|
|
|
|
key = (Provider.ANTHROPIC, ModelDatabase.normalize_model_name("claude-opus-4-7"))
|
|
assert ModelDatabase._PROVIDER_WIRE_MODEL_NAMES[key] == "anthropic.claude-opus-4-7"
|
|
|
|
|
|
# ── _strip_tool_use_caller ───────────────────────────────────────────────────
|
|
|
|
def test_strip_tool_use_caller_removes_caller_key() -> None:
|
|
blocks = [
|
|
{"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None},
|
|
{"type": "text", "text": "hello"},
|
|
{"type": "tool_use", "id": "t2", "name": "bar", "input": {}}, # no caller
|
|
]
|
|
result = mantle_shims._strip_tool_use_caller(blocks)
|
|
|
|
assert "caller" not in result[0]
|
|
assert result[1] == {"type": "text", "text": "hello"}
|
|
assert "caller" not in result[2]
|
|
|
|
|
|
def test_strip_tool_use_caller_is_idempotent() -> None:
|
|
blocks = [{"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None}]
|
|
mantle_shims._strip_tool_use_caller(blocks)
|
|
mantle_shims._strip_tool_use_caller(blocks) # second pass must be a no-op
|
|
assert "caller" not in blocks[0]
|
|
|
|
|
|
def test_strip_tool_use_caller_ignores_non_dict_blocks() -> None:
|
|
# Anthropic SDK model objects are sometimes passed instead of dicts;
|
|
# the helper must leave those untouched.
|
|
class Sentinel:
|
|
type = "tool_use"
|
|
|
|
s = Sentinel()
|
|
blocks = [s]
|
|
mantle_shims._strip_tool_use_caller(blocks)
|
|
assert blocks[0] is s # unchanged
|
|
|
|
|
|
# ── install_tool_use_caller_strip ────────────────────────────────────────────
|
|
|
|
def test_install_tool_use_caller_strip_patches_converter() -> None:
|
|
from fast_agent.llm.provider.anthropic.multipart_converter_anthropic import (
|
|
AnthropicConverter,
|
|
)
|
|
|
|
mantle_shims.install_tool_use_caller_strip()
|
|
|
|
# The patched deserialize must strip caller from replayed tool_use dicts.
|
|
channels: dict[str, list[dict]] = {"assistant_raw": []}
|
|
# We can't easily stub the original's internal behaviour, so just
|
|
# verify the patch is in place by round-tripping a destination list
|
|
# through _append_server_tool_channel_blocks, which we *know* ends by
|
|
# calling our strip helper on `destination`.
|
|
destination: list[dict] = [
|
|
{"type": "tool_use", "id": "t1", "name": "foo", "input": {}, "caller": None},
|
|
]
|
|
AnthropicConverter._append_server_tool_channel_blocks(None, destination)
|
|
assert "caller" not in destination[0]
|
|
|
|
|
|
def test_install_tool_use_caller_strip_is_idempotent() -> None:
|
|
mantle_shims.install_tool_use_caller_strip()
|
|
mantle_shims.install_tool_use_caller_strip() # must not raise or re-wrap
|
|
|
|
|
|
# ── maybe_install ────────────────────────────────────────────────────────────
|
|
|
|
def test_maybe_install_installs_when_mantle(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
calls: list[str] = []
|
|
monkeypatch.setattr(
|
|
mantle_shims, "install_wire_name_prefix",
|
|
lambda: calls.append("wire"),
|
|
)
|
|
monkeypatch.setattr(
|
|
mantle_shims, "install_tool_use_caller_strip",
|
|
lambda: calls.append("tool_use"),
|
|
)
|
|
|
|
installed = mantle_shims.maybe_install("https://bedrock-mantle.us-east-1.api.aws/anthropic")
|
|
assert installed is True
|
|
assert calls == ["wire", "tool_use"]
|
|
|
|
|
|
def test_maybe_install_noop_for_non_mantle(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
calls: list[str] = []
|
|
monkeypatch.setattr(
|
|
mantle_shims, "install_wire_name_prefix",
|
|
lambda: calls.append("wire"),
|
|
)
|
|
monkeypatch.setattr(
|
|
mantle_shims, "install_tool_use_caller_strip",
|
|
lambda: calls.append("tool_use"),
|
|
)
|
|
|
|
assert mantle_shims.maybe_install("https://api.anthropic.com") is False
|
|
assert mantle_shims.maybe_install(None) is False
|
|
assert mantle_shims.maybe_install("") is False
|
|
assert calls == []
|