security_demo/switch_model.py

#!/usr/bin/env python3
"""
快速切换本地LLM模型
"""
import yaml
import sys

def switch_model(provider_name):
    """切换到指定的模型提供商"""
    config_path = "configs/llm.yaml"
    
    try:
        # 读取配置
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        # 检查提供商是否存在
        if provider_name not in config.get('providers', {}):
            print(f"❌ 提供商 '{provider_name}' 不存在")
            print("可用的提供商:")
            for p in config.get('providers', {}):
                model = config['providers'][p].get('model', '未知')
                print(f"  - {p}: {model}")
            return False
        
        # 更新配置
        config['provider'] = provider_name
        
        # 写回配置
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
        
        model = config['providers'][provider_name].get('model', '未知')
        print(f"✅ 已切换到 {provider_name}: {model}")
        print("请重启后端服务生效: python3 run_backend.py")
        return True
        
    except Exception as e:
        print(f"❌ 切换失败: {e}")
        return False

def main():
    print("🔄 AI SOC分析平台 模型切换工具")
    print("=" * 40)
    
    if len(sys.argv) != 2:
        print("用法: python3 switch_model.py <provider_name>")
        print()
        print("快速选择:")
        print("  python3 switch_model.py local          # qwen3:latest (快速)")
        print("  python3 switch_model.py local_powerful # qwen3:32b (高质量)")
        print("  python3 switch_model.py local_deepseek # deepseek-r1:70b (最强推理)")
        print("  python3 switch_model.py local_coder    # qwen2.5-coder:32b (代码分析)")
        print("  python3 switch_model.py claude         # Claude 4 (需API key)")
        return 1
    
    provider = sys.argv[1]
    success = switch_model(provider)
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main())