akaAKR47/agent.py

from pydantic import BaseModel, Field
from typing import Dict, Any, List
import polars as pl
import joblib
from data_processing import preprocess_data

# 定义输入输出模型
class CustomerData(BaseModel):
    """客户数据模型"""
    gender: str = Field(..., description="性别: Male 或 Female")
    SeniorCitizen: int = Field(..., description="是否为老年人: 0 或 1")
    Partner: str = Field(..., description="是否有伴侣: Yes 或 No")
    Dependents: str = Field(..., description="是否有家属: Yes 或 No")
    tenure: int = Field(..., description="在网时长，单位为月")
    PhoneService: str = Field(..., description="是否开通电话服务: Yes 或 No")
    MultipleLines: str = Field(..., description="是否开通多条线路: Yes、No 或 No phone service")
    InternetService: str = Field(..., description="网络服务类型: DSL、Fiber optic 或 No")
    OnlineSecurity: str = Field(..., description="是否开通在线安全服务: Yes、No 或 No internet service")
    OnlineBackup: str = Field(..., description="是否开通在线备份服务: Yes、No 或 No internet service")
    DeviceProtection: str = Field(..., description="是否开通设备保护服务: Yes、No 或 No internet service")
    TechSupport: str = Field(..., description="是否开通技术支持服务: Yes、No 或 No internet service")
    StreamingTV: str = Field(..., description="是否开通流媒体电视服务: Yes、No 或 No internet service")
    StreamingMovies: str = Field(..., description="是否开通流媒体电影服务: Yes、No 或 No internet service")
    Contract: str = Field(..., description="合同类型: Month-to-month、One year 或 Two year")
    PaperlessBilling: str = Field(..., description="是否使用无纸化账单: Yes 或 No")
    PaymentMethod: str = Field(..., description="支付方式")
    MonthlyCharges: float = Field(..., description="月费用")
    TotalCharges: float = Field(..., description="总费用")

class ChurnPrediction(BaseModel):
    """客户流失预测结果"""
    prediction: int = Field(..., description="预测结果: 0 表示不流失，1 表示流失")
    probability: float = Field(..., description="流失概率")
    model_used: str = Field(..., description="使用的模型")

class ActionSuggestion(BaseModel):
    """基于预测结果的行动建议"""
    customer_id: str = Field(..., description="客户ID")
    prediction: int = Field(..., description="预测结果: 0 表示不流失，1 表示流失")
    probability: float = Field(..., description="流失概率")
    suggestions: List[str] = Field(..., description="可执行的行动建议")

