package ai import ( "bufio" "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strings" ) type OpenAIClient struct { apiKey string baseURL string model string } func NewOpenAIClient(cfg *Config) *OpenAIClient { baseURL := cfg.BaseURL if baseURL == "" { baseURL = "https://api.openai.com/v1" } model := cfg.Model if model == "" { model = "gpt-3.5-turbo" } return &OpenAIClient{ apiKey: cfg.APIKey, baseURL: baseURL, model: model, } } type openAIRequest struct { Model string `json:"model"` Messages []Message `json:"messages"` Stream bool `json:"stream"` } type openAIResponse struct { Choices []struct { Message struct { Content string `json:"content"` } `json:"message"` Delta struct { Content string `json:"content"` } `json:"delta"` } `json:"choices"` Error *struct { Message string `json:"message"` } `json:"error"` } func (c *OpenAIClient) Chat(ctx context.Context, messages []Message) (string, error) { if c.apiKey == "" { return "", fmt.Errorf("OpenAI API Key 未配置,请在 config.yaml 中设置 ai.openai.api_key") } reqBody := openAIRequest{ Model: c.model, Messages: messages, Stream: false, } body, _ := json.Marshal(reqBody) req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/chat/completions", bytes.NewReader(body)) if err != nil { return "", err } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.apiKey) resp, err := http.DefaultClient.Do(req) if err != nil { return "", fmt.Errorf("调用AI服务失败: %v", err) } defer resp.Body.Close() var result openAIResponse if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { return "", fmt.Errorf("解析AI响应失败: %v", err) } if result.Error != nil { return "", fmt.Errorf("AI服务错误: %s", result.Error.Message) } if len(result.Choices) == 0 { return "", fmt.Errorf("AI未返回有效响应") } return result.Choices[0].Message.Content, nil } func (c *OpenAIClient) ChatStream(ctx context.Context, messages []Message, writer io.Writer) error { if c.apiKey == "" { return fmt.Errorf("OpenAI API Key 未配置") } reqBody := openAIRequest{ Model: c.model, Messages: messages, Stream: true, } body, _ := json.Marshal(reqBody) req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/chat/completions", bytes.NewReader(body)) if err != nil { return err } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.apiKey) resp, err := http.DefaultClient.Do(req) if err != nil { return err } defer resp.Body.Close() // 解析 SSE 流 reader := bufio.NewReader(resp.Body) for { line, err := reader.ReadString('\n') if err == io.EOF { break } if err != nil { return err } line = strings.TrimSpace(line) if strings.HasPrefix(line, "data:") { data := strings.TrimPrefix(line, "data:") data = strings.TrimSpace(data) if data == "[DONE]" { break } var streamResp openAIResponse if err := json.Unmarshal([]byte(data), &streamResp); err != nil { continue } if len(streamResp.Choices) > 0 { content := streamResp.Choices[0].Delta.Content writer.Write([]byte(content)) } } } return nil }