Update main.py

main.py

@@ -44,14 +44,14 @@ async def index():
     """
     return HTMLResponse(content=html_content, status_code=200)
 
-class ChatCompletionRequest(BaseModel):
+class ChatCompletionRequestV0(BaseModel):
     prompt: str
 
 class Message(BaseModel):
     role: str
     content: str
 
-class ChatCompletionRequestV2(BaseModel):
+class ChatCompletionRequest(BaseModel):
     messages: List[Message]
     max_tokens: int = 100
 
@@ -89,7 +89,7 @@ async def chat(request: ChatCompletionRequest):
     return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
 
 @app.post("/v0/chat/completions")
-async def chat(request: ChatCompletionRequest, response_mode=None):
+async def chat(request: ChatCompletionRequestV0, response_mode=None):
     tokens = llm.tokenize(request.prompt)
     async def server_sent_events(chat_chunks, llm):
         for chat_chunk in llm.generate(chat_chunks):