Update app.py

app.py

@@ -1,14 +1,29 @@
 import time
 import json
+from typing import List, Literal
 
 from fastapi import FastAPI
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel
 
 from gradio_client import Client
 
 app = FastAPI()
 client = Client("AWeirdDev/mistral-7b-instruct-v0.2")
 
+class Message(BaseModel):
+    role: Literal["user", "assistant", "system"]
+    content: str
+
+class Payload(BaseModel):
+    stream: bool = False
+    model: Literal["mistral-7b-instruct-v0.2"] = "mistral-7b-instruct-v0.2"
+    messages: List[Message]
+    temperature: float
+    presence_penalty: float
+    frequency_penalty: float
+    top_p: float
+
 async def stream(iter):
     while True:
         try:
@@ -37,18 +52,46 @@ def make_chunk_obj(i, delta, fr):
 
 @app.get('/')
 async def index():
-    return { "message": "hello" }
+    return JSONResponse({ "message": "hello" })
 
 @app.post('/chat/completions')
-async def c_cmp():
+async def c_cmp(payload: Payload):
+    if not payload.stream:
+        return JSONResponse(
+            {
+                "id": str(time.time_ns()),
+                "object": "chat.completion",
+                "created": round(time.time()),
+                "model": payload.model,
+                "system_fingerprint": "wtf",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": client.predict(
+                                payload.messages.model_dump_json(),
+                                payload.temperature,
+                                4096,
+                                payload.top_p,
+                                payload.presence_penalty,
+                                api_name="/chat"
+                            )
+                        }
+                    }
+                ]
+            }
+        )
+        
+        
     def streamer():
         text = ""
         result = client.submit(
-            "Hello!!",
-            0.9,	# float (numeric value between 0.0 and 1.0) in 'Temperature' Slider component
+            payload.messages.model_dump_json(),
+            payload.temperature,	# float (numeric value between 0.0 and 1.0) in 'Temperature' Slider component
             4096,	# float (numeric value between 0 and 1048) in 'Max new tokens' Slider component
-            .9,	# float (numeric value between 0.0 and 1) in 'Top-p (nucleus sampling)' Slider component
-            1,	# float (numeric value between 1.0 and 2.0) in 'Repetition penalty' Slider component
+            payload.top_p,	# float (numeric value between 0.0 and 1) in 'Top-p (nucleus sampling)' Slider component
+            payload.presence_penalty,	# float (numeric value between 1.0 and 2.0) in 'Repetition penalty' Slider component
             api_name="/chat"
         )
         for i, item in enumerate(result):
@@ -58,7 +101,7 @@ async def c_cmp():
             )
             text = item
 
-        yield "data: " + json.dumps(make_chunk_obj(i, delta, "stop"))
+        yield "data: " + json.dumps(make_chunk_obj(i, "", "stop"))
         yield "data: [END]"
 
     return StreamingResponse(streamer())