Update app.py

app.py

@@ -46,6 +46,7 @@ app = FastAPI(docs_url='/')
 use_gpu = False
 output_dir = 'output'
 
+
 @app.on_event("startup")
 def startup_db():
     try:
@@ -54,6 +55,7 @@ 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")
@@ -70,6 +72,7 @@ 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:
@@ -80,6 +83,8 @@ def fetch_file_from_s3(file_key):
     except Exception as e:
         raise Exception(f"Failed to fetch file from S3: {str(e)}")
 
+
+# Function to summarize text using OpenAI GPT
 # Updated extraction function that handles PDF and image files differently
 def extract_invoice_data(file_data, content_type, json_schema):
     """
@@ -87,27 +92,14 @@ def extract_invoice_data(file_data, content_type, json_schema):
     For Images: Pass the Base64-encoded image to OpenAI (assuming a multimodal model)
     """
     system_prompt = "You are an expert in document data extraction."
-    base64_encoded_images = []  # To store Base64-encoded image data
-
-    extracted_data = {}
 
     if content_type == "application/pdf":
         # Use PyMuPDF to extract text directly from the PDF
         try:
             doc = fitz.open(stream=file_data, filetype="pdf")
-            num_pages = doc.page_count
-            
-            # Check if the number of pages exceeds 2
-            if num_pages > 2:
-                raise ValueError("The PDF contains more than 2 pages, extraction not supported.")
-            
             extracted_text = ""
             for page in doc:
                 extracted_text += page.get_text()
-
-            # Store the extracted text in the dictionary
-            extracted_data["text"] = extracted_text
-
         except Exception as e:
             logger.error(f"Error extracting text from PDF: {e}")
             raise
@@ -120,38 +112,18 @@ def extract_invoice_data(file_data, content_type, json_schema):
         )
 
     elif content_type.startswith("image/"):
-        # For images, determine if more than 2 images are provided
-        try:
-            img = Image.open(io.BytesIO(file_data))  # Open the image file
-            num_images = img.n_frames  # Get number of images (pages in the image file)
-            
-            if num_images > 2:
-                raise ValueError("The image file contains more than 2 pages, extraction not supported.")
-            
-            # Process each image page if there are 1 or 2 pages
-            for page_num in range(num_images):
-                img.seek(page_num)  # Move to the current page
-                img_bytes = io.BytesIO()
-                img.save(img_bytes, format="PNG")  # Save each page as a PNG image in memory
-                base64_encoded = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
-                base64_encoded_images.append(base64_encoded)
-            
-            # Add Base64 image data to the extracted data dictionary
-            extracted_data["base64_images"] = base64_encoded_images
-
-            # Build a prompt containing the image data for OpenAI
-            prompt = f"Extract the invoice data from the following images (Base64 encoded). Return only valid JSON that adheres to this schema:\n\n{json.dumps(json_schema, indent=2)}\n\n"
-            for base64_image in base64_encoded_images:
-                prompt += f"Image Data URL: data:{content_type};base64,{base64_image}\n"
-        
-        except Exception as e:
-            logger.error(f"Error handling images: {e}")
-            raise
-
+        # For images, encode as Base64 and pass to OpenAI
+        base64_encoded = base64.b64encode(file_data).decode('utf-8')
+        # In this example we assume the model accepts image inputs via a Base64 data URL.
+        # (This requires access to a multimodal model.)
+        prompt = (
+            f"Extract the invoice data from the following image. "
+            f"Return only valid JSON that adheres to this schema:\n\n{json.dumps(json_schema, indent=2)}\n\n"
+            f"Image Data URL:\n data:{content_type};base64,{base64_encoded}"
+        )
     else:
         raise ValueError(f"Unsupported content type: {content_type}")
 
-    # Send request to OpenAI for data extraction
     try:
         response = openai.ChatCompletion.create(
             model="gpt-4o-mini",
@@ -164,20 +136,21 @@ def extract_invoice_data(file_data, content_type, json_schema):
         )
 
         content = response.choices[0].message.content.strip()
+
+        # Clean and parse JSON output (remove markdown formatting if present)
         cleaned_content = content.strip().strip('```json').strip('```')
-        
         try:
             parsed_content = json.loads(cleaned_content)
-            extracted_data["extracted_json"] = parsed_content  # Store the parsed JSON data
-            return extracted_data
+            return parsed_content
         except json.JSONDecodeError as e:
             logger.error(f"JSON Parse Error: {e}")
-            return {"error": f"JSON Parse Error: {str(e)}"}
+            return None
 
     except Exception as e:
         logger.error(f"Error in data extraction: {e}")
         return {"error": str(e)}
 
+
 def get_content_type_from_s3(file_key):
     """Fetch the content type (MIME type) of a file stored in S3."""
     try:
@@ -186,21 +159,24 @@ def get_content_type_from_s3(file_key):
     except Exception as 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:
         raise HTTPException(status_code=401, detail="Invalid API Key")
 
+
 @app.get("/")
 def read_root():
     return {"message": "Welcome to the Invoice Summarization API!"}
 
+
 @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"),
-    entity_ref_key: str = Query(..., description="Entity Reference Key")
+        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 text from a PDF or Image stored in S3 and process it based on document size."""
     try:
@@ -218,9 +194,9 @@ def extract_text_from_file(
 
         json_schema = schema_doc.get("json_schema")
         if not json_schema:
-            raise ValueError("Schema is empty or not properly defined.")    
-            
-        # Retrieve file from S3 and determine content type
+            raise ValueError("Schema is empty or not properly defined.")
+
+            # Retrieve file from S3 and determine content type
         content_type = get_content_type_from_s3(file_key)
         file_data, _ = fetch_file_from_s3(file_key)
         extracted_data = extract_invoice_data(file_data, content_type, json_schema)
@@ -256,7 +232,8 @@ 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")