Update app.py

app.py

@@ -77,36 +77,33 @@ def fetch_file_from_s3(file_key):
     except Exception as e:
         raise Exception(f"Failed to fetch file from S3: {str(e)}")
 
-def extract_text_from_pdf(file_data):
-    try:
-        doc = fitz.open(stream=file_data, filetype="pdf")
-        return "\n".join([page.get_text("text") for page in doc]).strip()
-    except Exception as e:
-        logger.error(f"Failed to extract text from PDF: {e}")
-        return ""
-
 # Function to summarize text using OpenAI GPT
 def extract_invoice_data(file_data, content_type):
-    if content_type == "application/pdf":
-        text = extract_text_from_pdf(file_data)
-        if len(text.split()) > 500:  # Large document handling
-            return {"extracted_text": text}
-    
+    system_prompt = "You are an expert in document data extraction."
+
+    # Convert file to Base64
     base64_encoded = base64.b64encode(file_data).decode('utf-8')
-    mime_type = content_type if content_type.startswith("image/") else "application/pdf"
-    
+
+    # 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": "You are an expert in document data extraction."},
+                {"role": "system", "content": system_prompt},
                 {
                     "role": "user",
                     "content": [
                         {
                             "type": "image_url",
                             "image_url": {
-                                "url": f"data:{mime_type};base64,{base64_encoded}"
+                                "url": f"data:image/{mime_type};base64,{base64_encoded}"
                             }
                         }
                     ]