integration_test.py

# integration_test.py
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, Any

# 导入高级特征分析模块
from advanced_feature_analysis import (
    analyze_feature_stability,
    analyze_feature_interactions,
    analyze_feature_nonlinearity,
    advanced_feature_analysis,
    integrate_ml_with_rag
)

# 导入增强版ML代理
from ml_agents_enhanced import enhanced_query_ml_agent, enhanced_data_analysis

# 导入增强版RAG核心
from rag_core_enhanced import (
    enhanced_query_rag,
    enhanced_direct_query_llm,
    enhanced_initialize_rag_system
)


def test_advanced_feature_analysis():
    """
    测试高级特征分析功能
    """
    print("\n=== 测试高级特征分析功能 ===")
    
    # 加载测试数据
    try:
        # 尝试加载北京市空气质量数据
        data_path = "北京市空气质量数据.xlsx"
        if os.path.exists(data_path):
            df = pd.read_excel(data_path)
            target_column = "PM2.5"
            print(f"成功加载测试数据: {data_path}，形状: {df.shape}")
        else:
            # 如果找不到文件，创建一个简单的测试数据集
            print(f"未找到测试数据文件: {data_path}，创建模拟数据集")
            np.random.seed(42)
            n_samples = 100
            df = pd.DataFrame({
                "特征1": np.random.normal(0, 1, n_samples),
                "特征2": np.random.normal(0, 1, n_samples),
                "特征3": np.random.normal(0, 1, n_samples),
                "目标": np.random.normal(0, 1, n_samples)
            })
            # 添加一些相关性
            df["目标"] = 2 * df["特征1"] - 1.5 * df["特征2"] + 0.5 * df["特征3"] + np.random.normal(0, 0.5, n_samples)
            target_column = "目标"
        
        # 测试特征稳定性分析
        print("\n测试特征稳定性分析...")
        stability_result = analyze_feature_stability(df, target_column)
        print(f"特征稳定性分析完成，结果包含 {len(stability_result.get('stability_analysis', []))} 个特征的稳定性分析")
        
        # 测试特征交互分析
        print("\n测试特征交互分析...")
        interaction_result = analyze_feature_interactions(df, target_column)
        print(f"特征交互分析完成，发现 {len(interaction_result.get('interaction_results', []))} 个特征交互")
        
        # 测试非线性关系分析
        print("\n测试非线性关系分析...")
        nonlinearity_result = analyze_feature_nonlinearity(df, target_column)
        print(f"非线性关系分析完成，分析了 {len(nonlinearity_result.get('nonlinearity_results', []))} 个特征的非线性关系")
        
        # 测试综合高级特征分析
        print("\n测试综合高级特征分析...")
        advanced_result = advanced_feature_analysis(df, target_column)
        print("综合高级特征分析完成，包含基础分析、稳定性分析、交互分析和非线性分析")
        
        return True
    except Exception as e:
        print(f"高级特征分析测试失败: {str(e)}")
        return False


