Modified based on JSON Schema

app.py

@@ -19,6 +19,7 @@ from dotenv import load_dotenv
 import base64
 from bson.objectid import ObjectId
 
+db_client = None
 load_dotenv()
 
 # Set up logging
@@ -30,28 +31,28 @@ MONGODB_URI = os.getenv("MONGODB_URI")
 DATABASE_NAME = os.getenv("DATABASE_NAME")
 COLLECTION_NAME = os.getenv("COLLECTION_NAME", "invoice_collection")
 
-app = FastAPI(docs_url='/')
-use_gpu = False
-output_dir = 'output'
+# 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.")
 
 # Initialize MongoDB Connection
-client = MongoClient(MONGODB_URI)
-db = client[DATABASE_NAME]
+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='/')
 
 @app.on_event("startup")
 def startup_db():
     try:
-        # Create a MongoDB client and fetch server info to verify connection
-        client = MongoClient(MONGODB_URI)
-        server_info = client.server_info()  # This is a blocking call
-        print("MongoDB connection successful:", server_info)
+        db_client.server_info()
+        logger.info("MongoDB connection successful")
     except Exception as e:
-        print("MongoDB connection failed:", str(e))
+        logger.error(f"MongoDB connection failed: {str(e)}")
 
 # AWS S3 Configuration
 API_KEY = os.getenv("API_KEY")
@@ -70,7 +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
@@ -80,146 +81,153 @@ def fetch_file_from_s3_file(file_key):
         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 into a standardized JSON format. The invoice headers should include the following details:
-
-- **Vendor Information**:
-  - Vendor Name
-  - Vendor Address
-  - Vendor GST No.
-
-- **Invoice Details**:
-  - Invoice No.
-  - Invoice Date (formatted as dd-MMM-yyyy)
-  - Invoice Currency (e.g., USD, INR, etc.)
-  - 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 the following JSON format:
-
-{
-  "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": {
-                "Amount": {
-                  "type": "number",
-                  "description": "Amount for the line item"
-                },
-                "Description": {
-                  "type": "string",
-                  "description": "Description of the line item"
-                },
-                "ProductCode": {
-                  "type": "string",
-                  "description": "Product or service code for the line item"
-                }
-              },
-              "additionalProperties": false
-            },
-            "title": "Line Items"
-          },
-          "TaxAmount": {
-            "type": "number",
-            "description": "Total tax amount for the invoice"
-          },
-          "VendorGST": {
-            "type": "string",
-            "description": "Vendor's GST number"
-          },
-          "VendorName": {
-            "type": "string",
-            "description": "Name of the vendor"
-          },
-          "InvoiceDate": {
-            "type": "string",
-            "format": "date",
-            "description": "Invoice date in dd-MMM-yyyy format"
-          },
-          "TotalAmount": {
-            "type": "number",
-            "description": "Total amount for the invoice"
-          },
-          "InvoiceNumber": {
-            "type": "string",
-            "description": "Invoice number"
-          },
-          "VendorAddress": {
-            "type": "string",
-            "description": "Vendor's address"
-          },
-          "InvoiceCurrency": {
-            "type": "string",
-            "description": "Currency used in the invoice"
-          }
-        },
-        "required": [
-          "LineItems", "TaxAmount", "VendorGST", "VendorName",
-          "InvoiceDate", "TotalAmount", "InvoiceNumber", "VendorAddress", "InvoiceCurrency"
-        ],
-        "additionalProperties": false,
-        "description": "Schema for extracting specific information from invoices"
-      }
-    }
-  }
-}
-
-### Guidelines for Processing:
-- Extract all relevant data from the invoice using naming conventions such as "Bill to" or "Taxpayer Name" for Vendor and Customer info.
-- Convert the **Invoice Date** to dd-MMM-yyyy format.
-- Ensure correct handling of amounts (e.g., **Amount**, **Tax Amount**, **Total Invoice Amount**).
-- For line items, include descriptions, tax percentages, tax amounts, and amounts.
-- If certain values are missing, leave them empty or set them as null.
-- 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_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": f"{text}"}
+                {
+                    "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
         )
+
+        # Clean and parse JSON output
         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
+        cleaned_content = content.strip().strip('```json').strip('```')
 
-        # Step 2: Parse the cleaned content as JSON
         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)
+            logger.error(f"JSON Parse Error: {e}")
             return None
+
     except Exception as e:
-        return f"Error in summarization: {str(e)}"
+        logger.error(f"Error in data extraction: {e}")
+        return {"error": str(e)}
 
 # Dependency to check API Key
 def verify_api_key(api_key: str = Header(...)):
@@ -239,50 +247,27 @@ def extract_text_from_file(
 ):
     """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})
+        existing_document = invoice_collection.find_one({"entityrefkey": entity_ref_key})
 
-        if(existing_document):
+        if existing_document:
             existing_document["_id"] = str(existing_document["_id"])
             return {
                 "message": "Document Retrieved from MongoDB.",
                 "document": existing_document
             }
-        # Fetch file from S3
-        file_data, content_type = fetch_file_from_s3_file(file_key)
-
-        extracted_text = []
-        base64DataResp = None
-        summary = None
-
-        if content_type.startswith("image/"):  # Image file
-            image = Image.open(io.BytesIO(file_data)).convert("RGB")
-            extracted_text.append(pytesseract.image_to_string(image))
-            
-            # If single image, store Base64
-            base64Data = base64.b64encode(file_data).decode('utf-8')
-            base64DataResp = f"data:image/{content_type.lower()};base64,{base64Data}"
 
-        elif content_type == "application/pdf":  # PDF file
-            pdf_document = fitz.open(stream=io.BytesIO(file_data), filetype="pdf")
-            num_pages = len(pdf_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
 
-            for page_number in range(num_pages):
-                page = pdf_document[page_number]
-                extracted_text.append(page.get_text("text"))
+        # Extract text (Ensure Extraction function is implemented)
+        extracted_text, num_pages = extract_text_from_s3(file_key, content_type)
 
-            pdf_document.close()
-
-            # If 2 pages or less, store Base64
-            if num_pages <= 2:
-                base64Data = base64.b64encode(file_data).decode('utf-8')
-                base64DataResp = f"data:application/pdf;base64,{base64Data}"
-
-            # If 2 pages or less, generate summary
-            if num_pages <= 2:
-                full_text = " ".join(extracted_text)
-                summary = summarize_text(full_text)
-        else:
-            return {"error": f"Unsupported file type: {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
         document = {
@@ -290,20 +275,20 @@ def extract_text_from_file(
             "file_type": content_type,
             "document_type": document_type,
             "entityrefkey": entity_ref_key,
-            "num_pages": len(extracted_text),  # Store page count
+            "num_pages": num_pages,
             "base64DataResp": base64DataResp,  # Only for small files
-            "extracted_text": " ".join(extracted_text),
+            "extracted_text": extracted_text,
             "summary": summary,  # Only for small files
         }
 
         inserted_doc = invoice_collection.insert_one(document)
-        document_id = str(inserted_doc.inserted_id)        
+        document_id = str(inserted_doc.inserted_id)
 
         return {
             "message": "Document successfully stored in MongoDB",
             "document_id": document_id,
             "file_key": file_key,
-            "num_pages": len(extracted_text),
+            "num_pages": num_pages,
             "summary": summary if summary else "Skipped for large documents"
         }
 
@@ -314,7 +299,6 @@ def extract_text_from_file(
             "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")