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

aidial_adapter_anthropic/adapter/_claude/citations.py

@@ -0,0 +1,63 @@
+from typing import Callable, assert_never
+
+from anthropic.types.beta import (
+    BetaCitationCharLocation as CitationCharLocation,
+)
+from anthropic.types.beta import (
+    BetaCitationContentBlockLocation as CitationContentBlockLocation,
+)
+from anthropic.types.beta import (
+    BetaCitationPageLocation as CitationPageLocation,
+)
+from anthropic.types.beta import (
+    BetaCitationSearchResultLocation as CitationSearchResultLocation,
+)
+from anthropic.types.beta import (
+    BetaCitationsWebSearchResultLocation as CitationsWebSearchResultLocation,
+)
+from anthropic.types.beta import BetaTextCitation as TextCitation
+
+from aidial_adapter_anthropic.dial.consumer import Consumer
+from aidial_adapter_anthropic.dial.resource import DialResource
+
+
+def _add_document_citation(
+    consumer: Consumer,
+    get_document: Callable[[int], DialResource | None],
+    document_index: int,
+):
+    resource = get_document(document_index)
+    document = None if resource is None else resource.to_attachment()
+
+    # NOTE: multiple citations to the same document are merged into one citation
+    # until we find a better API to handle citations embedded in text.
+    display_index = consumer.add_citation_attachment(
+        document_id=document_index, document=document
+    )
+
+    # NOTE: avoid adding citation URLs into the generated content,
+    # since such references aren't easily portable (e.g. when a conversion is duplicated).
+    consumer.append_content(f"[{display_index}]")
+
+
+def create_citations(
+    consumer: Consumer,
+    get_document: Callable[[int], DialResource | None],
+    citation: TextCitation,
+):
+    match citation:
+        case CitationCharLocation(
+            document_index=document_index
+        ) | CitationPageLocation(document_index=document_index):
+            _add_document_citation(consumer, get_document, document_index)
+
+        # Custom document aren't supported yet
+        case CitationContentBlockLocation():
+            pass
+        # web search isn't supported yet
+        case CitationsWebSearchResultLocation():
+            pass
+        case CitationSearchResultLocation():
+            pass
+        case _:
+            assert_never(citation)

aidial_adapter_anthropic/adapter/_claude/config.py

@@ -0,0 +1,39 @@
+from typing import List, Literal
+
+from anthropic.types.anthropic_beta_param import AnthropicBetaParam
+from anthropic.types.beta import BetaThinkingConfigParam as ThinkingConfigParam
+from pydantic import Field
+
+from aidial_adapter_anthropic._utils.pydantic import ExtraForbidModel
+
+
+class ThinkingConfigEnabled(ExtraForbidModel):
+    type: Literal["enabled"]
+    budget_tokens: int
+
+    def to_claude(self) -> ThinkingConfigParam:
+        return {"type": "enabled", "budget_tokens": self.budget_tokens}
+
+
+class ThinkingConfigDisabled(ExtraForbidModel):
+    type: Literal["disabled"]
+
+    def to_claude(self) -> ThinkingConfigParam:
+        return {"type": "disabled"}
+
+
+class ClaudeConfiguration(ExtraForbidModel):
+    betas: List[AnthropicBetaParam] | None = Field(
+        default=None,
+        description="List of beta features to enable. Make sure to check if the given feature is supported by the Claude deployment you are using.",
+    )
+    enable_citations: bool = False
+
+
+class ClaudeConfigurationWithThinking(ClaudeConfiguration):
+    # NOTE: once migrated to Pydantic v2 we can use TypeAdapter over
+    # the anthropic's ThinkingConfigParam class directly.
+    thinking: ThinkingConfigEnabled | ThinkingConfigDisabled | None = None
+
+
+Configuration = ClaudeConfiguration | ClaudeConfigurationWithThinking

aidial_adapter_anthropic/adapter/_claude/converters.py

