feat(infra/llm): OpenAI client + base_url override + reasoning_content compat

This commit is contained in:
2026-05-04 09:40:41 +08:00
parent 027b5d5eba
commit be751e76f6
4 changed files with 306 additions and 0 deletions
+3
View File
@@ -30,6 +30,7 @@ require (
github.com/cpuguy83/dockercfg v0.3.2 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/distribution/reference v0.6.0 // indirect
github.com/dlclark/regexp2 v1.10.0 // indirect
github.com/docker/go-connections v0.6.0 // indirect
github.com/docker/go-units v0.5.0 // indirect
github.com/ebitengine/purego v0.10.0 // indirect
@@ -58,9 +59,11 @@ require (
github.com/moby/sys/user v0.4.0 // indirect
github.com/moby/sys/userns v0.1.0 // indirect
github.com/moby/term v0.5.2 // indirect
github.com/openai/openai-go v1.12.0 // indirect
github.com/opencontainers/go-digest v1.0.0 // indirect
github.com/opencontainers/image-spec v1.1.1 // indirect
github.com/pelletier/go-toml/v2 v2.2.4 // indirect
github.com/pkoukk/tiktoken-go v0.1.8 // indirect
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect
github.com/sagikazarmark/locafero v0.11.0 // indirect
+6
View File
@@ -37,6 +37,8 @@ github.com/dhui/dktest v0.4.6 h1:+DPKyScKSEp3VLtbMDHcUq6V5Lm5zfZZVb0Sk7Ahom4=
github.com/dhui/dktest v0.4.6/go.mod h1:JHTSYDtKkvFNFHJKqCzVzqXecyv+tKt8EzceOmQOgbU=
github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/docker/docker v28.3.3+incompatible h1:Dypm25kh4rmk49v1eiVbsAtpAsYURjYkaKubwuBdxEI=
github.com/docker/docker v28.3.3+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
github.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94=
@@ -122,6 +124,8 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
github.com/oklog/ulid/v2 v2.1.1 h1:suPZ4ARWLOJLegGFiZZ1dFAkqzhMjL3J1TzI+5wHz8s=
github.com/oklog/ulid/v2 v2.1.1/go.mod h1:rcEKHmBBKfef9DhnvX7y1HZBYxjXb0cP5ExxNsTT1QQ=
github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0=
github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y=
github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
@@ -131,6 +135,8 @@ github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0
github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pkoukk/tiktoken-go v0.1.8 h1:85ENo+3FpWgAACBaEUVp+lctuTcYUO7BtmfhlN/QTRo=
github.com/pkoukk/tiktoken-go v0.1.8/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+201
View File
@@ -0,0 +1,201 @@
package llm
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"github.com/openai/openai-go"
"github.com/openai/openai-go/option"
"github.com/openai/openai-go/packages/param"
tiktoken "github.com/pkoukk/tiktoken-go"
)
// Compile-time assertion: OpenAIClient must satisfy LLMClient.
var _ LLMClient = (*OpenAIClient)(nil)
// OpenAIClient 通过官方 SDK 调用 Chat Completions API。
// baseURL 改写后兼容 DeepSeek / Qwen / Ollama / vLLM / Moonshot 等 OpenAI-compatible 端点。
// 流式响应中识别 reasoning_content 扩展字段以支持 DeepSeek / Qwen 的 thinking 输出。
type OpenAIClient struct {
sdk openai.Client // SDK v1.12: NewClient returns value, not pointer
}
// NewOpenAI 构造客户端。baseURL 可空(用 SDK 默认)。
func NewOpenAI(apiKey, baseURL string) (*OpenAIClient, error) {
opts := []option.RequestOption{option.WithAPIKey(apiKey)}
if baseURL != "" {
opts = append(opts, option.WithBaseURL(baseURL))
}
c := openai.NewClient(opts...)
return &OpenAIClient{sdk: c}, nil
}
func (c *OpenAIClient) Provider() string { return "openai" }
// toOpenAIMessages 把中立 Message 切片转换为 SDK 类型。
// document/PDF 块由上层能力检查拦截,此处不翻译(OpenAI Chat Completions 不原生支持 PDF)。
// image 块编码为 base64 data URI 传入。
func toOpenAIMessages(msgs []Message) []openai.ChatCompletionMessageParamUnion {
out := make([]openai.ChatCompletionMessageParamUnion, 0, len(msgs))
for _, m := range msgs {
switch m.Role {
case RoleUser:
parts := make([]openai.ChatCompletionContentPartUnionParam, 0, len(m.Blocks))
for _, b := range m.Blocks {
switch b.Type {
case "text":
parts = append(parts, openai.TextContentPart(b.Text))
case "image":
// SDK v1.12: ImageContentPart 接受 ChatCompletionContentPartImageImageURLParam
encoded := base64.StdEncoding.EncodeToString(b.Data)
dataURI := fmt.Sprintf("data:%s;base64,%s", b.MimeType, encoded)
parts = append(parts, openai.ImageContentPart(openai.ChatCompletionContentPartImageImageURLParam{
URL: dataURI,
}))
// document/PDF 不翻译,由上层能力检查拦截
}
}
out = append(out, openai.UserMessage(parts))
case RoleAssistant:
// 助手消息只取第一个 text block(assistant 输出 PDF 场景不存在)
text := ""
for _, b := range m.Blocks {
if b.Type == "text" {
text = b.Text
break
}
}
out = append(out, openai.AssistantMessage(text))
}
}
return out
}
// Stream 启动流式请求,将 SDK chunk 转换为中立 StreamEvent。
// 使用 ctx-aware send 辅助函数防止下游 reader 退出后 goroutine 永久阻塞(同 anthropic.go 模式)。
func (c *OpenAIClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
msgs := make([]openai.ChatCompletionMessageParamUnion, 0)
if req.SystemPrompt != "" {
msgs = append(msgs, openai.SystemMessage(req.SystemPrompt))
}
msgs = append(msgs, toOpenAIMessages(req.Messages)...)
params := openai.ChatCompletionNewParams{
Model: openai.ChatModel(modelID),
Messages: msgs,
// SDK v1.12: MaxCompletionTokens 替代已弃用的 MaxTokens
StreamOptions: openai.ChatCompletionStreamOptionsParam{
IncludeUsage: param.NewOpt(true),
},
}
if req.MaxOutputTokens > 0 {
params.MaxCompletionTokens = param.NewOpt(int64(req.MaxOutputTokens))
}
stream := c.sdk.Chat.Completions.NewStreaming(ctx, params)
out := make(chan StreamEvent, 16)
// send 在 ctx 取消时立即放弃发送,防止下游 reader 已退出导致 goroutine 永久阻塞。
send := func(ev StreamEvent) bool {
select {
case <-ctx.Done():
return false
case out <- ev:
return true
}
}
go func() {
defer close(out)
var usageSent bool
for stream.Next() {
chunk := stream.Current()
// 处理 usage(最后一个含 usage 的 chunk,choices 为空)
if chunk.JSON.Usage.Valid() && chunk.Usage.PromptTokens > 0 {
if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
PromptTokens: int(chunk.Usage.PromptTokens),
CompletionTokens: int(chunk.Usage.CompletionTokens),
}}) {
return
}
usageSent = true
}
for _, choice := range chunk.Choices {
delta := choice.Delta
// reasoning_content 扩展字段(DeepSeek-Reasoner / Qwen-Thinking)
// SDK v1.12: 非标准字段通过 delta.JSON.ExtraFields["reasoning_content"] 访问。
// ExtraFields 中的条目 status=invalid(不可 Valid()),但 Raw() 非空时表示字段存在。
// Raw() 返回带引号的原始 JSON 字符串,需 json.Unmarshal 解析为 Go string。
if rcField, ok := delta.JSON.ExtraFields["reasoning_content"]; ok && rcField.Raw() != "" {
var rcText string
if err := json.Unmarshal([]byte(rcField.Raw()), &rcText); err == nil && rcText != "" {
if !send(StreamEvent{Type: EventThinkingDelta, Text: rcText}) {
return
}
}
}
// 普通文本 delta
if delta.Content != "" {
if !send(StreamEvent{Type: EventTextDelta, Text: delta.Content}) {
return
}
}
// 流结束
if choice.FinishReason != "" {
// usage 可能在同一 chunk 的 choices 之后,也可能在独立 chunk;
// 此处仅在尚未发送 usage 时检查本 chunk 的 usage
if !usageSent && chunk.JSON.Usage.Valid() && chunk.Usage.PromptTokens > 0 {
if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
PromptTokens: int(chunk.Usage.PromptTokens),
CompletionTokens: int(chunk.Usage.CompletionTokens),
}}) {
return
}
usageSent = true
}
if !send(StreamEvent{Type: EventDone, Reason: choice.FinishReason}) {
return
}
}
}
}
if err := stream.Err(); err != nil {
send(StreamEvent{Type: EventError, Err: fmt.Errorf("openai stream: %w", err)})
}
}()
return out, nil
}
// CountTokens 使用本地 tiktoken-go 统计 token 数,不发起远程 API 调用。
// 当模型不可识别时,自动回退到 cl100k_base 编码。
func (c *OpenAIClient) CountTokens(_ context.Context, modelID string, req Request) (int, error) {
enc, err := tiktoken.EncodingForModel(modelID)
if err != nil {
// SDK v1.12 / tiktoken-go v0.1.x: 未知模型时回退到 cl100k_base
enc, err = tiktoken.GetEncoding("cl100k_base")
if err != nil {
return 0, fmt.Errorf("tiktoken fallback cl100k_base: %w", err)
}
}
total := 0
if req.SystemPrompt != "" {
total += len(enc.Encode(req.SystemPrompt, nil, nil))
}
for _, m := range req.Messages {
for _, b := range m.Blocks {
if b.Type == "text" {
total += len(enc.Encode(b.Text, nil, nil))
}
}
}
return total, nil
}
+96
View File
@@ -0,0 +1,96 @@
package llm
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
const openaiStreamFixture = `data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}
data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}
data: [DONE]
`
const openaiReasoningFixture = `data: {"id":"c1","object":"chat.completion.chunk","model":"deepseek-reasoner","choices":[{"index":0,"delta":{"reasoning_content":"Let me think"},"finish_reason":null}]}
data: {"id":"c1","object":"chat.completion.chunk","model":"deepseek-reasoner","choices":[{"index":0,"delta":{"content":"Done"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":1}}
data: [DONE]
`
func TestOpenAI_Stream_TextDelta(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(openaiStreamFixture))
}))
defer srv.Close()
c, _ := NewOpenAI("k", srv.URL)
ch, err := c.Stream(context.Background(), "gpt-4o", Request{
Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}},
MaxOutputTokens: 100,
})
require.NoError(t, err)
var deltas []string
var sawDone, sawUsage bool
for ev := range ch {
switch ev.Type {
case EventTextDelta:
deltas = append(deltas, ev.Text)
case EventUsage:
sawUsage = true
require.Equal(t, 10, ev.Usage.PromptTokens)
require.Equal(t, 2, ev.Usage.CompletionTokens)
case EventDone:
sawDone = true
}
}
require.Equal(t, "Hello world", strings.Join(deltas, ""))
require.True(t, sawUsage)
require.True(t, sawDone)
}
func TestOpenAI_Stream_ReasoningContent_Compatible(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte(openaiReasoningFixture))
}))
defer srv.Close()
c, _ := NewOpenAI("k", srv.URL)
ch, _ := c.Stream(context.Background(), "deepseek-reasoner", Request{
Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}},
})
var thinking, text string
for ev := range ch {
switch ev.Type {
case EventThinkingDelta:
thinking += ev.Text
case EventTextDelta:
text += ev.Text
}
}
require.Equal(t, "Let me think", thinking)
require.Equal(t, "Done", text)
}
func TestOpenAI_CountTokens_LocalTiktoken(t *testing.T) {
c, _ := NewOpenAI("k", "")
got, err := c.CountTokens(context.Background(), "gpt-4o", Request{
SystemPrompt: "You are helpful.",
Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hello world"}}}},
})
require.NoError(t, err)
require.Greater(t, got, 0)
}