Update app.py

app.py

@@ -201,6 +201,10 @@ class ImageScraper:
         ]
 
     def get_images(self, query: str, num_images: int = 15) -> Dict[str, List[Dict[str, str]]]:
+        """
+        Get images from Pixabay API with enhanced error handling and categorization
+        Returns a dictionary with categorized image results
+        """
         try:
             # Ensure query is not empty
             if not query.strip():
@@ -216,67 +220,91 @@ class ImageScraper:
                 'lang': 'en'
             }
             
-            print(f"Querying Pixabay API with params: {params}")  # Debug log
-            response = requests.get(base_url, params=params)
-            print(f"API Response status: {response.status_code}")  # Debug log
+            print(f"Querying Pixabay API with params: {params}")
+            response = requests.get(base_url, params=params, headers=self.headers)
+            print(f"API Response status: {response.status_code}")
             
             if response.status_code == 200:
                 data = response.json()
                 hits = data.get('hits', [])
-                print(f"Found {len(hits)} images")  # Debug log
+                print(f"Found {len(hits)} images")
                 
                 if not hits:
                     # If no results, try backup keywords
                     backup_keywords = ["digital security", "technology", "cyber security", "business technology"]
                     for keyword in backup_keywords:
                         params['q'] = keyword
-                        backup_response = requests.get(base_url, params=params)
+                        backup_response = requests.get(base_url, params=params, headers=self.headers)
                         if backup_response.status_code == 200:
                             backup_data = backup_response.json()
                             hits.extend(backup_data.get('hits', []))
                             if hits:
                                 break
-    
-                # Structure response
+                
+                # Structure response with categories
                 result = {
                     'primary': [],
                     'secondary': [],
                     'general': []
                 }
-    
-                # Ensure we have at least some images
+                
+                # If no hits found even after backup keywords, return stock images
                 if not hits:
-                    result['general'] = [{'url': url, 'keyword': 'technology', 'relevance': 'Fallback', 'tags': 'technology'} 
-                                       for url in self.get_stock_images()]
+                    stock_images = self.get_stock_images()
+                    result['general'] = [
+                        {
+                            'url': url,
+                            'keyword': 'technology',
+                            'relevance': 'Fallback',
+                            'tags': 'technology'
+                        } for url in stock_images
+                    ]
                     return result
-    
-                # Categorize found images
+                
+                # Process and categorize the found images
                 for hit in hits[:num_images]:
                     image_data = {
                         'url': hit['largeImageURL'],
                         'keyword': query,
-                        'relevance': 'Primary match',
+                        'relevance': 'Primary match' if query.lower() in hit.get('tags', '').lower() else 'Related',
                         'tags': hit.get('tags', '')
                     }
                     
-                    if len(result['primary']) < num_images // 3:
+                    # Distribute images across categories
+                    if len(result['primary']) < num_images // 3 and query.lower() in hit.get('tags', '').lower():
                         result['primary'].append(image_data)
                     elif len(result['secondary']) < num_images // 3:
                         result['secondary'].append(image_data)
                     else:
                         result['general'].append(image_data)
     
+                # Ensure we have some results in each category
+                if not result['primary'] and result['secondary']:
+                    result['primary'] = result['secondary'][:num_images//3]
+                    result['secondary'] = result['secondary'][num_images//3:]
+                
+                if not result['secondary'] and result['general']:
+                    result['secondary'] = result['general'][:num_images//3]
+                    result['general'] = result['general'][num_images//3:]
+    
                 return result
     
         except Exception as e:
-            print(f"Error in get_images: {str(e)}")  # Debug log
+            print(f"Error in get_images: {str(e)}")
             # Return stock images as fallback
+            stock_images = self.get_stock_images()
             return {
-                'general': [{'url': url, 'keyword': 'technology', 'relevance': 'Fallback', 'tags': 'technology'} 
-                           for url in self.get_stock_images()]
+                'general': [
+                    {
+                        'url': url,
+                        'keyword': 'technology',
+                        'relevance': 'Fallback',
+                        'tags': 'technology'
+                    } for url in stock_images
+                ]
             }
-        
-        
+            
+            
     def get_fallback_keywords(self) -> List[Dict[str, str]]:
         """Return fallback keywords if AI extraction fails"""
         return [