aidial-adapter-anthropic 0.1.0__py3-none-any.whl

aidial_adapter_anthropic/adapter/_claude/tokenizer/approximate.py

@@ -0,0 +1,260 @@
+"""
+An attempt to approximate the tokenizer for Claude models.
+
+This tokenizer doesn't provide the precise token count,
+because Anthropic doesn't provide the exact tokenization algorithm.
+
+This tokenizer provides an *overestimation* of the request token count.
+We need to be conservative, since the tokenizer is used in the prompt
+truncation algorithm. So we are choosing to be unable to pack the request with tokens
+as tightly as possible over making an additional chat completion request,
+which is going to fail with a token overflow error.
+
+1. For the text parts of request we count every byte in their UTF-8 encoding.
+Note that the official Claude 2 tokenizer couldn't be used
+for anything more than a very rough estimate:
+https://github.com/anthropics/anthropic-sdk-python/blob/246a2978694b584429d4bbd5b44245ff8eac2ac2/src/anthropic/_client.py#L270-L283
+
+2. For the image parts we use the official approximation:
+> tokens = (width px * height px)/750
+https://docs.anthropic.com/en/docs/build-with-claude/vision#calculate-image-costs
+
+3. For the tool usage we use the official approximation:
+https://docs.anthropic.com/en/docs/build-with-claude/tool-use#pricing
+    a. tool-related components of the request are serialized to strings and tokenized as such,
+    b. the hidden tool-enabling system prompt is accounted as per the documentation.
+
+TODO: use the official tokenizer:
+    https://docs.anthropic.com/en/docs/build-with-claude/token-counting
+    once it's supported in Bedrock:
+    https://github.com/anthropics/anthropic-sdk-python/blob/599f2b9a9501b8c98fb3132043c3ec71e3026f84/src/anthropic/lib/bedrock/_client.py#L61-L62
+"""
+
+import base64
+import io
+import json
+import logging
+import math
+from typing import List, Literal, Tuple, assert_never
+
+from anthropic._types import Base64FileInput
+from anthropic.types.beta import (
+    BetaBashCodeExecutionToolResultBlock as BashCodeExecutionToolResultBlock,
+)
+from anthropic.types.beta import (
+    BetaCodeExecutionToolResultBlock as CodeExecutionToolResultBlock,
+)
+from anthropic.types.beta import (
+    BetaContainerUploadBlock as ContainerUploadBlock,
+)
+from anthropic.types.beta import BetaContentBlockParam as ContentBlockParam
+from anthropic.types.beta import BetaMCPToolResultBlock as MCPToolResultBlock
+from anthropic.types.beta import BetaMCPToolUseBlock as MCPToolUseBlock
+from anthropic.types.beta import BetaMessageParam as ClaudeMessage
+from anthropic.types.beta import (
+    BetaRedactedThinkingBlock as RedactedThinkingBlock,
+)
+from anthropic.types.beta import BetaServerToolUseBlock as ServerToolUseBlock
+from anthropic.types.beta import BetaTextBlock as TextBlock
+from anthropic.types.beta import (
+    BetaTextEditorCodeExecutionToolResultBlock as TextEditorCodeExecutionToolResultBlock,
+)
+from anthropic.types.beta import BetaThinkingBlock as ThinkingBlock
+from anthropic.types.beta import BetaToolParam as ToolParam
+from anthropic.types.beta import (
+    BetaToolResultBlockParam as ToolResultBlockParam,
+)
+from anthropic.types.beta import (
+    BetaToolSearchToolResultBlock as ToolSearchToolResultBlock,
+)
+from anthropic.types.beta import BetaToolUseBlock as ToolUseBlock
+from anthropic.types.beta import (
+    BetaWebFetchToolResultBlock as WebFetchToolResultBlock,
+)
+from anthropic.types.beta import (
+    BetaWebSearchToolResultBlock as WebSearchToolResultBlock,
+)
+from anthropic.types.beta.beta_image_block_param import Source
+from anthropic.types.beta.beta_tool_result_block_param import (
+    Content as ToolResultBlockParamContent,
+)
+from PIL import Image
+
+from aidial_adapter_anthropic.adapter._claude.params import ClaudeParameters
+from aidial_adapter_anthropic.adapter._tokenize import default_tokenize_string
+
+_log = logging.getLogger(__name__)
+
+
+class ApproximateTokenizer:
+    def tokenize_text(self, text: str) -> int:
+        return default_tokenize_string(text)
+
+    def _get_image_size(
+        self, image_data: str | Base64FileInput
+    ) -> Tuple[int, int]:
+        try:
+            if not isinstance(image_data, str):
+                raise ValueError("Images as files aren't yet supported.")
+
+            image_bytes = base64.b64decode(image_data)
+            with Image.open(io.BytesIO(image_bytes)) as img:
+                return img.size
+        except Exception:
+            _log.exception("Cannot compute image size, assuming 1000x1000")
+            return 1000, 1000
+
+    def _tokenize_image(self, source: Source) -> int:
+        match source["type"]:
+            case "url" | "file":
+                return 0
+            case "base64":
+                width, height = self._get_image_size(source["data"])
+                return math.ceil((width * height) / 750.0)
+            case _:
+                assert_never(source)
+
+    def _tokenize_tool_use(self, id: str, input: object, name: str) -> int:
+        return self.tokenize_text(f"{id} {name} {json.dumps(input)}")
+
+    def _tokenize_tool_result(self, message: ToolResultBlockParam) -> int:
+        tokens: int = self.tokenize_text(message["tool_use_id"])
+        if (content := message.get("content")) is not None:
+            if isinstance(content, str):
+                tokens += self.tokenize_text(content)
+            else:
+                for sub_message in content:
+                    tokens += self._tokenize_sub_message(sub_message)
+        return tokens
+
+    def _tokenize_sub_message(
+        self,
+        message: ContentBlockParam | ToolResultBlockParamContent,
+    ) -> int:
+        if isinstance(message, dict):
+            match message["type"]:
+                case "text":
+                    return self.tokenize_text(message["text"])
+                case "image":
+                    return self._tokenize_image(message["source"])
+                case "tool_use":
+                    return self._tokenize_tool_use(
+                        message["id"], message["input"], message["name"]
+                    )
+                case "tool_result":
+                    return self._tokenize_tool_result(message)
+                case "document":
+                    return self.tokenize_text(json.dumps(message))
+                case "thinking":
+                    return self.tokenize_text(message["thinking"])
+                case "redacted_thinking":
+                    return self.tokenize_text(message["data"])
+                case "server_tool_use":
+                    return self.tokenize_text(json.dumps(message["input"]))
+                case "web_search_tool_result":
+                    return self.tokenize_text(json.dumps(message["content"]))
+                case (
+                    "search_result"
+                    | "code_execution_tool_result"
+                    | "mcp_tool_use"
+                    | "mcp_tool_result"
+                    | "container_upload"
+                    | "bash_code_execution_tool_result"
+                    | "text_editor_code_execution_tool_result"
+                    | "web_fetch_tool_result"
+                    | "tool_search_tool_result"
+                    | "tool_reference"
+                ):
+                    return 0
+                case _:
+                    assert_never(message["type"])
+        else:
+            match message:
+                case TextBlock():
+                    return self.tokenize_text(message.text)
+                case ToolUseBlock():
+                    return self._tokenize_tool_use(
+                        message.id, message.input, message.name
+                    )
+                case ThinkingBlock(thinking=thinking):
+                    return self.tokenize_text(thinking)
+                case RedactedThinkingBlock(data=data):
+                    return self.tokenize_text(data)
+                case ServerToolUseBlock(input=input):
+                    return self.tokenize_text(json.dumps(input))
+                case WebSearchToolResultBlock(content=content):
+                    return self.tokenize_text(json.dumps(content))
+                case (
+                    CodeExecutionToolResultBlock()
+                    | MCPToolUseBlock()
+                    | MCPToolResultBlock()
+                    | ContainerUploadBlock()
+                    | BashCodeExecutionToolResultBlock()
+                    | TextEditorCodeExecutionToolResultBlock()
+                    | WebFetchToolResultBlock()
+                    | ToolSearchToolResultBlock()
+                ):
+                    return 0
+                case _:
+                    assert_never(message)
+
+    def _tokenize_message(self, message: ClaudeMessage) -> int:
+        tokens: int = 0
+        content = message["content"]
+
+        match content:
+            case str():
+                tokens += self.tokenize_text(content)
+            case _:
+                for item in content:
+                    tokens += self._tokenize_sub_message(item)
+
+        return tokens
+
+    def _tokenize_messages(self, messages: List[ClaudeMessage]) -> int:
+        # A rough estimation
+        per_message_tokens = 5
+
+        tokens: int = 0
+        for message in messages:
+            tokens += self._tokenize_message(message) + per_message_tokens
+        return tokens
+
+    def _tokenize_tool_param(self, tool: ToolParam) -> int:
+        return self.tokenize_text(json.dumps(tool))
+
+    def tokenize_tool_system_message(
+        self,
+        tool_choice: Literal["none", "auto", "any", "tool"],
+    ) -> int:
+        # Different models has different pricing for the tool use:
+        # https://platform.claude.com/docs/en/agents-and-tools/tool-use/overview#pricing
+        # Here we provide a default for Claude Opus 3
+        return 530 if tool_choice in ("auto", "none") else 281
+
+    async def tokenize(
+        self, params: ClaudeParameters, messages: List[ClaudeMessage]
+    ) -> int:
+        tokens: int = 0
+
+        if system := params["system"]:
+            if isinstance(system, str):
+                tokens += self.tokenize_text(system)
+            else:
+                for item in system:
+                    tokens += self._tokenize_sub_message(item)
+
+        if tools := params["tools"]:
+            if tool_choice := params["tool_choice"]:
+                choice = tool_choice["type"]
+            else:
+                choice = "auto"
+
+            tokens += self.tokenize_tool_system_message(choice)
+
+            for tool in tools:
+                tokens += self._tokenize_tool_param(tool)
+
+        tokens += self._tokenize_messages(messages)
+
+        return tokens

