113 lines
2.5 KiB
Go
113 lines
2.5 KiB
Go
package ai
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"resume-platform/pkg/config"
|
|
"time"
|
|
)
|
|
|
|
// DoubaoProvider 豆包 AI 供应商实现
|
|
type DoubaoProvider struct {
|
|
cfg *config.DoubaoConfig // 豆包配置
|
|
httpClient *http.Client // HTTP 客户端
|
|
}
|
|
|
|
// NewDoubaoProvider 创建豆包供应商实例
|
|
// @param cfg 豆包配置
|
|
// @return *DoubaoProvider 供应商实例
|
|
// @author sunct
|
|
func NewDoubaoProvider(cfg *config.DoubaoConfig) *DoubaoProvider {
|
|
return &DoubaoProvider{
|
|
cfg: cfg,
|
|
httpClient: &http.Client{
|
|
Timeout: time.Duration(cfg.Timeout) * time.Second,
|
|
},
|
|
}
|
|
}
|
|
|
|
// GetName 获取供应商名称
|
|
// @return string 供应商标识
|
|
// @author sunct
|
|
func (p *DoubaoProvider) GetName() string {
|
|
return "doubao"
|
|
}
|
|
|
|
// Generate 调用豆包 API 生成文本
|
|
// @param ctx 上下文
|
|
// @param prompt 提示词
|
|
// @param opts 可选参数
|
|
// @return string 生成结果
|
|
// @return error 调用错误
|
|
// @author sunct
|
|
func (p *DoubaoProvider) Generate(ctx context.Context, prompt string, opts ...Option) (string, error) {
|
|
options := &Options{
|
|
MaxTokens: p.cfg.MaxTokens,
|
|
Temperature: 0.7,
|
|
TopP: 0.9,
|
|
}
|
|
for _, opt := range opts {
|
|
opt(options)
|
|
}
|
|
|
|
body := map[string]interface{}{
|
|
"model": p.getModelName(),
|
|
"messages": []map[string]interface{}{
|
|
{
|
|
"role": "user",
|
|
"content": prompt,
|
|
},
|
|
},
|
|
"max_tokens": options.MaxTokens,
|
|
"temperature": options.Temperature,
|
|
"top_p": options.TopP,
|
|
}
|
|
|
|
jsonBody, _ := json.Marshal(body)
|
|
|
|
req, err := http.NewRequest("POST", p.cfg.BaseURL+"/chat/completions", bytes.NewBuffer(jsonBody))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+p.cfg.APIKey)
|
|
|
|
resp, err := p.httpClient.Do(req)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
respBody, _ := io.ReadAll(resp.Body)
|
|
|
|
var result map[string]interface{}
|
|
if err := json.Unmarshal(respBody, &result); err != nil {
|
|
return string(respBody), nil
|
|
}
|
|
|
|
if choices, ok := result["choices"].([]interface{}); ok && len(choices) > 0 {
|
|
if choice, ok := choices[0].(map[string]interface{}); ok {
|
|
if message, ok := choice["message"].(map[string]interface{}); ok {
|
|
if content, ok := message["content"].(string); ok {
|
|
return content, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return string(respBody), nil
|
|
}
|
|
|
|
// getModelName 获取模型名称，未配置时使用默认 Doubao
|
|
// @return string 模型名
|
|
// @author sunct
|
|
func (p *DoubaoProvider) getModelName() string {
|
|
if p.cfg.Model != "" {
|
|
return p.cfg.Model
|
|
}
|
|
return "Doubao"
|
|
} |