Update api/models.py

api/models.py

@@ -1,3 +1,5 @@
+# api/models.py
+
 from typing import List, Optional, Union
 from pydantic import BaseModel
 
@@ -15,13 +17,18 @@ class ChatRequest(BaseModel):
     max_tokens: Optional[int] = 99999999
 
 class ImageResponseModel(BaseModel):
-    images: str
+    images: List[str]
     alt: str
 
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: dict
+    finish_reason: Optional[str]
+
 class ChatCompletionResponse(BaseModel):
     id: str
     object: str
     created: int
     model: str
-    choices: List[dict]
+    choices: List[ChatCompletionChoice]
     usage: Optional[dict] = None