142 lines
3.9 KiB
Go
142 lines
3.9 KiB
Go
package apimart
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"juhe-factory/api/internal/provider"
|
|
)
|
|
|
|
type ChatResult struct {
|
|
Content string
|
|
Raw json.RawMessage
|
|
FinishReason string
|
|
InputTokens int64
|
|
OutputTokens int64
|
|
TotalTokens int64
|
|
UsageRaw json.RawMessage
|
|
}
|
|
|
|
func (c Client) Chat(ctx context.Context, baseURL, apiKey string, payload map[string]any) (string, json.RawMessage, error) {
|
|
result, err := c.ChatWithUsage(ctx, baseURL, apiKey, payload)
|
|
return result.Content, result.Raw, err
|
|
}
|
|
|
|
func (c Client) ChatWithUsage(ctx context.Context, baseURL, apiKey string, payload map[string]any) (ChatResult, error) {
|
|
var result ChatResult
|
|
requestPayload := make(map[string]any, len(payload)+1)
|
|
for key, value := range payload {
|
|
requestPayload[key] = value
|
|
}
|
|
requestPayload["stream"] = false
|
|
body, err := json.Marshal(requestPayload)
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(baseURL, "/")+"/chat/completions", bytes.NewReader(body))
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
request.Header.Set("Authorization", "Bearer "+apiKey)
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.HTTPClient.Do(request)
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
defer response.Body.Close()
|
|
data, err := io.ReadAll(io.LimitReader(response.Body, 16*1024*1024))
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
result.Raw = append(json.RawMessage(nil), data...)
|
|
return result, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
|
}
|
|
type chatChoice struct {
|
|
Message struct {
|
|
Content any `json:"content"`
|
|
} `json:"message"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}
|
|
var decoded struct {
|
|
Choices []chatChoice `json:"choices"`
|
|
Usage json.RawMessage `json:"usage"`
|
|
Data struct {
|
|
Choices []chatChoice `json:"choices"`
|
|
Usage json.RawMessage `json:"usage"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(data, &decoded); err != nil {
|
|
result.Raw = append(json.RawMessage(nil), data...)
|
|
return result, fmt.Errorf("文本模型响应不是有效 JSON: %w；上游响应: %s", err, safeBody(data))
|
|
}
|
|
result.Raw = append(json.RawMessage(nil), data...)
|
|
result.UsageRaw = decoded.Usage
|
|
choices := decoded.Choices
|
|
if len(choices) == 0 {
|
|
choices = decoded.Data.Choices
|
|
if len(result.UsageRaw) == 0 {
|
|
result.UsageRaw = decoded.Data.Usage
|
|
}
|
|
}
|
|
if len(choices) == 0 {
|
|
return result, fmt.Errorf("文本模型响应缺少 choices；上游响应: %s", safeBody(data))
|
|
}
|
|
result.FinishReason = choices[0].FinishReason
|
|
if len(result.UsageRaw) > 0 {
|
|
var usage map[string]any
|
|
if json.Unmarshal(result.UsageRaw, &usage) == nil {
|
|
result.InputTokens = intValue(usage["prompt_tokens"], usage["input_tokens"])
|
|
result.OutputTokens = intValue(usage["completion_tokens"], usage["output_tokens"])
|
|
result.TotalTokens = intValue(usage["total_tokens"])
|
|
if result.TotalTokens == 0 {
|
|
result.TotalTokens = result.InputTokens + result.OutputTokens
|
|
}
|
|
}
|
|
}
|
|
switch content := choices[0].Message.Content.(type) {
|
|
case string:
|
|
result.Content = content
|
|
return result, nil
|
|
case []any:
|
|
parts := make([]string, 0)
|
|
for _, item := range content {
|
|
if block, ok := item.(map[string]any); ok {
|
|
if text, ok := block["text"].(string); ok {
|
|
parts = append(parts, text)
|
|
}
|
|
}
|
|
}
|
|
result.Content = strings.Join(parts, "\n")
|
|
return result, nil
|
|
default:
|
|
return result, errors.New("文本模型响应内容无效")
|
|
}
|
|
}
|
|
|
|
func intValue(values ...any) int64 {
|
|
for _, value := range values {
|
|
switch number := value.(type) {
|
|
case float64:
|
|
if number > 0 {
|
|
return int64(number)
|
|
}
|
|
case int64:
|
|
if number > 0 {
|
|
return number
|
|
}
|
|
case json.Number:
|
|
if parsed, err := number.Int64(); err == nil && parsed > 0 {
|
|
return parsed
|
|
}
|
|
}
|
|
}
|
|
return 0
|
|
}
|