Update app.py

app.py

@@ -6,11 +6,12 @@ import logging
 from datetime import datetime
 from typing import List, Dict, Optional
 from fastapi import FastAPI, HTTPException, UploadFile, File
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import markdown
 import PyPDF2
+import asyncio
 
 # Setup logging
 logging.basicConfig(
@@ -102,31 +103,44 @@ async def startup_event():
     except Exception as e:
         logger.error(f"Startup error: {str(e)}")
 
-@app.post("/chat")
-async def chat_endpoint(request: ChatRequest):
-    try:
-        raw_response = agent.chat(
-            message=request.message,
-            history=request.history,
-            temperature=request.temperature,
-            max_new_tokens=request.max_new_tokens
-        )
-        formatted_response = {
-            "raw": raw_response,
-            "clean": clean_text_response(raw_response),
-            "structured": structure_medical_response(raw_response),
-            "html": markdown.markdown(raw_response)
-        }
-        return JSONResponse({
-            "status": "success",
-            "format": request.format,
-            "response": formatted_response.get(request.format, formatted_response["clean"]),
-            "timestamp": datetime.now().isoformat(),
-            "available_formats": list(formatted_response.keys())
-        })
-    except Exception as e:
-        logger.error(f"Chat error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+@app.post("/chat-stream")
+async def chat_stream_endpoint(request: ChatRequest):
+    async def token_stream():
+        try:
+            conversation = []
+            conversation.append({"role": "system", "content": agent.chat_prompt})
+            if request.history:
+                for msg in request.history:
+                    conversation.append({"role": msg["role"], "content": msg["content"]})
+            conversation.append({"role": "user", "content": request.message})
+
+            input_ids = agent.tokenizer.apply_chat_template(
+                conversation,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(agent.device)
+
+            streamer = agent.model.generate(
+                input_ids,
+                do_sample=True,
+                temperature=request.temperature,
+                max_new_tokens=request.max_new_tokens,
+                pad_token_id=agent.tokenizer.eos_token_id,
+                return_dict_in_generate=True,
+                output_scores=False
+            )
+
+            output = agent.tokenizer.decode(streamer["sequences"][0][input_ids.shape[1]:], skip_special_tokens=True)
+
+            for chunk in output.split():
+                yield chunk + " "
+                await asyncio.sleep(0.05)
+
+        except Exception as e:
+            logger.error(f"Streaming chat error: {str(e)}")
+            yield f"\n⚠️ Error: {str(e)}"
+
+    return StreamingResponse(token_stream(), media_type="text/plain")
 
 @app.post("/upload")
 async def upload_file(file: UploadFile = File(...)):