Update main.py

main.py

@@ -1,6 +1,7 @@
 import os
 import logging
 import socket
+import io
 import requests
 import dns.resolver
 import pandas as pd
@@ -11,11 +12,10 @@ 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__)
 
-
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -23,24 +23,21 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Load .env file
+# Load environment variables
 load_dotenv()
 
-# Configure DNS resolver with environment variables from Spaces
+# DNS Configuration
 nameserver1 = os.getenv('nameserver1', '8.8.8.8')
 nameserver2 = os.getenv('nameserver2', '8.8.4.4')
 
-# Configure global DNS resolution
 def setup_dns():
     """Configure DNS resolution globally"""
     resolver = dns.resolver.Resolver()
     resolver.nameservers = [nameserver1, nameserver2]
     
-    # Create a custom getaddrinfo function
     _orig_getaddrinfo = socket.getaddrinfo
     def new_getaddrinfo(*args, **kwargs):
         try:
-            # If the host is graph.facebook.com, use our custom resolver
             if args and args[0] == 'graph.facebook.com':
                 answers = resolver.resolve('graph.facebook.com', 'A')
                 ip = str(answers[0])
@@ -49,10 +46,8 @@ def setup_dns():
             logger.error(f"DNS resolution failed: {e}")
         return _orig_getaddrinfo(*args, **kwargs)
     
-    # Replace the global getaddrinfo function
     socket.getaddrinfo = new_getaddrinfo
 
-# Setup DNS resolution
 setup_dns()
 
 # Initialize WhatsApp
@@ -61,136 +56,18 @@ try:
         os.getenv("whatsapp_token"),
         phone_number_id=os.getenv("phone_number_id")
     )
-    if not os.getenv("whatsapp_token") or not os.getenv("phone_number_id"):
-        raise ValueError("Missing required environment variables")
 except Exception as e:
-    logger.error(f"Failed to initialize WhatsApp messenger: {str(e)}")
+    logger.error(f"Failed to initialize WhatsApp: {str(e)}")
     raise
 
 VERIFY_TOKEN = "30cca545-3838-48b2-80a7-9e43b1ae8ce4"
 
-@app.get("/health")
-def health_check():
-    """Health check endpoint to verify DNS resolution"""
-    try:
-        resolver = dns.resolver.Resolver()
-        resolver.nameservers = [nameserver1, nameserver2]
-        ip = resolver.resolve('graph.facebook.com', 'A')[0]
-        return {
-            "status": "healthy",
-            "dns_resolution": str(ip),
-            "nameserver1": nameserver1,
-            "nameserver2": nameserver2
-        }
-    except Exception as e:
-        return {
-            "status": "unhealthy",
-            "error": str(e)
-        }
-
-
-def create_reply_button_message(text: str, buttons: List[Dict]) -> Dict:
-    """Create a reply button message format compatible with WhatsApp API"""
-    return {
-        "type": "button",
-        "body": {
-            "text": text
-        },
-        "action": {
-            "buttons": [
-                {
-                    "type": "reply",
-                    "reply": {
-                        "id": btn["id"],
-                        "title": btn["title"]
-                    }
-                }
-                for btn in buttons
-            ]
-        }
-    }
-
-
-def handle_media_message(messenger, data, message_type, mobile):
-    """
-    Handle media messages (images, documents, audio)
-    Retrieve and process media content
-    """
-    try:
-        # Download media content based on type
-        if message_type == "image":
-            media_url = messenger.get_image_url(data)
-        elif message_type == "document":
-            media_url = messenger.get_document_url(data)
-        elif message_type == "audio":
-            media_url = messenger.get_audio_url(data)
-        else:
-            logger.error(f"Unsupported media type: {message_type}")
-            return None
-
-        # Download media content
-        response = requests.get(media_url, headers={
-            "Authorization": f"Bearer {os.getenv('whatsapp_token')}"
-        })
-        
-        if response.status_code != 200:
-            logger.error(f"Failed to download media: {response.status_code}")
-            messenger.send_message("Sorry, I couldn't process your media.", mobile)
-            return None
-
-        # 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)}")
-        messenger.send_message("Sorry, I couldn't process your media.", mobile)
-        return None
-
-
-@app.get("/")
-def verify_token():
-    """Webhook verification endpoint"""
-    try:
-        if request.args.get("hub.verify_token") == VERIFY_TOKEN:
-            logger.info("Verified webhook")
-            response = make_response(request.args.get("hub.challenge"), 200)
-            response.mimetype = "text/plain"
-            return response
-        logger.error("Webhook Verification failed")
-        return "Invalid verification token"
-    except Exception as e:
-        logger.error(f"Error in verify_token: {str(e)}")
-        return "Internal server error", 500
-
-# Load .env file
-load_dotenv()
-
-# Initialize Gemini
+# Gemini Configuration
 genai.configure(api_key=os.getenv("google_api_key"))
