优化代码
This commit is contained in:
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
|
||||
}
|
||||
+41
-20
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user