++

2026-02-05 10:18:53 +08:00
parent 717f8c1747
commit e8a4b2d318
364 changed files with 453102 additions and 0 deletions
@@ -0,0 +1,174 @@
+import os
+import sys
+import shutil
+import subprocess
+import json
+import platform
+import time
+from pathlib import Path
+
+# --- Helpers ---
+def run_cmd(cmd, shell=False):
+    print(f"Running: {' '.join(cmd) if isinstance(cmd, list) else cmd}")
+    subprocess.check_call(cmd, shell=shell)
+
+def get_home():
+    return Path.home()
+
+def set_unix_env_var(name, value):
+    """Appends export command to shell config files"""
+    if not value: return
+    
+    shell = os.environ.get("SHELL", "/bin/bash")
+    rc_files = set()
+    
+    # Detect shell config files
+    home = get_home()
+    if "zsh" in shell:
+        rc_files.add(home / ".zshrc")
+        rc_files.add(home / ".zprofile")
+    elif "bash" in shell:
+        rc_files.add(home / ".bashrc")
+        rc_files.add(home / ".bash_profile")
+    else:
+        rc_files.add(home / ".profile")
+    
+    entry = f'\nexport {name}="{value}"'
+    
+    for rc in rc_files:
+        if not rc.exists(): continue # Only append to existing files
+        try:
+            content = rc.read_text(encoding='utf-8')
+            if f'export {name}=' in content:
+                print(f"Update: {name} already exists in {rc.name}, skipping append (manual check recommended).")
+            else:
+                with open(rc, "a", encoding='utf-8') as f:
+                    f.write(entry)
+                print(f"✅ Added {name} to {rc.name}")
+        except Exception as e:
+            print(f"⚠️ Could not write to {rc}: {e}")
+
+# --- Steps ---
+
+def step1_init_files():
+    print("\n[Api/Files] Running Step 1: Project Initialization...")
+    try:
+        # Run Step1_init_project_files.py
+        run_cmd([sys.executable, "Step1_init_project_files.py"])
+    except Exception as e:
+        print(f"Step 1 failed: {e}")
+        sys.exit(1)
+
+def step2_brain_system_prompt():
+    print("\n[Config] Running Step 2: Configure Brain Consultant Persona...")
+    src = Path("brain-consultant.md").resolve()
+    if not src.exists():
+        print(f"Error: {src} not found.")
+        return
+
+    dest_dir = get_home() / ".claude" / "agents"
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    
+    dest = dest_dir / "brain-consultant.md"
+    shutil.copy2(src, dest)
+    print(f"✅ Copied Persona to: {dest}")
+
+def step3_setup_mcp():
+    print("\n[MCP] Running Step 3: Setup MCP Server...")
+    
+    script_path = Path("platform_functions.py").resolve()
+    if not script_path.exists():
+        print("Error: platform_functions.py not found.")
+        sys.exit(1)
+        
+    print("Installing Python dependencies for MCP...")
+    pip_cmd = [sys.executable, "-m", "pip", "install", "playwright", "beautifulsoup4", "pandas", "requests", "mcp"]
+    run_cmd(pip_cmd)
+    
+    print("Installing Playwright browsers...")
+    run_cmd([sys.executable, "-m", "playwright", "install", "chromium"])
+    
+    # Configure Claude MCP
+    print("Configuring Claude MCP...")
+    
+    # Remove old
+    try:
+        run_cmd(["claude", "mcp", "remove", "brain-mcp"])
+    except:
+        pass # Ignore if not exists
+        
+    # Add new
+    # On Unix we just call python3 directly
+    # Using sys.executable to ensure we use the same python env
+    mcp_cmd = ["claude", "mcp", "add", "--scope", "user", "brain-mcp", "--", sys.executable, str(script_path)]
+    try:
+        run_cmd(mcp_cmd)
+    except Exception as e:
+        print("Note: 'claude mcp add' might have failed if clauce CLI is not logged in or ready. Continuing to config edit...")
+
+    # Edit config for timeout
+    config_path = get_home() / ".claude.json"
+    if config_path.exists():
+        try:
+            with open(config_path, 'r', encoding='utf-8') as f:
+                config = json.load(f)
+            
+            # Ensure structure
+            if "mcpServers" not in config: config["mcpServers"] = {}
+            if "brain-mcp" not in config["mcpServers"]: config["mcpServers"]["brain-mcp"] = {}
+            
+            brain_mcp = config["mcpServers"]["brain-mcp"]
+            brain_mcp["timeout"] = 900
+            brain_mcp["description"] = "WorldQuant BRAIN Platform MCP Server - Comprehensive trading platform integration. Credentials in user_config.json. Docs/Forum support."
+            
+            with open(config_path, 'w', encoding='utf-8') as f:
+                json.dump(config, f, indent=2, ensure_ascii=False)
+            print("✅ Updated .claude.json with timeout/description")
+        except Exception as e:
+            print(f"Failed to update config json: {e}")
+
+def step4_api_config():
+    print("\n[API] Running Step 4: Configure API Model...")
+    
+    # We will import Step4 module to use its interactive logic, but we need to patch the env var setter
+    # Because Step4 currently uses PowerShell
+    
+    try:
+        # Load Step4 as a module
+        import importlib.util
+        spec = importlib.util.spec_from_file_location("Step4", "Step4_SetAPI_And_Check_LLMModel.py")
+        step4_module = importlib.util.module_from_spec(spec)
+        
+        # Monkey patch the windows env var function
+        step4_module.set_windows_env_var = set_unix_env_var
+        
+        # Execute
+        spec.loader.exec_module(step4_module)
+        
+        # Run main
+        if hasattr(step4_module, "main"):
+            step4_module.main()
+            
+    except Exception as e:
+        print(f"Error running Step 4 logic: {e}")
+        print("Fallback: Please manually set ANTHROPIC_BASE_URL and ANTHROPIC_API_KEY in your shell.")
+
+def main():
+    print("=== WorldQuant BRAIN Setup (Linux/Mac) ===")
+    
+    # Prereq check (Runtime)
+    if not shutil.which("claude"):
+        print("❌ 'claude' command not found. Please install it via 'npm install -g @anthropic-ai/claude-code'")
+        sys.exit(1)
+        
+    step1_init_files()
+    step2_brain_system_prompt()
+    step3_setup_mcp()
+    step4_api_config()
+    
+    print("\n✅ Setup Complete!")
+    print("Please RESTART your terminal to ensure environment variables are loaded.")
+    print("Then run: claude --agent brain-consultant")
+
+if __name__ == "__main__":
+    main()