302 lines
7.8 KiB
Go
302 lines
7.8 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// 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: &http.Client{
|
|
// 不设总超时：流式响应可能持续很久，由 ctx 控制。
|
|
Transport: &http.Transport{
|
|
Proxy: http.ProxyFromEnvironment,
|
|
MaxIdleConns: 8,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
ForceAttemptHTTP2: true,
|
|
TLSHandshakeTimeout: 15 * time.Second,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
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"`
|
|
} `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
|
|
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 {
|
|
if c.Delta.Content == "" {
|
|
continue
|
|
}
|
|
got.WriteString(c.Delta.Content)
|
|
if fn != nil {
|
|
if derr := fn(c.Delta.Content); derr != nil {
|
|
return derr
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return usage, err
|
|
}
|
|
if got.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"`
|
|
} `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 || resp.Choices[0].Message.Content == "" {
|
|
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,
|
|
}
|
|
}
|
|
text := resp.Choices[0].Message.Content
|
|
if fn != nil {
|
|
if err := fn(text); 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
|
|
}
|