nexus-claude-api/tests/test_translators.py

from __future__ import annotations

import base64

from nexus_claude_api.models import AnthropicMessagesRequest
from nexus_claude_api.translators.anthropic_to_bedrock import anthropic_to_bedrock_request
from nexus_claude_api.translators.bedrock_to_anthropic import bedrock_to_anthropic_response
from nexus_claude_api.translators.stream import bedrock_stream_to_anthropic_events


def test_text_request_translation() -> None:
    payload = AnthropicMessagesRequest.model_validate(
        {
            "model": "claude-opus-4.6",
            "messages": [{"role": "user", "content": "Hello"}],
            "system": "You are helpful.",
            "max_tokens": 100,
            "temperature": 0.2,
        }
    )

    request = anthropic_to_bedrock_request(payload)

    assert request["modelId"] == "claude-opus-4.6"
    assert request["messages"] == [{"role": "user", "content": [{"text": "Hello"}]}]
    assert request["system"] == [{"text": "You are helpful."}]
    assert request["inferenceConfig"]["maxTokens"] == 100
    assert request["inferenceConfig"]["temperature"] == 0.2


def test_image_request_translation() -> None:
    image_data = base64.b64encode(b"image-bytes").decode("ascii")
    payload = AnthropicMessagesRequest.model_validate(
        {
            "model": "claude-opus-4.6",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Describe this"},
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": "image/png",
                                "data": image_data,
                            },
                        },
                    ],
                }
            ],
            "max_tokens": 100,
        }
    )

    request = anthropic_to_bedrock_request(payload)

    image_block = request["messages"][0]["content"][1]["image"]
    assert image_block["format"] == "png"
    assert image_block["source"]["bytes"] == b"image-bytes"


def test_tool_translation() -> None:
    payload = AnthropicMessagesRequest.model_validate(
        {
            "model": "claude-opus-4.6",
            "messages": [
                {"role": "user", "content": "Use a tool"},
                {
                    "role": "assistant",
                    "content": [
                        {
                            "type": "tool_use",
                            "id": "toolu_1",
                            "name": "get_weather",
                            "input": {"city": "Berlin"},
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": "toolu_1",
                            "content": "Sunny",
                        }
                    ],
                },
            ],
            "tools": [
                {
                    "name": "get_weather",
                    "description": "Get weather",
                    "input_schema": {
                        "type": "object",
                        "properties": {"city": {"type": "string"}},
                    },
                }
            ],
            "tool_choice": {"type": "auto"},
            "max_tokens": 100,
        }
    )

    request = anthropic_to_bedrock_request(payload)

    assert request["messages"][1]["content"][0]["toolUse"]["toolUseId"] == "toolu_1"
    assert request["messages"][2]["content"][0]["toolResult"]["toolUseId"] == "toolu_1"
    assert request["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather"
    assert request["toolConfig"]["toolChoice"] == {"auto": {}}


def test_bedrock_response_translation() -> None:
    response = {
        "ResponseMetadata": {"RequestId": "req-1"},
        "output": {"message": {"content": [{"text": "Hello there"}]}},
        "stopReason": "end_turn",
        "usage": {"inputTokens": 10, "outputTokens": 3},
    }

    translated = bedrock_to_anthropic_response(response, model="claude-opus-4.6")

    assert translated.id == "req-1"
    assert translated.content[0].type == "text"
    assert translated.content[0].text == "Hello there"
    assert translated.usage.input_tokens == 10
    assert translated.usage.output_tokens == 3


def test_stream_translation() -> None:
    events = [
        {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "Hi"}}},
        {"contentBlockStop": {"contentBlockIndex": 0}},
        {"metadata": {"usage": {"inputTokens": 4, "outputTokens": 1}}},
        {"messageStop": {"stopReason": "end_turn"}},
    ]

    translated = list(
        bedrock_stream_to_anthropic_events(events, model="claude-opus-4.6")
    )

    assert translated[0]["type"] == "message_start"
    assert any(event["type"] == "content_block_delta" for event in translated)
    assert translated[-2]["type"] == "message_delta"
    assert translated[-2]["usage"]["input_tokens"] == 4
    assert translated[-1]["type"] == "message_stop"