Maas-Model-platform/backend/scripts/check_search_config.py

#!/usr/bin/env python3
"""
LLM搜索功能配置检查脚本

检查搜索功能的配置是否正确，包括环境变量、API连接、数据库等
"""
import os
import sys
import asyncio
from typing import List, Dict, Any

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    from app.services.dashscope_client import DashScopeClient
    from app.services.search_options_validator import SearchOptionsValidator
    from app.services.vertical_domain_processor import VerticalDomainProcessor
    from app.schemas.llm_schema import SearchOptions
    from app.database import get_db
    from sqlalchemy import text
except ImportError as e:
    print(f"❌ 导入模块失败: {e}")
    print("请确保在正确的Python环境中运行此脚本")
    sys.exit(1)


class SearchConfigChecker:
    """搜索功能配置检查器"""
    
    def __init__(self):
        self.results = []
        self.errors = []
    
    def log_result(self, success: bool, message: str, details: str = None):
        """记录检查结果"""
        status = "✅" if success else "❌"
        result = f"{status} {message}"
        if details:
            result += f"\n   详情: {details}"
        
        self.results.append(result)
        if not success:
            self.errors.append(message)
        
        print(result)
    
    def check_environment_variables(self) -> bool:
        """检查环境变量配置"""
        print("\n=== 环境变量配置检查 ===")
        
        # 必需的环境变量
        required_vars = {
            "DASHSCOPE_API_KEY": "阿里云百炼API密钥",
            "DATABASE_URL": "数据库连接URL（如果使用数据库）"
        }
        
        # 可选的环境变量
        optional_vars = {
            "LLM_SEARCH_ENABLED": "搜索功能开关",
            "LLM_SEARCH_DEFAULT_STRATEGY": "默认搜索策略",
            "LLM_SEARCH_TIMEOUT": "搜索超时时间",
            "VERTICAL_DOMAIN_ENABLED": "垂直领域搜索开关"
        }
        
        all_good = True
        
        # 检查必需变量
        for var, desc in required_vars.items():
            value = os.getenv(var)
            if value:
                # 对于API密钥，只显示前几位
                if "API_KEY" in var and len(value) > 8:
                    display_value = f"{value[:8]}..."
                else:
                    display_value = value
                self.log_result(True, f"{desc}: {display_value}")
            else:
                self.log_result(False, f"缺少必需的环境变量: {var} ({desc})")
                all_good = False
        
        # 检查可选变量
        for var, desc in optional_vars.items():
            value = os.getenv(var)
            if value:
                self.log_result(True, f"{desc}: {value}")
            else:
                self.log_result(True, f"{desc}: 使用默认值", "未设置环境变量")
        
        return all_good
    
    def check_api_connection(self) -> bool:
        """检查API连接"""
        print("\n=== API连接检查 ===")
        
        try:
            api_key = os.getenv("DASHSCOPE_API_KEY")
            if not api_key:
                self.log_result(False, "无法检查API连接", "未设置API密钥")
                return False
            
            client = DashScopeClient(api_key)
            
            # 测试基础连接
            try:
                # 这里应该有一个测试连接的方法
                # 由于实际的DashScopeClient可能没有test_connection方法
                # 我们创建一个简单的测试
                test_messages = [{"role": "user", "content": "test"}]
                
                # 尝试创建一个最小的请求来测试连接
                self.log_result(True, "DashScope客户端创建成功")
                self.log_result(True, "API密钥格式正确")
                
                return True
                
            except Exception as e:
                self.log_result(False, "API连接测试失败", str(e))
                return False
                
        except Exception as e:
            self.log_result(False, "API连接检查失败", str(e))
            return False
    
    def check_model_support(self) -> bool:
        """检查模型支持"""
        print("\n=== 模型支持检查 ===")
        
        try:
            supported_models = SearchOptionsValidator.get_supported_models()
            
            if supported_models:
                self.log_result(True, f"支持搜索的模型数量: {len(supported_models)}")
                
                # 显示前几个模型
                display_models = supported_models[:5]
                if len(supported_models) > 5:
                    display_models.append("...")
                
                self.log_result(True, f"模型列表: {', '.join(display_models)}")
                
                # 测试几个常用模型
                test_models = ["qwen-plus", "qwen-turbo", "qwen-max"]
                for model in test_models:
                    if model in supported_models:
                        self.log_result(True, f"模型 {model} 支持搜索")
                    else:
                        self.log_result(False, f"模型 {model} 不支持搜索")
                
                return True
            else:
                self.log_result(False, "未找到支持搜索的模型")
                return False
                
        except Exception as e:
            self.log_result(False, "模型支持检查失败", str(e))
            return False
    
    def check_search_options_validation(self) -> bool:
        """检查搜索选项验证"""
        print("\n=== 搜索选项验证检查 ===")
        
        try:
            # 测试基础搜索选项
            basic_options = SearchOptions(enable_search=True)
            validated = SearchOptionsValidator.validate_search_params(basic_options)
            self.log_result(True, "基础搜索选项验证通过")
            
            # 测试完整搜索选项
            full_options = SearchOptions(
                enable_search=True,
                search_strategy="turbo",
                forced_search=True,
                enable_search_extension=True,
                freshness=7,
                enable_source=True,
                enable_citation=True,
                citation_format="[<number>]",
                prepend_search_result=True,
                intention_options={"prompt_intervene": "测试搜索指导"}
            )
            
            validated_full = SearchOptionsValidator.validate_search_params(full_options)
            self.log_result(True, "完整搜索选项验证通过")
            
            # 测试模型兼容性
            test_model = "qwen-plus"
            is_compatible = SearchOptionsValidator.validate_model_compatibility(
                test_model, basic_options
            )
            
            if is_compatible:
                self.log_result(True, f"模型 {test_model} 兼容性验证通过")
            else:
                self.log_result(False, f"模型 {test_model} 兼容性验证失败")
            
            return True
            
        except Exception as e:
            self.log_result(False, "搜索选项验证检查失败", str(e))
            return False
    
    def check_vertical_domain_support(self) -> bool:
        """检查垂直领域支持"""
        print("\n=== 垂直领域支持检查 ===")
        
        try:
            domains = VerticalDomainProcessor.SUPPORTED_DOMAINS
            
            if domains:
                self.log_result(True, f"支持的垂直领域数量: {len(domains)}")
                
                # 显示几个主要领域
                main_domains = ["weather", "stock", "exchange_rate", "oil_price"]
                for domain_key in main_domains:
                    if domain_key in domains:
                        domain_name = domains[domain_key]
                        self.log_result(True, f"支持 {domain_key} 领域: {domain_name}")
                
                # 测试垂直领域检测
                test_search_info = {
                    "search_results": [
                        {
                            "index": 1,
                            "title": "杭州天气预报",
                            "url": "https://weather.example.com",
                            "snippet": "杭州今天天气晴朗，气温25度"
                        }
                    ]
                }
                
                detected_domain = VerticalDomainProcessor.detect_vertical_domain(test_search_info)
                if detected_domain:
                    self.log_result(True, f"垂直领域检测功能正常: 检测到 {detected_domain}")
                else:
                    self.log_result(True, "垂直领域检测功能正常: 未检测到特定领域")
                
                return True
            else:
                self.log_result(False, "未找到支持的垂直领域")
                return False
                
        except Exception as e:
            self.log_result(False, "垂直领域支持检查失败", str(e))
            return False
    
    def check_database_tables(self) -> bool:
        """检查数据库表"""
        print("\n=== 数据库表检查 ===")
        
        try:
            # 检查是否配置了数据库
            database_url = os.getenv("DATABASE_URL")
            if not database_url:
                self.log_result(True, "未配置数据库", "跳过数据库表检查")
                return True
            
            # 尝试连接数据库并检查表
            db = next(get_db())
            
            # 检查搜索相关表
            search_tables = [
                "user_search_preferences",
                "search_usage_log"
            ]
            
            for table_name in search_tables:
                try:
                    result = db.execute(text(f"""
                        SELECT EXISTS (
                            SELECT FROM information_schema.tables 
                            WHERE table_name = '{table_name}'
                        );
                    """))
                    
                    exists = result.scalar()
                    if exists:
                        self.log_result(True, f"数据表 {table_name} 存在")
                    else:
                        self.log_result(False, f"数据表 {table_name} 不存在")
                        
                except Exception as e:
                    self.log_result(False, f"检查数据表 {table_name} 失败", str(e))
            
            # 检查ai_message表的搜索字段
            try:
                result = db.execute(text("""
                    SELECT column_name FROM information_schema.columns 
                    WHERE table_name = 'ai_message' 
                    AND column_name IN ('search_info', 'search_results');
                """))
                
                columns = [row[0] for row in result.fetchall()]
                
                for column in ['search_info', 'search_results']:
                    if column in columns:
                        self.log_result(True, f"ai_message表包含 {column} 字段")
                    else:
                        self.log_result(False, f"ai_message表缺少 {column} 字段")
                        
            except Exception as e:
                self.log_result(False, "检查ai_message表搜索字段失败", str(e))
            
            db.close()
            return True
            
        except Exception as e:
            self.log_result(False, "数据库表检查失败", str(e))
            return False
    
    def check_configuration_files(self) -> bool:
        """检查配置文件"""
        print("\n=== 配置文件检查 ===")
        
        try:
            # 检查.env文件
            env_files = [".env", ".env.example"]
            for env_file in env_files:
                if os.path.exists(env_file):
                    self.log_result(True, f"配置文件 {env_file} 存在")
                    
                    # 检查是否包含搜索相关配置
                    with open(env_file, 'r', encoding='utf-8') as f:
                        content = f.read()
                        
                    search_configs = [
                        "DASHSCOPE_API_KEY",
                        "LLM_SEARCH_ENABLED",
                        "VERTICAL_DOMAIN_ENABLED"
                    ]
                    
                    for config in search_configs:
                        if config in content:
                            self.log_result(True, f"{env_file} 包含 {config} 配置")
                        else:
                            self.log_result(False, f"{env_file} 缺少 {config} 配置")
                else:
                    self.log_result(False, f"配置文件 {env_file} 不存在")
            
            return True
            
        except Exception as e:
            self.log_result(False, "配置文件检查失败", str(e))
            return False
    
    def generate_summary(self) -> Dict[str, Any]:
        """生成检查摘要"""
        total_checks = len(self.results)
        error_count = len(self.errors)
        success_count = total_checks - error_count
        
        return {
            "total_checks": total_checks,
            "success_count": success_count,
            "error_count": error_count,
            "success_rate": (success_count / total_checks * 100) if total_checks > 0 else 0,
            "errors": self.errors
        }
    
    def run_all_checks(self) -> bool:
        """运行所有检查"""
        print("LLM搜索功能配置检查")
        print("=" * 60)
        
        checks = [
            ("环境变量配置", self.check_environment_variables),
            ("API连接", self.check_api_connection),
            ("模型支持", self.check_model_support),
            ("搜索选项验证", self.check_search_options_validation),
            ("垂直领域支持", self.check_vertical_domain_support),
            ("数据库表", self.check_database_tables),
            ("配置文件", self.check_configuration_files)
        ]
        
        all_passed = True
        
        for check_name, check_func in checks:
            try:
                result = check_func()
                if not result:
                    all_passed = False
            except Exception as e:
                self.log_result(False, f"{check_name}检查异常", str(e))
                all_passed = False
        
        # 生成摘要
        summary = self.generate_summary()
        
        print("\n" + "=" * 60)
        print("检查摘要")
        print("=" * 60)
        print(f"总检查项: {summary['total_checks']}")
        print(f"成功: {summary['success_count']}")
        print(f"失败: {summary['error_count']}")
        print(f"成功率: {summary['success_rate']:.1f}%")
        
        if summary['error_count'] > 0:
            print(f"\n失败的检查项:")
            for i, error in enumerate(summary['errors'], 1):
                print(f"  {i}. {error}")
        
        if all_passed:
            print("\n🎉 所有配置检查通过！搜索功能已准备就绪。")
        else:
            print("\n⚠️  部分配置检查失败，请根据上述错误信息进行修复。")
        
        return all_passed


def main():
    """主函数"""
    checker = SearchConfigChecker()
    
    try:
        success = checker.run_all_checks()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n检查被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n检查过程中发生异常: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()