Update app.py

app.py

@@ -1,20 +1,21 @@
 from fastapi import FastAPI, UploadFile, File, HTTPException
-from tools import update_retriever
+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()
 
-# Configure CORS - restrict in production
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # TODO: Restrict to specific origins in production
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -23,82 +24,46 @@ app.add_middleware(
 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'}
+    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)}")
+    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"):
-    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"
-    
+        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",
@@ -108,54 +73,32 @@ async def chat(message: str, session_id: str = "default"):
             "X-Accel-Buffering": "no",
         },
     )
+# ---------------- STT ---------------- #
 
 @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)}")
+        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")
-        
-        # 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
+            filename="speech.mp3"
         )
-    
-    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
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))