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

aidial_adapter_anthropic/dial/request.py

@@ -0,0 +1,170 @@
+from typing import (
+    List,
+    Literal,
+    Optional,
+    Type,
+    TypeGuard,
+    TypeVar,
+    assert_never,
+)
+
+from aidial_sdk.chat_completion import (
+    MessageContentImagePart,
+    MessageContentPart,
+    MessageContentTextPart,
+    Role,
+)
+from aidial_sdk.chat_completion.request import (
+    ChatCompletionRequest,
+    MessageContentRefusalPart,
+)
+from aidial_sdk.exceptions import RequestValidationError
+from pydantic import BaseModel
+from pydantic.v1 import ValidationError as PydanticValidationError
+
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+from aidial_adapter_anthropic.dial.tools import (
+    ToolsConfig,
+    ToolsMode,
+    validate_messages,
+)
+
+MessageContent = str | List[MessageContentPart] | None
+MessageContentSpecialized = (
+    MessageContent
+    | List[MessageContentTextPart]
+    | List[MessageContentImagePart]
+)
+
+_Model = TypeVar("_Model", bound=BaseModel)
+
+
+class ModelParameters(BaseModel):
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    n: int = 1
+    stop: List[str] = []
+    seed: Optional[int] = None
+    max_tokens: Optional[int] = None
+    max_prompt_tokens: Optional[int] = None
+    stream: bool = False
+    tool_config: Optional[ToolsConfig] = None
+    configuration: Optional[dict] = None
+
+    @classmethod
+    def create(cls, request: ChatCompletionRequest) -> "ModelParameters":
+        stop: List[str] = []
+        if request.stop is not None:
+            stop = (
+                [request.stop]
+                if isinstance(request.stop, str)
+                else request.stop
+            )
+
+        validate_messages(request)
+
+        configuration = (
+            cf.configuration
+            if (cf := request.custom_fields) is not None
+            else None
+        )
+
+        return cls(
+            temperature=request.temperature,
+            top_p=request.top_p,
+            n=request.n or 1,
+            stop=stop,
+            seed=request.seed,
+            max_tokens=request.max_tokens,
+            max_prompt_tokens=request.max_prompt_tokens,
+            stream=request.stream,
+            tool_config=ToolsConfig.from_request(request),
+            configuration=configuration,
+        )
+
+    def add_stop_sequences(self, stop: List[str]) -> "ModelParameters":
+        return self.copy(update={"stop": [*self.stop, *stop]})
+
+    @property
+    def tools_mode(self) -> ToolsMode | None:
+        if self.tool_config is not None:
+            return self.tool_config.tools_mode
+        return None
+
+    def parse_configuration(self, cls: Type[_Model]) -> _Model:
+        try:
+            return cls.parse_obj(self.configuration or {})
+        except PydanticValidationError as e:
+            if self.configuration is None:
+                msg = "The configuration at path 'custom_fields.configuration' is missing."
+            else:
+                error = e.errors()[0]
+                path = ".".join(map(str, error["loc"]))
+                msg = f"Invalid request. Path: 'custom_fields.configuration.{path}', error: {error['msg']}"
+
+            raise RequestValidationError(msg)
+
+
+def collect_text_content(
+    content: MessageContentSpecialized, delimiter: str = "\n\n"
+) -> str:
+    match content:
+        case None:
+            return ""
+        case str():
+            return content
+        case list():
+            texts: List[str] = []
+            for part in content:
+                match part:
+                    case MessageContentTextPart(text=text):
+                        texts.append(text)
+                    case MessageContentImagePart():
+                        raise ValidationError(
+                            "Can't extract text from an image content part"
+                        )
+                    case MessageContentRefusalPart():
+                        raise ValidationError(
+                            "Can't extract text from a refusal content part"
+                        )
+                    case _:
+                        assert_never(part)
+            return delimiter.join(texts)
+        case _:
+            assert_never(content)
+
+
+def to_message_content(content: MessageContentSpecialized) -> MessageContent:
+    match content:
+        case None | str():
+            return content
+        case list():
+            return [*content]
+        case _:
+            assert_never(content)
+
+
+def is_text_content(
+    content: MessageContent,
+) -> TypeGuard[str | List[MessageContentTextPart]]:
+    match content:
+        case None:
+            return False
+        case str():
+            return True
+        case list():
+            return all(
+                isinstance(part, MessageContentTextPart) for part in content
+            )
+        case _:
+            assert_never(content)
+
+
+def is_plain_text_content(content: MessageContent) -> TypeGuard[str | None]:
+    return content is None or isinstance(content, str)
+
+
+def is_system_role(
+    role: Role,
+) -> TypeGuard[Literal[Role.SYSTEM, Role.DEVELOPER]]:
+    return role in [Role.SYSTEM, Role.DEVELOPER]

