diff --git a/src/banks/types.py b/src/banks/types.py index 62eca9b..fb011b9 100644 --- a/src/banks/types.py +++ b/src/banks/types.py @@ -15,37 +15,26 @@ class ContentBlockType(str, Enum): text = "text" - image = "image" - - -class MediaTypeBlockType(str, Enum): - image_jpeg = "image/jpeg" - image_png = "image/png" - image_gif = "image/gif" - image_webp = "image/webp" + image_url = "image_url" class CacheControl(BaseModel): type: str = "ephemeral" -class Source(BaseModel): - type: str = "base64" - media_type: MediaTypeBlockType - data: str +class ImageUrl(BaseModel): + url: str - class Config: - use_enum_values = True + @classmethod + def from_base64(cls, media_type: str, base64_str: str): + return cls(url=f"data:{media_type};base64,{base64_str}") class ContentBlock(BaseModel): type: ContentBlockType cache_control: CacheControl | None = None text: str | None = None - source: Source | None = None - - class Config: - use_enum_values = True + image_url: ImageUrl | None = None ChatMessageContent = list[ContentBlock] | str diff --git a/tests/test_cache_control.py b/tests/test_cache_control.py index 0a6e378..100153b 100644 --- a/tests/test_cache_control.py +++ b/tests/test_cache_control.py @@ -5,4 +5,4 @@ def test_cache_control(): res = cache_control("foo", "ephemeral") res = res.replace("", "") res = res.replace("", "") - assert res == '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","source":null}' + assert res == '{"type":"text","cache_control":{"type":"ephemeral"},"text":"foo","image_url":null}'