@@ -0,0 +1,303 @@
+from typing import List, Literal, Optional, Sequence, Set, Tuple, assert_never
+
+from aidial_sdk.chat_completion import FinishReason, Tool
+from aidial_sdk.chat_completion import ToolChoice as DialToolChoice
+from anthropic.types.beta import (
+    BetaCacheControlEphemeralParam as CacheControlEphemeralParam,
+)
+from anthropic.types.beta import BetaContentBlockParam as ContentBlockParam
+from anthropic.types.beta import BetaMessageParam as MessageParam
+from anthropic.types.beta import BetaStopReason as ClaudeStopReason
+from anthropic.types.beta import BetaTextBlockParam as TextBlockParam
+from anthropic.types.beta import BetaToolChoiceAnyParam as ToolChoiceAnyParam
+from anthropic.types.beta import BetaToolChoiceAutoParam as ToolChoiceAutoParam
+from anthropic.types.beta import BetaToolChoiceNoneParam as ToolChoiceNoneParam
+from anthropic.types.beta import BetaToolChoiceParam as ToolChoice
+from anthropic.types.beta import BetaToolChoiceToolParam as ToolChoiceToolParam
+from anthropic.types.beta import BetaToolParam as ToolParam
+from anthropic.types.beta import BetaUsage as Usage
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic._utils.list import ListProjection, group_by
+from aidial_adapter_anthropic.adapter._claude.blocks import (
+    create_text_block,
+    create_tool_result_block,
+    create_tool_use_block,
+)
+from aidial_adapter_anthropic.adapter._claude.config import Configuration
+from aidial_adapter_anthropic.adapter._claude.state import (
+    get_message_content_from_state,
+)
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+from aidial_adapter_anthropic.dial._attachments import (
+    AttachmentProcessors,
+    WithResources,
+)
+from aidial_adapter_anthropic.dial._message import (
+    AIRegularMessage,
+    AIToolCallMessage,
+    BaseMessage,
+    HumanRegularMessage,
+    HumanToolResultMessage,
+    SystemMessage,
+)
+from aidial_adapter_anthropic.dial.token_usage import TokenUsage
+from aidial_adapter_anthropic.dial.tools import ToolsConfig, ToolsMode
+
+DialMessage = BaseMessage | HumanToolResultMessage | AIToolCallMessage
+
+ClaudeMessage = WithResources[ContentBlockParam]
+
+_claude_cache_breakpoint = CacheControlEphemeralParam(type="ephemeral")
+
+
+def _add_cache_control(
+    message: DialMessage, claude_messages: Sequence[ContentBlockParam]
+) -> None:
+    if message.cache_breakpoint is not None:
+        for block in reversed(claude_messages):
+            if (
+                isinstance(block, dict)
+                and block["type"] != "thinking"
+                and block["type"] != "redacted_thinking"
+            ):
+                block["cache_control"] = _claude_cache_breakpoint
+                break
+
+
+def _get_claude_message_role(
+    dial_message: (
+        AIRegularMessage
+        | AIToolCallMessage
+        | HumanRegularMessage
+        | HumanToolResultMessage
+    ),
+) -> Literal["assistant", "user"]:
+    match dial_message:
+        case AIRegularMessage() | AIToolCallMessage():
+            return "assistant"
+        case HumanRegularMessage() | HumanToolResultMessage():
+            return "user"
+        case _:
+            assert_never(dial_message)
+
+
+_Elem = Tuple[WithResources[MessageParam], Set[int]]
+
+
+def _merge_messages_with_same_role(
+    messages: ListProjection[WithResources[MessageParam]],
+) -> ListProjection[WithResources[MessageParam]]:
+
+    def _key(message: _Elem) -> str:
+        return message[0].payload["role"]
+
+    def _merge_message_param(
+        msg1: MessageParam, msg2: MessageParam
+    ) -> MessageParam:
+        content1 = msg1["content"]
+        content2 = msg2["content"]
+
+        if isinstance(content1, str):
+            content1 = [TextBlockParam(type="text", text=content1)]
+
+        if isinstance(content2, str):
+            content2 = [TextBlockParam(type="text", text=content2)]
+
+        return MessageParam(
+            role=msg1["role"],
+            content=list(content1) + list(content2),
+        )
+
+    def _merge(a: _Elem, b: _Elem) -> _Elem:
+        (msg1, set1), (msg2, set2) = a, b
+        payload = _merge_message_param(msg1.payload, msg2.payload)
+        resources = msg1.resources + msg2.resources
+        return (WithResources(payload, resources), set1 | set2)
+
+    return ListProjection(group_by(messages.list, _key, lambda x: x, _merge))
+
+
+async def _get_claude_blocks(
+    handlers: AttachmentProcessors[
+        TextBlockParam, ContentBlockParam, Configuration
+    ],
+    message: (
+        HumanRegularMessage
+        | AIRegularMessage
+        | AIToolCallMessage
+        | HumanToolResultMessage
+    ),
+    message_idx: int,
+) -> WithResources[Sequence[ContentBlockParam]]:
+
+    match message:
+        case HumanRegularMessage():
+            return await handlers.process_attachments(message)
+
+        case HumanToolResultMessage():
+            blocks = [create_tool_result_block(message)]
+            return WithResources(payload=blocks)
+
+        case AIRegularMessage():
+            content = await handlers.process_attachments(message)
+
+            # Take the message content from the state if possible,
+            # since it may include certain content blocks that
+            # are missing from the DIAL message itself,
+            # such as thinking signatures and redacted thinking blocks.
+            if state := get_message_content_from_state(message_idx, message):
+                content.payload = state
+
+            return content
+
+        case AIToolCallMessage():
+            blocks = [create_tool_use_block(call) for call in message.calls]
+            if text_content := message.content:
+                blocks.insert(0, create_text_block(text_content))
+
+            content = WithResources(payload=blocks)
+            if state := get_message_content_from_state(message_idx, message):
+                content.payload = state
+
+            return content
+
+        case _:
+            assert_never(message)
+
+
+async def to_claude_messages(
+    handlers: AttachmentProcessors[
+        TextBlockParam, ContentBlockParam, Configuration
+    ],
+    messages: List[DialMessage],
+) -> Tuple[List[TextBlockParam], ListProjection[WithResources[MessageParam]]]:
+
+    idx_offset: int = 0
+    system_messages: List[TextBlockParam] = []
+
+    for message in messages:
+        if not isinstance(message, SystemMessage):
+            break
+
+        idx_offset += 1
+        sys_content = await handlers.process_system_message(message)
+        _add_cache_control(message, sys_content)
+
+        system_messages.extend(sys_content)
+
+    claude_messages: ListProjection[WithResources[MessageParam]] = (
+        ListProjection()
+    )
+
+    for idx, message in enumerate(messages[idx_offset:], start=idx_offset):
+        if isinstance(message, SystemMessage):
+            raise ValidationError(
+                "System and developer messages are only allowed in the beginning of the conversation."
+            )
+
+        blocks = await _get_claude_blocks(handlers, message, idx)
+        _add_cache_control(message, blocks.payload)
+
+        role = _get_claude_message_role(message)
+        claude_message = WithResources(
+            payload=MessageParam(role=role, content=blocks.payload),
+            resources=blocks.resources,
+        )
+
+        claude_messages.append(claude_message, idx)
+
+    return system_messages, _merge_messages_with_same_role(claude_messages)
+
+
+def to_dial_finish_reason(
+    finish_reason: Optional[ClaudeStopReason],
+    tools_mode: ToolsMode | None,
+) -> FinishReason:
+    if finish_reason is None:
+        return FinishReason.STOP
+
+    match finish_reason:
+        case "end_turn":
+            return FinishReason.STOP
+        case "max_tokens" | "model_context_window_exceeded":
+            return FinishReason.LENGTH
+        case "stop_sequence" | "pause_turn" | "refusal":
+            return FinishReason.STOP
+        case "tool_use":
+            match tools_mode:
+                case ToolsMode.TOOLS:
+                    return FinishReason.TOOL_CALLS
+                case ToolsMode.FUNCTIONS:
+                    return FinishReason.FUNCTION_CALL
+                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)
+
+        case _:
+            assert_never(finish_reason)
+
+
+def to_dial_usage(usage: Usage) -> TokenUsage:
+    read = usage.cache_creation_input_tokens or 0
+    write = usage.cache_read_input_tokens or 0
+    return TokenUsage(
+        completion_tokens=usage.output_tokens,
+        prompt_tokens=usage.input_tokens + read + write,
+        cache_write_input_tokens=read,
+        cache_read_input_tokens=write,
+    )
+
+
+def _to_claude_tool(tool: Tool) -> ToolParam:
+    function = tool.function
+    tool_param = ToolParam(
+        input_schema=function.parameters
+        or {"type": "object", "properties": {}},
+        name=function.name,
+        description=function.description or "",
+    )
+
+    if tool.custom_fields and tool.custom_fields.cache_breakpoint:
+        tool_param["cache_control"] = _claude_cache_breakpoint
+
+    return tool_param
+
+
+def _to_claude_tool_choice(
+    tool_choice: Literal["auto", "none", "required"] | DialToolChoice,
+) -> ToolChoice:
+    # NOTE tool_choice.disable_parallel_tool_use=True option isn't supported
+    # by older Claude3 versions, so we limit the number of generated function calls
+    # to one in the adapter itself for the functions mode.
+
+    match tool_choice:
+        case DialToolChoice(function=function):
+            return ToolChoiceToolParam(type="tool", name=function.name)
+        case "required":
+            return ToolChoiceAnyParam(type="any")
+        case "auto":
+            return ToolChoiceAutoParam(type="auto")
+        case "none":
+            return ToolChoiceNoneParam(type="none")
+        case _:
+            assert_never(tool_choice)
+
+
+class ClaudeToolsConfig(BaseModel):
+    tools: List[ToolParam]
+    tool_choice: ToolChoice
+
+
+def to_claude_tool_config(
+    tools_config: ToolsConfig | None,
+) -> ClaudeToolsConfig | None:
+    if tools_config is None or not tools_config.tools:
+        return None
+
+    tools = [_to_claude_tool(tool) for tool in tools_config.tools]
+    tool_choice = _to_claude_tool_choice(tools_config.tool_choice)
+    return ClaudeToolsConfig(tools=tools, tool_choice=tool_choice)

