improve(relay):优化gemini流式输出效率
This commit is contained in:
chengqian committed
1 parent
03f6c1f304
commit
bc42e00ac9
4 files changed
+245
-10
No files matched your search
@@ -1369,11 +1369,9 @@ func HandleResponsesStreamToChat(ctx context.Context, resp *http.Response, info
|
||||
if chunk == nil {
|
||||
return true
|
||||
}
|
||||
data, err := json.Marshal(chunk)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return helper.WriteSSEData(writer, string(data)) == nil
|
||||
// 池化缓冲 + 绑定 encoder：序列化与拼帧零分配，整帧单次写出。
|
||||
// 序列化失败与写失败都返回 false（同旧行为），由调用方终止流。
|
||||
return helper.WriteSSEDataJSON(writer, chunk) == nil
|
||||
}
|
||||
|
||||
sendStartIfNeeded := func() bool {
|
||||
|
||||
@@ -157,8 +157,8 @@ func handleGeminiInboundStream(ctx context.Context, resp *http.Response, info *c
|
||||
}
|
||||
|
||||
if len(geminiChunk.Candidates) > 0 {
|
||||
chunkData, _ := json.Marshal(geminiChunk)
|
||||
_ = helper.WriteSSEData(writer, string(chunkData))
|
||||
// 传指针而非结构体值：值传进 any 会额外堆分配一份拷贝用于装箱
|
||||
_ = helper.WriteSSEDataJSON(writer, &geminiChunk)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -194,10 +194,9 @@ func handleGeminiInboundStream(ctx context.Context, resp *http.Response, info *c
|
||||
finalChunk.UsageMetadata.ThoughtsTokenCount = usage.CompletionTokenDetails.ReasoningTokens
|
||||
}
|
||||
}
|
||||
chunkData, _ := json.Marshal(finalChunk)
|
||||
_ = helper.WriteSSEData(writer, string(chunkData))
|
||||
_ = helper.WriteSSEDataJSON(writer, &finalChunk)
|
||||
|
||||
helper.WriteSSEData(writer, "[DONE]")
|
||||
_ = helper.WriteSSEData(writer, "[DONE]")
|
||||
info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil)
|
||||
|
||||
if err := scanner.Err(); err != nil && err != io.EOF && ctx.Err() == nil {
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/qianfree/team-api/relay/common"
|
||||
"github.com/qianfree/team-api/relay/constant"
|
||||
"github.com/qianfree/team-api/relay/dto"
|
||||
)
|
||||
|
||||
// 本文件覆盖两条改用池化帧写入器（helper.WriteSSEDataJSON）的转换路径：
|
||||
// - handleGeminiInboundStream：OpenAI 上游 → Gemini 客户端
|
||||
// - HandleResponsesStreamToChat：Responses 上游 → Chat Completions 客户端
|
||||
//
|
||||
// 两条路径此前无测试覆盖。断言两件事：
|
||||
// 1. 输出仍是合法且内容正确的 SSE（池化不能改变对外协议）；
|
||||
// 2. 每次 Write 都落在帧边界（整帧单次写出，并发保活 ping 插不进帧中间）。
|
||||
|
||||
// assertFramesAreSingleWrites 断言每次 Write 都是一个以空行收尾的完整帧
|
||||
func assertFramesAreSingleWrites(t *testing.T, chunks []string) {
|
||||
t.Helper()
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("没有任何写出")
|
||||
}
|
||||
for i, c := range chunks {
|
||||
if !strings.HasSuffix(c, "\n\n") {
|
||||
t.Errorf("第 %d 次 Write 未落在帧边界，保活 ping 可插入其中:\n%q", i, c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// collectSSEDataFrames 从 SSE 文本中取出所有 data 行的内容
|
||||
func collectSSEDataFrames(raw string) []string {
|
||||
var out []string
|
||||
for _, line := range strings.Split(raw, "\n") {
|
||||
line = strings.TrimSuffix(line, "\r")
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
out = append(out, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func geminiInboundInfo() *common.RelayInfo {
|
||||
return &common.RelayInfo{
|
||||
IsStream: true,
|
||||
InboundFormat: constant.RelayFormatGemini,
|
||||
ClientFormat: constant.RelayFormatGemini,
|
||||
OriginModelName: "gemini-3-pro",
|
||||
StreamStatus: common.NewStreamStatus(),
|
||||
ChannelMeta: &common.ChannelMeta{
|
||||
ChannelType: int(constant.ProviderOpenAI),
|
||||
BaseURL: "https://upstream.example.com",
|
||||
UpstreamModelName: "gpt-4o",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleGeminiInboundStream_PooledFramesValid OpenAI SSE → Gemini SSE：
|
||||
// 每个 chunk 都是完整可解析的 Gemini JSON，且整帧单次写出。
|
||||
func TestHandleGeminiInboundStream_PooledFramesValid(t *testing.T) {
|
||||
upstream := strings.Join([]string{
|
||||
`data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"你好"}}]}`,
|
||||
``,
|
||||
`data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"世界"}}]}`,
|
||||
``,
|
||||
`data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`,
|
||||
``,
|
||||
`data: [DONE]`,
|
||||
``,
|
||||
}, "\n")
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": {"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(upstream)),
|
||||
}
|
||||
|
||||
info := geminiInboundInfo()
|
||||
rec := &sseFrameRecorder{}
|
||||
|
||||
usage, err := handleGeminiInboundStream(context.Background(), resp, info, rec)
|
||||
if err != nil {
|
||||
t.Fatalf("handleGeminiInboundStream error: %v", err)
|
||||
}
|
||||
|
||||
assertFramesAreSingleWrites(t, rec.chunks)
|
||||
|
||||
frames := collectSSEDataFrames(rec.body.String())
|
||||
if len(frames) < 2 {
|
||||
t.Fatalf("data 帧数 = %d, 期望至少 2（内容 chunk + 终止 chunk）\n%s", len(frames), rec.body.String())
|
||||
}
|
||||
|
||||
var text strings.Builder
|
||||
sawDone := false
|
||||
for i, f := range frames {
|
||||
if f == "[DONE]" {
|
||||
sawDone = true
|
||||
if i != len(frames)-1 {
|
||||
t.Errorf("[DONE] 不在最后一帧（位置 %d/%d）", i, len(frames)-1)
|
||||
}
|
||||
continue
|
||||
}
|
||||
var chunk dto.GeminiChatResponse
|
||||
if err := json.Unmarshal([]byte(f), &chunk); err != nil {
|
||||
t.Fatalf("第 %d 帧不是合法 Gemini JSON（池化若出错会产生这种残片）: %v\ndata=%q", i, err, f)
|
||||
}
|
||||
for _, cand := range chunk.Candidates {
|
||||
if cand.Content == nil {
|
||||
continue
|
||||
}
|
||||
for _, p := range cand.Content.Parts {
|
||||
text.WriteString(p.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawDone {
|
||||
t.Error("缺少 [DONE] 终止帧")
|
||||
}
|
||||
if got := text.String(); got != "你好世界" {
|
||||
t.Errorf("转发文本 = %q, want %q", got, "你好世界")
|
||||
}
|
||||
if usage == nil || usage.PromptTokens != 11 || usage.CompletionTokens != 7 {
|
||||
t.Errorf("usage = %+v, want prompt=11 completion=7", usage)
|
||||
}
|
||||
}
|
||||
|
||||
func responsesToChatInfo() *common.RelayInfo {
|
||||
return &common.RelayInfo{
|
||||
IsStream: true,
|
||||
InboundFormat: constant.RelayFormatOpenAI,
|
||||
ClientFormat: constant.RelayFormatOpenAI,
|
||||
OriginModelName: "gpt-4o",
|
||||
StreamStatus: common.NewStreamStatus(),
|
||||
ChannelMeta: &common.ChannelMeta{
|
||||
ChannelType: int(constant.ProviderOpenAI),
|
||||
BaseURL: "https://upstream.example.com",
|
||||
UpstreamModelName: "gpt-4o",
|
||||
SupportsResponses: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleResponsesStreamToChat_PooledFramesValid Responses SSE → Chat SSE：
|
||||
// 每个 chunk 都是完整可解析的 chat.completion.chunk，且整帧单次写出。
|
||||
func TestHandleResponsesStreamToChat_PooledFramesValid(t *testing.T) {
|
||||
upstream := strings.Join([]string{
|
||||
"event: response.created",
|
||||
`data: {"type":"response.created","response":{"id":"resp_1","object":"response","status":"in_progress"}}`,
|
||||
``,
|
||||
"event: response.output_text.delta",
|
||||
`data: {"type":"response.output_text.delta","delta":"你好"}`,
|
||||
``,
|
||||
"event: response.output_text.delta",
|
||||
`data: {"type":"response.output_text.delta","delta":"世界"}`,
|
||||
``,
|
||||
"event: response.completed",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","status":"completed","usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}}`,
|
||||
``,
|
||||
`data: [DONE]`,
|
||||
``,
|
||||
}, "\n")
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": {"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(upstream)),
|
||||
}
|
||||
|
||||
info := responsesToChatInfo()
|
||||
rec := &sseFrameRecorder{}
|
||||
|
||||
if _, err := HandleResponsesStreamToChat(context.Background(), resp, info, rec); err != nil {
|
||||
t.Fatalf("HandleResponsesStreamToChat error: %v", err)
|
||||
}
|
||||
|
||||
assertFramesAreSingleWrites(t, rec.chunks)
|
||||
|
||||
frames := collectSSEDataFrames(rec.body.String())
|
||||
var text strings.Builder
|
||||
sawDone := false
|
||||
for i, f := range frames {
|
||||
if f == "[DONE]" {
|
||||
sawDone = true
|
||||
continue
|
||||
}
|
||||
var chunk dto.ChatCompletionStreamResponse
|
||||
if err := json.Unmarshal([]byte(f), &chunk); err != nil {
|
||||
t.Fatalf("第 %d 帧不是合法 chat.completion.chunk: %v\ndata=%q", i, err, f)
|
||||
}
|
||||
if chunk.Object != "chat.completion.chunk" {
|
||||
t.Errorf("第 %d 帧 object = %q, want chat.completion.chunk", i, chunk.Object)
|
||||
}
|
||||
for _, c := range chunk.Choices {
|
||||
if s, ok := c.Delta.Content.(string); ok {
|
||||
text.WriteString(s)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawDone {
|
||||
t.Error("缺少 [DONE] 终止帧")
|
||||
}
|
||||
if got := text.String(); got != "你好世界" {
|
||||
t.Errorf("转发文本 = %q, want %q", got, "你好世界")
|
||||
}
|
||||
}
|
||||
@@ -832,6 +832,33 @@ func TestWriteSSEEventJSON_RealWorldFrameSizes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkSSEData_PooledJSON_Pointer / _Value 验证「传指针而非结构体值」的理由：
|
||||
// 结构体值传进 any 形参需要装箱，会额外堆分配一份拷贝；传指针则不会。
|
||||
// 迁移调用点时据此统一传 &chunk。
|
||||
func BenchmarkSSEData_PooledJSON_Pointer(b *testing.B) {
|
||||
w := &discardWriter{}
|
||||
payload := benchClaudeDelta() // 已是指针
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := WriteSSEDataJSON(w, payload); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSSEData_PooledJSON_Value(b *testing.B) {
|
||||
w := &discardWriter{}
|
||||
payload := *benchClaudeDelta() // 解引用成值
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := WriteSSEDataJSON(w, payload); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSSEFrameWriter_MaxBufDegradeSuppressesPing 超大帧降级为分次写出时，写出的这批字节
|
||||
// 不含帧尾 —— 此刻放行 ping 依然会把帧劈成两个事件。降级期间必须抑制 ping，直到分隔空行到达。
|
||||
func TestSSEFrameWriter_MaxBufDegradeSuppressesPing(t *testing.T) {
|
||||
|
||||
Reference in new issue
Block a user