feat: update to handle newer messagesCHECKPOINTprovider-config-poc

1 files changed, 16 insertions, 1 deletions

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]
+