Update app.py

app.py

@@ -102,14 +102,15 @@ def extract_invoice_data(file_data, content_type, json_schema):
     system_prompt = "You are an expert in document data extraction."
     base64_images = []
     base64DataResp = []
-
+    
     if content_type == "application/pdf":
         try:
             extracted_text = extract_pdf_text(file_data)
+            
             # Store PDF as Base64
             base64_pdf = base64.b64encode(file_data).decode('utf-8')
             base64DataResp.append(f"data:application/pdf;base64,{base64_pdf}")
-
+    
             images = convert_from_bytes(file_data)  # Convert PDF to images
 
             if len(images) > 2:
@@ -121,19 +122,16 @@ def extract_invoice_data(file_data, content_type, json_schema):
                 base64_encoded = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
                 base64_images.append(f"data:image/png;base64,{base64_encoded}")
 
-            # Store all images as a single JSON object
-            # base64DataResp = json.dumps(base64_images)  
-
         except Exception as e:
             logger.error(f"Error converting PDF to image: {e}")
             return {"error": "Failed to process PDF"}, None
 
-    else:
+    elif content_type.startswith("image/"):
         # Handle direct image files
-        base64_encoded = base64.b64encode(file_data).decode('utf-8')
-        base64DataResp.append(f"data:{content_type};base64,{base64_encoded}")
-        # base64_images.append(f"data:{content_type};base64,{base64_encoded}")
-        # base64DataResp = json.dumps(base64_images)  # Store as a JSON object
+        base64_pdf = base64.b64encode(file_data).decode('utf-8')
+        base64DataResp.append(f"data:{content_type};base64,{base64_pdf}")
+    else:
+        return {"error": f"Unsupported file type: {content_type}"}
 
     # Prepare OpenAI request
     openai_content = [{"type": "image_url", "image_url": {"url": img_base64}} for img_base64 in base64_images]