commit 000000239

app.py

@@ -2,8 +2,7 @@ import os
 import time
 import torch
 import re
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from pydantic import BaseModel
@@ -24,15 +23,15 @@ app = FastAPI(
     title="ChatMate Real-Time API",
     description="LangChain + DuckDuckGo + Phi-4",
     version="1.0",
-    docs_url="/apidocs",  # Swagger UI at /apidocs
-    redoc_url="/redoc"    # ReDoc at /redoc
+    docs_url="/apidocs",
+    redoc_url="/redoc"
 )
 
 # ✅ Static + templates
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
-# Enable CORS (important for browser clients)
+# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -94,24 +93,38 @@ class ChatRequest(BaseModel):
     message: str
     history: list = []
 
-# ---------------- FastAPI route ----------------
 # ---------------- Routes ----------------
 @app.get("/", summary="Serve homepage")
 async def home(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
-@app.post("/chat-stream", summary="Stream assistant reply", tags=["Chat"])
-async def chat_stream(body: ChatRequest):
+@app.websocket("/chat-stream")
+async def chat_stream_ws(websocket: WebSocket):
     """
-    Stream the AI assistant's reply token-by-token.
+    WebSocket endpoint to stream tokens to the client.
     """
-    def generate():
-        reply = generate_full_reply(body.message, body.history)
-        for token in reply:
-            yield token
-            time.sleep(0.05)
+    await websocket.accept()
+    try:
+        while True:
+            data = await websocket.receive_text()
+            try:
+                req = ChatRequest.parse_raw(data)
+            except Exception as e:
+                await websocket.send_text(f"[Error] Invalid request: {str(e)}")
+                continue
 
-    return StreamingResponse(generate(), media_type="text/plain")
+            reply = generate_full_reply(req.message, req.history)
+
+            # Stream token by token
+            for token in reply:
+                await websocket.send_text(token)
+                await asyncio.sleep(0.05)
+
+            # Optional: send special marker to indicate completion
+            await websocket.send_text("[END]")
+
+    except WebSocketDisconnect:
+        print("Client disconnected from WebSocket.")
 
 # ---------------- Startup warm-up ----------------
 @app.on_event("startup")
@@ -120,10 +133,6 @@ async def warmup_model():
     _ = generate_full_reply("Hello", [])
 
 # ---------------- Run with Uvicorn ----------------
-# In Hugging Face Spaces, just run: uvicorn app:app --host 0.0.0.0 --port 7860
 if __name__ == "__main__":
-    # Hugging Face Spaces usually expects port 7860
     port = int(os.environ.get("PORT", 7860))
-
-    # Run using uvicorn for FastAPI/Flask with ASGI wrapper
-    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)
+    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)