aidial_adapter_anthropic/dial/resource.py

@@ -0,0 +1,189 @@
+import base64
+import mimetypes
+from abc import ABC, abstractmethod
+from typing import List
+
+from aidial_sdk.chat_completion import Attachment
+from pydantic import BaseModel, Field, root_validator, validator
+
+from aidial_adapter_anthropic._utils.resource import Resource
+from aidial_adapter_anthropic._utils.text import truncate_string
+from aidial_adapter_anthropic.dial.storage import FileStorage, download_file
+
+
+class ValidationError(Exception):
+    message: str
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class MissingContentType(ValidationError):
+    pass
+
+
+class UnsupportedContentType(ValidationError):
+    type: str
+    supported_types: List[str]
+
+    def __init__(self, *, message: str, type: str, supported_types: List[str]):
+        self.type = type
+        self.supported_types = supported_types
+        super().__init__(message)
+
+
+class DialResource(ABC, BaseModel):
+    entity_name: str = Field(default=None)
+    supported_types: List[str] | None = Field(default=None)
+
+    @abstractmethod
+    def to_attachment(self) -> Attachment: ...
+
+    @abstractmethod
+    async def download(self, storage: FileStorage | None) -> Resource: ...
+
+    @abstractmethod
+    async def guess_content_type(self) -> str | None: ...
+
+    @abstractmethod
+    async def get_resource_name(self, storage: FileStorage | None) -> str: ...
+
+    async def get_content_type(self) -> str:
+        type = await self.guess_content_type()
+
+        if not type:
+            raise MissingContentType(
+                f"Can't derive content type of the {self.entity_name}"
+            )
+
+        if (
+            self.supported_types is not None
+            and type not in self.supported_types
+        ):
+            raise UnsupportedContentType(
+                message=f"The {self.entity_name} is not one of the supported types",
+                type=type,
+                supported_types=self.supported_types,
+            )
+
+        return type
+
+
+class URLResource(DialResource):
+    url: str
+    content_type: str | None = None
+
+    def to_attachment(self) -> Attachment:
+        return Attachment(type=self.content_type, url=self.url)
+
+    @root_validator
+    def validator(cls, values):
+        values["entity_name"] = values.get("entity_name") or "URL"
+        return values
+
+    async def download(self, storage: FileStorage | None) -> Resource:
+        type = await self.get_content_type()
+        data = await _download_url(storage, self.url)
+        return Resource(type=type, data=data)
+
+    async def guess_content_type(self) -> str | None:
+        return (
+            self.content_type
+            or Resource.parse_data_url_content_type(self.url)
+            or mimetypes.guess_type(self.url)[0]
+        )
+
+    def is_data_url(self) -> bool:
+        return Resource.parse_data_url_content_type(self.url) is not None
+
+    async def get_resource_name(self, storage: FileStorage | None) -> str:
+        if self.is_data_url():
+            return f"data URL ({await self.guess_content_type()})"
+
+        name = self.url
+        if storage is not None:
+            name = await storage.get_human_readable_name(self.url)
+
+        return truncate_string(name, n=50)
+
+
+class AttachmentResource(DialResource):
+    attachment: Attachment
+
+    def to_attachment(self) -> Attachment:
+        return self.attachment
+
+    @validator("attachment", pre=True)
+    def parse_attachment(cls, value):
+        if isinstance(value, dict):
+            attachment = Attachment.parse_obj(value)
+            # Working around the issue of defaulting missing type to a markdown:
+            # https://github.com/epam/ai-dial-sdk/blob/2835107e950c89645a2b619fecba2518fa2d7bb1/aidial_sdk/chat_completion/request.py#L22
+            if "type" not in value:
+                attachment.type = None
+            return attachment
+        return value
+
+    @root_validator(pre=True)
+    def validator(cls, values):
+        values["entity_name"] = values.get("entity_name") or "attachment"
+        return values
+
+    async def download(self, storage: FileStorage | None) -> Resource:
+        type = await self.get_content_type()
+
+        if self.attachment.data:
+            data = base64.b64decode(self.attachment.data)
+        elif self.attachment.url:
+            data = await _download_url(storage, self.attachment.url)
+        else:
+            raise ValidationError(f"Invalid {self.entity_name}")
+
+        return Resource(type=type, data=data)
+
+    def create_url_resource(self, url: str) -> URLResource:
+        return URLResource(
+            url=url,
+            content_type=self.informative_content_type,
+            entity_name=self.entity_name,
+        )
+
+    @property
+    def informative_content_type(self) -> str | None:
+        if (
+            self.attachment.type is None
+            or "octet-stream" in self.attachment.type
+        ):
+            return None
+        return self.attachment.type
+
+    async def guess_content_type(self) -> str | None:
+        if url := self.attachment.url:
+            type = await self.create_url_resource(url).guess_content_type()
+            if type:
+                return type
+
+        return self.attachment.type
+
+    async def get_resource_name(self, storage: FileStorage | None) -> str:
+        if title := self.attachment.title:
+            return title
+
+        if self.attachment.data:
+            return f"data {self.entity_name}"
+        elif url := self.attachment.url:
+            return await self.create_url_resource(url).get_resource_name(
+                storage
+            )
+        else:
+            raise ValidationError(f"Invalid {self.entity_name}")
+
+
+async def _download_url(file_storage: FileStorage | None, url: str) -> bytes:
+    if (resource := Resource.from_data_url(url)) is not None:
+        return resource.data
+
+    if file_storage:
+        return await file_storage.download_file(url)
+    return await download_file(url)

