323 lines
8.6 KiB
Go
323 lines
8.6 KiB
Go
package llm
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
)
|
||
|
||
// openAI 是 OpenAI 兼容协议（/chat/completions + SSE）的实现。
|
||
type openAI struct {
|
||
opts Options
|
||
http *http.Client
|
||
}
|
||
|
||
func newOpenAI(opts Options) *openAI {
|
||
if opts.MaxTokens <= 0 {
|
||
opts.MaxTokens = 1024
|
||
}
|
||
if opts.Temperature <= 0 {
|
||
opts.Temperature = 0.3
|
||
}
|
||
return &openAI{opts: opts, http: newHTTPClient()}
|
||
}
|
||
|
||
type reqMessage struct {
|
||
Role string `json:"role"`
|
||
Content any `json:"content"`
|
||
}
|
||
|
||
type contentPart struct {
|
||
Type string `json:"type"`
|
||
Text string `json:"text,omitempty"`
|
||
ImageURL *imageURL `json:"image_url,omitempty"`
|
||
}
|
||
|
||
type imageURL struct {
|
||
URL string `json:"url"`
|
||
}
|
||
|
||
type streamOptions struct {
|
||
IncludeUsage bool `json:"include_usage"`
|
||
}
|
||
|
||
type chatRequest struct {
|
||
Model string `json:"model"`
|
||
Messages []reqMessage `json:"messages"`
|
||
Stream bool `json:"stream"`
|
||
MaxTokens int `json:"max_tokens,omitempty"`
|
||
Temperature float64 `json:"temperature,omitempty"`
|
||
StreamOptions *streamOptions `json:"stream_options,omitempty"`
|
||
}
|
||
|
||
type streamChunk struct {
|
||
Choices []struct {
|
||
Delta struct {
|
||
Content string `json:"content"`
|
||
// 思考内容：DeepSeek-R1 / Qwen3 等推理模型走 reasoning_content，
|
||
// 也有网关用 reasoning，两个字段都收。
|
||
Reasoning string `json:"reasoning_content"`
|
||
ReasoningAlt string `json:"reasoning"`
|
||
} `json:"delta"`
|
||
FinishReason string `json:"finish_reason"`
|
||
} `json:"choices"`
|
||
Usage *struct {
|
||
PromptTokens int `json:"prompt_tokens"`
|
||
CompletionTokens int `json:"completion_tokens"`
|
||
TotalTokens int `json:"total_tokens"`
|
||
} `json:"usage"`
|
||
Error *struct {
|
||
Message string `json:"message"`
|
||
Type string `json:"type"`
|
||
} `json:"error"`
|
||
}
|
||
|
||
func (o *openAI) endpoint() string {
|
||
base := strings.TrimRight(o.opts.BaseURL, "/")
|
||
if strings.HasSuffix(base, "/chat/completions") {
|
||
return base
|
||
}
|
||
return base + "/chat/completions"
|
||
}
|
||
|
||
func (o *openAI) setHeaders(req *http.Request, stream bool) {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
if stream {
|
||
req.Header.Set("Accept", "text/event-stream")
|
||
} else {
|
||
req.Header.Set("Accept", "application/json")
|
||
}
|
||
if o.opts.APIKey != "" {
|
||
req.Header.Set("Authorization", "Bearer "+o.opts.APIKey)
|
||
}
|
||
for k, v := range o.opts.ExtraHeaders {
|
||
req.Header.Set(k, v)
|
||
}
|
||
}
|
||
|
||
func (o *openAI) buildBody(system string, msgs []Message, stream bool) ([]byte, error) {
|
||
body := chatRequest{
|
||
Model: o.opts.Model,
|
||
Stream: stream,
|
||
MaxTokens: o.opts.MaxTokens,
|
||
Temperature: o.opts.Temperature,
|
||
}
|
||
if stream {
|
||
body.StreamOptions = &streamOptions{IncludeUsage: true}
|
||
}
|
||
if strings.TrimSpace(system) != "" {
|
||
body.Messages = append(body.Messages, reqMessage{Role: "system", Content: system})
|
||
}
|
||
for _, m := range msgs {
|
||
role := string(m.Role)
|
||
if role == "" {
|
||
role = "user"
|
||
}
|
||
if len(m.Image) == 0 {
|
||
body.Messages = append(body.Messages, reqMessage{Role: role, Content: m.Text})
|
||
continue
|
||
}
|
||
text := m.Text
|
||
if strings.TrimSpace(text) == "" {
|
||
text = "请看图并回答。"
|
||
}
|
||
parts := []contentPart{
|
||
{Type: "text", Text: text},
|
||
{
|
||
Type: "image_url",
|
||
ImageURL: &imageURL{
|
||
URL: "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(m.Image),
|
||
},
|
||
},
|
||
}
|
||
body.Messages = append(body.Messages, reqMessage{Role: role, Content: parts})
|
||
}
|
||
raw, err := json.Marshal(body)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("构造请求失败: %w", err)
|
||
}
|
||
return raw, nil
|
||
}
|
||
|
||
// Chat 发起一次流式对话。
|
||
func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error) {
|
||
if len(msgs) == 0 {
|
||
return Usage{}, fmt.Errorf("没有可发送的消息")
|
||
}
|
||
payload, err := o.buildBody(system, msgs, true)
|
||
if err != nil {
|
||
return Usage{}, err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
|
||
if err != nil {
|
||
return Usage{}, fmt.Errorf("创建请求失败: %w", err)
|
||
}
|
||
o.setHeaders(req, true)
|
||
|
||
resp, err := o.http.Do(req)
|
||
if err != nil {
|
||
return Usage{}, fmt.Errorf("请求 %s 失败: %w", o.opts.Name, err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode != http.StatusOK {
|
||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
|
||
return Usage{}, fmt.Errorf("模型返回 HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
|
||
}
|
||
|
||
// 部分服务端会忽略 stream 参数直接返回完整 JSON，这里做兼容。
|
||
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "event-stream") {
|
||
return o.readWhole(resp.Body, fn)
|
||
}
|
||
|
||
var usage Usage
|
||
var got strings.Builder
|
||
var gotReasoning strings.Builder
|
||
err = readSSE(resp.Body, func(ev sseEvent) error {
|
||
if len(ev.Data) == 0 || bytes.Equal(ev.Data, []byte("[DONE]")) {
|
||
return nil
|
||
}
|
||
var chunk streamChunk
|
||
if json.Unmarshal(ev.Data, &chunk) != nil {
|
||
return nil // 忽略无法解析的心跳/杂项
|
||
}
|
||
if chunk.Error != nil {
|
||
return fmt.Errorf("模型返回错误: %s", chunk.Error.Message)
|
||
}
|
||
if chunk.Usage != nil {
|
||
usage = Usage{
|
||
PromptTokens: chunk.Usage.PromptTokens,
|
||
CompletionTokens: chunk.Usage.CompletionTokens,
|
||
TotalTokens: chunk.Usage.TotalTokens,
|
||
}
|
||
}
|
||
for _, c := range chunk.Choices {
|
||
reasoning := c.Delta.Reasoning
|
||
if reasoning == "" {
|
||
reasoning = c.Delta.ReasoningAlt
|
||
}
|
||
if reasoning != "" {
|
||
gotReasoning.WriteString(reasoning)
|
||
if fn != nil {
|
||
if derr := fn(Delta{Reasoning: reasoning}); derr != nil {
|
||
return derr
|
||
}
|
||
}
|
||
}
|
||
if c.Delta.Content == "" {
|
||
continue
|
||
}
|
||
got.WriteString(c.Delta.Content)
|
||
if fn != nil {
|
||
if derr := fn(Delta{Text: c.Delta.Content}); derr != nil {
|
||
return derr
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return usage, err
|
||
}
|
||
if got.Len() == 0 && gotReasoning.Len() == 0 {
|
||
return usage, fmt.Errorf("模型没有返回任何内容，请确认 %s 是支持图片输入的视觉模型", o.opts.Model)
|
||
}
|
||
return usage, nil
|
||
}
|
||
|
||
// readWhole 处理服务端忽略 stream 参数的情况。
|
||
func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
|
||
raw, err := io.ReadAll(io.LimitReader(r, 8<<20))
|
||
if err != nil {
|
||
return Usage{}, fmt.Errorf("读取响应失败: %w", err)
|
||
}
|
||
var resp struct {
|
||
Choices []struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
// 非流式响应里思考内容通常挂在 message 上。
|
||
ReasoningContent string `json:"reasoning_content"`
|
||
Reasoning string `json:"reasoning"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
Usage *struct {
|
||
PromptTokens int `json:"prompt_tokens"`
|
||
CompletionTokens int `json:"completion_tokens"`
|
||
TotalTokens int `json:"total_tokens"`
|
||
} `json:"usage"`
|
||
Error *struct {
|
||
Message string `json:"message"`
|
||
} `json:"error"`
|
||
}
|
||
if err := json.Unmarshal(raw, &resp); err != nil {
|
||
return Usage{}, fmt.Errorf("解析响应失败: %w", err)
|
||
}
|
||
if resp.Error != nil {
|
||
return Usage{}, fmt.Errorf("模型返回错误: %s", resp.Error.Message)
|
||
}
|
||
if len(resp.Choices) == 0 {
|
||
return Usage{}, fmt.Errorf("模型返回内容为空")
|
||
}
|
||
msg := resp.Choices[0].Message
|
||
if msg.Content == "" && msg.ReasoningContent == "" && msg.Reasoning == "" {
|
||
return Usage{}, fmt.Errorf("模型返回内容为空")
|
||
}
|
||
var usage Usage
|
||
if resp.Usage != nil {
|
||
usage = Usage{
|
||
PromptTokens: resp.Usage.PromptTokens,
|
||
CompletionTokens: resp.Usage.CompletionTokens,
|
||
TotalTokens: resp.Usage.TotalTokens,
|
||
}
|
||
}
|
||
if fn != nil {
|
||
reasoning := msg.ReasoningContent
|
||
if reasoning == "" {
|
||
reasoning = msg.Reasoning
|
||
}
|
||
if reasoning != "" {
|
||
if err := fn(Delta{Reasoning: reasoning}); err != nil {
|
||
return usage, err
|
||
}
|
||
}
|
||
if msg.Content != "" {
|
||
if err := fn(Delta{Text: msg.Content}); err != nil {
|
||
return usage, err
|
||
}
|
||
}
|
||
}
|
||
return usage, nil
|
||
}
|
||
|
||
// Ping 做一次最小代价的连通性检查。
|
||
func (o *openAI) Ping(ctx context.Context) error {
|
||
payload, err := o.buildBody("", []Message{{Role: RoleUser, Text: "ping"}}, false)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
|
||
if err != nil {
|
||
return fmt.Errorf("创建请求失败: %w", err)
|
||
}
|
||
o.setHeaders(req, false)
|
||
|
||
resp, err := o.http.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("连接 %s 失败: %w", o.opts.BaseURL, err)
|
||
}
|
||
defer resp.Body.Close()
|
||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||
if resp.StatusCode != http.StatusOK {
|
||
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
|
||
}
|
||
if len(raw) == 0 {
|
||
return fmt.Errorf("服务返回空响应")
|
||
}
|
||
return nil
|
||
}
|