From 90fa1cc3b30253bed1ace217d64fbd7df542c987 Mon Sep 17 00:00:00 2001
From: Henri Wenlin <Henri.Wenlin@kone.com>
Date: Thu, 20 Jun 2024 12:15:24 +0300
Subject: [PATCH] Refactor message content handling in schemas.py to add
 capability for image content

---
 schemas.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/schemas.py b/schemas.py
index da4122a7..57517dfd 100644
--- a/schemas.py
+++ b/schemas.py
@@ -1,14 +1,23 @@
-from typing import List, Optional
-from pydantic import BaseModel, ConfigDict
+from typing import List, Union, Optional
+from pydantic import BaseModel, RootModel, ConfigDict
 
+class ImageContent(BaseModel):
+    type: str
+    image_url: dict
+
+class TextContent(BaseModel):
+    type: str
+    text: str
+
+class MessageContent(RootModel):
+    root: Union[TextContent, ImageContent]
 
 class OpenAIChatMessage(BaseModel):
     role: str
-    content: str | List
+    content: Union[str, List[MessageContent]]
 
     model_config = ConfigDict(extra="allow")
 
-
 class OpenAIChatCompletionForm(BaseModel):
     stream: bool = True
     model: str
@@ -16,8 +25,7 @@ class OpenAIChatCompletionForm(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
-
 class FilterForm(BaseModel):
     body: dict
     user: Optional[dict] = None
-    model_config = ConfigDict(extra="allow")
+    model_config = ConfigDict(extra="allow")
\ No newline at end of file