Update main.py

main.py

@@ -1,19 +1,24 @@
 import io
 import os
+import time
+from datetime import datetime
 from pydantic import BaseModel
 from fastapi import FastAPI, HTTPException, File, UploadFile
-from pdfminer.high_level import extract_text
-from datetime import datetime
 from fastapi.middleware.cors import CORSMiddleware
+from pdfminer.high_level import extract_text
 from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
 from databases import Database
-
 from langchain.chains.question_answering import load_qa_chain
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.docstore.document import Document as LangchainDocument
+from requests.exceptions import HTTPError
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 app = FastAPI()
 
@@ -23,7 +28,7 @@ app.add_middleware(
     CORSMiddleware,
     allow_origins=origins,
     allow_credentials=True,
-    allow_methods=["GET", "POST", "PUT", "DELETE"],
+    allow_methods=["*"],
     allow_headers=["*"],
 )
 
@@ -63,8 +68,11 @@ async def upload_pdf(file: UploadFile = File(...)):
     content = await file.read()
     
     # Extract text from the PDF
-    with io.BytesIO(content) as pdf_file:
-        text_content = extract_text(pdf_file)
+    try:
+        with io.BytesIO(content) as pdf_file:
+            text_content = extract_text(pdf_file)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to extract text from PDF: {str(e)}")
     
     # Create a document object
     doc = Document(filename=file.filename, upload_date=str(datetime.now()), content=text_content)
@@ -78,7 +86,7 @@ async def upload_pdf(file: UploadFile = File(...)):
     last_record_id = await database.execute(query)
     
     # Return the document object
-    return doc
+    return {"id": last_record_id, "filename": doc.filename, "upload_date": doc.upload_date}
 
 # Pydantic model for input data
 class DataInput(BaseModel):
@@ -94,6 +102,9 @@ async def process_data(data: DataInput):
     
     # Load required models and components from Langchain library
     HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+    if not HUGGINGFACEHUB_API_TOKEN:
+        raise HTTPException(status_code=500, detail="HuggingFace API token not found.")
+    
     os.environ['HUGGINGFACEHUB_API_TOKEN'] = HUGGINGFACEHUB_API_TOKEN
 
     dom = [LangchainDocument(page_content=response_data, metadata={"source": "local"})]
@@ -104,16 +115,28 @@ async def process_data(data: DataInput):
     embeddings = HuggingFaceEmbeddings()
     db = FAISS.from_documents(docs, embeddings)
     
+    # Configure the HuggingFaceEndpoint with an increased timeout
     llm = HuggingFaceEndpoint(
         repo_id="google/flan-t5-xxl",
-        temperature=0.5
+        temperature=0.5,
+        timeout=300  # Increase the timeout to 300 seconds
     )
 
     chain = load_qa_chain(llm, chain_type="stuff")
 
-    # Perform similarity search and question answering
-    dm = db.similarity_search(user_input)
-    result = chain.run(input_documents=dm, question=user_input)
+    # Implement a retry mechanism
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            # Perform similarity search and question answering
+            dm = db.similarity_search(user_input)
+            result = chain.run(input_documents=dm, question=user_input)
+            break  # Break out of the loop if successful
+        except HTTPError as e:
+            if attempt < max_retries - 1:
+                time.sleep(10)  # Wait for 10 seconds before retrying
+            else:
+                raise HTTPException(status_code=503, detail="Service Unavailable. Please try again later.") from e
     
     return {"result": result}