Update main.py

main.py

@@ -6,8 +6,8 @@ from ctransformers import AutoModelForCausalLM
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
-from pydantic import BaseModel
-
+from pydantic import BaseModel, Field
+from typing_extensions import Literal
 from dialogue import DialogueTemplate
 
 llm = AutoModelForCausalLM.from_pretrained("NeoDim/starchat-alpha-GGML",
@@ -116,12 +116,14 @@ async def chat(prompt = "<|user|> Write an express server with server sent event
     return EventSourceResponse(server_sent_events(tokens, llm))
 
 
-class ChatCompletion(BaseModel):
-    role: str
-    content: str
+class ChatCompletionRequestMessage(BaseModel):
+    role: Literal["system", "user", "assistant"] = Field(
+        default="user", description="The role of the message."
+    )
+    content: str = Field(default="", description="The content of the message.")
 
 class ChatCompletionRequest(BaseModel):
-    messages: List[ChatCompletion]
+    messages: List[ChatCompletionRequestMessage]
 
 system_message = "Below is a conversation between a human user and a helpful AI coding assistant."