package main import ( "context" "encoding/json" "fmt" "io" "log" "net/http" "os" "path/filepath" "strconv" "strings" "sync" "time" "golang.org/x/sync/semaphore" ) // AlphaConfig Alpha因子配置 type AlphaConfig struct { Name string `json:"name"` Expression string `json:"expression"` Settings map[string]interface{} `json:"settings,omitempty"` } // SimulationResult 模拟结果 type SimulationResult struct { Name string `json:"name"` Expression string `json:"expression"` AlphaID string `json:"alpha_id,omitempty"` Success bool `json:"success"` Error string `json:"error,omitempty"` ResultData interface{} `json:"result_data,omitempty"` } // WorldQuantBrainAPI WorldQuant Brain API客户端 type WorldQuantBrainAPI struct { BaseURL string CredentialsFile string HTTPClient *http.Client Username string Password string } // NewWorldQuantBrainAPI 创建新的API客户端 func NewWorldQuantBrainAPI(credentialsFile string) *WorldQuantBrainAPI { if credentialsFile == "" { credentialsFile = "brain_credentials.txt" } // 扩展文件路径 fullPath := expandUser(credentialsFile) return &WorldQuantBrainAPI{ BaseURL: "https://api.worldquantbrain.com", CredentialsFile: fullPath, HTTPClient: &http.Client{Timeout: 30 * time.Second}, } } // expandUser 扩展用户主目录路径 func expandUser(path string) string { if strings.HasPrefix(path, "~/") { home, err := os.UserHomeDir() if err != nil { return path } return filepath.Join(home, path[2:]) } return path } // Login 登录认证 func (wq *WorldQuantBrainAPI) Login() error { // 读取凭证文件 data, err := os.ReadFile(wq.CredentialsFile) if err != nil { return fmt.Errorf("读取凭证文件失败: %v", err) } content := strings.TrimSpace(string(data)) var credentials []string // 方法1: 尝试标准JSON解析 if err := json.Unmarshal([]byte(content), &credentials); err == nil && len(credentials) >= 2 { // JSON解析成功 } else { // 方法2: 尝试处理单引号JSON singleQuoteContent := strings.ReplaceAll(content, "'", "\"") if err := json.Unmarshal([]byte(singleQuoteContent), &credentials); err == nil && len(credentials) >= 2 { // 单引号JSON解析成功 } else { // 方法3: 尝试其他文本格式 credentials = wq.parseTextCredentials(content) } } if len(credentials) < 2 { return fmt.Errorf("无法解析凭证文件,需要用户名和密码。当前内容: %s", content) } wq.Username = strings.TrimSpace(credentials[0]) wq.Password = strings.TrimSpace(credentials[1]) // 创建认证请求 req, err := http.NewRequest("POST", wq.BaseURL+"/authentication", nil) if err != nil { return fmt.Errorf("创建请求失败: %v", err) } req.SetBasicAuth(wq.Username, wq.Password) // 发送认证请求 resp, err := wq.HTTPClient.Do(req) if err != nil { return fmt.Errorf("认证请求失败: %v", err) } defer resp.Body.Close() fmt.Printf("认证状态: %d\n", resp.StatusCode) // 修复:检查 200 和 201 状态码都表示成功 if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { body, _ := io.ReadAll(resp.Body) return fmt.Errorf("登录失败,状态码: %d, 响应: %s", resp.StatusCode, string(body)) } // 读取成功响应 body, _ := io.ReadAll(resp.Body) fmt.Printf("登录成功响应: %s\n", string(body)) fmt.Println("登录成功!") return nil } func (wq *WorldQuantBrainAPI) parseTextCredentials(content string) []string { panic("unimplemented") } // SimulateAlpha 模拟单个Alpha因子 func (wq *WorldQuantBrainAPI) SimulateAlpha(config AlphaConfig) SimulationResult { if wq.HTTPClient == nil { return SimulationResult{ Name: config.Name, Success: false, Error: "请先调用Login方法登录", } } name := config.Name if name == "" { name = "未命名Alpha" } fmt.Printf("\n🚀 开始模拟: %s\n", name) fmt.Printf("📝 表达式: %s\n", config.Expression) // 构建模拟数据 simulationData := map[string]interface{}{ "type": "REGULAR", "settings": config.Settings, "regular": config.Expression, } jsonData, err := json.Marshal(simulationData) if err != nil { return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("序列化模拟数据失败: %v", err), } } // 发送模拟请求 req, err := http.NewRequest("POST", wq.BaseURL+"/simulations", strings.NewReader(string(jsonData))) if err != nil { return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("创建请求失败: %v", err), } } req.SetBasicAuth(wq.Username, wq.Password) req.Header.Set("Content-Type", "application/json") resp, err := wq.HTTPClient.Do(req) if err != nil { return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("模拟请求失败: %v", err), } } defer resp.Body.Close() fmt.Printf("📤 模拟提交状态: %d\n", resp.StatusCode) if resp.StatusCode != http.StatusOK { // 尝试读取错误响应体 body, _ := io.ReadAll(resp.Body) return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("提交失败: %d, 响应: %s", resp.StatusCode, string(body)), } } // 获取进度URL progressURL := resp.Header.Get("Location") if progressURL == "" { return SimulationResult{ Name: name, Success: false, Error: "未找到进度URL", } } fmt.Printf("⏳ %s 进度URL获取成功\n", name) // 轮询进度 var resultData map[string]interface{} for { progressReq, err := http.NewRequest("GET", progressURL, nil) if err != nil { return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("创建进度请求失败: %v", err), } } progressReq.SetBasicAuth(wq.Username, wq.Password) progressResp, err := wq.HTTPClient.Do(progressReq) if err != nil { return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("进度请求失败: %v", err), } } retryAfter := progressResp.Header.Get("Retry-After") if retryAfter == "" || retryAfter == "0" { // 模拟完成 if err := json.NewDecoder(progressResp.Body).Decode(&resultData); err != nil { progressResp.Body.Close() return SimulationResult{ Name: name, Success: false, Error: fmt.Sprintf("解析结果失败: %v", err), } } progressResp.Body.Close() break } progressResp.Body.Close() // 等待指定时间 waitSeconds, _ := strconv.ParseFloat(retryAfter, 64) fmt.Printf("⏰ %s 等待 %.1f 秒...\n", name, waitSeconds) time.Sleep(time.Duration(waitSeconds * float64(time.Second))) } // 提取Alpha ID alphaID, ok := resultData["alpha"].(string) if !ok { return SimulationResult{ Name: name, Success: false, Error: "未找到Alpha ID", } } fmt.Printf("✅ %s 模拟完成! Alpha ID: %s\n", name, alphaID) return SimulationResult{ Name: name, Expression: config.Expression, AlphaID: alphaID, Success: true, ResultData: resultData, } } // BatchSimulateAlphas 批量测试多个Alpha因子 func (wq *WorldQuantBrainAPI) BatchSimulateAlphas(configFile string, maxWorkers int) []SimulationResult { // 读取Alpha配置 configs, err := wq.loadAlphaConfigs(configFile) if err != nil { log.Printf("❌ 读取配置文件失败: %v", err) return nil } fmt.Printf("📁 读取到 %d 个Alpha配置\n", len(configs)) // 登录 if err := wq.Login(); err != nil { log.Printf("❌ 登录失败: %v", err) return nil } var results []SimulationResult var mu sync.Mutex var wg sync.WaitGroup // 使用信号量控制并发数量 sem := semaphore.NewWeighted(int64(maxWorkers)) ctx := context.Background() for _, config := range configs { wg.Add(1) go func(c AlphaConfig) { defer wg.Done() // 获取信号量 if err := sem.Acquire(ctx, 1); err != nil { log.Printf("❌ 获取信号量失败: %v", err) return } defer sem.Release(1) // 执行模拟 result := wq.SimulateAlpha(c) // 安全地添加结果 mu.Lock() results = append(results, result) mu.Unlock() }(config) } wg.Wait() // 打印汇总结果 wq.printSummary(results) return results } // loadAlphaConfigs 加载Alpha配置 func (wq *WorldQuantBrainAPI) loadAlphaConfigs(configFile string) ([]AlphaConfig, error) { data, err := os.ReadFile(configFile) if err != nil { return nil, fmt.Errorf("读取配置文件 %s 失败: %v", configFile, err) } var configs []AlphaConfig if err := json.Unmarshal(data, &configs); err != nil { return nil, fmt.Errorf("解析配置文件 %s 失败: %v", configFile, err) } return configs, nil } // printSummary 打印测试结果汇总 func (wq *WorldQuantBrainAPI) printSummary(results []SimulationResult) { fmt.Println("\n" + strings.Repeat("=", 50)) fmt.Println("📊 Alpha测试结果汇总") fmt.Println(strings.Repeat("=", 50)) successCount := 0 for _, result := range results { if result.Success { successCount++ } } failedCount := len(results) - successCount fmt.Printf("✅ 成功: %d 个\n", successCount) fmt.Printf("❌ 失败: %d 个\n", failedCount) fmt.Println("\n成功详情:") for _, result := range results { if result.Success { fmt.Printf(" ✓ %s: %s\n", result.Name, result.AlphaID) } } if failedCount > 0 { fmt.Println("\n失败详情:") for _, result := range results { if !result.Success { fmt.Printf(" ✗ %s: %s\n", result.Name, result.Error) } } } } // Close 关闭客户端 func (wq *WorldQuantBrainAPI) Close() { // 如果有需要关闭的资源,可以在这里处理 } // 保存结果到文件 func saveResultsToFile(results []SimulationResult, filename string) error { data, err := json.MarshalIndent(results, "", " ") if err != nil { return err } return os.WriteFile(filename, data, 0644) } func main() { // 创建API实例 wqAPI := NewWorldQuantBrainAPI("") // 确保在程序结束时关闭连接 defer wqAPI.Close() fmt.Println("🎯 开始批量测试Alpha因子...") startTime := time.Now() // 批量测试所有Alpha results := wqAPI.BatchSimulateAlphas("alphas.json", 2) // 并发数量,根据API限制调整 duration := time.Since(startTime) fmt.Printf("\n⏱️ 总耗时: %.2f 秒\n", duration.Seconds()) fmt.Println(results) }