Update app.py

app.py

@@ -101,7 +101,7 @@ def extract_invoice_data(file_data, content_type, json_schema):
     """
     system_prompt = "You are an expert in document data extraction."
     base64_images = []
-    base64DataResp = []
+    #base64DataResp = []
 
     if content_type == "application/pdf":
         try:
@@ -115,10 +115,10 @@ def extract_invoice_data(file_data, content_type, json_schema):
                 img_byte_arr = io.BytesIO()
                 img.save(img_byte_arr, format="PNG", dpi=(300, 300))
                 base64_encoded = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
-                base64_images.append(f"data:image/png;base64,{base64_encoded}")
+                base64_images.append(f"data:{content_type};base64,{base64_encoded}")
 
             # Store all images as a single JSON object
-            base64DataResp = json.dumps(base64_images)  
+            #base64DataResp = json.dumps(base64_images)  
 
         except Exception as e:
             logger.error(f"Error converting PDF to image: {e}")
@@ -128,7 +128,7 @@ def extract_invoice_data(file_data, content_type, json_schema):
         # Handle direct image files
         base64_encoded = base64.b64encode(file_data).decode('utf-8')
         base64_images.append(f"data:{content_type};base64,{base64_encoded}")
-        base64DataResp = json.dumps(base64_images)  # Store as a JSON object
+        #base64DataResp = json.dumps(base64_images)  # Store as a JSON object
 
     # Prepare OpenAI request
     openai_content = [{"type": "image_url", "image_url": {"url": img_base64}} for img_base64 in base64_images]
@@ -146,11 +146,11 @@ def extract_invoice_data(file_data, content_type, json_schema):
         )
 
         parsed_content = json.loads(response.choices[0].message.content.strip())
-        return parsed_content, base64DataResp
+        return parsed_content, base64_images
 
     except Exception as e:
         logger.error(f"Error in OpenAI processing: {e}")
-        return {"error": str(e)}, base64DataResp
+        return {"error": str(e)}, base64_images
         
 
 def get_content_type_from_s3(file_key):
@@ -201,14 +201,14 @@ def extract_text_from_file(
         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)
+        extracted_data, base64_images = extract_invoice_data(file_data, content_type, json_schema)
 
         # Store document in MongoDB
         document = {
             "file_key": file_key,
             "file_type": content_type,
             "document_type": document_type,
-            "base64DataResp": base64DataResp,
+            "base64DataResp": base64_images,
             "entityrefkey": entity_ref_key,
             "extracted_data": extracted_data
         }
@@ -221,7 +221,7 @@ def extract_text_from_file(
             "message": "Document successfully stored in MongoDB",
             "document_id": document_id,
             "entityrefkey": entity_ref_key,
-            "base64DataResp": base64DataResp,
+            "base64DataResp": base64_images,
             "extracted_data": extracted_data
         }