You've already forked agentic-coding-workflow
feat(infra/llm): OpenAI client + base_url override + reasoning_content compat
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user