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

aidial_adapter_anthropic/dial/_message.py

@@ -0,0 +1,341 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Self, Union
+
+from aidial_sdk.chat_completion import (
+    Attachment,
+    CacheBreakpoint,
+    CustomContent,
+    FunctionCall,
+)
+from aidial_sdk.chat_completion import Message as DialMessage
+from aidial_sdk.chat_completion import (
+    MessageContentPart,
+    MessageContentTextPart,
+    MessageCustomFields,
+    Role,
+    ToolCall,
+)
+from pydantic import BaseModel
+
+from aidial_adapter_anthropic.adapter._errors import ValidationError
+from aidial_adapter_anthropic.dial.request import (
+    collect_text_content,
+    is_plain_text_content,
+    is_system_role,
+    is_text_content,
+    to_message_content,
+)
+
+
+class MessageABC(ABC, BaseModel):
+    cache_breakpoint: CacheBreakpoint | None = None
+
+    @property
+    def custom_fields(self) -> MessageCustomFields | None:
+        if self.cache_breakpoint:
+            return MessageCustomFields(cache_breakpoint=self.cache_breakpoint)
+        return None
+
+    @abstractmethod
+    def to_message(self) -> DialMessage: ...
+
+    @classmethod
+    @abstractmethod
+    def from_message(cls, message: DialMessage) -> Self | None: ...
+
+
+class BaseMessageABC(MessageABC):
+    @property
+    @abstractmethod
+    def text_content(self) -> str: ...
+
+
+def _get_cache_breakpoint(message: DialMessage) -> CacheBreakpoint | None:
+    if message.custom_fields is None:
+        return None
+    return message.custom_fields.cache_breakpoint
+
+
+class SystemMessage(BaseMessageABC):
+    content: str | List[MessageContentTextPart]
+    is_developer: bool = False
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.DEVELOPER if self.is_developer else Role.SYSTEM,
+            content=to_message_content(self.content),
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if not is_system_role(message.role):
+            return None
+
+        content = message.content
+
+        if not is_text_content(content):
+            raise ValidationError(
+                "System message is expected to be a string or a list of text content parts"
+            )
+
+        return cls(
+            is_developer=message.role == Role.DEVELOPER,
+            cache_breakpoint=_get_cache_breakpoint(message),
+            content=content,
+        )
+
+    @property
+    def text_content(self) -> str:
+        return collect_text_content(self.content)
+
+
+class HumanRegularMessage(BaseMessageABC):
+    content: str | List[MessageContentPart]
+    custom_content: CustomContent | None = None
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.USER,
+            content=self.content,
+            custom_content=self.custom_content,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.USER:
+            return None
+
+        content = message.content
+        if content is None:
+            raise ValidationError(
+                "User message is expected to have content field"
+            )
+
+        return cls(
+            content=content,
+            custom_content=message.custom_content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+    @property
+    def text_content(self) -> str:
+        return collect_text_content(self.content)
+
+    @property
+    def attachments(self) -> List[Attachment]:
+        return (
+            self.custom_content.attachments or [] if self.custom_content else []
+        )
+
+
+class HumanToolResultMessage(MessageABC):
+    id: str
+    content: str
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.TOOL,
+            tool_call_id=self.id,
+            content=self.content,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.TOOL:
+            return None
+
+        if not is_plain_text_content(message.content):
+            raise ValidationError(
+                "The tool message shouldn't contain content parts"
+            )
+
+        if message.content is None or message.tool_call_id is None:
+            raise ValidationError(
+                "The tool message is expected to have content and tool_call_id fields"
+            )
+
+        return cls(
+            id=message.tool_call_id,
+            content=message.content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+
+class HumanFunctionResultMessage(MessageABC):
+    name: str
+    content: str
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.FUNCTION,
+            name=self.name,
+            content=self.content,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.FUNCTION:
+            return None
+
+        if not is_plain_text_content(message.content):
+            raise ValidationError(
+                "The function message shouldn't contain content parts"
+            )
+
+        if message.content is None or message.name is None:
+            raise ValidationError(
+                "The function message is expected to have content and name fields"
+            )
+
+        return cls(
+            name=message.name,
+            content=message.content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+
+class AIRegularMessage(BaseMessageABC):
+    content: str | List[MessageContentPart]
+    """
+    According to Azure OpenAI API, the assistant message could only have textual content.
+    However, we leave a loophole to provide image content parts just in case
+    one day multi-modal Bedrock models will be able to accept images in assistant messages.
+    """
+
+    custom_content: Optional[CustomContent] = None
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.ASSISTANT,
+            content=self.content,  # type: ignore
+            custom_content=self.custom_content,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.ASSISTANT:
+            return None
+
+        if message.function_call is not None or message.tool_calls is not None:
+            return None
+
+        content = message.content
+        if content is None:
+            raise ValidationError(
+                "Assistant message is expected to have content field"
+            )
+
+        return cls(
+            content=content,
+            custom_content=message.custom_content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+    @property
+    def text_content(self) -> str:
+        return collect_text_content(self.content)
+
+    @property
+    def attachments(self) -> List[Attachment]:
+        return (
+            self.custom_content.attachments or [] if self.custom_content else []
+        )
+
+
+class AIToolCallMessage(MessageABC):
+    calls: List[ToolCall]
+    content: Optional[str] = None
+    custom_content: Optional[CustomContent] = None
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.ASSISTANT,
+            content=self.content,
+            tool_calls=self.calls,
+            custom_content=self.custom_content,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.ASSISTANT:
+            return None
+
+        if message.tool_calls is None or message.function_call is not None:
+            return None
+
+        if not is_plain_text_content(message.content):
+            raise ValidationError(
+                "The assistant message with tool calls shouldn't contain content parts"
+            )
+
+        return cls(
+            calls=message.tool_calls,
+            content=message.content,
+            custom_content=message.custom_content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+
+class AIFunctionCallMessage(MessageABC):
+    call: FunctionCall
+    content: Optional[str] = None
+
+    def to_message(self) -> DialMessage:
+        return DialMessage(
+            role=Role.ASSISTANT,
+            content=self.content,
+            function_call=self.call,
+            custom_fields=self.custom_fields,
+        )
+
+    @classmethod
+    def from_message(cls, message: DialMessage) -> Self | None:
+        if message.role != Role.ASSISTANT:
+            return None
+
+        if message.function_call is None or message.tool_calls is not None:
+            return None
+
+        if not is_plain_text_content(message.content):
+            raise ValidationError(
+                "The assistant message with function call shouldn't contain content parts"
+            )
+
+        return cls(
+            call=message.function_call,
+            content=message.content,
+            cache_breakpoint=_get_cache_breakpoint(message),
+        )
+
+
+BaseMessage = Union[SystemMessage, HumanRegularMessage, AIRegularMessage]
+
+ToolMessage = Union[
+    HumanToolResultMessage,
+    HumanFunctionResultMessage,
+    AIToolCallMessage,
+    AIFunctionCallMessage,
+]
+
+
+def parse_dial_message(msg: DialMessage) -> BaseMessage | ToolMessage:
+    message = (
+        SystemMessage.from_message(msg)
+        or HumanRegularMessage.from_message(msg)
+        or HumanToolResultMessage.from_message(msg)
+        or HumanFunctionResultMessage.from_message(msg)
+        or AIRegularMessage.from_message(msg)
+        or AIToolCallMessage.from_message(msg)
+        or AIFunctionCallMessage.from_message(msg)
+    )
+
+    if message is None:
+        raise ValidationError("Unknown message type or invalid message")
+
+    return message

aidial_adapter_anthropic/dial/consumer.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import dataclasses
+from abc import ABC, abstractmethod
+from types import TracebackType
+from typing import ContextManager, List, Optional, Protocol, Self, Tuple
+
+from aidial_sdk.chat_completion import (
+    Attachment,
+    Choice,
+    FinishReason,
+    FunctionCall,
+    Response,
+    ToolCall,
+)
+
+from aidial_adapter_anthropic.adapter._truncate_prompt import DiscardedMessages
+from aidial_adapter_anthropic.dial._lazy_stage import LazyStage
+from aidial_adapter_anthropic.dial.token_usage import TokenUsage
+
+
+class _ArgumentConsumer(Protocol):
+    def append_arguments(self, arguments: str) -> Self: ...
+
+
+@dataclasses.dataclass
+class ToolUseMessage:
+    call: _ArgumentConsumer
+    snapshot: str
+
+    def append_arguments(self, arguments: str) -> Self:
+        self.call.append_arguments(arguments)
+        self.snapshot += arguments
+        return self
+
+    def close(self) -> Self:
+        if not self.snapshot.strip():
+            self.append_arguments("{}")
+        return self
+
+
+class Consumer(ContextManager, ABC):
+    @abstractmethod
+    def fork(self) -> Consumer: ...
+
+    @property
+    @abstractmethod
+    def choice(self) -> Choice: ...
+
+    @abstractmethod
+    def close_content(self, finish_reason: FinishReason | None = None): ...
+
+    @abstractmethod
+    def append_content(self, content: str): ...
+
+    @abstractmethod
+    def add_attachment(self, attachment: Attachment): ...
+
+    @abstractmethod
+    def add_citation_attachment(
+        self, document_id: int, document: Attachment | None
+    ) -> int: ...
+
+    @abstractmethod
+    def add_usage(self, usage: TokenUsage): ...
+
+    @abstractmethod
+    def set_discarded_messages(
+        self, discarded_messages: Optional[DiscardedMessages]
+    ): ...
+
+    @abstractmethod
+    def get_discarded_messages(self) -> Optional[DiscardedMessages]: ...
+
+    @abstractmethod
+    def create_function_tool_call(self, call: ToolCall) -> ToolUseMessage: ...
+
+    @abstractmethod
+    def create_function_call(self, call: FunctionCall) -> ToolUseMessage: ...
+
+    @property
+    @abstractmethod
+    def has_function_call(self) -> bool: ...
+
+    def create_stage(self, title: str) -> LazyStage:
+        # NOTE: eta conversion to `factory = self.choice.create_stage`
+        # is invalid, since `self.choice` must be created lazily.
+        def factory(content: str):
+            return self.choice.create_stage(content)
+
+        return LazyStage(factory, title)
+
+
+class ChoiceConsumer(Consumer):
+    response: Response
+
+    usage: Optional[TokenUsage]
+    discarded_messages: Optional[DiscardedMessages]
+
+    _root: Optional[Consumer]
+    _choice: Optional[Choice]
+    _tool_calls: List[ToolUseMessage]
+    _citations: dict[int, Tuple[int, Attachment | None]]
+
+    def __init__(self, response: Response, root: Optional[Consumer] = None):
+        self.response = response
+
+        self.usage = None
+        self.discarded_messages = None
+
+        self._choice = None
+        self._root = root
+        self._tool_calls = []
+        self._citations = {}
+
+    def fork(self) -> Consumer:
+        return ChoiceConsumer(self.response, self._root or self)
+
+    @property
+    def choice(self) -> Choice:
+        if self._choice is None:
+            choice = self._choice = self.response.create_choice()
+            # Delay opening a choice to the very last moment
+            # so as to give opportunity for exceptions to bubble up to
+            # the level of HTTP response (instead of error objects in a stream).
+            choice.open()
+            return choice
+        else:
+            return self._choice
+
+    def __enter__(self) -> ChoiceConsumer:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        for tool_call in self._tool_calls:
+            tool_call.close()
+
+        if exc is None and self._choice is not None:
+            self._choice.close()
+
+        if self._root is None:
+            if self.usage is not None:
+                self.response.set_usage(
+                    prompt_tokens=self.usage.prompt_tokens,
+                    completion_tokens=self.usage.completion_tokens,
+                    prompt_tokens_details={
+                        "cached_tokens": self.usage.cache_read_input_tokens
+                    },
+                )
+
+            if self.discarded_messages is not None:
+                self.response.set_discarded_messages(self.discarded_messages)
+
+        return False
+
+    def close_content(self, finish_reason: FinishReason | None = None):
+        # Choice.close(finish_reason: Optional[FinishReason]) can be called only once
+        # Currently, there's no other way to explicitly set the finish reason
+        self.choice._last_finish_reason = finish_reason
+
+    def append_content(self, content: str):
+        self.choice.append_content(content)
+
+    def add_attachment(self, attachment: Attachment):
+        self.choice.add_attachment(attachment)
+
+    def add_citation_attachment(
+        self, document_id: int, document: Attachment | None
+    ) -> int:
+        if document_id in self._citations:
+            return self._citations[document_id][0]
+
+        display_index = len(self._citations) + 1
+        self._citations[document_id] = (display_index, document)
+
+        if document:
+            document = document.copy()
+            document.title = f"[{display_index}] {document.title or ''}".strip()
+            document.reference_type = document.reference_type or document.type
+            document.reference_url = document.reference_url or document.url
+            self.add_attachment(document)
+
+        return display_index
+
+    def add_usage(self, usage: TokenUsage):
+        if self._root:
+            self._root.add_usage(usage)
+        else:
+            self.usage = (self.usage or TokenUsage()).accumulate(usage)
+
+    def set_discarded_messages(
+        self, discarded_messages: Optional[DiscardedMessages]
+    ):
+        if self._root:
+            self._root.set_discarded_messages(discarded_messages)
+        else:
+            self.discarded_messages = discarded_messages
+
+    def get_discarded_messages(self) -> Optional[DiscardedMessages]:
+        if self._root:
+            return self._root.get_discarded_messages()
+        else:
+            return self.discarded_messages
+
+    def create_function_tool_call(self, call: ToolCall) -> ToolUseMessage:
+        tool_call = ToolUseMessage(
+            call=self.choice.create_function_tool_call(
+                id=call.id,
+                name=call.function.name,
+                arguments=call.function.arguments,
+            ),
+            snapshot=call.function.arguments,
+        )
+        self._tool_calls.append(tool_call)
+        return tool_call
+
+    def create_function_call(self, call: FunctionCall) -> ToolUseMessage:
+        tool_call = ToolUseMessage(
+            call=self.choice.create_function_call(
+                name=call.name,
+                arguments=call.arguments,
+            ),
+            snapshot=call.arguments,
+        )
+        self._tool_calls.append(tool_call)
+        return tool_call
+
+    @property
+    def has_function_call(self) -> bool:
+        return self._choice is not None and self._choice.has_function_call