102 lines
2.5 KiB
Go
102 lines
2.5 KiB
Go
// Package tokenhub 封装 Tokenhub 大模型 API 调用，用于生成每日星座运势与穿搭推荐。
|
|
package tokenhub
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"time"
|
|
|
|
"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
|
|
)
|
|
|
|
const baseURL = "https://tokenhub.tencentmaas.com/v1/chat/completions"
|
|
|
|
// Client Tokenhub API 客户端。
|
|
type Client struct {
|
|
apiKey string
|
|
model string
|
|
httpClient *http.Client
|
|
}
|
|
|
|
// NewClient 从全局配置创建客户端。
|
|
func NewClient() *Client {
|
|
return &Client{
|
|
apiKey: global.TokenhubSetting.APIKey,
|
|
model: global.TokenhubSetting.Model,
|
|
httpClient: &http.Client{
|
|
Timeout: 120 * time.Second,
|
|
},
|
|
}
|
|
}
|
|
|
|
// Message 聊天消息。
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// chatResponse Tokenhub chat/completions 外层响应结构（只取需要的字段）。
|
|
type chatResponse struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
// ResponseFormat JSON Schema 响应格式。
|
|
type ResponseFormat struct {
|
|
Type string `json:"type"`
|
|
JSONSchema *JSONSchema `json:"json_schema"`
|
|
}
|
|
|
|
// JSONSchema 结构化输出定义。
|
|
type JSONSchema struct {
|
|
Name string `json:"name"`
|
|
Schema interface{} `json:"schema"`
|
|
}
|
|
|
|
// ChatCompletion 调用混元 chat/completions，返回原始响应体。
|
|
func (c *Client) ChatCompletion(ctx context.Context, messages []Message, responseFormat *ResponseFormat) ([]byte, error) {
|
|
body := map[string]interface{}{
|
|
"model": c.model,
|
|
"messages": messages,
|
|
}
|
|
if responseFormat != nil {
|
|
body["response_format"] = responseFormat
|
|
}
|
|
|
|
jsonBody, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal request: %w", err)
|
|
}
|
|
|
|
request, err := http.NewRequestWithContext(ctx, "POST", baseURL, bytes.NewBuffer(jsonBody))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create request: %w", err)
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
|
|
response, err := c.httpClient.Do(request)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("do request: %w", err)
|
|
}
|
|
defer response.Body.Close()
|
|
|
|
data, err := io.ReadAll(response.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read response: %w", err)
|
|
}
|
|
|
|
if response.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("tokenhub api error: status=%d, body=%s", response.StatusCode, string(data))
|
|
}
|
|
|
|
return data, nil
|
|
}
|