Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks
+from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, HTTPException
 from fastapi.responses import JSONResponse
 from uuid import uuid4
 from vector_handler import create_vector_store, delete_vector_store, query_vector_store
@@ -6,6 +6,7 @@ from pdf_utils import extract_text_from_pdfs, split_text
 import os
 from dotenv import load_dotenv
 from datetime import datetime, timedelta
+from pinecone.openapi_support.exceptions import NotFoundException
 
 load_dotenv()
 
@@ -36,5 +37,21 @@ async def upload_pdfs(background_tasks: BackgroundTasks, files: list[UploadFile]
 
 @app.post("/query/")
 async def query_pdf(session_id: str = Form(...), question: str = Form(...)):
-    response = query_vector_store(session_id, question)
-    return JSONResponse(content={"answer": response})
+    try:
+        response = query_vector_store(session_id, question)
+        return JSONResponse(content={"answer": response})
+    except NotFoundException:
+        raise HTTPException(status_code=404, detail=f"Session '{session_id}' does not exist or has expired.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/clear/")
+async def clear_session(session_id: str = Form(...)):
+    try:
+        delete_vector_store(session_id)
+        return {"message": f"Session '{session_id}' has been cleared."}
+    except NotFoundException:
+        raise HTTPException(status_code=404, detail=f"Session '{session_id}' does not exist.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))