From 21f54bd0d623563436d506b6f855ba39cb7ea078 Mon Sep 17 00:00:00 2001
From: Max Resnick <max@ofmax.li>
Date: Tue, 4 Feb 2025 22:22:55 -0800
Subject: feat: update to handle newer messages

---
 py/types.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

(limited to 'py')

diff --git a/py/types.py b/py/types.py
index 819afe5..4eb0ccc 100644
--- a/py/types.py
+++ b/py/types.py
@@ -1,5 +1,5 @@
 from collections.abc import Sequence, Mapping
-from typing import TypedDict, Protocol
+from typing import TypedDict, Literal, Union, List
 
 
 class AIProvider(Protocol):
@@ -15,3 +15,18 @@ class Message(TypedDict):
     content: str
     type: str
 
+
+class TextContent(TypedDict):
+    type: Literal['text']
+    text: str
+
+class ImageUrlContent(TypedDict):
+    type: Literal['image_url']
+    image_url: dict[str, str]  # {'url': str}
+
+MessageContent = Union[TextContent, ImageUrlContent]
+
+class Message(TypedDict):
+    role: Literal['system', 'user', 'assistant']
+    content: List[MessageContent]
+
-- 
cgit v1.2.3