Update app.py

app.py

@@ -1,5 +1,4 @@
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import EventSourceResponse
+from fastapi import FastAPI, HTTPException, Request, StreamingResponse
 from llama_cpp import Llama
 import time
 import os
@@ -14,6 +13,19 @@ def typewrite(text, delay=0.01):
         time.sleep(delay)
     print(end='', flush=True)  # Print newline to move to the next line
 
+async def generate_text_stream(ask):
+    prompt = f"Llama-2-Chat [INST] <<SYS>>You're an assistant named Tusti. You are Developed by Aritra Roy. Don't share any false information.<</SYS>> {ask} [/INST]"
+
+    output_stream = llm(prompt, max_tokens=1024, echo=False, temperature=0.2, top_p=0.1, stream=True)
+
+    while True:
+        try:
+            chunk = next(output_stream)
+            if chunk.get('choices') and chunk['choices'][0].get('text'):
+                yield chunk['choices'][0]['text']
+        except StopIteration:
+            break
+
 @app.post("/chat")
 async def chat(request: Request):
     try:
@@ -28,31 +40,13 @@ async def chat(request: Request):
 
         ask = user_input
 
-        prompt = f"Llama-2-Chat [INST] <<SYS>>You're an assistant named Tusti. You are Developed by Aritra Roy. Don't share any false information.<</SYS>> {ask} [/INST]"
-
-        output_stream = llm(prompt, max_tokens=1024, echo=False, temperature=0.2, top_p=0.1, stream=True)
-
-        async def generate_chunks():
-            nonlocal output_stream
-            response_text_chunk = ''
-            try:
-                while True:
-                    chunk = next(output_stream)
-                    if chunk.get('choices') and chunk['choices'][0].get('text'):
-                        response_text_chunk += chunk['choices'][0]['text']
-                        typewrite(response_text_chunk, delay=0.00)
-                        yield f"data: {response_text_chunk}\n\n"
-            except StopIteration:
-                pass
-
         if ask == 'clear':
             os.system("cls")
 
-        return EventSourceResponse(generate_chunks())
+        return StreamingResponse(generate_text_stream(ask))
 
     except Exception as e:
-        print(f"Exception: {e}")
-        raise HTTPException(status_code=500, detail="Internal Server Error")
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn