Update main.py

main.py

@@ -22,10 +22,24 @@ with open("model_map.json", "r", encoding="utf-8") as f:
     MODEL_PROMPTS = json.load(f)
 
 # Request schema
+
+# Define ContentType for vision
+class ContentImage(BaseModel):
+    type: str  # must be "image_url"
+    image_url: dict  # {"url": "https://..." or "data:image/...;base64,..."}
+
+class ContentText(BaseModel):
+    type: str  # must be "text"
+    text: str
+
+ContentType = Union[ContentText, ContentImage]
+
+# Message model allows BOTH old and new formats
 class Message(BaseModel):
     role: str
-    content: str
+    content: Union[str, List[ContentType]]  # str (legacy) or list of ContentType
 
+# ChatRequest model
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
@@ -41,9 +55,38 @@ class ChatRequest(BaseModel):
 def build_payload(chat: ChatRequest):
     system_prompt = MODEL_PROMPTS.get(chat.model, "You are a helpful assistant.")
     filtered_messages = [msg for msg in chat.messages if msg.role != "system"]
-    payload_messages = [{"role": "system", "content": system_prompt}] + [
-        {"role": msg.role, "content": msg.content} for msg in filtered_messages
-    ]
+
+    payload_messages = [{"role": "system", "content": system_prompt}]
+
+    for msg in filtered_messages:
+        # Legacy format: string
+        if isinstance(msg.content, str):
+            payload_messages.append({"role": msg.role, "content": msg.content})
+
+        # Multimodal format
+        elif isinstance(msg.content, list):
+            content_payload = []
+            for content_item in msg.content:
+                # ContentText
+                if content_item.type == "text":
+                    content_payload.append({
+                        "type": "text",
+                        "text": content_item.text
+                    })
+                # ContentImage
+                elif content_item.type == "image_url":
+                    content_payload.append({
+                        "type": "image_url",
+                        "image_url": content_item.image_url
+                    })
+                else:
+                    logger.warning(f"Unknown content type: {content_item.type}, skipping.")
+
+            payload_messages.append({"role": msg.role, "content": content_payload})
+
+        else:
+            logger.warning(f"Unknown message content format: {msg.content}")
+
     return {
         "model": BACKEND_MODEL,
         "messages": payload_messages,