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

aidial_adapter_anthropic/adapter/_truncate_prompt.py

@@ -0,0 +1,168 @@
+from abc import ABC, abstractmethod
+from typing import Awaitable, Callable, List, Optional, Set, Tuple, TypeVar
+
+from aidial_sdk.exceptions import ContextLengthExceededError
+from aidial_sdk.exceptions import HTTPException as DialException
+from aidial_sdk.exceptions import (
+    InvalidRequestError,
+    TruncatePromptSystemAndLastUserError,
+)
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic._utils.list import (
+    omit_by_indices,
+    select_by_indices,
+)
+
+
+class TruncatePromptError(ABC, BaseModel):
+    @abstractmethod
+    def to_dial_exception(self) -> DialException:
+        pass
+
+    def print(self) -> str:
+        return self.to_dial_exception().message
+
+
+class InconsistentLimitsError(TruncatePromptError):
+    user_limit: int
+    model_limit: int
+
+    def to_dial_exception(self) -> DialException:
+        return InvalidRequestError(
+            f"The request maximum prompt tokens is {self.user_limit}. "
+            f"However, the model's maximum context length is {self.model_limit} tokens."
+        )
+
+
+class ModelLimitOverflow(TruncatePromptError):
+    model_limit: int
+    token_count: int
+
+    def to_dial_exception(self) -> DialException:
+        return ContextLengthExceededError(self.model_limit, self.token_count)
+
+
+class UserLimitOverflow(TruncatePromptError):
+    user_limit: int
+    token_count: int
+
+    def to_dial_exception(self) -> DialException:
+        return TruncatePromptSystemAndLastUserError(
+            self.user_limit, self.token_count
+        )
+
+
+def _partition_indexer(chunks: List[int]) -> Callable[[int], List[int]]:
+    """
+    Returns a function that maps an index to indices of its partition.
+    """
+    mapping: dict[int, List[int]] = {}
+    offset = 0
+    for size in chunks:
+        chunk = list(range(offset, offset + size))
+        for idx in range(size):
+            mapping[offset + idx] = chunk
+        offset += size
+
+    return mapping.__getitem__
+
+
+_T = TypeVar("_T")
+DiscardedMessages = List[int]
+
+
+async def truncate_prompt(
+    messages: List[_T],
+    tokenizer: Callable[[List[_T]], Awaitable[int]],
+    keep_message: Callable[[List[_T], int], bool],
+    partitioner: Callable[[List[_T]], List[int]],
+    model_limit: Optional[int],
+    user_limit: Optional[int],
+) -> Tuple[DiscardedMessages, List[_T]]:
+    """
+    Returns a list of indices of discarded messages and a list of preserved messages
+    """
+
+    result = await compute_discarded_messages(
+        messages,
+        tokenizer,
+        keep_message,
+        partitioner,
+        model_limit,
+        user_limit,
+    )
+
+    if isinstance(result, TruncatePromptError):
+        raise result.to_dial_exception()
+
+    return (list(result), omit_by_indices(messages, result))
+
+
+async def compute_discarded_messages(
+    messages: List[_T],
+    tokenizer: Callable[[List[_T]], Awaitable[int]],
+    keep_message: Callable[[List[_T], int], bool],
+    partitioner: Callable[[List[_T]], List[int]],
+    model_limit: Optional[int],
+    user_limit: Optional[int],
+) -> DiscardedMessages | TruncatePromptError:
+    if (
+        user_limit is not None
+        and model_limit is not None
+        and user_limit > model_limit
+    ):
+        return InconsistentLimitsError(
+            user_limit=user_limit, model_limit=model_limit
+        )
+
+    if user_limit is None:
+        if model_limit is None:
+            return []
+
+        token_count = await tokenizer(messages)
+        if token_count <= model_limit:
+            return []
+
+        return ModelLimitOverflow(
+            model_limit=model_limit, token_count=token_count
+        )
+
+    partition_sizes = partitioner(messages)
+    if sum(partition_sizes) != len(messages):
+        raise ValueError(
+            "Partition sizes must add up to the number of messages."
+        )
+
+    async def _tokenize_selected(indices: Set[int]) -> int:
+        return await tokenizer(select_by_indices(messages, indices))
+
+    get_partition_indices = _partition_indexer(partition_sizes)
+
+    n = len(messages)
+    kept_indices: Set[int] = {
+        j
+        for i in range(n)
+        for j in get_partition_indices(i)
+        if keep_message(messages, i)
+    }
+
+    token_count = await _tokenize_selected(kept_indices)
+    if token_count > user_limit:
+        return UserLimitOverflow(user_limit=user_limit, token_count=token_count)
+
+    for idx in reversed(range(n)):
+        if idx in kept_indices:
+            continue
+
+        chunk_indices = get_partition_indices(idx)
+        new_token_count = await _tokenize_selected(
+            {*kept_indices, *chunk_indices}
+        )
+        if new_token_count > user_limit:
+            break
+
+        kept_indices.update(chunk_indices)
+
+    all_indices = set(range(n))
+    return sorted(list(all_indices - kept_indices))

