Update app.py

app.py

@@ -5,6 +5,7 @@ from time import time
 import logging
 import llama_cpp
 import llama_cpp.llama_tokenizer
+from pydantic import BaseModel
 
 llama = llama_cpp.Llama.from_pretrained(
     repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF",
@@ -57,23 +58,25 @@ def health():
     return {"status": "ok"}
 
 
+class GenModel(BaseModel):
+    question: str
+    system: str = "You are a story writing assistant."
+    temperature: float = 0.7
+    seed: int = 42
+    
 # Chat Completion API
 @app.get("/generate_stream")
-async def complete(
-    question: str,
-    system: str = "You are a story writing assistant.",
-    temperature: float = 0.7,
-    seed: int = 42,
+async def complete(gen:GenModel
 ) -> dict:
     try:
         st = time()
         output = llama.create_chat_completion(
             messages=[
                 {"role": "system", "content": system},
-                {"role": "user", "content": question},
+                {"role": "user", "content": gen['question']},
             ],
-            temperature=temperature,
-            seed=seed,
+            temperature=gen['temperature'],
+            seed=gen['seed'],
             #stream=True
         )
         """