Update agent_tools/ml_tools.py

agent_tools/ml_tools.py

@@ -10,7 +10,6 @@ import json
 from pathlib import Path
 from datetime import datetime
 import duckdb
-import streamlit as st
 
 # Global model cache for HF Spaces
 _model_cache = {}
@@ -41,7 +40,7 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
         # Load model
         model_data = load_model_with_cache()
         if model_data is None:
-            return json.dumps({"error": "Model not found. Please wait for training to complete."})
+            return json.dumps({"error": "Model not found. Please train the model first."})
         
         model = model_data['model']
         label_encoders = model_data['label_encoders']
@@ -54,20 +53,22 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
             CREATE TABLE customers AS 
             SELECT * FROM 'hf://datasets/SAP/SALT/I_Customer.parquet'
             LIMIT 2000
-        """)  # Limit for performance
+        """)
         
         conn.execute("""
             CREATE TABLE sales_docs AS 
             SELECT * FROM 'hf://datasets/SAP/SALT/I_SalesDocument.parquet'
             LIMIT 5000
-        """)  # Limit for performance
+        """)
         
         # Filter customers if specified
         if customer_ids:
             customer_list = [f"'{cid.strip()}'" for cid in customer_ids.split(',')]
             where_clause = f"WHERE c.Customer IN ({','.join(customer_list)})"
+            limit_clause = ""
         else:
-            where_clause = "LIMIT 500"  # Further limit for demo
+            where_clause = ""
+            limit_clause = "LIMIT 500"  # Limit for demo
         
         # Get customer data
         customer_data = conn.execute(f"""
@@ -81,13 +82,13 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
                 MIN(s.CreationDate) as first_order_date
             FROM customers c
             LEFT JOIN sales_docs s ON c.Customer = s.SoldToParty
-            {where_clause if not customer_ids else ""}
+            {where_clause}
             GROUP BY c.Customer, c.CustomerName, c.Country, c.CustomerGroup
-            {where_clause if customer_ids else ""}
+            {limit_clause}
         """).df()
         
         if len(customer_data) == 0:
-            return json.dumps({"error": "No customers found"})
+            return json.dumps({"error": "No customers found for analysis"})
         
         # Feature engineering (same as training)
         reference_date = pd.to_datetime('2024-12-31')
@@ -114,7 +115,6 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
                         customer_data[col].fillna('Unknown')
                     )
                 except:
-                    # Handle unseen categories
                     customer_data[f'{col}_encoded'] = 0
         
         # Make predictions
@@ -133,7 +133,7 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
             # High risk customers
             high_risk = results[results['churn_probability'] >= risk_threshold].sort_values(
                 'churn_probability', ascending=False
-            ).head(20)  # Limit results for HF Spaces
+            ).head(20)
             
             # Generate recommendations
             recommendations = []
@@ -153,8 +153,8 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
                 "high_risk_count": len(high_risk),
                 "churn_rate_predicted": round(len(high_risk) / len(results) * 100, 2) if len(results) > 0 else 0,
                 "urgent_actions": recommendations,
-                "model_performance": f"Accuracy: {model_data.get('accuracy', 'N/A')}",
-                "hf_spaces_note": "Results limited for demo performance"
+                "model_performance": "Model ready and operational",
+                "note": "Results limited for demo performance"
             })
             
         except Exception as e:
@@ -163,7 +163,7 @@ def predict_customer_churn_hf(customer_ids: str = None, risk_threshold: float =
     except Exception as e:
         return json.dumps({
             "error": f"Churn analysis failed: {str(e)}",
-            "suggestion": "Please ensure model is trained"
+            "suggestion": "Please ensure model is trained and data is available"
         })
 
 @tool