Update main.py

main.py

@@ -10,7 +10,7 @@ from dotenv import load_dotenv
 from flask import Flask, request, make_response
 from typing import Dict, List, Union
 import json
-
+from PIL import Image
 # Initialize Flask App
 app = Flask(__name__)
 
@@ -120,7 +120,6 @@ def handle_media_message(messenger, data, message_type, mobile):
         # Download media content based on type
         if message_type == "image":
             media_url = messenger.get_image_url(data)
-            print(media_url)
         elif message_type == "document":
             media_url = messenger.get_document_url(data)
         elif message_type == "audio":
@@ -139,9 +138,28 @@ def handle_media_message(messenger, data, message_type, mobile):
             messenger.send_message("Sorry, I couldn't process your media.", mobile)
             return None
 
-        # Base64 encode the image for Gemini
-        base64_image = base64.b64encode(response.content).decode('utf-8')
-        return base64_image
+        # Ensure image is in a compatible format
+        try:
+            img = Image.open(io.BytesIO(response.content))
+            
+            # Convert to RGB if needed
+            if img.mode != 'RGB':
+                img = img.convert('RGB')
+            
+            # Resize large images to reduce processing time
+            max_size = (1024, 1024)
+            img.thumbnail(max_size)
+            
+            # Save to a BytesIO object
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='JPEG')
+            img_byte_arr = img_byte_arr.getvalue()
+            
+            return img_byte_arr
+        except Exception as img_error:
+            logger.error(f"Image processing error: {img_error}")
+            messenger.send_message("Sorry, I had trouble processing the image.", mobile)
+            return None
 
     except Exception as e:
         logger.error(f"Error handling media message: {str(e)}")
@@ -186,8 +204,9 @@ carts: Dict[str, List[Dict]] = {}
 class ShoppingAssistant:
     def __init__(self):
         self.product_data = product_data
+        self.last_analyzed_products = {}
         
-    def analyze_content(self, content: Union[str, bytes], content_type: str, budget: float = None) -> dict:
+    def analyze_content(self, content: Union[str, bytes], content_type: str, mobile: str, budget: float = None) -> dict:
         try:
             if content_type == "text":
                 response = model.generate_content(
@@ -195,14 +214,10 @@ class ShoppingAssistant:
                     specific products needed, and any budget constraints. Request: {content}"""
                 )
             else:  # For images, documents, etc.
-                response = vision_model.generate_content(
-                    [
-                        "Analyze this content and extract shopping-related information, products needed, and any specifications.",
-                        content
-                    ]
-                )
+                response = vision_model.generate_content(content)
             
             analysis = response.text
+            logger.info(f"Image analysis result: {analysis}")
             
             # Second pass to find specific products
             products_prompt = f"""Based on this analysis: {analysis}
@@ -211,7 +226,15 @@ class ShoppingAssistant:
             Return only product names and prices that match the request."""
             
             products_response = model.generate_content(products_prompt)
-            return self.format_recommendations(products_response.text, budget)
+            logger.info(f"Products response: {products_response.text}")
+            
+            result = self.format_recommendations(products_response.text, budget)
+            
+            # Store last analyzed products for this mobile number
+            if mobile:
+                self.last_analyzed_products[mobile] = result["products"]
+            
+            return result
             
         except Exception as e:
             logger.error(f"Error in content analysis: {str(e)}")
@@ -266,10 +289,10 @@ def hook():
 
         if message_type == "text":
             message = messenger.get_message(data)
-            result = shopping_assistant.analyze_content(message, "text")
+            result = shopping_assistant.analyze_content(message, "text", mobile)
             
             if "error" in result:
-                messenger.send_message(f"Sorry, I couldn't process your request.", mobile)
+                messenger.send_message("Sorry, I couldn't process your request.", mobile)
                 return "OK", 200
 
             if not result["products"]:
@@ -283,7 +306,6 @@ def hook():
             products_text += f"\nTotal: ${result['total']:.2f}"
 
             # Send interactive message with buttons
-            # Create and send reply buttons
             buttons = [
                 {"id": "add_to_cart", "title": "Add to Cart"},
                 {"id": "cancel", "title": "Cancel"}
@@ -295,29 +317,33 @@ def hook():
             )
 
         elif message_type == "interactive":
-            handle_interactive_message(messenger, data, mobile, carts)
+            handle_interactive_message(messenger, data, mobile, carts, shopping_assistant)
 
         elif message_type in ["image", "document", "audio"]:
             media_content = handle_media_message(messenger, data, message_type, mobile)
             if media_content:
-                result = shopping_assistant.analyze_content(media_content, message_type)
-
-                products_text= "Here are the suggested products:\n\n"
-                for product in result["products"]:
-                    products_text += f"- {product['product']}: ${product['price']}\n"
-                products_text += f"\nTotal: ${result['total']:.2f}"
-                # Create and send reply buttons
-                buttons = [
-                {"id": "add_to_cart", "title": "Add to Cart"},
-                {"id": "cancel", "title": "Cancel"}
-                 ]
-                button_message = create_reply_button_message(products_text, buttons)
-                messenger.send_reply_button(
-                   recipient_id=mobile,
-                   button=button_message
-                   )
-            else:
-                messenger.send_message("I couldn't find any matching products in the image.", mobile)
+                result = shopping_assistant.analyze_content(media_content, message_type, mobile)
+                
+                # Check if products were found
+                if result.get("products"):
+                    # Format product message
+                    products_text = "Here are the suggested products:\n\n"
+                    for product in result["products"]:
+                        products_text += f"- {product['product']}: ${product['price']}\n"
+                    products_text += f"\nTotal: ${result['total']:.2f}"
+
+                    # Create and send reply buttons
+                    buttons = [
+                        {"id": "add_to_cart", "title": "Add to Cart"},
+                        {"id": "cancel", "title": "Cancel"}
+                    ]
+                    button_message = create_reply_button_message(products_text, buttons)
+                    messenger.send_reply_button(
+                        recipient_id=mobile,
+                        button=button_message
+                    )
+                else:
+                    messenger.send_message("I couldn't find any matching products in the image.", mobile)
 
         return "OK", 200