aidial_adapter_anthropic/adapter/_claude/tokenizer/base.py

@@ -0,0 +1,26 @@
+from typing import List, Protocol, Set, Tuple, runtime_checkable
+
+from anthropic.types.beta import BetaMessageParam as ClaudeMessage
+
+from aidial_adapter_anthropic.adapter._claude.params import ClaudeParameters
+from aidial_adapter_anthropic.dial._attachments import WithResources
+
+
+@runtime_checkable
+class ClaudeTokenizer(Protocol):
+    def tokenize_text(self, text: str) -> int: ...
+
+    async def tokenize(
+        self, params: ClaudeParameters, messages: List[ClaudeMessage]
+    ) -> int: ...
+
+
+def create_tokenizer(tokenizer: ClaudeTokenizer, params: ClaudeParameters):
+    async def _tokenize(
+        messages: List[Tuple[WithResources[ClaudeMessage], Set[int]]],
+    ) -> int:
+        return await tokenizer.tokenize(
+            params, [msg[0].payload for msg in messages]
+        )
+
+    return _tokenize

aidial_adapter_anthropic/adapter/_claude/tools.py

@@ -0,0 +1,98 @@
+import json
+import logging
+from typing import assert_never
+
+from aidial_sdk.chat_completion import FunctionCall, ToolCall
+from anthropic.types.beta import BetaToolUseBlock as ToolUseBlock
+
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+from aidial_adapter_anthropic.dial._message import (
+    AIFunctionCallMessage,
+    AIRegularMessage,
+    AIToolCallMessage,
+    BaseMessage,
+    HumanFunctionResultMessage,
+    HumanRegularMessage,
+    HumanToolResultMessage,
+    SystemMessage,
+    ToolMessage,
+)
+from aidial_adapter_anthropic.dial.consumer import Consumer, ToolUseMessage
+from aidial_adapter_anthropic.dial.tools import ToolsMode
+
+_log = logging.getLogger(__name__)
+
+
+def to_dial_function_call(block: ToolUseBlock, streaming: bool) -> FunctionCall:
+    arguments = "" if streaming else json.dumps(block.input)
+    return FunctionCall(name=block.name, arguments=arguments)
+
+
+def to_dial_tool_call(block: ToolUseBlock, streaming: bool) -> ToolCall:
+    return ToolCall(
+        id=block.id,
+        type="function",
+        function=to_dial_function_call(block, streaming),
+    )
+
+
+def process_tools_block(
+    consumer: Consumer,
+    block: ToolUseBlock,
+    tools_mode: ToolsMode | None,
+    *,
+    streaming: bool,
+) -> ToolUseMessage | None:
+    match tools_mode:
+        case ToolsMode.TOOLS:
+            return consumer.create_function_tool_call(
+                to_dial_tool_call(block, streaming)
+            )
+        case ToolsMode.FUNCTIONS:
+            if consumer.has_function_call:
+                _log.warning(
+                    "The model generated more than one tool call. "
+                    "Only the first one will be taken in to account."
+                )
+                return None
+            else:
+                return consumer.create_function_call(
+                    to_dial_function_call(block, streaming)
+                )
+        case None:
+            raise ValidationError(
+                "A model has called a tool, but no tools were given to the model in the first place."
+            )
+        case _:
+            assert_never(tools_mode)
+
+
+def function_to_tool_messages(
+    message: BaseMessage | ToolMessage,
+) -> BaseMessage | HumanToolResultMessage | AIToolCallMessage:
+    match message:
+        case (
+            SystemMessage()
+            | HumanRegularMessage()
+            | AIRegularMessage()
+            | HumanToolResultMessage()
+            | AIToolCallMessage()
+        ):
+            return message
+        case AIFunctionCallMessage():
+            return AIToolCallMessage(
+                content=message.content,
+                calls=[
+                    ToolCall(
+                        id=message.call.name,
+                        type="function",
+                        function=message.call,
+                    )
+                ],
+            )
+        case HumanFunctionResultMessage():
+            return HumanToolResultMessage(
+                id=message.name, content=message.content
+            )
+        case _:
+            assert_never(message)