-model = genai.GenerativeModel('gemini-2.0-flash-exp')
 vision_model = genai.GenerativeModel('gemini-2.0-flash-exp')
+text_model = genai.GenerativeModel('gemini-2.0-flash-exp')
 
-# Load product data
+# Product Data and Cart
 try:
     df = pd.read_csv("supermarket_data1.csv")
     product_data = df.to_dict('records')
@@ -198,264 +75,204 @@ except Exception as e:
     logger.error(f"Failed to load product data: {str(e)}")
     raise
 
-# Cart management
 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, mobile: str, budget: float = None) -> dict:
+    
+    def process_input(self, content: Union[str, bytes], content_type: str, mobile: str) -> dict:
+        """Multimodal input processor"""
         try:
             if content_type == "text":
-                response = model.generate_content(
-                    f"""As a retail shopping assistant, analyze this request and extract shopping context, 
-                    specific products needed, pick only one item per product extracted 
-                    unless the user specifies the amount or number of items for a particular product. 
-                    If the user specifies budget make sure your suggest is within that budget. 
-
-                    If a user suggest a scenario planning a dinner, planning a party, I want to clean my car etc, be analytic and creative
-                    and suggest appropriate product
-                    Request: {content}"""
-                )
-            else:  # For images, documents, etc.
-                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}
-            Find matching products from our inventory considering these details:
-            {json.dumps(self.product_data)}
-            Return only product names and prices that match the request."""
-            
-            products_response = model.generate_content(products_prompt)
-            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
-            
+                return self._process_text(content, mobile)
+            elif content_type == "image":
+                return self._process_image(content, mobile)
+            else:
+                return {"error": "Unsupported content type"}
         except Exception as e:
-            logger.error(f"Error in content analysis: {str(e)}")
-            return {"error": "Failed to analyze content"}
-
-    def format_recommendations(self, recommendations: str, budget: float = None) -> dict:
-        try:
-            # Parse recommendations and filter by budget if specified
-            suggested_products = []
-            total = 0
-            
-            for product in self.product_data:
-                if product['product'].lower() in recommendations.lower():
-                    if budget and (total + product['price']) > budget:
-                        continue
-                    suggested_products.append(product)
-                    total += product['price']
+            logger.error(f"Processing error: {str(e)}")
+            return {"error": "Failed to process input"}
+
+    def _process_text(self, text: str, mobile: str) -> dict:
+        """Process text input using Gemini"""
+        response = text_model.generate_content(
+            f"""Analyze this shopping request: {text}
+            Identify products needed, quantities, and special requirements.
+            Respond with JSON format: {{products: [{{name: '', quantity: 1}}]}}"""
+        )
+        analysis = json.loads(response.text)
+        return self._match_products(analysis.get('products', []), mobile)
+
+    def _process_image(self, image_bytes: bytes, mobile: str) -> dict:
+        """Process image input using Gemini Vision"""
+        img = Image.open(io.BytesIO(image_bytes))
+        response = vision_model.generate_content([
+            "Analyze this shopping-related image and identify products",
+            img
+        ])
+        analysis = json.loads(response.text)
+        return self._match_products(analysis.get('products', []), mobile)
+
+    def _match_products(self, detected_items: List[dict], mobile: str) -> dict:
+        """Match detected items with inventory"""
+        matched_products = []
+        for item in detected_items:
+            product = next(
+                (p for p in self.product_data 
+                 if p['product'].lower() == item['name'].lower()),
+                None
+            )
+            if product:
+                matched_products.append({
+                    **product,
+                    'quantity': item.get('quantity', 1)
+                })
+        
+        if mobile:
+            self.last_analyzed_products[mobile] = matched_products
             
-            return {
-                "products": suggested_products,
-                "total": total,
-                "message": "No products found" if not suggested_products else "Products found"
-            }
-        except Exception as e:
-            logger.error(f"Error formatting recommendations: {str(e)}")
-            return {"error": "Failed to format recommendations"}
+        return {
+            "products": matched_products,
+            "total": sum(p['price']*p.get('quantity',1) for p in matched_products),
+            "message": f"Found {len(matched_products)} matching products"
+        }
+
+def create_reply_button_message(text: str, buttons: List[Dict]) -> Dict:
+    """Create interactive button message"""
+    return {
+        "type": "button",
+        "body": {"text": text},
+        "action": {
+            "buttons": [
+                {
+                    "type": "reply",
+                    "reply": {
+                        "id": btn["id"],
+                        "title": btn["title"]
+                    }
+                } for btn in buttons
+            ]
+        }
+    }
 
-# Initialize WhatsApp and DNS setup (keeping your original configuration)
+@app.get("/")
+def verify_token():
+    """Webhook verification"""
+    if request.args.get("hub.verify_token") == VERIFY_TOKEN:
+        return request.args.get("hub.challenge")
+    return "Invalid verification token"
 
 @app.post("/")
-def hook():
+def webhook():
     """Main webhook handler"""
     try:
         data = request.get_json()
-        logger.info("Received webhook data: %s", data)
-        
-        changed_field = messenger.changed_field(data)
-        if changed_field != "messages":
-            return "OK", 200
-            
-        if not messenger.is_message(data):
-            delivery = messenger.get_delivery(data)
-            if delivery:
-                logger.info(f"Message delivery status: {delivery}")
-            return "OK", 200
-
         mobile = messenger.get_mobile(data)
-        name = messenger.get_name(data)
         message_type = messenger.get_message_type(data)
-        
-        shopping_assistant = ShoppingAssistant()
+        assistant = ShoppingAssistant()
 
+        # Handle text input
         if message_type == "text":
             message = messenger.get_message(data)
-            result = shopping_assistant.analyze_content(message, "text", mobile)
+            result = assistant.process_input(message, "text", mobile)
             
-            if "error" in result:
-                messenger.send_message("Sorry, I couldn't process your request.", mobile)
-                return "OK", 200
-
-            if not result["products"]:
-                messenger.send_message("I couldn't find any matching products.", mobile)
-                return "OK", 200
-
-            # 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}"
-
-            # Send interactive message with 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
-            )
-
-        elif message_type == "interactive":
-            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, mobile)
+            if result.get('products'):
+                products_text = "🛒 *Suggested Products:*\n\n"
+                for p in result['products']:
+                    products_text += f"➤ {p['product']} x{p.get('quantity',1)}\n   Price: ${p['price']}\n\n"
+                products_text += f"💵 Total: ${result['total']:.2f}"
                 
-                # 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"}
+                ]
+                messenger.send_reply_button(
+                    mobile,
+                    create_reply_button_message(products_text, buttons)
+                )
+            else:
+                messenger.send_message("❌ No matching products found", mobile)
+
+        # Handle image input
+        elif message_type == "image":
+            image_url = messenger.get_image_url(data)
+            response = requests.get(image_url, headers={
+                "Authorization": f"Bearer {os.getenv('whatsapp_token')}"
+            })
+            
+            if response.status_code == 200:
+                result = assistant.process_input(response.content, "image", mobile)
+                
+                if result.get('products'):
+                    products_text = "📸 *Products Found in Your Image:*\n\n"
+                    for p in result['products']:
+                        products_text += f"➤ {p['product']} x{p.get('quantity',1)}\n   Price: ${p['price']}\n\n"
+                    products_text += f"💵 Total: ${result['total']:.2f}"
+                    
                     buttons = [
-                        {"id": "add_to_cart", "title": "Add to Cart"},
-                        {"id": "cancel", "title": "Cancel"}
+                        {"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
+                        mobile,
+                        create_reply_button_message(products_text, buttons)
                     )
                 else:
-                    messenger.send_message("I couldn't find any matching products in the image.", mobile)
-
-        return "OK", 200
-
-    except Exception as e:
-        logger.error(f"Error in webhook handler: {str(e)}")
-        return "Internal server error", 500
-
-def handle_interactive_message(messenger, data, mobile, carts, shopping_assistant):
-    """Handle interactive message responses with enhanced cart functionality"""
-    try:
-        response = messenger.get_interactive_response(data)
-        button_id = response.get("id")
-        
-        if button_id == "add_to_cart":
-            # Add last analyzed products to cart
-            if mobile not in carts:
-                carts[mobile] = []
-            
-            # Add products from last analysis
-            if mobile in shopping_assistant.last_analyzed_products:
-                carts[mobile].extend(shopping_assistant.last_analyzed_products[mobile])
-                del shopping_assistant.last_analyzed_products[mobile]
-            
-            buttons = [
-                {"id": "continue_shopping", "title": "Continue Shopping"},
-                {"id": "checkout", "title": "Checkout"}
-            ]
-            button_message = create_reply_button_message(
-                "Products added to cart! What would you like to do next?", 
-                buttons
-            )
-            messenger.send_reply_button(
-                recipient_id=mobile,
-                button=button_message
-            )
-
-        elif button_id == "continue_shopping":
-            # Encourage further shopping
-            messenger.send_message(
-                "Great! You can continue shopping by sending me a text or image of what you're looking for.", 
-                mobile
-            )
-
-        elif button_id == "checkout":
-            if mobile not in carts or not carts[mobile]:
-                messenger.send_message("Your cart is empty!", mobile)
-                return
+                    messenger.send_message("❌ No products found in the image", mobile)
 
-            # Calculate total
-            total = sum(product['price'] for product in carts[mobile])
+        # Handle button interactions
+        elif message_type == "interactive":
+            response = messenger.get_interactive_response(data)
+            button_id = response.get("id")
             
-            # Generate receipt
-            receipt_text = "Your Purchase Receipt:\n\n"
-            for product in carts[mobile]:
-                receipt_text += f"- {product['product']}: ${product['price']:.2f}\n"
-            receipt_text += f"\nTotal: ${total:.2f}\n\n"
-            receipt_text += "Expected Delivery: Within 1 hour\n"
-            receipt_text += "Thank you for your purchase!"
-
-            buttons = [
-                {"id": "confirm_order", "title": "Confirm Order"},
-                {"id": "modify_order", "title": "Modify Order"}
-            ]
-            button_message = create_reply_button_message(receipt_text, buttons)
-            messenger.send_reply_button(
-                recipient_id=mobile,
-                button=button_message
-            )
+            if button_id == "add_to_cart":
+                if mobile in assistant.last_analyzed_products:
+                    if mobile not in carts:
+                        carts[mobile] = []
+                    carts[mobile].extend(assistant.last_analyzed_products[mobile])
+                    
+                    buttons = [
+                        {"id": "continue_shopping", "title": "🛍️ Continue Shopping"},
+                        {"id": "checkout", "title": "💳 Checkout"}
+                    ]
+                    messenger.send_reply_button(
+                        mobile,
+                        create_reply_button_message(
+                            "✅ Items added to cart!\nWhat would you like to do next?",
+                            buttons
+                        )
+                    )
+                else:
+                    messenger.send_message("❌ No products to add", mobile)
             
+            elif button_id == "checkout":
+                if mobile in carts and carts[mobile]:
+                    total = sum(p['price']*p.get('quantity',1) for p in carts[mobile])
+                    receipt = "📋 *Your Cart Summary:*\n\n"
+                    receipt += "\n".join(
+                        f"➤ {p['product']} x{p.get('quantity',1)} - ${p['price']}"
+                        for p in carts[mobile]
+                    )
+                    receipt += f"\n\n💳 Total: ${total:.2f}"
+                    
+                    buttons = [
+                        {"id": "confirm_checkout", "title": "✅ Confirm Order"},
+                        {"id": "edit_cart", "title": "✏️ Edit Cart"}
+                    ]
+                    messenger.send_reply_button(
+                        mobile,
+                        create_reply_button_message(receipt, buttons)
+                    )
+                else:
+                    messenger.send_message("🛒 Your cart is empty!", mobile)
 
-        elif button_id == "confirm_order":
-            # Finalize order
-            messenger.send_message(
-                "Order confirmed! Your items will be delivered within 1 hour. Thank you!", 
-                mobile
-            )
-            # Clear the cart
-            carts[mobile] = []
-
-        elif button_id == "modify_order":
-            # Allow user to modify order
-            messenger.send_message(
-                "Please tell us any changes you'd like to make to your order.", 
-                mobile
-            )
-
-    except Exception as e:
-        logger.error(f"Error handling interactive message: {str(e)}")
-        messenger.send_message("Sorry, something went wrong.", mobile)
-
-
+        return "OK", 200
 
-def test_connection():
-    """Test DNS resolution and connection to Graph API"""
-    try:
-        resolver = dns.resolver.Resolver()
-        resolver.nameservers = [nameserver1, nameserver2]
-        ip = resolver.resolve('graph.facebook.com', 'A')[0]
-        logger.info(f"DNS Resolution successful: {ip}")
-        
-        response = requests.get('https://graph.facebook.com')
-        logger.info(f"HTTPS Connection successful: {response.status_code}")
-        
     except Exception as e:
-        logger.error(f"Connection test failed: {e}")
+        logger.error(f"Webhook error: {str(e)}")
+        return "Internal Server Error", 500
 
 if __name__ == "__main__":
-    test_connection()
-    app.run(debug=True, host="0.0.0.0", port=7860)
+    app.run(host="0.0.0.0", port=7860)