Update app.py

app.py

@@ -26,7 +26,7 @@ model = configure_gemini(os.environ['GOOGLE_API_KEY'])
 # Initialize Gemini models
 llm_flash_exp = ChatGoogleGenerativeAI(
     model="gemini-2.0-flash-exp",
-    max_retries=2  # Limit retries to 2
+    max_retries=2
 )
 
 class SmartShoppingAssistant:
@@ -36,38 +36,82 @@ class SmartShoppingAssistant:
         self.setup_agent()
     
     def find_closest_product(self, product_name, threshold=0.6):
-        """Find the closest matching product name using fuzzy matching"""
         matches = get_close_matches(
             product_name.upper(),
             self.df['ProductName'].str.upper().tolist(),
-            n=1,
+            n=3,  # Increased to get more potential matches
             cutoff=threshold
         )
-        return matches[0] if matches else None
+        return matches if matches else []
     
-    def search_products_fuzzy(self, product_names):
-        """Search for products using fuzzy matching"""
+    def match_products_with_catalogue(self, extracted_items):
+        """Match extracted items with catalogue products using Gemini"""
+        product_list = self.df['ProductName'].tolist()
+        product_string = ", ".join(product_list)
+        
+        prompt = f"""
+        Given these extracted items and quantities:
+        {extracted_items}
+        
+        And this product catalogue:
+        {product_string}
+
+        Match each item with the most appropriate product from the catalogue.
+        For each item, provide:
+        1. The exact product name from the catalogue
+        2. The quantity (if specified, otherwise assume 1)
+        3. Any specific requirements (brand, size, etc.)
+
+        Format the response as:
+        ProductName == "MATCHED_PRODUCT" quantity: NUMBER or ProductName == "MATCHED_PRODUCT" quantity: NUMBER
+
+        Only include products that have good matches in the catalogue.
+        """
+        
+        try:
+            matches = llm_flash_exp.predict(prompt)
+            return matches.strip()
+        except Exception as e:
+            return f"Error matching products: {str(e)}"
+    
+    def search_products_fuzzy(self, product_names_with_quantities):
+        """Search for products using fuzzy matching with quantity information"""
         results = pd.DataFrame()
-        for name in product_names:
-            closest_match = self.find_closest_product(name)
-            if closest_match:
-                matched_products = self.df[self.df['ProductName'].str.upper() == closest_match.upper()]
-                results = pd.concat([results, matched_products])
+        for item in product_names_with_quantities:
+            product_info = item.split('quantity:')
+            product_name = product_info[0].strip()
+            quantity = int(product_info[1].strip()) if len(product_info) > 1 else 1
+            
+            # Clean up product name
+            if 'ProductName ==' in product_name:
+                product_name = product_name.split('==')[1].strip(' "\'')
+            
+            closest_matches = self.find_closest_product(product_name)
+            for match in closest_matches:
+                matched_products = self.df[self.df['ProductName'].str.upper() == match.upper()]
+                if not matched_products.empty:
+                    matched_products['Quantity'] = quantity
+                    results = pd.concat([results, matched_products])
+                    break
+        
         return results
         
     def setup_agent(self):
         """Set up the LangChain agent with necessary tools"""
         def search_products(query):
             try:
-                # Extract product names from the query
-                product_names = [name.strip('"\'') for name in query.split(' or ')]
-                product_names = [name.split('==')[1].strip() if '==' in name else name for name in product_names]
+                # Split into individual product entries
+                product_entries = [entry.strip() for entry in query.split('or')]
                 
-                if 'RetailPrice' in self.df.columns:
-                    self.df['RetailPrice'] = pd.to_numeric(self.df['RetailPrice'].str.replace('$', ''), errors='coerce')
-                
-                results = self.search_products_fuzzy(product_names)
-                return results.to_string() if not results.empty else "No products found matching your criteria."
+                results = self.search_products_fuzzy(product_entries)
+                if not results.empty:
+                    # Format results with quantity
+                    formatted_results = results.apply(
+                        lambda x: f"{x['ProductName']} (Quantity: {x['Quantity']})\nPrice: ${x['RetailPrice']:.2f}\n",
+                        axis=1
+                    )
+                    return "\n".join(formatted_results)
+                return "No products found matching your criteria."
             except Exception as e:
                 return f"Error executing query: {str(e)}"
         
@@ -85,39 +129,40 @@ class SmartShoppingAssistant:
             llm=llm_flash_exp,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=True,
-            max_iterations=3  # Limit the number of iterations
+            max_iterations=3
         )
     
     def process_natural_language_query(self, query):
-        """Process natural language query and return relevant products"""
-
-        product_list = self.df['ProductName'].tolist()
-        product_string = ", ".join(product_list)
+        """Process natural language query with two-step matching"""
         try:
-            prompt = f"""
-            Convert this shopping request into a list of product names to search for. The response should be in the format:
-            ProductName == "PRODUCT1" or ProductName == "PRODUCT2" or ProductName == "PRODUCT3"
+            # First step: Extract items and quantities
+            extraction_prompt = f"""
+            Extract the products and their quantities from this shopping request.
+            Include any specific requirements mentioned.
             
             Shopping request: {query}
-
-            Available products will be matched approximately or use you intelligent to look at product and list and see the ones that match
-            , so focus on the main product names.
-            Return only the search string, nothing else.
-
-            This is the list of products: {product_string}
+            
+            Format each item with its quantity (assume 1 if not specified).
             """
             
-            response = llm_flash_exp.predict(prompt)
-            structured_query = response.strip()
+            extracted_items = llm_flash_exp.predict(extraction_prompt)
+            
+            # Second step: Match with catalogue
+            matched_products = self.match_products_with_catalogue(extracted_items)
             
-            result = self.agent.run(f"Search for products matching the specified names: {structured_query}")
+            # Third step: Search and return results
+            result = self.agent.run(f"Search for products matching the specified names: {matched_products}")
             return result
         except Exception as e:
             return f"Error processing query: {str(e)}"
     
     def extract_text_from_image(self, image):
         """Extract text from uploaded image using Gemini"""
-        prompt = """Analyze this image and extract products text from the image"""
+        prompt = """
+        Analyze this image and extract products and their quantities.
+        If quantities aren't specified, make reasonable assumptions based on typical shopping patterns.
+        List each item with its quantity.
+        """
         try:
             response = model.generate_content([prompt, image])
             return response.text
@@ -147,7 +192,6 @@ def main():
     df = load_product_data()
     assistant = SmartShoppingAssistant(df)
     
-    # Sidebar for file uploads
     with st.sidebar:
         st.header("Upload Shopping List")
         uploaded_file = st.file_uploader(
@@ -169,15 +213,14 @@ def main():
             except Exception as e:
                 st.error(f"Error processing file: {str(e)}")
 
-    # Main content area
     col1, col2 = st.columns([2, 1])
     
     with col1:
         st.header("Search Products")
         query = st.text_area(
-            "Describe what you're looking for:",
+            "Describe what you're looking for (include quantities if needed):",
             height=100,
-            placeholder="Example: I need healthy breakfast cereals under $5",
+            placeholder="Example: 2 boxes of healthy breakfast cereals under $5, 1 gallon of milk",
             value=st.session_state.get('query', '')
         )