Update api/models.py

api/models.py

@@ -1,14 +1,14 @@
-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.7
-    top_p: Optional[float] = 0.9
-    max_tokens: Optional[int] = 8192
+from typing import List, Optional, Union
+from pydantic import BaseModel, Field
+
+class MessageContent(BaseModel):
+    role: str
+    content: Union[str, List[dict]]  # Handle both text and complex structures like images
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[MessageContent]
+    stream: Optional[bool] = False
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 8192