63 lines
1.6 KiB
Go
63 lines
1.6 KiB
Go
package ai
|
|
|
|
import (
|
|
"fmt"
|
|
)
|
|
|
|
type Client struct {
|
|
enableSearch bool // 是否开启搜索
|
|
searchStrategy string // 搜索策略（开了搜索才生效）
|
|
temperature float32 // 采样温度
|
|
}
|
|
|
|
func NewAIClient(enableSearch bool, searchStrategy string, temperature float32) *Client {
|
|
return &Client{
|
|
enableSearch: enableSearch,
|
|
searchStrategy: searchStrategy,
|
|
temperature: temperature,
|
|
}
|
|
}
|
|
|
|
func (c *Client) RequestAI(pre []string, prompt []string, cb func(content string, i int) error) ([]Message, error) {
|
|
inToken, outToken := int32(0), int32(0)
|
|
messages := make([]Message, 0)
|
|
for _, s := range pre {
|
|
messages = append(messages, Message{
|
|
Role: "user",
|
|
Content: s,
|
|
})
|
|
}
|
|
for i, step := range prompt {
|
|
messages = append(messages, Message{
|
|
Role: "user",
|
|
Content: step,
|
|
})
|
|
r, err := request("qwen-plus", Parameters{
|
|
//r, err := request("deepseek-v3.2", Parameters{
|
|
EnableSearch: c.enableSearch,
|
|
SearchOptions: SearchOptions{
|
|
SearchStrategy: c.searchStrategy,
|
|
},
|
|
ResultFormat: "message",
|
|
Temperature: c.temperature,
|
|
}, messages)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
inToken += r.Usage.InputTokens
|
|
outToken += r.Usage.OutputTokens
|
|
if len(r.Output.Choices) == 0 {
|
|
fmt.Printf("AI response err: %+v\n", r)
|
|
return nil, err
|
|
} else {
|
|
messages = append(messages, r.Output.Choices[0].Message)
|
|
}
|
|
fmt.Printf("AI 回答 %v: %v\n", i, r.Output.Choices[0].Message.Content)
|
|
if err = cb(r.Output.Choices[0].Message.Content, i); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
fmt.Printf("AI finished, use token: in-%v out-%v\n", inToken, outToken)
|
|
return messages, nil
|
|
}
|