|
1 | 1 | from enum import Enum |
2 | 2 | from typing import Any, Dict, List, Literal, Mapping, Optional, Union |
3 | 3 |
|
| 4 | +from typing_extensions import assert_never |
| 5 | + |
4 | 6 | from aidial_sdk.chat_completion.enums import Status |
5 | 7 | from aidial_sdk.deployment.from_request_mixin import FromRequestDeploymentMixin |
| 8 | +from aidial_sdk.exceptions import InvalidRequestError |
6 | 9 | from aidial_sdk.pydantic_v1 import ( |
7 | 10 | ConstrainedFloat, |
8 | 11 | ConstrainedInt, |
@@ -58,15 +61,51 @@ class Role(str, Enum): |
58 | 61 | TOOL = "tool" |
59 | 62 |
|
60 | 63 |
|
| 64 | +class ImageURL(ExtraForbidModel): |
| 65 | + url: StrictStr |
| 66 | + detail: Optional[Literal["auto", "low", "high"]] = None |
| 67 | + |
| 68 | + |
| 69 | +class MessageContentImagePart(ExtraForbidModel): |
| 70 | + type: Literal["image_url"] |
| 71 | + image_url: ImageURL |
| 72 | + |
| 73 | + |
| 74 | +class MessageContentTextPart(ExtraForbidModel): |
| 75 | + type: Literal["text"] |
| 76 | + text: StrictStr |
| 77 | + |
| 78 | + |
| 79 | +MessageContentPart = Union[MessageContentTextPart, MessageContentImagePart] |
| 80 | + |
| 81 | + |
61 | 82 | class Message(ExtraForbidModel): |
62 | 83 | role: Role |
63 | | - content: Optional[StrictStr] = None |
| 84 | + content: Optional[Union[StrictStr, List[MessageContentPart]]] = None |
64 | 85 | custom_content: Optional[CustomContent] = None |
65 | 86 | name: Optional[StrictStr] = None |
66 | 87 | tool_calls: Optional[List[ToolCall]] = None |
67 | 88 | tool_call_id: Optional[StrictStr] = None |
68 | 89 | function_call: Optional[FunctionCall] = None |
69 | 90 |
|
| 91 | + def text(self) -> str: |
| 92 | + """ |
| 93 | + Returns content of the message only if it's present as a string. |
| 94 | + Otherwise, throws an invalid request exception. |
| 95 | + """ |
| 96 | + |
| 97 | + def _error_message(actual: str) -> str: |
| 98 | + return f"Unable to retrieve text content of the message: the actual content is {actual}." |
| 99 | + |
| 100 | + if self.content is None: |
| 101 | + raise InvalidRequestError(_error_message("null or missing")) |
| 102 | + elif isinstance(self.content, str): |
| 103 | + return self.content |
| 104 | + elif isinstance(self.content, list): |
| 105 | + raise InvalidRequestError(_error_message("a list of content parts")) |
| 106 | + else: |
| 107 | + assert_never(self.content) |
| 108 | + |
70 | 109 |
|
71 | 110 | class Addon(ExtraForbidModel): |
72 | 111 | name: Optional[StrictStr] = None |
|
0 commit comments