优化代码

This commit is contained in:
hero920103 committed 2026-10-28 08:47:09 +08:00
1 parent 3553790475
commit eda7e8bc96
23 files changed
+3113 -342

No files matched your search

+341
View File
@@ -0,0 +1,341 @@
package llm
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
)
// DefaultAnthropicVersion 是 anthropic-version 请求头的默认值。
const DefaultAnthropicVersion = "2023-06-01"
// anthropic 是 Anthropic Messages 协议(/v1/messages + SSE)的实现。
// 与 OpenAI 兼容格式的三处关键差异:鉴权用 x-api-key 而非 Bearer;
// system 是请求体顶层字段而非一条消息;图片走 source.base64 而非 data URL。
type anthropic struct {
opts Options
http *http.Client
}
func newAnthropic(opts Options) *anthropic {
if opts.MaxTokens <= 0 {
opts.MaxTokens = 1024
}
if opts.Temperature <= 0 {
opts.Temperature = 0.3
}
return &anthropic{opts: opts, http: newHTTPClient()}
}
type anthropicRequest struct {
Model string `json:"model"`
Messages []anthropicMessage `json:"messages"`
System string `json:"system,omitempty"`
Stream bool `json:"stream"`
MaxTokens int `json:"max_tokens"`
Temperature float64 `json:"temperature,omitempty"`
}
type anthropicMessage struct {
Role string `json:"role"`
Content any `json:"content"`
}
type anthropicBlock struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
Source *anthropicImage `json:"source,omitempty"`
}
type anthropicImage struct {
Type string `json:"type"`
MediaType string `json:"media_type"`
Data string `json:"data"`
}
type anthropicUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
}
// anthropicEvent 覆盖流式与非流式两种响应体,用不到的字段留空即可。
type anthropicEvent struct {
Type string `json:"type"`
Message *struct {
Usage *anthropicUsage `json:"usage"`
Content []struct {
Type string `json:"type"`
Text string `json:"text"`
Thinking string `json:"thinking"`
} `json:"content"`
} `json:"message"`
Delta *struct {
Type string `json:"type"`
Text string `json:"text"`
// thinking_delta 里思考文本在 thinking 字段。
Thinking string `json:"thinking"`
} `json:"delta"`
Usage *anthropicUsage `json:"usage"`
Error *struct {
Type string `json:"type"`
Message string `json:"message"`
} `json:"error"`
}
// endpoint 兼容三种写法:官方根地址、已带 /v1、以及直接给全 /messages 的代理。
func (a *anthropic) endpoint() string {
base := strings.TrimRight(a.opts.BaseURL, "/")
switch {
case strings.HasSuffix(base, "/messages"):
return base
case strings.HasSuffix(base, "/v1"):
return base + "/messages"
default:
return base + "/v1/messages"
}
}
func (a *anthropic) 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 a.opts.APIKey != "" {
req.Header.Set("x-api-key", a.opts.APIKey)
}
if a.opts.ExtraHeaders["anthropic-version"] == "" {
req.Header.Set("anthropic-version", DefaultAnthropicVersion)
}
for k, v := range a.opts.ExtraHeaders {
req.Header.Set(k, v)
}
}
func (a *anthropic) buildBody(system string, msgs []Message, stream bool) ([]byte, error) {
body := anthropicRequest{
Model: a.opts.Model,
Stream: stream,
MaxTokens: a.opts.MaxTokens,
Temperature: a.opts.Temperature,
System: strings.TrimSpace(system),
}
for _, m := range msgs {
role := "user"
if m.Role == RoleAssistant {
role = "assistant"
}
if len(m.Image) == 0 {
body.Messages = append(body.Messages, anthropicMessage{Role: role, Content: m.Text})
continue
}
text := m.Text
if strings.TrimSpace(text) == "" {
text = "请看图并回答。"
}
blocks := []anthropicBlock{
{Type: "text", Text: text},
{
Type: "image",
Source: &anthropicImage{
Type: "base64",
MediaType: "image/jpeg",
Data: base64.StdEncoding.EncodeToString(m.Image),
},
},
}
body.Messages = append(body.Messages, anthropicMessage{Role: role, Content: blocks})
}
if len(body.Messages) == 0 {
return nil, fmt.Errorf("没有可发送的消息")
}
raw, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("构造请求失败: %w", err)
}
return raw, nil
}
// Chat 发起一次流式对话。
func (a *anthropic) Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error) {
if len(msgs) == 0 {
return Usage{}, fmt.Errorf("没有可发送的消息")
}
payload, err := a.buildBody(system, msgs, true)
if err != nil {
return Usage{}, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(payload))
if err != nil {
return Usage{}, fmt.Errorf("创建请求失败: %w", err)
}
a.setHeaders(req, true)
resp, err := a.http.Do(req)
if err != nil {
return Usage{}, fmt.Errorf("请求 %s 失败: %w", a.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 a.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 {
return nil
}
var e anthropicEvent
if json.Unmarshal(ev.Data, &e) != nil {
return nil // 忽略无法解析的心跳/杂项
}
switch e.Type {
case "error":
msg := "未知错误"
if e.Error != nil && e.Error.Message != "" {
msg = e.Error.Message
}
return fmt.Errorf("模型返回错误: %s", msg)
case "message_start":
if e.Message != nil && e.Message.Usage != nil {
usage.PromptTokens = e.Message.Usage.InputTokens
usage.CompletionTokens = e.Message.Usage.OutputTokens
}
case "content_block_delta":
if e.Delta == nil {
return nil
}
// thinking_delta 把思考文本放在 thinking 字段,text_delta 才是正式回答。
if e.Delta.Thinking != "" {
gotReasoning.WriteString(e.Delta.Thinking)
if fn != nil {
if derr := fn(Delta{Reasoning: e.Delta.Thinking}); derr != nil {
return derr
}
}
return nil
}
// input_json_delta 等其它类型走别的字段,直接跳过。
if e.Delta.Text == "" {
return nil
}
got.WriteString(e.Delta.Text)
if fn != nil {
if derr := fn(Delta{Text: e.Delta.Text}); derr != nil {
return derr
}
}
case "message_delta":
// 结束前的这一帧才带最终的 output_tokens。
if e.Usage != nil && e.Usage.OutputTokens > 0 {
usage.CompletionTokens = e.Usage.OutputTokens
}
}
return nil
})
if err != nil {
return usage, err
}
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
if got.Len() == 0 && gotReasoning.Len() == 0 {
return usage, fmt.Errorf("模型没有返回任何内容,请确认 %s 是支持图片输入的视觉模型", a.opts.Model)
}
return usage, nil
}
// readWhole 处理服务端忽略 stream 参数的情况。
func (a *anthropic) 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 e anthropicEvent
if err := json.Unmarshal(raw, &e); err != nil {
return Usage{}, fmt.Errorf("解析响应失败: %w", err)
}
if e.Error != nil {
return Usage{}, fmt.Errorf("模型返回错误: %s", e.Error.Message)
}
var usage Usage
var text strings.Builder
var thinking strings.Builder
if e.Message != nil {
if e.Message.Usage != nil {
usage.PromptTokens = e.Message.Usage.InputTokens
usage.CompletionTokens = e.Message.Usage.OutputTokens
}
for _, b := range e.Message.Content {
switch b.Type {
case "text":
text.WriteString(b.Text)
case "thinking":
thinking.WriteString(b.Thinking)
}
}
}
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
if text.Len() == 0 && thinking.Len() == 0 {
return usage, fmt.Errorf("模型返回内容为空")
}
if fn != nil {
if thinking.Len() > 0 {
if err := fn(Delta{Reasoning: thinking.String()}); err != nil {
return usage, err
}
}
if text.Len() > 0 {
if err := fn(Delta{Text: text.String()}); err != nil {
return usage, err
}
}
}
return usage, nil
}
// Ping 做一次最小代价的连通性检查。
func (a *anthropic) Ping(ctx context.Context) error {
payload, err := a.buildBody("", []Message{{Role: RoleUser, Text: "ping"}}, false)
if err != nil {
return err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(payload))
if err != nil {
return fmt.Errorf("创建请求失败: %w", err)
}
a.setHeaders(req, false)
resp, err := a.http.Do(req)
if err != nil {
return fmt.Errorf("连接 %s 失败: %w", a.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("服务返回空响应")
}
// 有些代理会把错误塞进 200 的响应体里。
var probe anthropicEvent
if json.Unmarshal(raw, &probe) == nil && probe.Error != nil {
return fmt.Errorf("模型返回错误: %s", probe.Error.Message)
}
return nil
}
+41 -20
View File
@@ -9,7 +9,6 @@ import (
"io"
"net/http"
"strings"
"time"
)
// openAI 是 OpenAI 兼容协议(/chat/completions + SSE)的实现。
@@ -25,19 +24,7 @@ func newOpenAI(opts Options) *openAI {
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,
},
},
}
return &openAI{opts: opts, http: newHTTPClient()}
}
type reqMessage struct {
@@ -72,6 +59,10 @@ 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"`
@@ -186,6 +177,7 @@ func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn Del
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
@@ -205,12 +197,24 @@ func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn Del
}
}
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(c.Delta.Content); derr != nil {
if derr := fn(Delta{Text: c.Delta.Content}); derr != nil {
return derr
}
}
@@ -220,7 +224,7 @@ func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn Del
if err != nil {
return usage, err
}
if got.Len() == 0 {
if got.Len() == 0 && gotReasoning.Len() == 0 {
return usage, fmt.Errorf("模型没有返回任何内容,请确认 %s 是支持图片输入的视觉模型", o.opts.Model)
}
return usage, nil
@@ -236,6 +240,9 @@ func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
Choices []struct {
Message struct {
Content string `json:"content"`
// 非流式响应里思考内容通常挂在 message 上。
ReasoningContent string `json:"reasoning_content"`
Reasoning string `json:"reasoning"`
} `json:"message"`
} `json:"choices"`
Usage *struct {
@@ -253,7 +260,11 @@ func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
if resp.Error != nil {
return Usage{}, fmt.Errorf("模型返回错误: %s", resp.Error.Message)
}
if len(resp.Choices) == 0 || resp.Choices[0].Message.Content == "" {
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
@@ -264,10 +275,20 @@ func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
TotalTokens: resp.Usage.TotalTokens,
}
}
text := resp.Choices[0].Message.Content
if fn != nil {
if err := fn(text); err != nil {
return usage, err
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
+62 -9
View File
@@ -1,11 +1,36 @@
// Package llm 提供与 OpenAI 兼容接口的大模型通信能力(含视觉与流式)。
// Package llm 提供大模型通信能力(含视觉与流式),
// 支持 OpenAI Chat Completions 与 Anthropic Messages 两种接口格式。
package llm
import (
"context"
"fmt"
"net/http"
"time"
)
// 接口格式。它决定请求体结构、鉴权头与流式事件的解析方式,
// 也是配置里唯一需要用户先选的东西。
const (
// FormatOpenAI 是 OpenAI Chat Completions 兼容格式(POST {base}/chat/completions)。
FormatOpenAI = "openai"
// FormatAnthropic 是 Anthropic Messages 兼容格式(POST {base}/v1/messages)。
FormatAnthropic = "anthropic"
)
// FormatLabel 返回接口格式的中文名,供界面显示。
func FormatLabel(format string) string {
switch format {
case FormatAnthropic:
return "Anthropic 兼容"
default:
return "OpenAI 兼容"
}
}
// Formats 返回全部受支持的接口格式,顺序即界面下拉顺序。
func Formats() []string { return []string{FormatOpenAI, FormatAnthropic} }
// Role 是消息角色。
type Role string
@@ -36,12 +61,21 @@ func (u *Usage) Add(o Usage) {
u.TotalTokens += o.TotalTokens
}
// DeltaFunc 接收流式增量文本,返回 error 可中断请求。
type DeltaFunc func(delta string) error
// Delta 是一次流式增量。
// Text 是正式回答,Reasoning 是模型吐出来的思考过程;两者都按原样透传,
// 服务端吐什么、界面就显示什么。
type Delta struct {
Text string
Reasoning string
}
// Options 描述一个模型接入点。
// DeltaFunc 接收流式增量,返回 error 可中断请求。
type DeltaFunc func(delta Delta) error
// Options 描述一个模型的接入参数。
type Options struct {
Name string
Format string
BaseURL string
APIKey string
Model string
@@ -51,7 +85,6 @@ type Options struct {
}
// Provider 是大模型服务的最小抽象。
// 目前只有 OpenAI 兼容实现,其它私有协议按此接口新增即可。
type Provider interface {
// Chat 发起一次流式对话,system 为系统提示词。
Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error)
@@ -59,13 +92,33 @@ type Provider interface {
Ping(ctx context.Context) error
}
// New 根据接入点创建 Provider。
// New 根据模型配置创建 Provider,按 Format 分派到对应协议实现。
func New(opts Options) (Provider, error) {
if opts.BaseURL == "" {
return nil, fmt.Errorf("模型接入点缺少 BaseURL")
return nil, fmt.Errorf("模型配置缺少 BaseURL")
}
if opts.Model == "" {
return nil, fmt.Errorf("模型接入点 %s 缺少 model", opts.Name)
return nil, fmt.Errorf("模型 %s 缺少模型名", opts.Name)
}
switch opts.Format {
case "", FormatOpenAI:
return newOpenAI(opts), nil
case FormatAnthropic:
return newAnthropic(opts), nil
default:
return nil, fmt.Errorf("未知的接口格式 %q", opts.Format)
}
}
// newHTTPClient 返回不设总超时的客户端:流式响应可能持续很久,由 ctx 控制。
func newHTTPClient() *http.Client {
return &http.Client{
Transport: &http.Transport{
Proxy: http.ProxyFromEnvironment,
MaxIdleConns: 8,
IdleConnTimeout: 90 * time.Second,
ForceAttemptHTTP2: true,
TLSHandshakeTimeout: 15 * time.Second,
},
}
return newOpenAI(opts), nil
}
+139
View File
@@ -0,0 +1,139 @@
package llm
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
// sseServer 起一个假的 SSE 服务端,按脚本逐条吐数据。
func sseServer(t *testing.T, chunks []string, gap time.Duration) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
flusher, _ := w.(http.Flusher)
for _, c := range chunks {
fmt.Fprint(w, c)
if flusher != nil {
flusher.Flush()
}
time.Sleep(gap)
}
}))
}
func openAIChunk(content, reasoning string) string {
delta := fmt.Sprintf(`{"choices":[{"delta":{"content":%q,"reasoning_content":%q}}]}`, content, reasoning)
return "data: " + delta + "\n\n"
}
func TestOpenAIStreamsIncrementally(t *testing.T) {
srv := sseServer(t, []string{
openAIChunk("第一段", ""),
openAIChunk("第二段", ""),
openAIChunk("第三段", ""),
}, 60*time.Millisecond)
defer srv.Close()
prov, err := New(Options{Format: FormatOpenAI, BaseURL: srv.URL, Model: "m", APIKey: "k"})
if err != nil {
t.Fatal(err)
}
type ev struct {
text string
at time.Duration
}
var evs []ev
start := time.Now()
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
evs = append(evs, ev{text: d.Text, at: time.Since(start)})
return nil
})
if err != nil {
t.Fatal(err)
}
if len(evs) != 3 {
t.Fatalf("应收到 3 段增量,实际 %d", len(evs))
}
// 第一段必须在服务端还没吐完时就到,否则就是"一次性返回"。
if evs[0].at > 50*time.Millisecond {
t.Fatalf("第一段增量到达太晚(%v),说明没有真正流式", evs[0].at)
}
if got := evs[0].text + evs[1].text + evs[2].text; got != "第一段第二段第三段" {
t.Fatalf("增量内容不对: %q", got)
}
}
func TestOpenAIParsesReasoning(t *testing.T) {
srv := sseServer(t, []string{
openAIChunk("", "先想一下…"),
openAIChunk("", "再想一下…"),
openAIChunk("正式回答", ""),
}, 10*time.Millisecond)
defer srv.Close()
prov, err := New(Options{Format: FormatOpenAI, BaseURL: srv.URL, Model: "m"})
if err != nil {
t.Fatal(err)
}
var think, text strings.Builder
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
think.WriteString(d.Reasoning)
text.WriteString(d.Text)
return nil
})
if err != nil {
t.Fatal(err)
}
if think.String() != "先想一下…再想一下…" {
t.Fatalf("思考内容解析不对: %q", think.String())
}
if text.String() != "正式回答" {
t.Fatalf("正文解析不对: %q", text.String())
}
}
func TestAnthropicStreamsThinkingAndText(t *testing.T) {
srv := sseServer(t, []string{
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"我在想…\"}}\n\n",
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"答案\"}}\n\n",
"event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":12}}\n\n",
}, 60*time.Millisecond)
defer srv.Close()
prov, err := New(Options{Format: FormatAnthropic, BaseURL: srv.URL, Model: "m"})
if err != nil {
t.Fatal(err)
}
var think, text strings.Builder
start := time.Now()
var firstAt time.Duration
_, err = prov.Chat(context.Background(), "", []Message{{Role: RoleUser, Text: "hi"}}, func(d Delta) error {
if firstAt == 0 {
firstAt = time.Since(start)
}
think.WriteString(d.Reasoning)
text.WriteString(d.Text)
return nil
})
if err != nil {
t.Fatal(err)
}
if think.String() != "我在想…" {
t.Fatalf("思考内容解析不对: %q", think.String())
}
if text.String() != "答案" {
t.Fatalf("正文解析不对: %q", text.String())
}
if firstAt > 50*time.Millisecond {
t.Fatalf("第一段增量到达太晚(%v),说明没有真正流式", firstAt)
}
}