Update app.py

app.py

@@ -1,23 +1,17 @@
 import uvicorn
 from fastapi.staticfiles import StaticFiles
 import hashlib
-from enum import Enum
 from fastapi import FastAPI, Header, Query, Depends, HTTPException
-from PIL import Image
-import io
-import fitz  # PyMuPDF for PDF handling
-import logging
 from pymongo import MongoClient
-
 import boto3
 import openai
 import os
-import traceback  # For detailed traceback of errors
-import re
+import traceback  
 import json
 from dotenv import load_dotenv
 import base64
 from bson.objectid import ObjectId
+import logging
 
 db_client = None
 load_dotenv()
@@ -26,26 +20,26 @@ load_dotenv()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# MongoDB Configuration
+# Validate and Load Environment Variables
 MONGODB_URI = os.getenv("MONGODB_URI")
 DATABASE_NAME = os.getenv("DATABASE_NAME")
 COLLECTION_NAME = os.getenv("COLLECTION_NAME", "invoice_collection")
+AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
+AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+API_KEY = os.getenv("API_KEY")
 
-# use_gpu = False
-# output_dir = 'output'
-
-# Check if environment variables are set
-if not MONGODB_URI:
-    raise ValueError("MONGODB_URL is not set. Please add it to Hugging Face secrets.")
+if not all([MONGODB_URI, DATABASE_NAME, AWS_ACCESS_KEY, AWS_SECRET_KEY, S3_BUCKET_NAME, OPENAI_API_KEY, API_KEY]):
+    raise ValueError("One or more required environment variables are missing.")
 
 # Initialize MongoDB Connection
 db_client = MongoClient(MONGODB_URI)
 db = db_client[DATABASE_NAME]
 invoice_collection = db[COLLECTION_NAME]
+openai.api_key = OPENAI_API_KEY
 
 app = FastAPI(docs_url='/')
-use_gpu = False
-output_dir = 'output'
 
 @app.on_event("startup")
 def startup_db():
@@ -55,15 +49,6 @@ def startup_db():
     except Exception as e:
         logger.error(f"MongoDB connection failed: {str(e)}")
 
-# AWS S3 Configuration
-API_KEY = os.getenv("API_KEY")
-AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
-AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
-S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
-
-# OpenAI Configuration
-openai.api_key = os.getenv("OPENAI_API_KEY")
-
 # S3 Client
 s3_client = boto3.client(
     's3',
@@ -71,247 +56,59 @@ s3_client = boto3.client(
     aws_secret_access_key=AWS_SECRET_KEY
 )
 
-# Function to fetch file from S3
 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
-        file_data = response['Body'].read()
-        return file_data, content_type  # Return file data as BytesIO
+        return response['Body'].read(), response['ContentType']
     except Exception as e:
-        raise Exception(f"Failed to fetch file from S3: {str(e)}")
-
-# Function to summarize text using OpenAI GPT
-def extract_invoice_data(file_data, content_type):
-    system_prompt = "You are an expert in document data extraction."
-
-    # Convert file to Base64
-    base64_encoded = base64.b64encode(file_data).decode('utf-8')
-
-    # Determine the correct MIME type for OpenAI
-    if content_type.startswith("image/"):
-        mime_type = content_type  # e.g., image/png, image/jpeg
-    elif content_type == "application/pdf":
-        mime_type = "application/pdf"
-    else:
-        raise ValueError(f"Unsupported content type: {content_type}")
-
-    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:{mime_type};base64,{base64_encoded}"
-                            }
-                        }
-                    ]
-                }
-            ],
-            response_format={
-                "type": "json_schema",
-                "json_schema": {
-                    "name": "invoice",
-                    "strict": True,
-                    "schema": {
-                        "type": "object",
-                        "title": "Invoice Information Extractor",
-                        "$schema": "http://json-schema.org/draft-07/schema#",
-                        "properties": {
-                            "LineItems": {
-                                "type": "array",
-                                "items": {
-                                    "type": "object",
-                                    "required": [
-                                        "ProductCode",
-                                        "Description",
-                                        "Amount"
-                                    ],
-                                    "properties": {
-                                        "ProductCode": {
-                                            "type": "string",
-                                            "title": "Product Code",
-                                            "description": "The code of the product"
-                                        },
-                                        "Description": {
-                                            "type": "string",
-                                            "title": "Description",
-                                            "description": "Description of the product"
-                                        },
-                                        "Amount": {
-                                            "type": "number",
-                                            "title": "Amount",
-                                            "description": "The amount of the product"
-                                        }
-                                    },
-                                    "additionalProperties": False
-                                },
-                                "title": "Line Items",
-                                "description": "List of line items on the invoice"
-                            },
-                            "TaxAmount": {
-                                "type": "number",
-                                "title": "Tax Amount",
-                                "description": "The tax amount on the invoice"
-                            },
-                            "VendorGST": {
-                                "type": "string",
-                                "title": "Vendor GST",
-                                "description": "The GST number of the vendor"
-                            },
-                            "VendorName": {
-                                "type": "string",
-                                "title": "Vendor Name",
-                                "description": "The name of the vendor"
-                            },
-                            "InvoiceDate": {
-                                "type": "string",
-                                "title": "Invoice Date",
-                                "description": "The date of the invoice (format: dd-MMM-yyyy)"
-                            },
-                            "TotalAmount": {
-                                "type": "number",
-                                "title": "Total Amount",
-                                "description": "The total amount on the invoice"
-                            },
-                            "InvoiceNumber": {
-                                "type": "string",
-                                "title": "Invoice Number",
-                                "description": "The number of the invoice"
-                            },
-                            "VendorAddress": {
-                                "type": "string",
-                                "title": "Vendor Address",
-                                "description": "The address of the vendor"
-                            },
-                            "InvoiceCurrency": {
-                                "type": "string",
-                                "title": "Invoice Currency",
-                                "description": "The currency used in the invoice (e.g., USD, INR, AUD)"
-                            }
-                        },
-                        "required": [
-                            "LineItems",
-                            "TaxAmount",
-                            "VendorGST",
-                            "VendorName",
-                            "InvoiceDate",
-                            "TotalAmount",
-                            "InvoiceNumber",
-                            "VendorAddress",
-                            "InvoiceCurrency"
-                        ],
-                        "additionalProperties": False,
-                        "description": "Schema for extracting structured invoice data"
-                    }
-                }
-            },
-            temperature=0.5,
-            max_tokens=16384
-        )
+        logger.error(f"Failed to fetch file from S3: {e}")
+        raise HTTPException(status_code=500, detail=f"S3 fetch error: {str(e)}")
 
