Update main.py

main.py

@@ -1,94 +1,77 @@
-from flask import Flask, request, jsonify
-import re
-import pandas as pd
-import google.generativeai as genai
-import pypdf
+import os
+import io
 import json
+import hashlib
 from datetime import datetime
-import os
+from PIL import Image
+
+from flask import Flask, request, jsonify
 from flask_cors import CORS
+import google.generativeai as genai
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
 
-api_key = os.environ['Gemini']
+# Use the Gemini API key from the environment, or set it here for testing.
+api_key = os.environ.get('Gemini', 'YOUR_GEMINI_API_KEY')
 
 def configure_gemini(api_key):
     genai.configure(api_key=api_key)
-    return genai.GenerativeModel('gemini-2.0-flash-exp')
-
-def read_pdf(file_path):
-    text_content = []
-    with open(file_path, 'rb') as file:
-        pdf_reader = pypdf.PdfReader(file)  
-        for page in pdf_reader.pages:
-            text = page.extract_text()
-            if text:
-                text_content.append(text)
-    return "\n".join(text_content)
+    return genai.GenerativeModel('gemini-2.0-flash-thinking-exp')
 
-def process_with_gemini(model, text):
-    prompt = """Analyze this bank statement and extract transactions in JSON format with these fields:
-    - Date (format DD/MM/YYYY)
-    - Description
-    - Amount (just the integer value)
-    - Type (is 'income' if 'credit amount', else 'expense')
-    - Customer Name (Only If Type is 'income' and if no name is extracted write 'general income' and if type is not 'income' write 'expense')
-    - City (In address of bank statement)
-    
-    Return ONLY valid JSON with this structure:
-    {
-        "transactions": [
-            {
-                "Date": "string",
-                "Description": "string",
-                "Customer_name": "string",
-                "City": "string",
-                "Amount": number,
-                "Type": "string"
-            }
-        ]
-    }"""
-    
-    response = model.generate_content([prompt, text])
+def process_receipt(model, image):
+    prompt = (
+        "Analyze this image and determine if it's a receipt. If it is a receipt, extract:\n"
+        " - Total amount (as float)\n"
+        " - List of items purchased (array of strings)\n"
+        " - Date of transaction (DD/MM/YYYY format)\n"
+        " - Receipt number (as string)\n"
+        "Return JSON format with keys: is_receipt (boolean), total, items, date, receipt_number.\n"
+        "If not a receipt, return {\"is_receipt\": false}"
+    )
+    response = model.generate_content([prompt, image])
     return response.text
 
-@app.route('/process-pdf', methods=['POST'])
-def process_pdf():
+@app.route('/process-receipt', methods=['POST'])
+def process_receipt_endpoint():
     try:
-        if 'file' not in request.files:
+        if 'receipt' not in request.files:
             return jsonify({'error': 'No file uploaded'}), 400
-            
-        file = request.files['file']
+
+        file = request.files['receipt']
         if file.filename == '':
             return jsonify({'error': 'No file selected'}), 400
-            
-        if not file.filename.endswith('.pdf'):
-            return jsonify({'error': 'File must be a PDF'}), 400
 
-        # Save uploaded file temporarily
-        temp_path = "temp.pdf"
-        file.save(temp_path)
-        
-        # Configure Gemini and process PDF
+        # Read file bytes and compute a hash (for duplicate checking or logging)
+        image_bytes = file.read()
+        file_hash = hashlib.md5(image_bytes).hexdigest()
+
+        # Open the image using Pillow
+        image = Image.open(io.BytesIO(image_bytes))
+
+        # Configure Gemini and process the receipt image
         model = configure_gemini(api_key)
-        pdf_text = read_pdf(temp_path)
-        json_response = process_with_gemini(model, pdf_text)
-        
-        # Clean JSON response
-        json_str = json_response[json_response.find('{'):json_response.rfind('}')+1]
+        result_text = process_receipt(model, image)
+
+        # Attempt to extract JSON from the response text
+        json_start = result_text.find('{')
+        json_end = result_text.rfind('}')
+        if json_start == -1 or json_end == -1:
+            return jsonify({'error': 'Invalid response format', 'raw': result_text}), 500
+
+        json_str = result_text[json_start:json_end+1]
+        # Clean up any markdown formatting if necessary
         json_str = json_str.replace('```json', '').replace('```', '')
-        
-        # Parse JSON and return
         data = json.loads(json_str)
-        
-        # Clean up temp file
-        os.remove(temp_path)
-        
+
+        # Optionally, add metadata to the response
+        data['file_hash'] = file_hash
+        data['timestamp'] = datetime.now().isoformat()
+
         return jsonify(data)
-        
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':
-    app.run(debug=True,host="0.0.0.0", port=7860)
+    # The server listens on all interfaces at port 7860.
+    app.run(debug=True, host="0.0.0.0", port=7860)