Improve(interface): replace some struct with interface for testing (#994)

* fix(trader): get peakPnlPct using posKey * fix(docs): keep readme at the same page * improve(interface): replace with interface * refactor mcp --------- Co-authored-by: zbhan <zbhan@freewheel.tv>
2026-06-06 05:51:19 +08:00 · 2025-11-13 22:22:05 -05:00
parent 5577d284f6
commit b96c86fce4
8 changed files with 164 additions and 93 deletions
--- a/mcp/client.go
+++ b/mcp/client.go
@@ -13,18 +13,17 @@ import (
 	"time"
 )

-// Provider AI提供商类型
-type Provider string
-
 const (
-	ProviderDeepSeek Provider = "deepseek"
-	ProviderQwen     Provider = "qwen"
-	ProviderCustom   Provider = "custom"
+	ProviderCustom = "custom"
+)
+
+var (
+	DefaultTimeout = 120 * time.Second
 )

 // Client AI API配置
 type Client struct {
-	Provider   Provider
+	Provider   string
 	APIKey     string
 	BaseURL    string
 	Model      string
@@ -33,7 +32,7 @@ type Client struct {
 	MaxTokens  int  // AI响应的最大token数
 }

-func New() *Client {
+func New() AIClient {
 	// 从环境变量读取 MaxTokens，默认 2000
 	maxTokens := 2000
 	if envMaxTokens := os.Getenv("AI_MAX_TOKENS"); envMaxTokens != "" {
@@ -48,65 +47,15 @@ func New() *Client {
 	// 默认配置
 	return &Client{
 		Provider:  ProviderDeepSeek,
-		BaseURL:   "https://api.deepseek.com/v1",
-		Model:     "deepseek-chat",
-		Timeout:   120 * time.Second, // 增加到120秒，因为AI需要分析大量数据
+		BaseURL:   DefaultDeepSeekBaseURL,
+		Model:     DefaultDeepSeekModel,
+		Timeout:   DefaultTimeout,
 		MaxTokens: maxTokens,
 	}
 }

-// SetDeepSeekAPIKey 设置DeepSeek API密钥
-// customURL 为空时使用默认URL，customModel 为空时使用默认模型
-func (client *Client) SetDeepSeekAPIKey(apiKey string, customURL string, customModel string) {
-	client.Provider = ProviderDeepSeek
-	client.APIKey = apiKey
-	if customURL != "" {
-		client.BaseURL = customURL
-		log.Printf("🔧 [MCP] DeepSeek 使用自定义 BaseURL: %s", customURL)
-	} else {
-		client.BaseURL = "https://api.deepseek.com/v1"
-		log.Printf("🔧 [MCP] DeepSeek 使用默认 BaseURL: %s", client.BaseURL)
-	}
-	if customModel != "" {
-		client.Model = customModel
-		log.Printf("🔧 [MCP] DeepSeek 使用自定义 Model: %s", customModel)
-	} else {
-		client.Model = "deepseek-chat"
-		log.Printf("🔧 [MCP] DeepSeek 使用默认 Model: %s", client.Model)
-	}
-	// 打印 API Key 的前后各4位用于验证
-	if len(apiKey) > 8 {
-		log.Printf("🔧 [MCP] DeepSeek API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
-	}
-}
-
-// SetQwenAPIKey 设置阿里云Qwen API密钥
-// customURL 为空时使用默认URL，customModel 为空时使用默认模型
-func (client *Client) SetQwenAPIKey(apiKey string, customURL string, customModel string) {
-	client.Provider = ProviderQwen
-	client.APIKey = apiKey
-	if customURL != "" {
-		client.BaseURL = customURL
-		log.Printf("🔧 [MCP] Qwen 使用自定义 BaseURL: %s", customURL)
-	} else {
-		client.BaseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
-		log.Printf("🔧 [MCP] Qwen 使用默认 BaseURL: %s", client.BaseURL)
-	}
-	if customModel != "" {
-		client.Model = customModel
-		log.Printf("🔧 [MCP] Qwen 使用自定义 Model: %s", customModel)
-	} else {
-		client.Model = "qwen3-max"
-		log.Printf("🔧 [MCP] Qwen 使用默认 Model: %s", client.Model)
-	}
-	// 打印 API Key 的前后各4位用于验证
-	if len(apiKey) > 8 {
-		log.Printf("🔧 [MCP] Qwen API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
-	}
-}
-
 // SetCustomAPI 设置自定义OpenAI兼容API
-func (client *Client) SetCustomAPI(apiURL, apiKey, modelName string) {
+func (client *Client) SetAPIKey(apiKey, apiURL, customModel string) {
 	client.Provider = ProviderCustom
 	client.APIKey = apiKey

@@ -119,22 +68,14 @@ func (client *Client) SetCustomAPI(apiURL, apiKey, modelName string) {
 		client.UseFullURL = false
 	}

-	client.Model = modelName
+	client.Model = customModel
 	client.Timeout = 120 * time.Second
 }

-// SetClient 设置完整的AI配置（高级用户）
-func (client *Client) SetClient(Client Client) {
-	if Client.Timeout == 0 {
-		Client.Timeout = 30 * time.Second
-	}
-	client = &Client
-}
-
 // CallWithMessages 使用 system + user prompt 调用AI API（推荐）
 func (client *Client) CallWithMessages(systemPrompt, userPrompt string) (string, error) {
 	if client.APIKey == "" {
-		return "", fmt.Errorf("AI API密钥未设置，请先调用 SetDeepSeekAPIKey() 或 SetQwenAPIKey()")
+		return "", fmt.Errorf("AI API密钥未设置，请先调用 SetAPIKey")
 	}

 	// 重试配置
@@ -171,6 +112,10 @@ func (client *Client) CallWithMessages(systemPrompt, userPrompt string) (string,
 	return "", fmt.Errorf("重试%d次后仍然失败: %w", maxRetries, lastErr)
 }

+func (client *Client) setAuthHeader(reqHeader http.Header) {
+	reqHeader.Set("Authorization", fmt.Sprintf("Bearer %s", client.APIKey))
+}
+
 // callOnce 单次调用AI API（内部使用）
 func (client *Client) callOnce(systemPrompt, userPrompt string) (string, error) {
 	// 打印当前 AI 配置
@@ -234,17 +179,7 @@ func (client *Client) callOnce(systemPrompt, userPrompt string) (string, error)

 	req.Header.Set("Content-Type", "application/json")

-	// 根据不同的Provider设置认证方式
-	switch client.Provider {
-	case ProviderDeepSeek:
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", client.APIKey))
-	case ProviderQwen:
-		// 阿里云Qwen使用API-Key认证
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", client.APIKey))
-		// 注意：如果使用的不是兼容模式，可能需要不同的认证方式
-	default:
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", client.APIKey))
-	}
+	client.setAuthHeader(req.Header)

 	// 发送请求
 	httpClient := &http.Client{Timeout: client.Timeout}
--- a/mcp/deepseek_client.go
+++ b/mcp/deepseek_client.go
@@ -0,0 +1,53 @@
+package mcp
+
+import (
+	"log"
+	"net/http"
+)
+
+const (
+	ProviderDeepSeek       = "deepseek"
+	DefaultDeepSeekBaseURL = "https://api.deepseek.com/v1"
+	DefaultDeepSeekModel   = "deepseek-chat"
+)
+
+type DeepSeekClient struct {
+	*Client
+}
+
+func NewDeepSeekClient() AIClient {
+	client := New().(*Client)
+	client.Provider = ProviderDeepSeek
+	client.Model = DefaultDeepSeekModel
+	client.BaseURL = DefaultDeepSeekBaseURL
+	return &DeepSeekClient{
+		Client: client,
+	}
+}
+
+func (dsClient *DeepSeekClient) SetAPIKey(apiKey string, customURL string, customModel string) {
+	if dsClient.Client == nil {
+		dsClient.Client = New().(*Client)
+	}
+	dsClient.Client.APIKey = apiKey
+
+	if len(apiKey) > 8 {
+		log.Printf("🔧 [MCP] DeepSeek API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
+	}
+	if customURL != "" {
+		dsClient.Client.BaseURL = customURL
+		log.Printf("🔧 [MCP] DeepSeek 使用自定义 BaseURL: %s", customURL)
+	} else {
+		log.Printf("🔧 [MCP] DeepSeek 使用默认 BaseURL: %s", dsClient.Client.BaseURL)
+	}
+	if customModel != "" {
+		dsClient.Client.Model = customModel
+		log.Printf("🔧 [MCP] DeepSeek 使用自定义 Model: %s", customModel)
+	} else {
+		log.Printf("🔧 [MCP] DeepSeek 使用默认 Model: %s", dsClient.Client.Model)
+	}
+}
+
+func (dsClient *DeepSeekClient) setAuthHeader(reqHeaders http.Header) {
+	dsClient.Client.setAuthHeader(reqHeaders)
+}
--- a/mcp/interface.go
+++ b/mcp/interface.go
@@ -0,0 +1,12 @@
+package mcp
+
+import "net/http"
+
+// AIClient AI客户端接口
+type AIClient interface {
+	SetAPIKey(apiKey string, customURL string, customModel string)
+	// CallWithMessages 使用 system + user prompt 调用AI API
+	CallWithMessages(systemPrompt, userPrompt string) (string, error)
+
+	setAuthHeader(reqHeaders http.Header)
+}
--- a/mcp/qwen_client.go
+++ b/mcp/qwen_client.go
@@ -0,0 +1,53 @@
+package mcp
+
+import (
+	"log"
+	"net/http"
+)
+
+const (
+	ProviderQwen       = "qwen"
+	DefaultQwenBaseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+	DefaultQwenModel   = "qwen3-max"
+)
+
+type QwenClient struct {
+	*Client
+}
+
+func NewQwenClient() AIClient {
+	client := New().(*Client)
+	client.Provider = ProviderQwen
+	client.Model = DefaultQwenModel
+	client.BaseURL = DefaultQwenBaseURL
+	return &QwenClient{
+		Client: client,
+	}
+}
+
+func (qwenClient *QwenClient) SetAPIKey(apiKey string, customURL string, customModel string) {
+	if qwenClient.Client == nil {
+		qwenClient.Client = New().(*Client)
+	}
+	qwenClient.Client.APIKey = apiKey
+
+	if len(apiKey) > 8 {
+		log.Printf("🔧 [MCP] Qwen API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
+	}
+	if customURL != "" {
+		qwenClient.Client.BaseURL = customURL
+		log.Printf("🔧 [MCP] Qwen 使用自定义 BaseURL: %s", customURL)
+	} else {
+		log.Printf("🔧 [MCP] Qwen 使用默认 BaseURL: %s", qwenClient.Client.BaseURL)
+	}
+	if customModel != "" {
+		qwenClient.Client.Model = customModel
+		log.Printf("🔧 [MCP] Qwen 使用自定义 Model: %s", customModel)
+	} else {
+		log.Printf("🔧 [MCP] Qwen 使用默认 Model: %s", qwenClient.Client.Model)
+	}
+}
+
+func (qwenClient *QwenClient) setAuthHeader(reqHeaders http.Header) {
+	qwenClient.Client.setAuthHeader(reqHeaders)
+}