script-build-and-framework-tmp/agent/agent/llm/anthropic_protocol.py

"""Provider-neutral helpers for Anthropic Messages API adapters."""

from __future__ import annotations

import base64
import json
import logging
import mimetypes
import struct
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import httpx

from .pricing import calculate_cost
from .usage import TokenUsage


ANTHROPIC_RETRYABLE_EXCEPTIONS = (
    httpx.RemoteProtocolError,
    httpx.ConnectError,
    httpx.ReadTimeout,
    httpx.WriteTimeout,
    httpx.ConnectTimeout,
    httpx.PoolTimeout,
    ConnectionError,
)

ANTHROPIC_MODEL_EXACT = {
    "claude-sonnet-4-6": "claude-sonnet-4-6",
    "claude-sonnet-4.6": "claude-sonnet-4-6",
    "claude-sonnet-4-5-20250929": "claude-sonnet-4-5-20250929",
    "claude-sonnet-4-5": "claude-sonnet-4-5-20250929",
    "claude-sonnet-4.5": "claude-sonnet-4-5-20250929",
    "claude-opus-4-6": "claude-opus-4-6",
    "claude-opus-4-5-20251101": "claude-opus-4-5-20251101",
    "claude-opus-4-5": "claude-opus-4-5-20251101",
    "claude-opus-4-1-20250805": "claude-opus-4-1-20250805",
    "claude-opus-4-1": "claude-opus-4-1-20250805",
    "claude-haiku-4-5-20251001": "claude-haiku-4-5-20251001",
    "claude-haiku-4-5": "claude-haiku-4-5-20251001",
}

ANTHROPIC_MODEL_FUZZY: List[Tuple[str, str]] = [
    ("sonnet-4-6", "claude-sonnet-4-6"),
    ("sonnet-4.6", "claude-sonnet-4-6"),
    ("sonnet-4-5", "claude-sonnet-4-5-20250929"),
    ("sonnet-4.5", "claude-sonnet-4-5-20250929"),
    ("opus-4-6", "claude-opus-4-6"),
    ("opus-4.6", "claude-opus-4-6"),
    ("opus-4-5", "claude-opus-4-5-20251101"),
    ("opus-4.5", "claude-opus-4-5-20251101"),
    ("opus-4-1", "claude-opus-4-1-20250805"),
    ("opus-4.1", "claude-opus-4-1-20250805"),
    ("haiku-4-5", "claude-haiku-4-5-20251001"),
    ("haiku-4.5", "claude-haiku-4-5-20251001"),
    ("sonnet", "claude-sonnet-4-6"),
    ("opus", "claude-opus-4-6"),
    ("haiku", "claude-haiku-4-5-20251001"),
]


def resolve_anthropic_model(
    model: str,
    *,
    provider_prefix: str = "",
    preserve_unknown_prefix: bool = False,
    logger: Optional[logging.Logger] = None,
) -> str:
    """Resolve framework aliases while preserving each Provider's fallback."""

    original = model
    bare = model.split("/", 1)[1] if "/" in model else model
    target = ANTHROPIC_MODEL_EXACT.get(bare)
    if target is None:
        bare_lower = bare.lower()
        target = next(
            (
                candidate
                for keyword, candidate in ANTHROPIC_MODEL_FUZZY
                if keyword in bare_lower
            ),
            None,
        )
    if target is not None:
        resolved = f"{provider_prefix}{target}"
        if logger and bare != target:
            logger.info("Anthropic model alias resolved: %s -> %s", model, resolved)
        return resolved

    fallback = original if preserve_unknown_prefix else bare
    if logger:
        logger.warning("Could not resolve Anthropic model alias: %s", model)
    return fallback


def normalize_tool_call_ids(
    messages: List[Dict[str, Any]], target_prefix: str
) -> List[Dict[str, Any]]:
    """Rewrite cross-Provider tool call IDs without mutating input messages."""

    id_map: Dict[str, str] = {}
    for message in messages:
        if message.get("role") != "assistant":
            continue
        for tool_call in message.get("tool_calls") or ():
            old_id = tool_call.get("id", "")
            if old_id and not old_id.startswith(f"{target_prefix}_"):
                id_map.setdefault(old_id, f"{target_prefix}_{len(id_map):06x}")

    if not id_map:
        return messages

    normalized = []
    for message in messages:
        if message.get("role") == "assistant" and message.get("tool_calls"):
            tool_calls = [
                {**tool_call, "id": id_map[tool_call["id"]]}
                if tool_call.get("id") in id_map
                else tool_call
                for tool_call in message["tool_calls"]
            ]
            normalized.append({**message, "tool_calls": tool_calls})
        elif message.get("role") == "tool" and message.get("tool_call_id") in id_map:
            normalized.append(
                {**message, "tool_call_id": id_map[message["tool_call_id"]]}
            )
        else:
            normalized.append(message)
    return normalized


