Update app.py

app.py

@@ -1,104 +1,161 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from tools import create_rag_tool, update_retriever
-from chatbot import app as app_graph
-from langchain_core.messages import HumanMessage
-import os
-from fastapi.responses import StreamingResponse, FileResponse
-from langchain_core.messages import AIMessage
-from fastapi.middleware.cors import CORSMiddleware
-import asyncio
-from pydantic import BaseModel
-from utils import TTS, STT
-
-
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-class TTSRequest(BaseModel):
-    text: str
-
-
-UPLOAD_DIR = "uploads"
-
-@app.get("/")
-def health():
-    return {'Status' : 'The api is live and running'}
-
-@app.post("/upload")
-async def upload_file(file: UploadFile = File(...)):
-    os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-    file_path = os.path.join(UPLOAD_DIR, file.filename)
-
-    with open(file_path, "wb") as f:
-        f.write(await file.read())
-
-    update_retriever(file_path)
-
-    return {
-        "status": "success",
-        "filename": file.filename
-    }
-
-
-@app.post("/chat")
-async def chat(message: str, session_id: str = "default"):
-
-    async def event_generator():
-        async for chunk in app_graph.astream(
-            {"messages": [HumanMessage(content=message)]},
-            config={"configurable": {"thread_id": session_id}},
-            stream_mode="messages"
-        ):
-            if len(chunk) >= 1:
-                message_chunk = chunk[0] if isinstance(chunk, tuple) else chunk
-                if hasattr(message_chunk, 'content') and message_chunk.content:
-                    data = str(message_chunk.content).replace("\n", "\\n")
-                    yield f"data: {data}\n\n"
-                    await asyncio.sleep(0.01)
-
-    return StreamingResponse(
-        event_generator(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "X-Accel-Buffering": "no",
-        },
-    )
-# ---------------- STT ---------------- #
-
-@app.post("/stt")
-async def transcribe_audio(file: UploadFile = File(...)):
-    try:
-        return await STT(file)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-
-@app.post("/tts")
-async def generate_tts(request: TTSRequest):
-    try:
-        if not request.text.strip():
-            raise HTTPException(status_code=400, detail="Text is empty")
-
-        audio_path = await TTS(text=request.text)
-
-        if not os.path.exists(audio_path):
-            raise HTTPException(status_code=500, detail="Audio file not created")
-
-        return FileResponse(
-            path=audio_path,
-            media_type="audio/mpeg",
-            filename="speech.mp3"
-        )
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from tools import update_retriever
+from chatbot import app as app_graph
+from langchain_core.messages import HumanMessage
+import os
+from fastapi.responses import StreamingResponse, FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+import asyncio
+from pydantic import BaseModel
+from utils import TTS, STT
+
+app = FastAPI()
+
+# Configure CORS - restrict in production
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # TODO: Restrict to specific origins in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class TTSRequest(BaseModel):
+    text: str
+
+UPLOAD_DIR = "uploads"
+ALLOWED_EXTENSIONS = {".pdf", ".txt", ".docx", ".md"}  # Add allowed file types
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB limit
+
+@app.get("/")
+def health():
+    return {'status': 'The api is live and running'}
+
+@app.post("/upload")
+async def upload_file(file: UploadFile = File(...)):
+    try:
+        # Validate file extension
+        file_ext = os.path.splitext(file.filename)[1].lower()
+        if file_ext not in ALLOWED_EXTENSIONS:
+            raise HTTPException(
+                status_code=400, 
+                detail=f"File type not allowed. Allowed types: {', '.join(ALLOWED_EXTENSIONS)}"
+            )
+        
+        # Create upload directory
+        os.makedirs(UPLOAD_DIR, exist_ok=True)
+        
+        # Read file content
+        content = await file.read()
+        
+        # Validate file size
+        if len(content) > MAX_FILE_SIZE:
+            raise HTTPException(
+                status_code=400,
+                detail=f"File too large. Maximum size: {MAX_FILE_SIZE / (1024*1024)}MB"
+            )
+        
+        # Save file
+        file_path = os.path.join(UPLOAD_DIR, file.filename)
+        with open(file_path, "wb") as f:
+            f.write(content)
+        
+        # Update retriever
+        update_retriever(file_path)
+        
+        return {
+            "status": "success",
+            "filename": file.filename,
+            "size": len(content)
+        }
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Upload failed: {str(e)}")
+
+@app.post("/chat")
+async def chat(message: str, session_id: str = "default"):
+    if not message.strip():
+        raise HTTPException(status_code=400, detail="Message cannot be empty")
+    
+    # Validate session_id (basic validation)
+    if not session_id or len(session_id) > 100:
+        raise HTTPException(status_code=400, detail="Invalid session_id")
+    
+    async def event_generator():
+        try:
+            async for chunk in app_graph.astream(
+                {"messages": [HumanMessage(content=message)]},
+                config={"configurable": {"thread_id": session_id}},
+                stream_mode="messages"
+            ):
+                if len(chunk) >= 1:
+                    message_chunk = chunk[0] if isinstance(chunk, tuple) else chunk
+                    if hasattr(message_chunk, 'content') and message_chunk.content:
+                        data = str(message_chunk.content).replace("\n", "\\n")
+                        yield f"data: {data}\n\n"
+                        await asyncio.sleep(0.01)
+        except Exception as e:
+            yield f"data: Error: {str(e)}\n\n"
+    
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+@app.post("/stt")
+async def transcribe_audio(file: UploadFile = File(...)):
+    try:
+        # Validate audio file type
+        if not file.content_type or not file.content_type.startswith("audio/"):
+            raise HTTPException(status_code=400, detail="Invalid audio file")
+        
+        return await STT(file)
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
+
+@app.post("/tts")
+async def generate_tts(request: TTSRequest):
+    try:
+        if not request.text.strip():
+            raise HTTPException(status_code=400, detail="Text is empty")
+        
+        # Limit text length to prevent abuse
+        if len(request.text) > 5000:
+            raise HTTPException(status_code=400, detail="Text too long (max 5000 characters)")
+        
+        audio_path = await TTS(text=request.text)
+        
+        if not os.path.exists(audio_path):
+            raise HTTPException(status_code=500, detail="Audio file not created")
+        
+        return FileResponse(
+            path=audio_path,
+            media_type="audio/mpeg",
+            filename="speech.mp3",
+            background=None  # Consider adding cleanup after response
+        )
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"TTS failed: {str(e)}")
+
+# Optional: Add cleanup endpoint for old files
+@app.on_event("startup")
+async def startup_event():
+    os.makedirs(UPLOAD_DIR, exist_ok=True)
+
+# Optional: Graceful shutdown
+@app.on_event("shutdown")
+async def shutdown_event():
+    # Clean up temporary files if needed
+    pass