Update app.py

app.py

@@ -2,13 +2,12 @@ import uvicorn
 from fastapi.staticfiles import StaticFiles
 import hashlib
 from enum import Enum
-from fastapi import FastAPI,Header, Query,Depends,HTTPException
-from paddleocr import PaddleOCR, PPStructure, save_structure_res
+from fastapi import FastAPI, Header, Query, Depends, HTTPException
 from PIL import Image
 import io
-import numpy as np
 import fitz  # PyMuPDF for PDF handling
 import logging
+from pymongo import MongoClient
 
 import boto3
 import openai
@@ -17,29 +16,43 @@ import traceback  # For detailed traceback of errors
 import re
 import json
 from dotenv import load_dotenv
-import uvicorn
 import base64
+from bson.objectid import ObjectId
 
+db_client = None
 load_dotenv()
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-#Set up MongoDB
+# MongoDB Configuration
 MONGODB_URI = os.getenv("MONGODB_URI")
 DATABASE_NAME = os.getenv("DATABASE_NAME")
-COLLECTION_NAME = os.getenv("COLLECTION_NAME",DocumentExtractionSchema)
+COLLECTION_NAME = os.getenv("COLLECTION_NAME")
+SCHEMA = os.getenv("SCHEMA")
+
+# Check if environment variables are set
+if not MONGODB_URI:
+    raise ValueError("MONGODB_URI is not set. Please add it to your secrets.")
+
+# Initialize MongoDB Connection
+db_client = MongoClient(MONGODB_URI)
+db = db_client[DATABASE_NAME]
+invoice_collection = db[COLLECTION_NAME]
+schema_collection = db[SCHEMA]
 
 app = FastAPI(docs_url='/')
 use_gpu = False
 output_dir = 'output'
 
-if not MONGODB_URI:
-    raise ValueError("MONGODB_URL is not set. Please add it to Hugging Face secrets.")
-
-# Initialize PaddleOCR
-ocr = PaddleOCR(use_angle_cls=True, lang='en')
+@app.on_event("startup")
+def startup_db():
+    try:
+        db_client.server_info()
+        logger.info("MongoDB connection successful")
+    except Exception as e:
+        logger.error(f"MongoDB connection failed: {str(e)}")
 
 # AWS S3 Configuration
 API_KEY = os.getenv("API_KEY")
@@ -58,8 +71,7 @@ s3_client = boto3.client(
 )
 
 # Function to fetch file from S3
-
-def fetch_file_from_s3_file(file_key):
+def fetch_file_from_s3(file_key):
     try:
         response = s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=file_key)
         content_type = response['ContentType']  # Retrieve MIME type
@@ -68,98 +80,95 @@ def fetch_file_from_s3_file(file_key):
     except Exception as e:
         raise Exception(f"Failed to fetch file from S3: {str(e)}")
 
-# Function to summarize text using OpenAI GPT
-def summarize_text(text):
-    system_prompt = """You are tasked with extracting and structuring all relevant information from an invoice in a standardized JSON format for storing invoice headers and line items. The invoice headers should include the following details:
-
-Vendor Information:
-
-Vendor Name
-Vendor Address
-Vendor GST No.
-Invoice Details:
-
-Invoice No.
-Invoice Date → Considered as InvoiceDate (formatted as dd-MMM-yyyy).
-Invoice Currency/Currency
-Base Amount/Amount
-Tax Amount
-Total Invoice Amount
-Type of Invoice (e.g., "Tax Invoice", "Proforma Invoice", etc.)
-Customer Information:
-
-Customer Name
-Customer Address
-Customer GST No.
-Shipping and References:
-
-MBL No./HBL No./Container No./Shipping Bill No./Shipper Invoice No./Manifest No./MAWB/HAWB/OBL No./Bill of Lading Number/REF/Ocean Bill of Lading/House Bill of Lading/BL No./Job No. → Considered as RefNo.
-Shipping Order
-You should extract this data and structure it into a table-like format in the following JSON format:
-{
-  "invoice_headers": {
-    "VendorName": "",
-    "VendorAddress": "",
-    "VendorGSTNo": "",
-    "InvoiceNo": "",
-    "InvoiceDate": "",
-    "InvoiceCurrency": "",
-    "BaseAmount": "",
-    "TaxAmount": "",
-    "TotalInvoiceAmt": "",
-    "TypeofInvoice": "",
-    "CustomerName": "",
-    "CustomerAddress": "",
-    "CustomerGSTNO": "",
-    "RefNo": "",
-    "ShippingOrder": ""
-  },
-  "line_items": [
-    {    
-      "Description": "",
-      "TaxPercentage": "",
-      "TaxAmount": "",
-      "Amount": 0
-    }
-  ]
-}
-Guidelines for Processing:
-
-Ensure accurate extraction of data from the invoice by recognizing alternative naming conventions (e.g., Bill to, Taxpayer Name, etc.).
-Convert the Invoice Date to the specified dd-MMM-yyyy format.
-Use the correct currency and amounts for each invoice field.
-For each line item, provide the Description, Tax Percentage, Tax Amount, and Amount.
-If certain values are missing or not applicable, leave them empty or set them as null where necessary.
-This JSON format will be used to store and manage invoices in a structured and uniform way. Please ensure only return JSON format. No extra content should not provide."""
+def extract_pdf_text(file_data):
+    """
+    Extracts text from a PDF file using PyMuPDF (fitz).
+    """
     try:
