package ai import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "regexp" "strings" "time" ) const ( defaultOpenAIBaseURL = "https://api.openai.com/v1" defaultOpenAIModel = "gpt-4o-mini" ) // OpenAIProvider OpenAI 兼容协议实现(DeepSeek/通义/本地 vLLM 同协议可用)。 type OpenAIProvider struct { baseURL string apiKey string model string hc *http.Client } func NewOpenAIProvider(cfg Config) *OpenAIProvider { base := strings.TrimRight(strings.TrimSpace(cfg.BaseURL), "/") if base == "" { base = defaultOpenAIBaseURL } model := strings.TrimSpace(cfg.Model) if model == "" { model = defaultOpenAIModel } return &OpenAIProvider{ baseURL: base, apiKey: cfg.APIKey, model: model, hc: &http.Client{Timeout: 60 * time.Second}, } } func (p *OpenAIProvider) Name() string { return "openai" } type chatCompletionReq struct { Model string `json:"model"` Messages []Message `json:"messages"` } type chatCompletionRes struct { Choices []struct { Message Message `json:"message"` } `json:"choices"` Error *struct { Message string `json:"message"` } `json:"error"` } // Complete 单次非流式调用;ctx 全透传,HTTP 层 60s 超时兜底。 func (p *OpenAIProvider) Complete(ctx context.Context, req ChatReq) (ChatRes, error) { model := p.model if strings.TrimSpace(req.Model) != "" { model = strings.TrimSpace(req.Model) } payload := chatCompletionReq{ Model: model, Messages: BuildMessages(req), } body, err := json.Marshal(payload) if err != nil { return ChatRes{}, err } httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body)) if err != nil { return ChatRes{}, err } httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Authorization", "Bearer "+p.apiKey) resp, err := p.hc.Do(httpReq) if err != nil { return ChatRes{}, fmt.Errorf("AI 请求失败: %w", err) } defer resp.Body.Close() raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) if err != nil { return ChatRes{}, fmt.Errorf("读取 AI 响应失败: %w", err) } var out chatCompletionRes if err := json.Unmarshal(raw, &out); err != nil { return ChatRes{}, fmt.Errorf("解析 AI 响应失败(HTTP %d)", resp.StatusCode) } if out.Error != nil { return ChatRes{}, fmt.Errorf("AI 服务返回错误: %s", out.Error.Message) } if resp.StatusCode != http.StatusOK { return ChatRes{}, fmt.Errorf("AI 服务 HTTP %d", resp.StatusCode) } if len(out.Choices) == 0 { return ChatRes{}, fmt.Errorf("AI 服务未返回候选结果") } content := strings.TrimSpace(out.Choices[0].Message.Content) return ChatRes{Content: content, SQL: extractSQL(content)}, nil } var sqlFence = regexp.MustCompile("(?s)```sql\\s*\n(.+?)```") // extractSQL 从模型回复中提取 ```sql 代码块内容。 func extractSQL(content string) string { m := sqlFence.FindStringSubmatch(content) if len(m) < 2 { return "" } return strings.TrimSpace(m[1]) }