aidial_adapter_anthropic/adapter/claude.py

@@ -0,0 +1,17 @@
+from aidial_adapter_anthropic.adapter._claude.adapter import create_adapter
+from aidial_adapter_anthropic.adapter._claude.state import MessageState
+from aidial_adapter_anthropic.adapter._claude.tokenizer.approximate import (
+    ApproximateTokenizer,
+)
+from aidial_adapter_anthropic.adapter._claude.tokenizer.base import (
+    ClaudeTokenizer,
+    create_tokenizer,
+)
+
+__all__ = [
+    "create_adapter",
+    "MessageState",
+    "create_tokenizer",
+    "ApproximateTokenizer",
+    "ClaudeTokenizer",
+]

aidial_adapter_anthropic/dial/_attachments.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass, field
+from typing import (
+    AsyncIterator,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Protocol,
+    Sequence,
+    Set,
+    TypeVar,
+    assert_never,
+    runtime_checkable,
+)
+
+from aidial_sdk.chat_completion import (
+    MessageContentImagePart,
+    MessageContentRefusalPart,
+    MessageContentTextPart,
+)
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic._utils.list import aiter_to_list
+from aidial_adapter_anthropic._utils.resource import Resource
+from aidial_adapter_anthropic.adapter._errors import UserError, ValidationError
+from aidial_adapter_anthropic.dial._message import BaseMessage, SystemMessage
+from aidial_adapter_anthropic.dial.resource import (
+    AttachmentResource,
+    DialResource,
+    UnsupportedContentType,
+    URLResource,
+)
+from aidial_adapter_anthropic.dial.storage import FileStorage
+
+_T = TypeVar("_T", covariant=True)
+_Txt = TypeVar("_Txt", covariant=True)
+_Config = TypeVar("_Config", bound=BaseModel, contravariant=True)
+
+
+@runtime_checkable
+class Handler(Protocol, Generic[_T]):
+    def __call__(self, resource: Resource) -> _T: ...
+
+
+@runtime_checkable
+class HandlerWithConfig(Protocol, Generic[_T, _Config]):
+    def __call__(self, resource: Resource, config: _Config | None) -> _T: ...
+
+
+class AttachmentProcessor(BaseModel, Generic[_T, _Config]):
+    class Config:
+        arbitrary_types_allowed = True
+
+    supported_types: Dict[str, Set[str]]
+    """MIME type to file extensions mapping"""
+
+    handler: Handler[_T] | HandlerWithConfig[_T, _Config]
+
+    def handle(self, resource: Resource, config: _Config | None) -> _T:
+        sig = inspect.signature(self.handler)
+        params = list(sig.parameters.values())
+
+        with_config = (
+            len(params) >= 2
+            and params[1].kind
+            in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            )
+        ) or any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params)
+
+        if with_config:
+            return self.handler(resource, config)  # type: ignore
+        return self.handler(resource)  # type: ignore
+
+
+@dataclass
+class WithResources(Generic[_T]):
+    payload: _T
+    resources: List[DialResource] = field(default_factory=list)
+
+    @staticmethod
+    def transpose(xs: List[WithResources[_T]]) -> WithResources[List[_T]]:
+        resources = [r for x in xs for r in x.resources]
+        payload = [x.payload for x in xs]
+        return WithResources(payload=payload, resources=resources)
+
+
+class AttachmentProcessors(BaseModel, Generic[_Txt, _T, _Config]):
+    config: _Config | None = None
+    attachment_processors: Sequence[AttachmentProcessor[_T, _Config]]
+    text_handler: Callable[[str], _Txt]
+    file_storage: FileStorage | None
+
+    @property
+    def supported_types(self) -> Dict[str, Set[str]]:
+        ret: Dict[str, Set[str]] = {}
+        for processor in self.attachment_processors:
+            for mime_type, file_exts in processor.supported_types.items():
+                ret.setdefault(mime_type, set()).update(file_exts)
+        return ret
+
+    @property
+    def supported_mime_types(self) -> List[str]:
+        return list(self.supported_types)
+
+    @property
+    def supported_image_types(self) -> List[str]:
+        return [t for t in self.supported_mime_types if t.startswith("image/")]
+
+    def _text_handler(self, text: str) -> WithResources[_Txt]:
+        return WithResources(self.text_handler(text))
+
+    async def process_system_message(
+        self, message: SystemMessage
+    ) -> List[_Txt]:
+        def _gen():
+            match (content := message.content):
+                case str():
+                    if content:
+                        yield self.text_handler(content)
+                case list():
+                    for part in content:
+                        match part:
+                            case MessageContentTextPart(text=text):
+                                if text:
+                                    yield self.text_handler(text)
+                            case _:
+                                assert_never(part)
+                case _:
+                    assert_never(content)
+
+        return [x for x in _gen()]
+
+    async def process_attachments(
+        self, message: BaseMessage
+    ) -> WithResources[List[_T | _Txt]]:
+        ret = await aiter_to_list(self._process_attachments_iter(message)) or [
+            self._text_handler("")
+        ]
+        return WithResources.transpose(ret)
+
+    async def _process_attachments_iter(
+        self, message: BaseMessage
+    ) -> AsyncIterator[WithResources[_T | _Txt]]:
+        if not isinstance(message, SystemMessage):
+            for attachment in message.attachments:
+                yield await self._handle_dial_resource(
+                    AttachmentResource(
+                        attachment=attachment,
+                        entity_name="attachment",
+                        supported_types=self.supported_mime_types,
+                    ),
+                )
+
+        content = message.content
+
+        match content:
+            case str():
+                if content:
+                    yield self._text_handler(content)
+            case list():
+                for part in content:
+                    match part:
+                        case MessageContentTextPart(text=text):
+                            if text:
+                                yield self._text_handler(text)
+                        case MessageContentImagePart(image_url=image_url):
+                            yield await self._handle_dial_resource(
+                                URLResource(
+                                    url=image_url.url,
+                                    entity_name="image url",
+                                    supported_types=self.supported_image_types,
+                                ),
+                            )
+                        case MessageContentRefusalPart():
+                            raise ValidationError(
+                                "Refuse content parts aren't supported"
+                            )
+                        case _:
+                            assert_never(part)
+            case _:
+                assert_never(content)
+
+    async def _download_resource(self, dial_resource: DialResource) -> Resource:
+        try:
+            return await dial_resource.download(self.file_storage)
+        except UnsupportedContentType as e:
+            raise UserError(
+                f"Unsupported media type: {e.type}",
+                _get_usage_message(self.get_file_exts(e.supported_types)),
+            )
+
+    async def _handle_resource(self, resource: Resource) -> _T:
+        for processor in self.attachment_processors:
+            if resource.type in processor.supported_types:
+                return processor.handle(resource, self.config)
+
+        raise UserError(
+            f"Unsupported media type: {resource.type}",
+            _get_usage_message(self.get_file_exts(self.supported_mime_types)),
+        )
+
+    async def _handle_dial_resource(
+        self, dial_resource: DialResource
+    ) -> WithResources[_T]:
+        resource = await self._download_resource(dial_resource)
+        message = await self._handle_resource(resource)
+        return WithResources(message, resources=[dial_resource])
+
+    def get_file_exts(self, mime_types: List[str]) -> List[str]:
+        return [
+            file_ext
+            for mime_type, file_exts in self.supported_types.items()
+            if mime_type in mime_types
+            for file_ext in file_exts
+        ]
+
+
+def _get_usage_message(supported_exts: List[str]) -> str:
+    document_hint = ""
+    if "pdf" in supported_exts:
+        document_hint = '- "Summarize the document" for a PDF document'
+
+    return f"""
+The application answers queries about attached files.
+Attach file(s) and ask questions about them in the same message.
+
+Supported attachment types: {', '.join(supported_exts)}.
+
+Examples of queries:
+- "Describe this picture" for an image
+- "What are in these images? Is there any difference between them?" for multiple images
+{document_hint}
+""".strip()

aidial_adapter_anthropic/dial/_lazy_stage.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Callable
+
+from aidial_sdk.chat_completion import Stage
+
+_StageFactory = Callable[[str], Stage]
+
+
+class LazyStage:
+    title: str
+    stage_factory: _StageFactory
+
+    _stage: Stage | None = None
+
+    def __init__(self, stage_factory: _StageFactory, title: str):
+        self.stage_factory = stage_factory
+        self.title = title
+
+    def __enter__(self) -> LazyStage:
+        return self
+
+    async def __aenter__(self) -> LazyStage:
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+    def append_content(self, text: str) -> None:
+        if self._stage is None:
+            self._stage = self.stage_factory(self.title)
+            self._stage.open()
+        self._stage.append_content(text)
+
+    def close(self) -> None:
+        if self._stage is not None:
+            self._stage.close()