main.py

import io
import os
import time
from datetime import datetime
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException, File, UploadFile
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

app = FastAPI()

# Set up CORS (Cross-Origin Resource Sharing) for allowing requests from all origins
origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Define SQLAlchemy engine and metadata
DATABASE_URL = "sqlite:///./test.db"
engine = create_engine(DATABASE_URL)
metadata = MetaData()

# Define the document table schema
documents = Table(
    "documents",
    metadata,
    Column("id", Integer, primary_key=True),
    Column("filename", String),
    Column("upload_date", String),
    Column("content", String),
)

# Create the document table in the database
metadata.create_all(engine)
database = Database(DATABASE_URL)

# Define Pydantic model for the document
class Document(BaseModel):
    filename: str
    upload_date: str
    content: str

# Endpoint for uploading PDF files
@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...)):
    # Check if the uploaded file is a PDF
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed.")
    
    # Read content of the uploaded PDF file
    content = await file.read()
    
    # Extract text from the PDF
    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)
    
    # Insert the document data into the database
    query = documents.insert().values(
        filename=doc.filename,
        upload_date=doc.upload_date,
        content=doc.content
    )
    last_record_id = await database.execute(query)
    
    # Return the document object
    return {"id": last_record_id, "filename": doc.filename, "upload_date": doc.upload_date}

# Pydantic model for input data
class DataInput(BaseModel):
    responseData: str
    userInput: str

# Endpoint for processing user data
@app.post("/doc/")
async def process_data(data: DataInput):
    # Access responseData and userInput
    response_data = data.responseData
    user_input = data.userInput
    
    # 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"})]

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(dom)

    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,
        timeout=300  # Increase the timeout to 300 seconds
    )

    chain = load_qa_chain(llm, chain_type="stuff")

    # 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}

# To ensure the database connection is managed properly
@app.on_event("startup")
async def startup():
    await database.connect()

@app.on_event("shutdown")
async def shutdown():
    await database.disconnect()