-        # Clean and parse JSON output
-        content = response.choices[0].message.content.strip()
-        cleaned_content = content.strip().strip('```json').strip('```')
+def extract_text_from_s3(file_key, content_type):
+    return "Extracted text from file", 1  # Placeholder for real extraction logic
 
-        try:
-            parsed_content = json.loads(cleaned_content)
-            return parsed_content
-        except json.JSONDecodeError as e:
-            logger.error(f"JSON Parse Error: {e}")
-            return None
+def convert_to_base64(file_key):
+    return "Base64 encoded data"  # Placeholder
 
-    except Exception as e:
-        logger.error(f"Error in data extraction: {e}")
-        return {"error": str(e)}
+def generate_summary(extracted_text):
+    return "Summarized text"  # Placeholder
 
-# Dependency to check API Key
 def verify_api_key(api_key: str = Header(...)):
     if api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API Key")
 
-@app.get("/")
-def read_root():
-    return {"message": "Welcome to the Invoice Summarization API!"}
-
-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:
-        raise Exception(f"Failed to get content type from S3: {str(e)}")
-
-
 @app.get("/ocr/extraction")
 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"),
+    file_key: str = Query(..., description="S3 file key"),
+    document_type: str = Query(..., description="Document type"),
     entity_ref_key: str = Query(..., description="Entity Reference Key")
 ):
-    """Extract text from a PDF or Image stored in S3 and process it based on document size."""
     try:
         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
-            }
-
-        # Retrieve file from S3 and determine content type (Ensure this step is implemented)
-        content_type = get_content_type_from_s3(file_key)  # Implement this function
-
-        # Extract text (Ensure Extraction function is implemented)
+            return {"message": "Document Retrieved from MongoDB.", "document": existing_document}
+        
+        file_data, content_type = fetch_file_from_s3(file_key)
         extracted_text, num_pages = extract_text_from_s3(file_key, content_type)
-
-        # Define values for small/large files
         base64DataResp = None
         summary = None
         if num_pages <= 2:
-            base64DataResp = convert_to_base64(file_key)  # Implement this function
-            summary = generate_summary(extracted_text)  # Implement this function
-
-        # Store extracted data in MongoDB
+            base64DataResp = convert_to_base64(file_key)
+            summary = generate_summary(extracted_text)
+        
         document = {
             "file_key": file_key,
             "file_type": content_type,
             "document_type": document_type,
             "entityrefkey": entity_ref_key,
             "num_pages": num_pages,
-            "base64DataResp": base64DataResp,  # Only for small files
-            "extracted_text": extracted_text,
-            "summary": summary,  # Only for small files
+            "base64DataResp": base64DataResp,
+            "extracted_text": extracted_text if num_pages <= 2 else "Skipped",
+            "summary": summary,
         }
-
         inserted_doc = invoice_collection.insert_one(document)
-        document_id = str(inserted_doc.inserted_id)
-
-        return {
-            "message": "Document successfully stored in MongoDB",
-            "document_id": document_id,
-            "file_key": file_key,
-            "num_pages": num_pages,
-            "summary": summary if summary else "Skipped for large documents"
-        }
-
+        return {"message": "Document stored in MongoDB", "document_id": str(inserted_doc.inserted_id)}
     except Exception as e:
-        error_details = {
-            "error_type": type(e).__name__,
-            "error_message": str(e),
-            "traceback": traceback.format_exc()
-        }
-        return {"error": error_details}
-     
-# Serve the output folder as static files
-app.mount("/output", StaticFiles(directory="output", follow_symlink=True, html=True), name="output")
-
-if __name__ == '__main__':
-    uvicorn.run(app=app)
+        return {"error": {"type": type(e).__name__, "message": str(e), "traceback": traceback.format_exc()}}