def _image_dimensions(data: bytes) -> Optional[Tuple[int, int]]:
    """Read PNG/JPEG dimensions from headers without a Pillow dependency."""

    try:
        if data[:8] == b"\x89PNG\r\n\x1a\n" and len(data) >= 24:
            return struct.unpack(">II", data[16:24])
        if data[:2] == b"\xff\xd8":
            offset = 2
            while offset < len(data) - 9:
                if data[offset] != 0xFF:
                    break
                marker = data[offset + 1]
                if marker in (0xC0, 0xC2):
                    height, width = struct.unpack(">HH", data[offset + 5 : offset + 9])
                    return width, height
                length = struct.unpack(">H", data[offset + 2 : offset + 4])[0]
                offset += 2 + length
    except (IndexError, struct.error):
        return None
    return None


def to_anthropic_content(
    content: Any,
    *,
    resolve_local_files: bool = True,
    include_image_metadata: bool = True,
    logger: Optional[logging.Logger] = None,
) -> Any:
    """Convert OpenAI content blocks to Anthropic content blocks."""

    if not isinstance(content, list):
        return content

    converted = []
    for block in content:
        if not isinstance(block, dict) or block.get("type") != "image_url":
            converted.append(block)
            continue

        image_url = block.get("image_url", {})
        url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url)
        image_data: Optional[bytes] = None
        if url.startswith("data:"):
            header, _, encoded = url.partition(",")
            media_type = (
                header.split(":", 1)[1].split(";", 1)[0]
                if ":" in header
                else "image/png"
            )
            source = {"type": "base64", "media_type": media_type, "data": encoded}
            if include_image_metadata:
                image_data = base64.b64decode(encoded)
        elif resolve_local_files and Path(url).is_file():
            path = Path(url)
            image_data = path.read_bytes()
            media_type = mimetypes.guess_type(str(path))[0] or "image/png"
            source = {
                "type": "base64",
                "media_type": media_type,
                "data": base64.b64encode(image_data).decode("ascii"),
            }
            if logger:
                logger.info("Converted local image to base64: %s (%d bytes)", url, len(image_data))
        else:
            source = {"type": "url", "url": url}

        image_block: Dict[str, Any] = {"type": "image", "source": source}
        if include_image_metadata and image_data:
            dimensions = _image_dimensions(image_data)
            if dimensions:
                image_block["_image_meta"] = {
                    "width": dimensions[0],
                    "height": dimensions[1],
                }
        converted.append(image_block)
    return converted


def to_anthropic_messages(
    messages: List[Dict[str, Any]],
    *,
    resolve_local_files: bool = True,
    include_image_metadata: bool = True,
    split_tool_result_images: bool = True,
    logger: Optional[logging.Logger] = None,
) -> Tuple[Any, List[Dict[str, Any]]]:
    """Convert OpenAI message history to the Anthropic Messages format."""

    def convert(content: Any) -> Any:
        return to_anthropic_content(
            content,
            resolve_local_files=resolve_local_files,
            include_image_metadata=include_image_metadata,
            logger=logger,
        )

    system_prompt = None
    anthropic_messages: List[Dict[str, Any]] = []
    for message in messages:
        role = message.get("role", "")
        content = message.get("content", "")
        if role == "system":
            system_prompt = content
        elif role == "user":
            anthropic_messages.append({"role": "user", "content": convert(content)})
        elif role == "assistant":
            tool_calls = message.get("tool_calls")
            if not tool_calls:
                anthropic_messages.append({"role": "assistant", "content": content})
                continue
            content_blocks: List[Dict[str, Any]] = []
            if content:
                assistant_content = convert(content)
                if isinstance(assistant_content, list):
                    content_blocks.extend(assistant_content)
                elif isinstance(assistant_content, str) and assistant_content.strip():
                    content_blocks.append({"type": "text", "text": assistant_content})
            for tool_call in tool_calls:
                function = tool_call.get("function", {})
                arguments = function.get("arguments", "{}")
                try:
                    parsed_arguments = (
                        json.loads(arguments) if isinstance(arguments, str) else arguments
                    )
                except json.JSONDecodeError:
                    parsed_arguments = {}
                content_blocks.append(
                    {
                        "type": "tool_use",
                        "id": tool_call.get("id", ""),
                        "name": function.get("name", ""),
                        "input": parsed_arguments,
                    }
                )
            anthropic_messages.append(
                {"role": "assistant", "content": content_blocks}
            )
        elif role == "tool":
            tool_content = convert(content)
            if not split_tool_result_images:
                blocks = [
                    {
                        "type": "tool_result",
                        "tool_use_id": message.get("tool_call_id", ""),
                        "content": tool_content,
                    }
                ]
            else:
                text_parts: List[Dict[str, Any]] = []
                image_parts: List[Dict[str, Any]] = []
                if isinstance(tool_content, list):
                    for block in tool_content:
                        if isinstance(block, dict) and block.get("type") == "image":
                            image_parts.append(block)
                        else:
                            text_parts.append(block)
                elif isinstance(tool_content, str) and tool_content:
                    text_parts.append({"type": "text", "text": tool_content})
                tool_result: Dict[str, Any] = {
                    "type": "tool_result",
                    "tool_use_id": message.get("tool_call_id", ""),
                }
                if len(text_parts) == 1 and text_parts[0].get("type") == "text":
                    tool_result["content"] = text_parts[0]["text"]
                elif text_parts:
                    tool_result["content"] = text_parts
                blocks = [tool_result, *image_parts]

            if (
                anthropic_messages
                and anthropic_messages[-1].get("role") == "user"
                and isinstance(anthropic_messages[-1].get("content"), list)
                and anthropic_messages[-1]["content"]
                and anthropic_messages[-1]["content"][0].get("type") == "tool_result"
            ):
                anthropic_messages[-1]["content"].extend(blocks)
            else:
                anthropic_messages.append({"role": "user", "content": blocks})
    return system_prompt, anthropic_messages


