Update app.py

app.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 from llama_cpp import Llama
 import time
 import io
+import asyncio
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -29,15 +30,21 @@ class RequestModel(BaseModel):
     logprobs: int = None
     stop: list = None
 
-# Response model
+# Response model based on OpenAI API structure
+class Choice(BaseModel):
+    text: str
+    index: int
+    logprobs: int = None
+    finish_reason: str
+
 class ResponseModel(BaseModel):
     id: str
     object: str
     created: int
     model: str
-    choices: list
+    choices: list[Choice]
 
-@app.post("/v1/completions")
+@app.post("/v1/completions", response_model=ResponseModel)
 async def create_completion(request: RequestModel):
     try:
         if request.stream:
@@ -79,12 +86,12 @@ async def create_completion(request: RequestModel):
                 created=int(time.time()),
                 model=request.model,
                 choices=[
-                    {
-                        "text": generated_text,
-                        "index": 0,
-                        "logprobs": None,
-                        "finish_reason": "stop"
-                    }
+                    Choice(
+                        text=generated_text,
+                        index=0,
+                        logprobs=None,
+                        finish_reason="stop"
+                    )
                 ]
             )
 
@@ -93,5 +100,4 @@ async def create_completion(request: RequestModel):
 
 if __name__ == "__main__":
     import uvicorn
-    import asyncio
     uvicorn.run(app, host="0.0.0.0", port=8000)