/ pkg / aliyun / chat.go
chat.go
 1  package aliyun
 2  
 3  import (
 4  	"context"
 5  	goopenai "github.com/sashabaranov/go-openai"
 6  	"go.uber.org/zap"
 7  	"krillin-ai/log"
 8  )
 9  
10  type ChatClient struct {
11  	*goopenai.Client
12  }
13  
14  func NewChatClient(apiKey string) *ChatClient {
15  	cfg := goopenai.DefaultConfig(apiKey)
16  	cfg.BaseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1" // 使用阿里云的openai兼容模式调用
17  	return &ChatClient{
18  		Client: goopenai.NewClientWithConfig(cfg),
19  	}
20  }
21  
22  func (c ChatClient) ChatCompletion(query string) (string, error) {
23  	req := goopenai.ChatCompletionRequest{
24  		Model: "qwen-plus",
25  		Messages: []goopenai.ChatCompletionMessage{
26  			{
27  				Role:    goopenai.ChatMessageRoleSystem,
28  				Content: "You are an assistant that helps with subtitle translation.",
29  			},
30  			{
31  				Role:    goopenai.ChatMessageRoleUser,
32  				Content: query,
33  			},
34  		},
35  	}
36  
37  	resp, err := c.CreateChatCompletion(context.Background(), req)
38  	if err != nil {
39  		log.GetLogger().Error("aliyun openai create chat completion failed", zap.Error(err))
40  		return "", err
41  	}
42  
43  	resContent := resp.Choices[0].Message.Content
44  
45  	return resContent, nil
46  }