aidial_adapter_anthropic/dial/storage.py

@@ -0,0 +1,138 @@
+import base64
+import hashlib
+import io
+import logging
+import mimetypes
+from typing import Mapping, Optional, TypedDict
+from urllib.parse import unquote, urljoin
+
+import aiohttp
+from pydantic import BaseModel
+
+_log = logging.getLogger(__name__)
+
+
+class FileMetadata(TypedDict):
+    name: str
+    parentPath: str
+    bucket: str
+    url: str
+
+
+class Bucket(TypedDict):
+    bucket: str
+    appdata: str
+
+
+class FileStorage(BaseModel):
+    dial_url: str
+    api_key: str
+    bucket: Optional[Bucket] = None
+
+    @property
+    def auth_headers(self) -> Mapping[str, str]:
+        return {"api-key": self.api_key}
+
+    async def _get_bucket(self, session: aiohttp.ClientSession) -> Bucket:
+        if self.bucket is None:
+            async with session.get(
+                f"{self.dial_url}/v1/bucket",
+                headers=self.auth_headers,
+            ) as response:
+                response.raise_for_status()
+                self.bucket = bucket = await response.json()
+                _log.debug(f"bucket: {self.bucket}")
+                return bucket
+
+        return self.bucket
+
+    async def _get_user_bucket(self, session: aiohttp.ClientSession) -> str:
+        bucket = await self._get_bucket(session)
+        appdata = bucket.get("appdata")
+        if appdata is None:
+            raise ValueError(
+                "Can't retrieve user bucket because appdata isn't available"
+            )
+        return appdata.split("/", 1)[0]
+
+    @staticmethod
+    def _to_form_data(
+        filename: str, content_type: str, content: bytes
+    ) -> aiohttp.FormData:
+        data = aiohttp.FormData()
+        data.add_field(
+            "file",
+            io.BytesIO(content),
+            filename=filename,
+            content_type=content_type,
+        )
+        return data
+
+    async def upload(
+        self, filename: str, content_type: str, content: bytes
+    ) -> FileMetadata:
+        async with aiohttp.ClientSession() as session:
+            bucket = await self._get_bucket(session)
+
+            appdata = bucket["appdata"]
+            ext = mimetypes.guess_extension(content_type) or ""
+            url = f"{self.dial_url}/v1/files/{appdata}/{filename}{ext}"
+
+            data = FileStorage._to_form_data(filename, content_type, content)
+
+            async with session.put(
+                url=url,
+                data=data,
+                headers=self.auth_headers,
+            ) as response:
+                response.raise_for_status()
+                meta = await response.json()
+                _log.debug(f"Uploaded file: url={url}, metadata={meta}")
+                return meta
+
+    async def upload_file_as_base64(
+        self, upload_dir: str, data: str, content_type: str
+    ) -> FileMetadata:
+        filename = f"{upload_dir}/{compute_hash_digest(data)}"
+        content: bytes = base64.b64decode(data)
+        return await self.upload(filename, content_type, content)
+
+    def attachment_link_to_url(self, link: str) -> str:
+        return urljoin(f"{self.dial_url}/v1/", link)
+
+    def _url_to_attachment_link(self, url: str) -> str:
+        return url.removeprefix(f"{self.dial_url}/v1/")
+
+    async def download_file(self, link: str) -> bytes:
+        url = self.attachment_link_to_url(link)
+        headers: Mapping[str, str] = {}
+        if url.lower().startswith(self.dial_url.lower()):
+            headers = self.auth_headers
+        return await download_file(url, headers)
+
+    async def get_human_readable_name(self, link: str) -> str:
+        url = self.attachment_link_to_url(link)
+        link = self._url_to_attachment_link(url)
+
+        link = link.removeprefix("files/")
+
+        if link.startswith("public/"):
+            bucket = "public"
+        else:
+            async with aiohttp.ClientSession() as session:
+                bucket = await self._get_user_bucket(session)
+
+        link = link.removeprefix(f"{bucket}/")
+        decoded_link = unquote(link)
+        return link if link == decoded_link else repr(decoded_link)
+
+
+async def download_file(url: str, headers: Mapping[str, str] = {}) -> bytes:
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url, headers=headers) as response:
+            response.raise_for_status()
+            return await response.read()
+
+
+def compute_hash_digest(file_content: str) -> str:
+    return hashlib.sha256(file_content.encode()).hexdigest()

