342 lines
9.5 KiB
Go
342 lines
9.5 KiB
Go
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
|
||
}
|