-        response = openai.ChatCompletion.create(
-            model="gpt-4o-mini",
-            messages=[
-               {"role": "system", "content": system_prompt},
-               {"role": "user", "content": f"{text}"}
-             ],
-             temperature=0.5,
-             max_tokens=16384
-        )
-        content = response.choices[0].message.content.strip()
-        print("Before content:", content)
-        cleaned_content = re.sub(r'^.*```json\n', '', content)  # Remove '```json\n' at the beginning
-        cleaned_content = re.sub(r'\n```$', '', cleaned_content)  # Remove '\n```' at the end
-
-        # Step 2: Parse the cleaned content as JSON
-        #parsed_content = json.loads(cleaned_content)
-
-        # Step 3: Print the parsed JSON object
+        pdf_document = fitz.open(stream=file_data, filetype="pdf")
+        text = "\n".join([page.get_text("text") for page in pdf_document])
+        return text
+    except Exception as e:
+        logger.error(f"PDF Extraction Error: {e}")
+        return None
+
+# Function to summarize text using OpenAI GPT
+def extract_invoice_data(file_data, content_type, json_schema):
+    """
+    Extracts data from a PDF or image and returns structured JSON based on the provided schema.
+    """
+    system_prompt = "You are an expert in document data extraction. Extract relevant fields from the document and return structured JSON based on the provided schema."
+
+    # Convert file to Base64
+    base64_encoded = base64.b64encode(file_data).decode('utf-8')
+    base64dataresp = f"data:{content_type};base64,{base64_encoded}"
+
+    # Handle PDF Extraction & Format to JSON Schema
+    if content_type == "application/pdf":
+        extracted_text = extract_pdf_text(file_data)
+        if not extracted_text:
+            return {"error": "Failed to extract text from PDF"}, base64dataresp
+
         try:
-            parsed_content = json.loads(cleaned_content)
-            return parsed_content
-        except json.JSONDecodeError as e:
-            print("Error parsing JSON:", e)
-            # Optionally, print the cleaned content to debug
-            print("Cleaned content:", cleaned_content)
-            return None
+            # Send extracted text to OpenAI for structured JSON conversion
+            response = openai.ChatCompletion.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": extracted_text}
+                ],
+                response_format={"type": "json_schema", "json_schema": json_schema},
+                temperature=0.5,
+                max_tokens=16384
+            )
+
+            parsed_content = json.loads(response.choices[0].message.content.strip())
+            return parsed_content, base64dataresp  # Return structured JSON
+        except Exception as e:
+            logger.error(f"Error in OpenAI text-to-JSON conversion: {e}")
+            return {"error": str(e)}, base64dataresp
+
+    # Handle Image Extraction using OpenAI Vision API
+    elif content_type.startswith("image/"):
+        try:
+            response = openai.ChatCompletion.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": f"data:{content_type};base64,{base64_encoded}"
+                                }
+                            }
+                        ]
+                    }
+                ],
+                response_format={"type": "json_schema", "json_schema": json_schema},
+                temperature=0.5,
+                max_tokens=16384
+            )
+
+            parsed_content = json.loads(response.choices[0].message.content.strip())
+            return parsed_content, base64dataresp  # Return structured JSON
+        except Exception as e:
+            logger.error(f"Error in OpenAI image processing: {e}")
+            return {"error": str(e)}, base64dataresp
+
+    else:
+        raise ValueError(f"Unsupported content type: {content_type}")
+
+def get_content_type_from_s3(file_key):
+    """Fetch the content type (MIME type) of a file stored in S3."""
+    try:
+        response = s3_client.head_object(Bucket=S3_BUCKET_NAME, Key=file_key)
+        return response.get('ContentType', 'application/octet-stream')  # Default to binary if not found
     except Exception as e:
-        return f"Error in summarization: {str(e)}"
+        raise Exception(f"Failed to get content type from S3: {str(e)}")
+
 # Dependency to check API Key
 def verify_api_key(api_key: str = Header(...)):
     if api_key != API_KEY:
@@ -167,76 +176,68 @@ def verify_api_key(api_key: str = Header(...)):
 
 @app.get("/")
 def read_root():
-    return {"message": "Welcome to the PaddleOCR with S3 and GPT Summarization API!"}
+    return {"message": "Welcome to the Invoice Summarization API!"}
 
 @app.get("/ocr/extraction")
