feat: migrated to Attachment from DIAL SDK

epam · Nov 29, 2024 · f15bb2d · f15bb2d
1 parent ad77191
commit f15bb2d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 13 deletions.
diff --git a/aidial_adapter_bedrock/llm/consumer.py b/aidial_adapter_bedrock/llm/consumer.py
@@ -3,13 +3,13 @@
 from typing import Optional, assert_never
 
 from aidial_sdk.chat_completion import (
+    Attachment,
     Choice,
     FinishReason,
     FunctionCall,
     Response,
     ToolCall,
 )
-from pydantic import BaseModel
 
 from aidial_adapter_bedrock.dial_api.token_usage import TokenUsage
 from aidial_adapter_bedrock.llm.message import (
@@ -20,15 +20,6 @@
 from aidial_adapter_bedrock.llm.truncate_prompt import DiscardedMessages
 
 
-class Attachment(BaseModel):
-    type: str | None = None
-    title: str | None = None
-    data: str | None = None
-    url: str | None = None
-    reference_url: str | None = None
-    reference_type: str | None = None
-
-
 class Consumer(ABC):
     @abstractmethod
     def append_content(self, content: str):
@@ -147,7 +138,7 @@ def append_content(self, content: str):
         self._process_content(content)
 
     def add_attachment(self, attachment: Attachment):
-        self.choice.add_attachment(**attachment.dict())
+        self.choice.add_attachment(attachment)
 
     def add_usage(self, usage: TokenUsage):
         self.usage.accumulate(usage)

diff --git a/aidial_adapter_bedrock/llm/model/stability/v1.py b/aidial_adapter_bedrock/llm/model/stability/v1.py
@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
+from aidial_sdk.chat_completion import Attachment
 from pydantic import BaseModel, Field
 
 from aidial_adapter_bedrock.bedrock import Bedrock
@@ -14,7 +15,7 @@
     TextCompletionAdapter,
     TextCompletionPrompt,
 )
-from aidial_adapter_bedrock.llm.consumer import Attachment, Consumer
+from aidial_adapter_bedrock.llm.consumer import Consumer
 from aidial_adapter_bedrock.llm.errors import ValidationError
 from aidial_adapter_bedrock.llm.message import BaseMessage
 from aidial_adapter_bedrock.llm.model.stability.storage import save_to_storage

diff --git a/aidial_adapter_bedrock/llm/model/stability/v2.py b/aidial_adapter_bedrock/llm/model/stability/v2.py
@@ -2,6 +2,7 @@
 from typing import List, Optional, Tuple, assert_never
 
 from aidial_sdk.chat_completion import (
+    Attachment,
     Message,
     MessageContentImagePart,
     MessageContentTextPart,
@@ -26,7 +27,7 @@
 )
 from aidial_adapter_bedrock.dial_api.token_usage import TokenUsage
 from aidial_adapter_bedrock.llm.chat_model import ChatCompletionAdapter
-from aidial_adapter_bedrock.llm.consumer import Attachment, Consumer
+from aidial_adapter_bedrock.llm.consumer import Consumer
 from aidial_adapter_bedrock.llm.errors import UserError, ValidationError
 from aidial_adapter_bedrock.llm.model.stability.storage import save_to_storage
 from aidial_adapter_bedrock.llm.truncate_prompt import DiscardedMessages