import os
import tempfile
from typing import List, Annotated
from fastapi import FastAPI, File, UploadFile, Header, HTTPException
from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex, Document, Settings
from llama_index.vector_stores.postgres import PGVectorStore
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
import psycopg2
from sqlalchemy import make_url
from pdf2image import convert_from_path
from paddleocr import PaddleOCR
import io

app = FastAPI()


# Global variables
index = None
vector_store = None


# Initialize PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang='en')



# Helper function to perform OCR on a PDF
def perform_ocr_on_pdf(pdf_path):
    print(f"Processing PDF: {pdf_path}")

    try:
        images = convert_from_path(pdf_path, fmt="jpeg")
    except Exception as e:
        print(f"Error converting {pdf_path} to images: {e}")
        return ""

    full_text = ''

    for i, image in enumerate(images):
        try:
            # Convert PIL Image to bytes
            img_byte_arr = io.BytesIO()
            image.save(img_byte_arr, format='JPEG')
            img_byte_arr = img_byte_arr.getvalue()

            # Perform OCR directly on the image bytes
            result = ocr.ocr(img_byte_arr, cls=True)

            if result is None or len(result) == 0:
                print(f"Warning: OCR returned no results for page {i+1} of {pdf_path}")
                continue

            page_text = ''
            for line in result:
                if line is None:
                    print(f"Warning: Empty line in OCR result for page {i+1} of {pdf_path}")
                    continue
                for word_info in line:
                    if word_info is None or len(word_info) < 2:
                        print(f"Warning: Invalid word_info in OCR result for page {i+1} of {pdf_path}")
                        continue
                    page_text += word_info[1][0] + ' '
                page_text += '\n'

            full_text += f'--- Page {i+1} ---\n'
            full_text += page_text + '\n'

        except Exception as e:
            print(f"Error processing page {i+1} of {pdf_path}: {e}")

        print(f'Processed page {i+1} of {pdf_path}')

    if not full_text:
        print(f"Warning: No text extracted from {pdf_path}")

    return full_text

#ollama

ollama_llm = Ollama(base_url='https://artglobal-lisa.hf.space', model="mistral-nemo", request_timeout=120.0)
# Ollama embedding setup
ollama_embedding = OllamaEmbedding(base_url='https://artglobal-embed.hf.space', model_name="mxbai-embed-large", request_timeout=5000)

Settings.embed_model = ollama_embedding
Settings.llm = ollama_llm

# Database setup
connection_string = "postgresql://postgres:Tuvoc_2324@82.165.4.46:5432"
db_name = "vector_db"
def setup_database():
    global vector_store
    conn = psycopg2.connect(connection_string)
    conn.autocommit = True

    with conn.cursor() as c:
        c.execute(f"DROP DATABASE IF EXISTS {db_name}")
        c.execute(f"CREATE DATABASE {db_name}")

    url = make_url(connection_string)
    vector_store = PGVectorStore.from_params(
        database=db_name,
        host=url.host,
        password=url.password,
        port=url.port,
        user=url.username,
        table_name="lease_db",
        embed_dim=1024,
        hybrid_search=True,
        hnsw_kwargs={
            "hnsw_m": 16,
            "hnsw_ef_construction": 64,
            "hnsw_ef_search": 40,
            "hnsw_dist_method": "vector_cosine_ops",
        },
    )

# Call setup_database at startup
setup_database()

# Endpoint to upload and process PDF files
@app.post("/upload")
async def upload_files(files: List[UploadFile] = File(...)):
    global index
    try:
        documents = []
        for file in files:
            if not file.filename.lower().endswith('.pdf'):
                raise HTTPException(status_code=400, detail=f"File {file.filename} is not a PDF.")
            
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_pdf:
                temp_pdf.write(await file.read())
                temp_pdf_path = temp_pdf.name

            extracted_text = perform_ocr_on_pdf(temp_pdf_path)
            print(extracted_text)
            if extracted_text:
                documents.append(Document(
                    text=extracted_text,
                    metadata={"source": file.filename}
                ))

            os.unlink(temp_pdf_path)

        if documents:
            storage_context = StorageContext.from_defaults(vector_store=vector_store)
            index = VectorStoreIndex.from_documents(
                documents, storage_context=storage_context, show_progress=True, embed_model=ollama_embedding, llm=ollama_llm)
            return {"message": "Files uploaded, processed, and indexed successfully."}
        else:
            raise HTTPException(status_code=400, detail="No text extracted from uploaded files.")

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Query endpoint
@app.post("/query")
async def query_documents(query: str):
    global index
    if index is None:
        raise HTTPException(status_code=400, detail="No documents have been indexed. Please upload documents first.")
    
    try:
        query_engine = index.as_query_engine(llm=ollama_llm)
        response = query_engine.query(query)
        return {"answer": str(response)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

i am using this code to do following things


1. parse pdf
2. save pdf data in postgres with pgvector
3. perform rag to retrive data

now the issue i am facing is when i upload two pdfs and i ask give me list of documents it happily gives me list of 2 document name but when it increases like 5 or 10 it still gives me two name now the point of this is that once i ask give me list of document i can move forward and ask things like summarize this document etc etc. 

You understand what i want to do?