aidial_adapter_anthropic/adapter/_decorator/base.py

@@ -0,0 +1,53 @@
+from typing import Callable, List
+
+from aidial_sdk.chat_completion import Message
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic.adapter._base import ChatCompletionAdapter
+from aidial_adapter_anthropic.adapter._truncate_prompt import DiscardedMessages
+from aidial_adapter_anthropic.dial.consumer import Consumer
+from aidial_adapter_anthropic.dial.request import ModelParameters
+
+
+class ChatCompletionDecorator(ChatCompletionAdapter):
+    adapter: ChatCompletionAdapter
+
+    async def chat(
+        self,
+        consumer: Consumer,
+        params: ModelParameters,
+        messages: List[Message],
+    ) -> None:
+        await self.adapter.chat(consumer, params, messages)
+
+    async def configuration(self) -> type[BaseModel]:
+        return await self.adapter.configuration()
+
+    async def count_prompt_tokens(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> int:
+        return await self.adapter.count_prompt_tokens(params, messages)
+
+    async def count_completion_tokens(self, string: str) -> int:
+        return await self.adapter.count_completion_tokens(string)
+
+    async def compute_discarded_messages(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> DiscardedMessages | None:
+        return await self.adapter.compute_discarded_messages(params, messages)
+
+
+ChatCompletionTransformer = Callable[
+    [ChatCompletionAdapter], ChatCompletionAdapter
+]
+
+
+def compose_decorators(
+    *decorators: ChatCompletionTransformer,
+) -> ChatCompletionTransformer:
+    def compose(adapter: ChatCompletionAdapter) -> ChatCompletionAdapter:
+        for decorator in reversed(decorators):
+            adapter = decorator(adapter)
+        return adapter
+
+    return compose

aidial_adapter_anthropic/adapter/_decorator/preprocess.py

@@ -0,0 +1,63 @@
+from typing import Callable, List
+
+from aidial_sdk.chat_completion import Message
+
+from aidial_adapter_anthropic._utils.list import ListProjection
+from aidial_adapter_anthropic.adapter._decorator.base import (
+    ChatCompletionDecorator,
+    ChatCompletionTransformer,
+)
+from aidial_adapter_anthropic.adapter._truncate_prompt import DiscardedMessages
+from aidial_adapter_anthropic.dial.consumer import Consumer
+from aidial_adapter_anthropic.dial.request import ModelParameters
+
+
+def preprocess_messages_decorator(
+    on_messages: Callable[[List[Message]], ListProjection[Message]],
+) -> ChatCompletionTransformer:
+    return lambda adapter: PreprocessMessagesDecorator(
+        on_messages=on_messages, adapter=adapter
+    )
+
+
+class PreprocessMessagesDecorator(ChatCompletionDecorator):
+    on_messages: Callable[[List[Message]], ListProjection[Message]]
+
+    async def chat(
+        self,
+        consumer: Consumer,
+        params: ModelParameters,
+        messages: List[Message],
+    ) -> None:
+        new_messages = self.on_messages(messages)
+        await self.adapter.chat(consumer, params, new_messages.raw_list)
+        if (
+            discarded_messages := consumer.get_discarded_messages()
+        ) is not None:
+            discarded_messages = list(
+                new_messages.to_original_indices(discarded_messages)
+            )
+            consumer.set_discarded_messages(discarded_messages)
+
+    async def count_prompt_tokens(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> int:
+        new_messages = self.on_messages(messages)
+        return await self.adapter.count_prompt_tokens(
+            params, new_messages.raw_list
+        )
+
+    async def compute_discarded_messages(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> DiscardedMessages | None:
+        new_messages = self.on_messages(messages)
+        discarded_messages = await self.adapter.compute_discarded_messages(
+            params, new_messages.raw_list
+        )
+
+        if discarded_messages is not None:
+            discarded_messages = list(
+                new_messages.to_original_indices(discarded_messages)
+            )
+
+        return discarded_messages

aidial_adapter_anthropic/adapter/_decorator/replicator.py

@@ -0,0 +1,32 @@
+import asyncio
+from typing import List
+
+from aidial_sdk.chat_completion import Message
+
+from aidial_adapter_anthropic.adapter._decorator.base import (
+    ChatCompletionDecorator,
+    ChatCompletionTransformer,
+)
+from aidial_adapter_anthropic.dial.consumer import Consumer
+from aidial_adapter_anthropic.dial.request import ModelParameters
+
+
+def replicator_decorator() -> ChatCompletionTransformer:
+    return lambda adapter: ReplicatorDecorator(adapter=adapter)
+
+
+class ReplicatorDecorator(ChatCompletionDecorator):
+    async def chat(
+        self,
+        consumer: Consumer,
+        params: ModelParameters,
+        messages: List[Message],
+    ) -> None:
+        params1 = params.copy()
+        params1.n = 1
+
+        async def _chat(root_consumer: Consumer):
+            with root_consumer.fork() as consumer:
+                await self.adapter.chat(consumer, params1, messages)
+
+        await asyncio.gather(*(_chat(consumer) for _ in range(params.n)))

aidial_adapter_anthropic/adapter/_errors.py

@@ -0,0 +1,71 @@
+from typing import Optional
+
+from aidial_sdk.chat_completion import Choice
+from aidial_sdk.exceptions import HTTPException as DialException
+from aidial_sdk.exceptions import RequestValidationError
+
+
+class UserError(Exception):
+    """
+    The user errors are aimed to a DIAL chat user.
+    So whenever an exceptional situation arises that could be handled by a chat user themselves,
+    we should raise a UserError with a `display_message` explaining the error and
+    an optional `usage` message to help the user understand how to use the application correctly:
+
+    * `error_message` is what the chat user will be shown as an error message,
+    * `usage_message` is reported in a `Usage` dialog stage to educate the chat user.
+
+    A typical example of a user error is validation of supported input data attachments.
+    The chat user has full control over the list of attachments, so they can fix the issue themselves.
+    """
+
+    error_message: str
+    usage_message: Optional[str]
+
+    def __init__(self, error_message: str, usage_message: Optional[str] = None):
+        self.error_message = error_message
+        self.usage_message = usage_message
+        super().__init__(self.error_message)
+
+    async def report_usage(self, choice: Choice) -> None:
+        if self.usage_message is not None:
+            with choice.create_stage("Usage") as stage:
+                stage.append_content(self.usage_message)
+
+    def to_dial_exception(self) -> DialException:
+        return RequestValidationError(
+            message=self.error_message,
+            display_message=self.error_message,
+            code="invalid_argument",
+        )
+
+
+class ValidationError(Exception):
+    """
+    The validation errors are aimed to a DIAL API client (e.g. DIAL application developer).
+    They report in which way the request to the application is invalid.
+
+    Typically the validation errors are raised when the request not semantically valid but syntactically well-formed.
+
+    For example, an application doesn't support tools/functions feature, but the request contains it.
+    It's of no use to report such an error to a chat user, because they can't fix it themselves in the chat.
+    But the DIAL application developer who has a finer control over the request can fix the issue by modifying the request.
+    """
+
+    message: str
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
+    def to_dial_exception(self) -> DialException:
+        return RequestValidationError(
+            message=self.message,
+            code="invalid_argument",
+        )
+
+
+# The third category of errors is everything else, including standard Python exceptions, like ValueError or KeyError.
+# These kind of errors are internal to the DIAL application and thus highlight bugs in the application code itself.
+# Neither the chat user nor the DIAL application developer can fix the issue, because there is nothing wrong with the request.
+# Thus, such errors are simply reported as internal server errors with HTTP code 500.

aidial_adapter_anthropic/adapter/_tokenize.py

@@ -0,0 +1,12 @@
+def default_tokenize_string(string: str) -> int:
+    """
+    The number of bytes is a proxy for the number of tokens for
+    models which do not provide any means to count tokens.
+
+    Any token number estimator should satisfy the following requirements:
+    1. Overestimation of number of tokens is allowed.
+    It's ok to truncate the chat history more than necessary.
+    2. Underestimation of number of tokens is prohibited.
+    It's wrong to leave the chat history as is when the truncation was actually required.
+    """
+    return len(string.encode("utf-8"))