aidial_adapter_anthropic/dial/token_usage.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+
+class TokenUsage(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    cache_read_input_tokens: int = 0
+    cache_write_input_tokens: int = 0
+
+    @property
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+    def accumulate(self, other: "TokenUsage") -> "TokenUsage":
+        self.prompt_tokens += other.prompt_tokens
+        self.completion_tokens += other.completion_tokens
+        self.cache_read_input_tokens += other.cache_read_input_tokens
+        self.cache_write_input_tokens += other.cache_write_input_tokens
+        return self

aidial_adapter_anthropic/dial/tools.py

@@ -0,0 +1,180 @@
+import logging
+from enum import Enum
+from typing import Dict, List, Literal, Self
+
+from aidial_sdk.chat_completion import (
+    Function,
+    FunctionChoice,
+    Message,
+    Role,
+    Tool,
+    ToolChoice,
+)
+from aidial_sdk.chat_completion.request import (
+    AzureChatCompletionRequest,
+    StaticTool,
+)
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+
+_log = logging.getLogger(__name__)
+
+
+class ToolsMode(Enum):
+    TOOLS = "TOOLS"
+    FUNCTIONS = "FUNCTIONS"
+    """
+    Functions are deprecated instrument that came before tools
+    """
+
+
+class ToolsConfig(BaseModel):
+    tools: List[Tool]
+    """
+    List of functions/tools.
+    """
+
+    tools_mode: ToolsMode
+
+    tool_choice: Literal["auto", "none", "required"] | ToolChoice
+
+    tool_ids: Dict[str, str]
+    """
+    Mapping from tool call IDs to corresponding tool names.
+    Empty when there are no tool calls in the messages.
+    """
+
+    def not_supported(self) -> None:
+        if not self.tools:
+            return
+        if self.tools_mode == ToolsMode.TOOLS:
+            raise ValidationError("The tools aren't supported")
+        raise ValidationError("The functions aren't supported")
+
+    def create_fresh_tool_call_id(self, tool_name: str) -> str:
+        idx = 1
+        while True:
+            tool_id = f"{tool_name}_{idx}"
+            if tool_id not in self.tool_ids:
+                self.tool_ids[tool_id] = tool_name
+                return tool_id
+            idx += 1
+
+    def get_tool_name(self, tool_call_id: str) -> str:
+        tool_name = self.tool_ids.get(tool_call_id)
+        if tool_name is None:
+            raise ValidationError(f"Tool call ID not found: {self.tool_ids}")
+        return tool_name
+
+    @staticmethod
+    def _function_call_to_tool_choice(
+        function_call: Literal["auto", "none"] | FunctionChoice | None,
+    ) -> Literal["auto", "none", "required"] | ToolChoice | None:
+        match function_call:
+            case FunctionChoice():
+                return ToolChoice(type="function", function=function_call)
+            case _:
+                return function_call
+
+    @staticmethod
+    def _get_tool_from_function(tool: Function | Tool | StaticTool) -> Tool:
+        if isinstance(tool, StaticTool):
+            raise ValidationError("Static tools aren't supported")
+        if isinstance(tool, Function):
+            return Tool(type="function", function=tool)
+        else:
+            return tool
+
+    @staticmethod
+    def _get_tools_from_functions(
+        tools: List[Function] | List[Tool | StaticTool],
+    ) -> List[Tool]:
+        return [ToolsConfig._get_tool_from_function(tool) for tool in tools]
+
+    @classmethod
+    def from_request(cls, request: AzureChatCompletionRequest) -> Self | None:
+        validate_messages(request)
+
+        tool_ids = _collect_tool_ids(request.messages)
+
+        if request.functions is not None:
+            tools_mode = ToolsMode.FUNCTIONS
+            tools = cls._get_tools_from_functions(request.functions)
+            tool_choice = cls._function_call_to_tool_choice(
+                request.function_call
+            )
+        elif request.tools is not None:
+            tools_mode = ToolsMode.TOOLS
+            tools = cls._get_tools_from_functions(request.tools)
+            tool_choice = request.tool_choice
+        elif tool_ids:
+            tools_mode = ToolsMode.TOOLS
+            tools = []
+            tool_choice = None
+        else:
+            return None
+
+        return cls(
+            tools=tools,
+            tools_mode=tools_mode,
+            tool_choice=tool_choice or "auto",
+            tool_ids=tool_ids,
+        )
+
+
+def validate_messages(request: AzureChatCompletionRequest) -> None:
+    decl_tools = request.tools is not None
+    decl_functions = request.functions is not None
+
+    if decl_functions and decl_tools:
+        raise ValidationError("Both functions and tools are not allowed")
+
+    def warn(msg: str):
+        _log.warning(
+            f"The request is incomplete: {msg}. The model may misbehave."
+        )
+
+    tool_defs_are_missing = (
+        "the request is missing tool definitions in the 'tools' field"
+    )
+    func_defs_are_missing = (
+        "the request is missing function definitions in the 'functions' field"
+    )
+
+    for idx, message in enumerate(request.messages):
+        if (
+            message.role == Role.ASSISTANT
+            and message.tool_calls is not None
+            and not decl_tools
+        ):
+            warn(
+                f"'messages[{idx}]' is an Assistant message with a tool call, but {tool_defs_are_missing}"
+            )
+        if (
+            message.role == Role.ASSISTANT
+            and message.function_call is not None
+            and not decl_functions
+        ):
+            warn(
+                f"'messages[{idx}]' is an Assistant messages with a function call, but {func_defs_are_missing}"
+            )
+        if message.role == Role.FUNCTION and not decl_functions:
+            warn(
+                f"'messages[{idx}]' is a Function message, but {func_defs_are_missing}"
+            )
+        if message.role == Role.TOOL and not decl_tools:
+            warn(
+                f"'messages[{idx}]' is a Tool message, but {tool_defs_are_missing}"
+            )
+
+
+def _collect_tool_ids(messages: List[Message]) -> Dict[str, str]:
+    ret: Dict[str, str] = {}
+
+    for message in messages:
+        if message.role == Role.ASSISTANT and message.tool_calls is not None:
+            for tool_call in message.tool_calls:
+                ret[tool_call.id] = tool_call.function.name
+
+    return ret