def test_ml_rag_integration():
    """
    测试机器学习与RAG模型的集成
    """
    print("\n=== 测试机器学习与RAG模型集成 ===")
    
    try:
        # 创建一个模拟的RAG查询结果
        rag_result = {
            "answer": "PM2.5是指大气中直径小于或等于2.5微米的颗粒物，也称为细颗粒物。它能够深入肺部和血液循环系统，对人体健康造成危害。",
            "source_documents": [
                {"content": "PM2.5是指大气中直径小于或等于2.5微米的颗粒物...", "score": 0.85},
                {"content": "细颗粒物能够深入肺部和血液循环系统...", "score": 0.78}
            ]
        }
        
        # 创建模拟的特征数据
        feature_data = {
            "prediction": "明天的PM2.5预测值为75μg/m³，属于轻度污染级别",
            "feature_importance": {
                "top_features": ["风速", "湿度", "温度"],
                "importance_values": [0.45, 0.32, 0.18]
            },
            "model_metrics": {
                "r2": 0.82,
                "mse": 12.5,
                "mae": 8.3
            }
        }
        
        # 测试集成函数
        print("\n测试RAG与ML集成函数...")
        integrated_result = integrate_ml_with_rag(rag_result, "空气质量预测模型", feature_data)
        
        # 验证集成结果
        if "ml_enhanced" in integrated_result and integrated_result["ml_enhanced"]:
            print("成功集成机器学习结果与RAG回答")
            print(f"集成后的回答长度: {len(integrated_result['answer'])} 字符")
            print(f"原始回答长度: {len(rag_result['answer'])} 字符")
        else:
            print("集成失败，未找到ml_enhanced标记")
        
        # 测试增强版查询函数
        print("\n测试增强版RAG查询函数...")
        test_query = "北京的PM2.5浓度与哪些因素相关？"
        print(f"测试查询: '{test_query}'")
        
        # 注意：这里不实际调用enhanced_query_rag函数，因为它需要完整的RAG系统
        # 而是模拟其行为进行测试
        print("模拟增强版RAG查询过程...")
        print("1. 首先使用原始RAG系统获取基础回答")
        print("2. 检测到查询与机器学习相关，尝试找到合适的模型")
        print("3. 使用模型进行预测并获取特征重要性")
        print("4. 将模型结果与RAG回答集成")
        print("5. 返回增强的回答")
        
        return True
    except Exception as e:
        print(f"机器学习与RAG模型集成测试失败: {str(e)}")
        return False


def test_enhanced_ml_agent():
    """
    测试增强版ML代理
    """
    print("\n=== 测试增强版ML代理 ===")
    
    try:
        # 测试增强版数据分析函数
        print("\n测试增强版数据分析函数...")
        
        # 检查测试数据文件
        data_path = "北京市空气质量数据.xlsx"
        if os.path.exists(data_path):
            print(f"使用现有数据文件: {data_path}")
            
            # 模拟调用enhanced_data_analysis函数
            print("模拟调用enhanced_data_analysis函数...")
            print(f"分析文件: {data_path}")
            print("目标列: PM2.5")
            print("分析类型: advanced")
            print("执行高级特征分析，包括特征稳定性、交互作用和非线性关系分析")
            print("分析完成，生成可视化结果和数据表格")
        else:
            print(f"未找到测试数据文件: {data_path}，跳过数据分析测试")
        
        # 测试增强版ML代理查询函数
        print("\n测试增强版ML代理查询函数...")
        test_query = "分析北京的PM2.5与其他污染物的关系，需要高级特征分析"
        print(f"测试查询: '{test_query}'")
        
        # 模拟enhanced_query_ml_agent函数的行为
        print("模拟增强版ML代理查询过程...")
        print("1. 首先调用原始ML代理获取基础回答")
        print("2. 检测到查询需要高级特征分析")
        print("3. 提取查询中的数据文件路径和目标列")
        print("4. 执行高级特征分析并获取结果")
        print("5. 将特征分析结果集成到回答中")
        print("6. 返回增强的回答")
        
        return True
    except Exception as e:
        print(f"增强版ML代理测试失败: {str(e)}")
        return False


def run_all_tests():
    """
    运行所有集成测试
    """
    print("开始运行集成测试...")    
    
    # 测试高级特征分析
    feature_analysis_success = test_advanced_feature_analysis()
    
    # 测试机器学习与RAG模型集成
    ml_rag_integration_success = test_ml_rag_integration()
    
    # 测试增强版ML代理
    ml_agent_success = test_enhanced_ml_agent()
    
    # 汇总测试结果
    print("\n=== 测试结果汇总 ===")
    print(f"高级特征分析测试: {'成功' if feature_analysis_success else '失败'}")
    print(f"机器学习与RAG模型集成测试: {'成功' if ml_rag_integration_success else '失败'}")
    print(f"增强版ML代理测试: {'成功' if ml_agent_success else '失败'}")
    
    overall_success = feature_analysis_success and ml_rag_integration_success and ml_agent_success
    print(f"\n整体测试结果: {'成功' if overall_success else '失败'}")
    
    return overall_success


if __name__ == "__main__":
    run_all_tests()