Update api/models.py

api/models.py

@@ -1,14 +1,16 @@
-from typing import List, Optional
-from pydantic import BaseModel
-
-class Message(BaseModel):
-    role: str
-    content: str | list
-
-class ChatRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    stream: Optional[bool] = False
-    temperature: Optional[float] = 0.5
-    top_p: Optional[float] = 0.9
-    max_tokens: Optional[int] = 1024
+# api/models.py
+
+from typing import List, Optional, Union
+from pydantic import BaseModel
+
+class Message(BaseModel):
+    role: str
+    content: Union[str, List[dict]]  # Allow content to be a string or a list of dictionaries for images
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    stream: Optional[bool] = False
+    temperature: Optional[float] = 0.5
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 1024