aidial_adapter_anthropic/adapter/_claude/params.py

@@ -0,0 +1,25 @@
+from typing import List, TypedDict
+
+from anthropic import Omit
+from anthropic.types.anthropic_beta_param import AnthropicBetaParam
+from anthropic.types.beta import BetaTextBlockParam as TextBlockParam
+from anthropic.types.beta import BetaThinkingConfigParam as ThinkingConfigParam
+from anthropic.types.beta import BetaToolChoiceParam as ToolChoice
+from anthropic.types.beta import BetaToolParam as ToolParam
+
+
+class ClaudeParameters(TypedDict):
+    """
+    Subset of parameters to Anthropic Messages API request:
+    https://github.com/anthropics/anthropic-sdk-python/blob/ff83982c44db0920f435916aadb37c3523083079/src/anthropic/resources/messages.py#L1827-L1847
+    """
+
+    max_tokens: int
+    stop_sequences: List[str] | Omit
+    system: str | List[TextBlockParam] | Omit
+    temperature: float | Omit
+    top_p: float | Omit
+    tools: List[ToolParam] | Omit
+    tool_choice: ToolChoice | Omit
+    thinking: ThinkingConfigParam | Omit
+    betas: List[AnthropicBetaParam] | Omit

aidial_adapter_anthropic/adapter/_claude/state.py