-def ocr_from_s3(api_key: str = Depends(verify_api_key),file_key: str = Query(..., description="S3 file key for the file"),document_type: str = Query(...,description="JobDocument Doc_Type"),entityrefkey: str = Query(..., description="JobDocument JOD_PK")):
-    """
-    Perform OCR on a file (PDF or Image) stored in S3 and summarize the text using GPT.
-    """
+def extract_text_from_file(
+    api_key: str = Depends(verify_api_key), 
+    file_key: str = Query(..., description="S3 file key for the file"),
+    document_type: str = Query(..., description="Type of document"),
+    entity_ref_key: str = Query(..., description="Entity Reference Key")
+):
+    """Extract structured data from a PDF or Image stored in S3."""
     try:
-        # Fetch file from S3
-        file_data, content_type = fetch_file_from_s3_file(file_key)
-
-        extracted_text = []
-        base64Data = base64.b64encode(file_data).decode('utf-8')
-        # Determine file type based on MIME type
-        if content_type.startswith("image/"):  # Image file
-            image = Image.open(io.BytesIO(file_data)).convert("RGB")  # Use BytesIO stream directly
-            image_np = np.array(image)  # Convert to NumPy array
-            result = ocr.ocr(image_np, cls=True)
-            base64DataResp = f"data:image/{content_type.lower()};base64,{base64Data}"
-            # Extract text from OCR results
-            for line in result:
-                for word_info in line:
-                    extracted_text.append(word_info[1][0])
-
-        elif content_type == "application/pdf":  # PDF file
-            # Open PDF using PyMuPDF
-            pdf_document = fitz.open(stream=io.BytesIO(file_data), filetype="pdf")
-
-            extracted_text = []
-
-            # Process each page in the PDF
-            for page_number in range(len(pdf_document)):
-                page = pdf_document[page_number]
-
-                # Render the page as an image
-                pix = page.get_pixmap()
-                image = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")
-
-                # Convert Pillow image to NumPy array (for PaddleOCR compatibility)
-                image_np = np.array(image)
-
-                # Run OCR on the image
-                result = ocr.ocr(image_np, cls=True)
-                for line in result:
-                    for word_info in line:
-                        extracted_text.append(word_info[1][0])
-
-            pdf_document.close()
-            base64DataResp = f"data:application/pdf;base64,{base64Data}"
-        else:
-            return {"error": f"Unsupported file type: {content_type}"}
-
-        # Combine extracted text
-        full_text = " ".join(extracted_text)
-
-        # Summarize the extracted text
-        summary = summarize_text(full_text)
-
-        return {
+        existing_document = invoice_collection.find_one({"entityrefkey": entity_ref_key})
+        if existing_document:
+            existing_document["_id"] = str(existing_document["_id"])
+            return {
+                "message": "Document Retrieved from MongoDB.",
+                "document": existing_document
+            }
+
+        # Fetch JSON schema for the document type
+        schema_doc = schema_collection.find_one({"document_type": document_type})
+        if not schema_doc:
+            raise ValueError("No schema found for the given document type")
+
+        json_schema = schema_doc.get("json_schema")
+        if not json_schema:
+            raise ValueError("Schema is empty or not properly defined.")    
+
+        # Retrieve file from S3
+        content_type = get_content_type_from_s3(file_key)
+        file_data, _ = fetch_file_from_s3(file_key)
+
+        # Extract structured data from the document
+        extracted_data, base64dataresp = extract_invoice_data(file_data, content_type, json_schema)
+
+        # Build and store document in MongoDB
+        document = {
             "file_key": file_key,
             "file_type": content_type,
-            "document_type":document_type,
-            "entityrefkey":entityrefkey,
-            "base64DataResp":base64DataResp,
-            "extracted_text": full_text,
-            "summary": summary
+            "document_type": document_type,
+            "base64dataResp": base64dataresp,
+            "entityrefkey": entity_ref_key,
+            "extracted_data": extracted_data
+        }
+
+        try:
+            inserted_doc = invoice_collection.insert_one(document)
+            document_id = str(inserted_doc.inserted_id)
+            logger.info(f"Document inserted with ID: {document_id}")
+        except Exception as e:
+            logger.error(f"Error inserting document: {str(e)}")
+            raise HTTPException(status_code=500, detail="Error inserting document into MongoDB")
+
+        return {
+            "message": "Document successfully stored in MongoDB",
+            "document_id": document_id,
+            "entityrefkey": entity_ref_key,
+            "base64dataResp": base64dataresp,
+            "extracted_data": extracted_data
         }
 
     except Exception as e:
-        # Detailed error information
         error_details = {
             "error_type": type(e).__name__,
             "error_message": str(e),
@@ -248,4 +249,4 @@ def ocr_from_s3(api_key: str = Depends(verify_api_key),file_key: str = Query(...
 app.mount("/output", StaticFiles(directory="output", follow_symlink=True, html=True), name="output")
 
 if __name__ == '__main__':
-    uvicorn.run(app=app)
+    uvicorn.run(app=app)