|
|
@@ -0,0 +1,422 @@
|
|
|
+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)
|
|
|
+}
|