@@ -0,0 +1,45 @@
+import logging
+from typing import List
+
+import pydantic
+from anthropic.types.beta import BetaContentBlock as ContentBlock
+from anthropic.types.beta import BetaContentBlockParam as ContentBlockParam
+from anthropic.types.beta.parsed_beta_message import (
+    ParsedBetaContentBlock as ParsedContentBlock,
+)
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic.dial._message import (
+    AIRegularMessage,
+    AIToolCallMessage,
+)
+
+_log = logging.getLogger(__name__)
+
+
+class MessageState(BaseModel):
+    claude_message_content: List[ParsedContentBlock] | List[ContentBlock]
+
+    def to_dict(self) -> dict:
+        return self.dict(
+            # FIXME: a hack to exclude the private __json_buf field
+            exclude={"claude_message_content": {"__all__": {"__json_buf"}}},
+            # Excluding `citations: null`, since they could not be even parsed
+            # currently by the Bedrock.
+            exclude_none=True,
+        )
+
+
+def get_message_content_from_state(
+    idx: int, message: AIRegularMessage | AIToolCallMessage
+) -> List[ContentBlockParam] | None:
+    if (cc := message.custom_content) and (state_dict := cc.state):
+        try:
+            state = MessageState.parse_obj(state_dict)
+            return [block.to_dict() for block in state.claude_message_content]  # type: ignore
+        except pydantic.ValidationError as e:
+            _log.error(
+                f"Invalid state at the path 'messages[{idx}].custom_content.state': {e}"
+            )
+
+    return None

