cancel signal handling

.gitignore

@@ -2,4 +2,6 @@
 cache
 users.db
 .env
-__pycache__
+__pycache__
+-d
+-H

main.py

@@ -4,7 +4,7 @@ from fastapi.staticfiles import StaticFiles
 from jose import JWTError
 from schemas import UserRegister, TokenResponse, RefreshTokenRequest, QueryInput
 from auth import  register_user, get_db, authenticate_user, create_token, verify_token, verify_access_token, Session
-from utils import generate_stream, generate_response
+from utils import generate_stream, generate_response, shutdown_event
 from fastapi.security import OAuth2PasswordRequestForm
 from pathlib import Path
 from datetime import timedelta
@@ -103,11 +103,19 @@ async def generate(
     username: str = Depends(verify_access_token),
     stream: bool = Query(False, description="Enable streaming response"),
 ):
-    """API endpoint that supports both streaming and non-streaming responses."""
-    if stream:
-        return StreamingResponse(generate_stream(query_input.query), media_type="text/event-stream")
-    else:
-        return JSONResponse(await generate_response(query_input.query))
+    """Handles both streaming and non-streaming responses, with shutdown detection."""
+    if shutdown_event.is_set():
+        return JSONResponse({"message": "Server shutting down..."})
+
+    try:
+        if stream:
+            return StreamingResponse(generate_stream(query_input.query), media_type="text/event-stream")
+        else:
+            return JSONResponse(await generate_response(query_input.query))
+    except Exception as e:
+        logging.error(f"Error in generate endpoint: {e}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
 # WebSocket endpoint for streaming
 @app.on_event("startup")
 async def startup_event():

utils.py

@@ -1,9 +1,21 @@
 import asyncio
 import ollama
 import json
+import signal
+
+shutdown_event = asyncio.Event()
+
+def shutdown_handler(sig, frame):
+    """Handles shutdown signals like Ctrl + C."""
+    print("\n⛔ Shutdown requested! Stopping API...")
+    shutdown_event.set()
+
+# Attach signal handlers for graceful termination
+signal.signal(signal.SIGINT, shutdown_handler)
+signal.signal(signal.SIGTERM, shutdown_handler)
 
 async def generate_stream(query: str):
-    """Generates streamed responses from Ollama using LLaMA 3 in JSON format."""
+    """Generates streamed responses with cancellation support."""
     try:
         stream = ollama.chat(
             model="llama3.2",
@@ -12,9 +24,18 @@ async def generate_stream(query: str):
         )
 
         for chunk in stream:
+            if shutdown_event.is_set():
+                print("⛔ Stopping content generation...")
+                break  # Exit loop when shutdown is requested
+
             if "message" in chunk and "content" in chunk["message"]:
                 response_data = json.dumps({"content": chunk["message"]["content"]})
-                yield f"data: {response_data}\n\n"  # SSE format
+                yield f"data: {response_data}\n\n"
+                await asyncio.sleep(0.1)  # Allow graceful processing
+
+    except asyncio.CancelledError:
+        print("⛔ Stream cancelled by user.")
+        raise  # Propagate cancellation
 
     except Exception as e:
         error_data = json.dumps({"error": str(e)})