Update app.py

app.py

@@ -199,34 +199,58 @@ import pandas as pd
 from typing import Dict, List, Tuple, Union
 from smolagents import CodeAgent, HfApiModel, tool
 
-def process_agent_response(response: Union[Dict, str]) -> Tuple[pd.DataFrame, str]:
+def process_agent_response(response: Union[Dict, List, str]) -> Tuple[pd.DataFrame, str]:
     """
     Process the agent's response and convert it to a DataFrame
     Returns DataFrame and error message (if any)
     """
-    if isinstance(response, dict):
-        if 'error' in response:
-            return pd.DataFrame(), f"Error: {response['error']}"
-        elif 'products' in response:
+    def extract_products_from_response(resp):
+        """Helper function to extract product data from various response formats"""
+        if isinstance(resp, list):
+            # Response is already a list of products
+            return resp
+        elif isinstance(resp, dict):
+            # Check if it's wrapped in a 'products' key
+            if 'products' in resp:
+                return resp['products']
+            # Check if it's a single product
+            elif 'title' in resp:
+                return [resp]
+            elif 'error' in resp:
+                return None
+        elif isinstance(resp, str):
             try:
-                df = pd.DataFrame(response['products'])
-                return df, ""
-            except Exception as e:
-                return pd.DataFrame(), f"Error processing data: {str(e)}"
-    elif isinstance(response, str):
-        try:
-            # Try to evaluate string response as dict if needed
-            import ast
-            data = ast.literal_eval(response)
-            if isinstance(data, dict):
-                if 'products' in data:
-                    df = pd.DataFrame(data['products'])
-                else:
-                    df = pd.DataFrame([data])
-                return df, ""
-        except:
-            pass
-    return pd.DataFrame(), "Unexpected response format"
+                # Try to parse the string response
+                import ast
+                # Remove common prefixes
+                if "Final answer:" in resp:
+                    resp = resp.split("Final answer:", 1)[1].strip()
+                elif "Out - Final answer:" in resp:
+                    resp = resp.split("Out - Final answer:", 1)[1].strip()
+                
+                parsed = ast.literal_eval(resp)
+                # Recursively process the parsed result
+                return extract_products_from_response(parsed)
+            except:
+                pass
+        return None
+
+    try:
+        products = extract_products_from_response(response)
+        
+        if products is None:
+            return pd.DataFrame(), "No valid product data found"
+        
+        # Convert to DataFrame
+        df = pd.DataFrame(products)
+        
+        # Clean up column names
+        df.columns = [col.lower().strip() for col in df.columns]
+        
+        return df, ""
+            
+    except Exception as e:
+        return pd.DataFrame(), f"Error processing data: {str(e)}"
 
 def search_products(keyword: str, max_products: int) -> Tuple[pd.DataFrame, str, str]:
     """
@@ -238,13 +262,24 @@ def search_products(keyword: str, max_products: int) -> Tuple[pd.DataFrame, str,
         df, error_msg = process_agent_response(result)
         
         if not df.empty:
-            # Select and reorder relevant columns
+            # Select and reorder relevant columns, using lowercase names
             display_columns = [
                 'title', 'price', 'rating', 'reviews', 'description', 
                 'bullet_points', 'average_rating', 'total_reviews'
             ]
-            display_columns = [col for col in display_columns if col in df.columns]
-            df = df[display_columns]
+            # Filter for columns that actually exist in the DataFrame
+            available_columns = [col for col in display_columns if col in df.columns]
+            df = df[available_columns]
+            
+            # Clean up the display
+            if 'price' in df.columns:
+                df['price'] = df['price'].apply(lambda x: f"${str(x).strip('.')}")
+            
+            # Truncate long text fields for better display
+            if 'description' in df.columns:
+                df['description'] = df['description'].apply(lambda x: x[:200] + '...' if len(x) > 200 else x)
+            if 'bullet_points' in df.columns:
+                df['bullet_points'] = df['bullet_points'].apply(lambda x: x[:200] + '...' if len(x) > 200 else x)
             
             status_msg = f"Found {len(df)} products"
             return df, status_msg, ""