security-platform/setup_llm.py

#!/usr/bin/env python3
"""
LLM集成配置助手
"""
import os
import shutil
import yaml

def main():
    print("🤖 AI SOC分析平台 LLM集成配置助手")
    print("=" * 50)
    
    # 创建.env文件
    if not os.path.exists(".env"):
        if os.path.exists(".env.example"):
            shutil.copy(".env.example", ".env")
            print("✅ 已创建 .env 文件")
        else:
            print("❌ .env.example 文件不存在")
            return
    else:
        print("ℹ️  .env 文件已存在")
    
    # 显示配置选项
    print("\n🎯 选择你的LLM提供商:")
    print("1. Claude 4 (推荐，需API密钥)")
    print("2. 本地模型 (免费，需安装Ollama)")
    print("3. OpenAI (需API密钥)")
    print("4. 暂不配置 (使用模板生成)")
    
    try:
        choice = input("\n请输入选项 (1-4): ").strip()
        
        if choice == "1":
            configure_claude()
        elif choice == "2":
            configure_local()
        elif choice == "3":
            configure_openai()
        elif choice == "4":
            print("✅ 将使用模板生成，你可以稍后配置LLM")
        else:
            print("❌ 无效选项")
            return
            
        print(f"\n🎉 配置完成！请重启后端服务生效")
        print("重启命令: python3 run_backend.py")
        
    except KeyboardInterrupt:
        print("\n\n👋 配置已取消")

def configure_claude():
    print("\n🎨 配置Claude 4")
    print("请访问 https://console.anthropic.com/ 获取API密钥")
    
    api_key = input("请输入Claude API Key (sk-ant-...): ").strip()
    if not api_key.startswith("sk-ant-"):
        print("⚠️  API密钥格式可能不正确")
    
    # 更新.env文件
    update_env_file("CLAUDE_API_KEY", api_key)
    
    # 更新llm.yaml配置
    update_llm_config("claude")
    
    print("✅ Claude 4 配置完成")

def configure_local():
    print("\n🏠 配置本地模型")
    print("请确保已安装Ollama: https://ollama.ai/")
    
    print("\n推荐模型:")
    print("- qwen2.5:7b (中文优秀, ~4.4GB)")
    print("- qwen2.5:14b (性能更强, ~8.2GB)")
    print("- llama3.1:8b (英文强, ~4.7GB)")
    
    model = input("请输入模型名称 (默认: qwen2.5:7b): ").strip()
    if not model:
        model = "qwen2.5:7b"
    
    port = input("请输入Ollama端口 (默认: 11434): ").strip()
    if not port:
        port = "11434"
    
    # 更新llm.yaml配置
    update_llm_config("local", model, f"http://localhost:{port}/v1")
    
    print(f"✅ 本地模型配置完成")
    print(f"请确保运行: ollama pull {model}")
    print("然后启动服务: ollama serve")

def configure_openai():
    print("\n🌐 配置OpenAI")
    
    api_key = input("请输入OpenAI API Key (sk-proj-...): ").strip()
    if not api_key.startswith("sk-"):
        print("⚠️  API密钥格式可能不正确")
    
    model = input("请输入模型名称 (默认: gpt-3.5-turbo): ").strip()
    if not model:
        model = "gpt-3.5-turbo"
    
    # 更新.env文件
    update_env_file("OPENAI_API_KEY", api_key)
    
    # 更新llm.yaml配置
    update_llm_config("openai", model)
    
    print("✅ OpenAI 配置完成")

def update_env_file(key, value):
    """更新.env文件"""
    try:
        # 读取现有内容
        lines = []
        if os.path.exists(".env"):
            with open(".env", "r") as f:
                lines = f.readlines()
        
        # 更新或添加键值
        found = False
        for i, line in enumerate(lines):
            if line.startswith(f"{key}="):
                lines[i] = f"{key}={value}\n"
                found = True
                break
        
        if not found:
            lines.append(f"{key}={value}\n")
        
        # 写回文件
        with open(".env", "w") as f:
            f.writelines(lines)
            
    except Exception as e:
        print(f"❌ 更新.env文件失败: {e}")

def update_llm_config(provider, model=None, api_base=None):
    """更新LLM配置文件"""
    try:
        config_path = "configs/llm.yaml"
        
        # 读取现有配置
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        
        # 更新提供商
        config["provider"] = provider
        
        # 更新具体配置
        if model and provider in config.get("providers", {}):
            config["providers"][provider]["model"] = model
        
        if api_base and provider in config.get("providers", {}):
            config["providers"][provider]["api_base"] = api_base
        
        # 写回文件
        with open(config_path, "w", encoding="utf-8") as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
            
    except Exception as e:
        print(f"❌ 更新配置文件失败: {e}")

if __name__ == "__main__":
    main()