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

aidial_adapter_anthropic/_utils/json.py

@@ -0,0 +1,116 @@
+"""
+Utilities for pretty-printing JSON in debug logs.
+These functions are useful for dumping large data structures,
+with options to trim long strings and lists to specified limits.
+"""
+
+import json
+from dataclasses import asdict, is_dataclass
+from enum import Enum
+from typing import Any
+
+from anthropic import Omit
+from pydantic import BaseModel
+
+
+def json_dumps_short(
+    obj: Any, *, string_limit: int = 100, list_len_limit: int = 10, **kwargs
+) -> str:
+    def default(obj) -> str:
+        return _truncate_strings(str(obj), string_limit)
+
+    return json.dumps(
+        _truncate_lists(
+            _truncate_strings(_to_dict(obj, **kwargs), string_limit),
+            list_len_limit,
+        ),
+        default=default,
+    )
+
+
+def json_dumps(obj: Any, **kwargs) -> str:
+    return json.dumps(_to_dict(obj, **kwargs))
+
+
+def _to_dict(obj: Any, **kwargs) -> Any:
+    def rec(val):
+        return _to_dict(val, **kwargs)
+
+    def dict_field(key: str, val: Any) -> Any:
+        if key in kwargs.get("excluded_keys", []):
+            return "<excluded>"
+        return val
+
+    if isinstance(obj, bytes):
+        return f"<bytes>({len(obj):_} B)"
+
+    if isinstance(obj, Enum):
+        return obj.value
+
+    if isinstance(obj, dict):
+        return {key: rec(dict_field(key, value)) for key, value in obj.items()}
+
+    if isinstance(obj, list):
+        return [rec(element) for element in obj]
+
+    if isinstance(obj, tuple):
+        return tuple(rec(element) for element in obj)
+
+    if isinstance(obj, BaseModel):
+        return rec(obj.dict())
+
+    if hasattr(obj, "to_dict"):
+        return rec(obj.to_dict())
+
+    if is_dataclass(type(obj)):
+        return rec(asdict(obj))
+
+    if isinstance(obj, Omit):
+        return "omit"
+
+    return obj
+
+
+def _truncate_strings(obj: Any, limit: int) -> Any:
+    def rec(val):
+        return _truncate_strings(val, limit)
+
+    if isinstance(obj, dict):
+        return {key: rec(value) for key, value in obj.items()}
+
+    if isinstance(obj, list):
+        return [rec(element) for element in obj]
+
+    if isinstance(obj, tuple):
+        return tuple(rec(element) for element in obj)
+
+    if isinstance(obj, str) and len(obj) > limit:
+        skip = len(obj) - limit
+        return (
+            obj[: limit // 2] + f"...({skip:_} skipped)..." + obj[-limit // 2 :]
+        )
+
+    return obj
+
+
+def _truncate_lists(obj: Any, limit: int) -> Any:
+    def rec(val):
+        return _truncate_lists(val, limit)
+
+    if isinstance(obj, dict):
+        return {key: rec(value) for key, value in obj.items()}
+
+    if isinstance(obj, list):
+        if len(obj) > limit:
+            skip = len(obj) - limit
+            obj = (
+                obj[: limit // 2]
+                + [f"...({skip:_} skipped)..."]
+                + obj[-limit // 2 :]
+            )
+        return [rec(element) for element in obj]
+
+    if isinstance(obj, tuple):
+        return tuple(rec(element) for element in obj)
+
+    return obj

aidial_adapter_anthropic/_utils/list.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass, field
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Container,
+    Generic,
+    Iterable,
+    List,
+    Self,
+    Set,
+    Tuple,
+    TypeVar,
+)
+
+_T = TypeVar("_T")
+_V = TypeVar("_V")
+
+
+def select_by_indices(lst: List[_T], indices: Container[int]) -> List[_T]:
+    return [elem for idx, elem in enumerate(lst) if idx in indices]
+
+
+def omit_by_indices(lst: List[_T], indices: Container[int]) -> List[_T]:
+    return [elem for idx, elem in enumerate(lst) if idx not in indices]
+
+
+def group_by(
+    lst: List[_T],
+    key: Callable[[_T], Any],
+    init: Callable[[_T], _V],
+    merge: Callable[[_V, _T], _V],
+) -> List[_V]:
+
+    def _gen():
+        if not lst:
+            return
+
+        prev_val = init(lst[0])
+        prev_key = key(lst[0])
+
+        for elem in lst[1:]:
+            if prev_key == key(elem):
+                prev_val = merge(prev_val, elem)
+            else:
+                yield prev_val
+                prev_val = init(elem)
+                prev_key = key(elem)
+
+        yield prev_val
+
+    return list(_gen())
+
+
+@dataclass
+class ListProjection(Generic[_T]):
+    """
+    The class represents a transformation of the original list which may
+    include merge, removal and addition of the original list elements.
+
+    Each derivative element is mapped onto a subset of original elements.
+    The subsets must be disjoint.
+    """
+
+    list: List[Tuple[_T, Set[int]]] = field(default_factory=list)
+
+    @property
+    def raw_list(self) -> List[_T]:
+        return [msg for msg, _ in self.list]
+
+    def to_original_indices(self, indices: Iterable[int]) -> Set[int]:
+        return {
+            orig_index
+            for index in indices
+            for orig_index in self.list[index][1]
+        }
+
+    def append(self, elem: _T, idx: int) -> Self:
+        self.list.append((elem, {idx}))
+        return self
+
+
+async def aiter_to_list(iterator: AsyncIterator[_T]) -> List[_T]:
+    return [item async for item in iterator]

aidial_adapter_anthropic/_utils/pydantic.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class ExtraForbidModel(BaseModel):
+    class Config:
+        extra = "forbid"

aidial_adapter_anthropic/_utils/resource.py

@@ -0,0 +1,54 @@
+import base64
+import re
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Resource(BaseModel):
+    type: str
+    data: bytes
+
+    @classmethod
+    def from_base64(cls, type: str, data_base64: str) -> "Resource":
+        try:
+            data = base64.b64decode(data_base64, validate=True)
+        except Exception:
+            raise ValueError("Invalid base64 data")
+
+        return cls(type=type, data=data)
+
+    @classmethod
+    def from_data_url(cls, data_url: str) -> Optional["Resource"]:
+        """
+        Parsing a resource encoded as a data URL.
+        See https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs for reference.
+        """
+
+        type = cls.parse_data_url_content_type(data_url)
+        if type is None:
+            return None
+
+        data_base64 = data_url.removeprefix(cls._to_data_url_prefix(type))
+
+        return cls.from_base64(type, data_base64)
+
+    @property
+    def data_base64(self) -> str:
+        return base64.b64encode(self.data).decode()
+
+    def to_data_url(self) -> str:
+        return f"{self._to_data_url_prefix(self.type)}{self.data_base64}"
+
+    @staticmethod
+    def parse_data_url_content_type(data_url: str) -> Optional[str]:
+        pattern = r"^data:([^;]+);base64,"
+        match = re.match(pattern, data_url)
+        return None if match is None else match.group(1)
+
+    @staticmethod
+    def _to_data_url_prefix(content_type: str) -> str:
+        return f"data:{content_type};base64,"
+
+    def __str__(self) -> str:
+        return self.to_data_url()[:100] + "..."

aidial_adapter_anthropic/_utils/text.py

@@ -0,0 +1,4 @@
+def truncate_string(s: str, n: int) -> str:
+    if len(s) <= n:
+        return s
+    return s[:n] + "..."

aidial_adapter_anthropic/adapter/__init__.py

@@ -0,0 +1,4 @@
+from aidial_adapter_anthropic.adapter._base import ChatCompletionAdapter
+from aidial_adapter_anthropic.adapter._errors import UserError, ValidationError
+
+__all__ = ["ChatCompletionAdapter", "UserError", "ValidationError"]

aidial_adapter_anthropic/adapter/_base.py

@@ -0,0 +1,95 @@
+from abc import ABC, abstractmethod
+from typing import Any, List, Set, Tuple, Type
+
+from aidial_sdk.chat_completion import Message
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic._utils.list import ListProjection
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+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,
+    collect_text_content,
+    is_system_role,
+)
+
+
+class ChatCompletionAdapter(ABC, BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    @abstractmethod
+    async def chat(
+        self,
+        consumer: Consumer,
+        params: ModelParameters,
+        messages: List[Message],
+    ) -> None:
+        pass
+
+    async def configuration(self) -> Type[BaseModel]:
+        raise NotImplementedError
+
+    async def count_prompt_tokens(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> int:
+        raise NotImplementedError
+
+    async def count_completion_tokens(self, string: str) -> int:
+        raise NotImplementedError
+
+    async def compute_discarded_messages(
+        self, params: ModelParameters, messages: List[Message]
+    ) -> DiscardedMessages | None:
+        """
+        The method truncates the list of messages to fit
+        into the token limit set in `params.max_prompt_tokens`.
+
+        If the limit isn't provided, then it returns None.
+        Otherwise, returns the indices of _discarded_ messages which should be
+        removed from the list to make the rest fit into the token limit.
+        """
+        raise NotImplementedError
+
+
+def default_preprocess_messages(
+    messages: List[Message],
+) -> ListProjection[Message]:
+    def _is_empty_system_message(msg: Message) -> bool:
+        return (
+            is_system_role(msg.role)
+            and collect_text_content(msg.content).strip() == ""
+        )
+
+    ret: List[Tuple[Message, Set[int]]] = []
+    idx: Set[int] = set()
+
+    for i, msg in enumerate(messages):
+        idx.add(i)
+        if _is_empty_system_message(msg):
+            continue
+        ret.append((msg, idx))
+        idx = set()
+
+    if len(ret) == 0:
+        raise ValidationError("List of messages must not be empty")
+
+    return ListProjection(ret)
+
+
+def keep_last(messages: List[Any], idx: int) -> bool:
+    return idx == len(messages) - 1
+
+
+def keep_last_and_system_messages(messages: List[Message], idx: int) -> bool:
+    return is_system_role(messages[idx].role) or keep_last(messages, idx)
+
+
+def trivial_partitioner(messages: List[Any]) -> List[int]:
+    return [1] * len(messages)
+
+
+def turn_based_partitioner(messages: List[Any]) -> List[int]:
+    n = len(messages)
+    return [2] * (n // 2) + [1] * (n % 2)