# Agent工具类
class ChurnPredictionAgent:
    def __init__(self):
        # 加载最佳模型（使用LightGBM，因为它通常表现更好）
        self.model = joblib.load("models/lightgbm_model.pkl")
        self.model_name = "lightgbm"
    
    # 工具1：ML预测工具
    def predict_churn(self, customer_data: CustomerData) -> ChurnPrediction:
        """
        预测客户是否会流失
        
        Args:
            customer_data: 客户数据
            
        Returns:
            ChurnPrediction: 预测结果
        """
        # 将客户数据转换为Polars DataFrame
        customer_dict = customer_data.model_dump()
        df = pl.DataFrame([customer_dict])
        
        # 数据预处理（使用专门的单个客户预处理函数）
        from data_processing import preprocess_single_customer
        X_np = preprocess_single_customer(df)
        
        # 预测
        probability = self.model.predict_proba(X_np)[0, 1]
        prediction = 1 if probability >= 0.5 else 0
        
        return ChurnPrediction(
            prediction=prediction,
            probability=probability,
            model_used=self.model_name
        )
    
    # 工具2：行动建议工具
    def get_action_suggestions(self, customer_id: str, prediction: int, 
                             probability: float, customer_data: CustomerData) -> ActionSuggestion:
        """
        基于预测结果给出可执行的行动建议
        
        Args:
            customer_id: 客户ID
            prediction: 预测结果
            probability: 流失概率
            customer_data: 客户数据
            
        Returns:
            ActionSuggestion: 行动建议
        """
        suggestions = []
        
        if prediction == 1:
            # 高流失风险客户
            suggestions.append(f"客户 {customer_id} 有 {probability:.2%} 的概率会流失，需要重点关注")
            
            # 基于客户特征给出具体建议
            if customer_data.Contract == "Month-to-month":
                suggestions.append("建议提供长期合同折扣，鼓励客户转为一年或两年合同")
            
            if customer_data.TechSupport == "No":
                suggestions.append("建议提供免费的技术支持服务，提高客户满意度")
            
            if customer_data.OnlineSecurity == "No":
                suggestions.append("建议提供免费的在线安全服务，增加客户粘性")
            
            if customer_data.tenure < 12:
                suggestions.append("建议提供新客户忠诚度奖励计划，鼓励客户继续使用服务")
            
            if customer_data.MonthlyCharges > 70:
                suggestions.append(f"客户月费用较高 ({customer_data.MonthlyCharges} 元)，建议提供费用优化方案")
        else:
            # 低流失风险客户
            suggestions.append(f"客户 {customer_id} 流失风险较低 ({probability:.2%})，可维持现有服务")
            
            # 基于客户特征给出具体建议
            if customer_data.Contract == "Month-to-month":
                suggestions.append("建议定期发送满意度调查，了解客户需求")
            
            if customer_data.tenure >= 24:
                suggestions.append("建议提供忠诚客户专属优惠，巩固客户关系")
        
        return ActionSuggestion(
            customer_id=customer_id,
            prediction=prediction,
            probability=probability,
            suggestions=suggestions
        )
    
    # 工具3：批量预测工具（额外工具）
    def batch_predict(self, customer_data_list: List[CustomerData]) -> List[ChurnPrediction]:
        """
        批量预测客户是否会流失
        
        Args:
            customer_data_list: 客户数据列表
            
        Returns:
            List[ChurnPrediction]: 预测结果列表
        """
        results = []
        for customer_data in customer_data_list:
            result = self.predict_churn(customer_data)
            results.append(result)
        return results

# 测试Agent
if __name__ == "__main__":
    # 创建Agent实例
    agent = ChurnPredictionAgent()
    
    # 测试数据
    test_customer = CustomerData(
        gender="Male",
        SeniorCitizen=0,
        Partner="Yes",
        Dependents="No",
        tenure=12,
        PhoneService="Yes",
        MultipleLines="No",
        InternetService="Fiber optic",
        OnlineSecurity="No",
        OnlineBackup="Yes",
        DeviceProtection="No",
        TechSupport="No",
        StreamingTV="Yes",
        StreamingMovies="Yes",
        Contract="Month-to-month",
        PaperlessBilling="Yes",
        PaymentMethod="Electronic check",
        MonthlyCharges=79.85,
        TotalCharges=977.6
    )
    
    # 1. 使用ML预测工具
    print("=== 使用ML预测工具 ===")
    prediction_result = agent.predict_churn(test_customer)
    print(f"预测结果: {'会流失' if prediction_result.prediction == 1 else '不会流失'}")
    print(f"流失概率: {prediction_result.probability:.2%}")
    print(f"使用模型: {prediction_result.model_used}")
    
    # 2. 使用行动建议工具
    print("\n=== 使用行动建议工具 ===")
    suggestions = agent.get_action_suggestions(
        customer_id="TEST-123",
        prediction=prediction_result.prediction,
        probability=prediction_result.probability,
        customer_data=test_customer
    )
    
    print(f"客户ID: {suggestions.customer_id}")
    print(f"预测结果: {'会流失' if suggestions.prediction == 1 else '不会流失'}")
    print(f"流失概率: {suggestions.probability:.2%}")
    print("行动建议:")
    for i, suggestion in enumerate(suggestions.suggestions, 1):
        print(f"  {i}. {suggestion}")