def to_anthropic_tools(tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Convert OpenAI function definitions to Anthropic tool definitions."""

    return [
        {
            "name": function.get("name", ""),
            "description": function.get("description", ""),
            "input_schema": function.get(
                "parameters", {"type": "object", "properties": {}}
            ),
        }
        for tool in tools
        if tool.get("type") == "function"
        for function in (tool["function"],)
    ]


def parse_anthropic_response(result: Dict[str, Any]) -> Dict[str, Any]:
    """Convert an Anthropic response to the framework's unified shape."""

    text_parts = []
    tool_calls = []
    for block in result.get("content", []):
        if block.get("type") == "text":
            text_parts.append(block.get("text", ""))
        elif block.get("type") == "tool_use":
            tool_calls.append(
                {
                    "id": block.get("id", ""),
                    "type": "function",
                    "function": {
                        "name": block.get("name", ""),
                        "arguments": json.dumps(
                            block.get("input", {}), ensure_ascii=False
                        ),
                    },
                }
            )
    stop_reason = result.get("stop_reason", "end_turn")
    finish_reason = {
        "end_turn": "stop",
        "tool_use": "tool_calls",
        "max_tokens": "length",
        "stop_sequence": "stop",
    }.get(stop_reason, stop_reason)
    raw_usage = result.get("usage", {})
    return {
        "content": "\n".join(text_parts),
        "tool_calls": tool_calls or None,
        "finish_reason": finish_reason,
        "usage": TokenUsage(
            input_tokens=raw_usage.get("input_tokens", 0),
            output_tokens=raw_usage.get("output_tokens", 0),
            cache_creation_tokens=raw_usage.get("cache_creation_input_tokens", 0),
            cache_read_tokens=raw_usage.get("cache_read_input_tokens", 0),
        ),
    }


def count_cache_controls(system_prompt: Any, messages: List[Dict[str, Any]]) -> int:
    """Count cache-control blocks for optional Provider diagnostics."""

    blocks = list(system_prompt) if isinstance(system_prompt, list) else []
    for message in messages:
        if isinstance(message.get("content"), list):
            blocks.extend(message["content"])
    return sum(
        1 for block in blocks if isinstance(block, dict) and "cache_control" in block
    )


def build_anthropic_result(parsed: Dict[str, Any], model: str) -> Dict[str, Any]:
    """Attach cost and token compatibility fields to a parsed response."""

    usage = parsed["usage"]
    return {
        "content": parsed["content"],
        "tool_calls": parsed["tool_calls"],
        "prompt_tokens": usage.input_tokens,
        "completion_tokens": usage.output_tokens,
        "reasoning_tokens": usage.reasoning_tokens,
        "cache_creation_tokens": usage.cache_creation_tokens,
        "cache_read_tokens": usage.cache_read_tokens,
        "finish_reason": parsed["finish_reason"],
        "cost": calculate_cost(model, usage),
        "usage": usage,
    }


__all__ = [
    "ANTHROPIC_MODEL_EXACT",
    "ANTHROPIC_MODEL_FUZZY",
    "ANTHROPIC_RETRYABLE_EXCEPTIONS",
    "build_anthropic_result",
    "count_cache_controls",
    "normalize_tool_call_ids",
    "parse_anthropic_response",
    "resolve_anthropic_model",
    "to_anthropic_content",
    "to_anthropic_messages",
    "to_anthropic_tools",
]