aidial_adapter_anthropic/adapter/_claude/tokenizer/__init__.py

@@ -0,0 +1,10 @@
+from .anthropic import AnthropicTokenizer
+from .approximate import ApproximateTokenizer
+from .base import ClaudeTokenizer, create_tokenizer
+
+__all__ = [
+    "ClaudeTokenizer",
+    "create_tokenizer",
+    "AnthropicTokenizer",
+    "ApproximateTokenizer",
+]

aidial_adapter_anthropic/adapter/_claude/tokenizer/anthropic.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+from typing import List
+
+from aidial_sdk.exceptions import InternalServerError
+from anthropic import (
+    AsyncAnthropic,
+    AsyncAnthropicBedrock,
+    AsyncAnthropicFoundry,
+    AsyncAnthropicVertex,
+)
+from anthropic._resource import AsyncAPIResource
+from anthropic.resources.beta import AsyncMessages as FirstPartyAsyncMessagesAPI
+from anthropic.types.beta import BetaMessageParam as ClaudeMessageParam
+
+from aidial_adapter_anthropic.adapter._claude.params import ClaudeParameters
+
+AnthropicClient = (
+    AsyncAnthropic
+    | AsyncAnthropicBedrock
+    | AsyncAnthropicVertex
+    | AsyncAnthropicFoundry
+)
+
+
+# Beta AsyncMessages doesn't provide the 'count_tokens' method,
+# so we enabled it via the adapter.
+class _AsyncMessagesAdapter(AsyncAPIResource):
+    count_tokens = FirstPartyAsyncMessagesAPI.count_tokens
+
+    def __init__(self, resource: AsyncAPIResource):
+        super().__init__(resource._client)
+
+
+@dataclass
+class AnthropicTokenizer:
+    deployment: str
+    client: AnthropicClient
+
+    def tokenize_text(self, text: str) -> int:
+        raise InternalServerError(
+            "Tokenization of strings is not supported by Anthropic API"
+        )
+
+    async def tokenize(
+        self, params: ClaudeParameters, messages: List[ClaudeMessageParam]
+    ) -> int:
+        return (
+            await _AsyncMessagesAdapter(self.client.beta.messages).count_tokens(
+                model=self.deployment,
+                messages=messages,
+                system=params["system"],
+                thinking=params["thinking"],
+                tools=params["tools"],
+                tool_choice=params["tool_choice"],
+                betas=params["betas"],
+            )
+        ).input_tokens