优化代码
This commit is contained in:
hero920103 committed
1 parent
3553790475
commit
eda7e8bc96
23 files changed
+3113
-342
No files matched your search
@@ -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
|
||||
}
|
||||
Reference in new issue
Block a user