Update main.py

main.py

@@ -11,7 +11,7 @@ from fastapi import FastAPI, HTTPException, Request, Depends, Header, status
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field, validator
-from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+from typing import List, Dict, Any, Optional, Union, AsyncGenerator, Literal
 from datetime import datetime
 from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
@@ -113,10 +113,11 @@ def count_tokens(messages: List[Dict[str, Any]], model: str) -> int:
     for message in messages:
         if isinstance(message['content'], list):
             for content_part in message['content']:
-                if content_part.get('type') == 'text':
-                    tokens += len(encoding.encode(content_part['text']))
-                elif content_part.get('type') == 'image_url':
-                    tokens += len(encoding.encode(content_part['image_url']['url']))
+                if isinstance(content_part, dict):
+                    if content_part.get('type') == 'text':
+                        tokens += len(encoding.encode(content_part['text']))
+                    elif content_part.get('type') == 'image_url':
+                        tokens += len(encoding.encode(content_part['image_url']['url']))
         else:
             tokens += len(encoding.encode(message['content']))
     return tokens
@@ -377,11 +378,11 @@ class Blackbox:
 
 # Pydantic Models
 class TextContent(BaseModel):
-    type: str = Field(..., description="Type of content, e.g., 'text'.")
+    type: Literal["text"] = Field(..., description="Type of content, e.g., 'text'.")
     text: str = Field(..., description="The text content.")
 
 class ImageURLContent(BaseModel):
-    type: str = Field(..., description="Type of content, e.g., 'image_url'.")
+    type: Literal["image_url"] = Field(..., description="Type of content, e.g., 'image_url'.")
     image_url: Dict[str, str] = Field(..., description="Dictionary containing the image URL.")
 
 Content = Union[TextContent, ImageURLContent]
@@ -393,7 +394,17 @@ class Message(BaseModel):
     @validator('content', pre=True)
     def validate_content(cls, v):
         if isinstance(v, list):
-            return [Content(**item) for item in v]
+            processed_content = []
+            for item in v:
+                if 'type' not in item:
+                    raise ValueError("Each content part must have a 'type' field.")
+                if item['type'] == 'text':
+                    processed_content.append(TextContent(**item))
+                elif item['type'] == 'image_url':
+                    processed_content.append(ImageURLContent(**item))
+                else:
+                    raise ValueError(f"Unsupported content type: {item['type']}")
+            return processed_content
         elif isinstance(v, str):
             return v
         else:
@@ -452,10 +463,10 @@ async def chat_completions(
                 # Convert list of content parts to a structured format
                 combined_content = []
                 for part in msg.content:
-                    if isinstance(part, dict) and part.get('type') == 'text':
-                        combined_content.append({"type": part['type'], "text": part['text']})
-                    elif isinstance(part, dict) and part.get('type') == 'image_url':
-                        combined_content.append({"type": part['type'], "image_url": part['image_url']})
+                    if isinstance(part, TextContent):
+                        combined_content.append({"type": part.type, "text": part.text})
+                    elif isinstance(part, ImageURLContent):
+                        combined_content.append({"type": part.type, "image_url": part.image_url})
                 processed_messages.append({"role": msg.role, "content": combined_content})
             else:
                 processed_messages.append({"role": msg.role, "content": msg.content})