Update app.py

app.py

@@ -96,6 +96,46 @@ async def complete(
             status_code=500, content={"message": "Internal Server Error"}
         )
 
+# Chat Completion API
+@app.post("/generate")
+async def complete(
+    question: str,
+    system: str = "You are a story writing assistant.",
+    temperature: float = 0.7,
+    seed: int = 42,
+) -> dict:
+    try:
+        st = time()
+        output = llama.create_chat_completion(
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user", "content": question},
+            ],
+            temperature=temperature,
+            seed=seed,
+            stream=True
+        )
+        """
+        for chunk in output:
+            
+            delta = chunk['choices'][0]['delta']
+            if 'role' in delta:
+                print(delta['role'], end=': ')
+            elif 'content' in delta:
+                print(delta['content'], end='')
+            
+            print(chunk)
+        """
+        et = time()
+        output["time"] = et - st
+        return output
+    except Exception as e:
+        logger.error(f"Error in /complete endpoint: {e}")
+        return JSONResponse(
+            status_code=500, content={"message": "Internal Server Error"}
+        )
+
+
 
 if __name__ == "__main__":
     import uvicorn