463 lines
13 KiB
Go
463 lines
13 KiB
Go
package relayconvert
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGeminiToOpenAIStatefulStreamTerminal(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
chunk *dto.GeminiChatResponse
|
|
wantFinishReason string
|
|
wantFinishOnFinalize bool
|
|
wantEmptyFinishDelta bool
|
|
}{
|
|
{
|
|
name: "stop",
|
|
chunk: terminalTestGeminiChunk("Hello", "STOP", false),
|
|
wantFinishReason: types.FinishReasonStop,
|
|
wantEmptyFinishDelta: true,
|
|
},
|
|
{
|
|
name: "tool call",
|
|
chunk: terminalTestGeminiChunk("", "STOP", true),
|
|
wantFinishReason: types.FinishReasonToolCalls,
|
|
wantEmptyFinishDelta: true,
|
|
},
|
|
{
|
|
name: "non stop finish reason",
|
|
chunk: terminalTestGeminiChunk("partial", "MAX_TOKENS", false),
|
|
wantFinishReason: types.FinishReasonLength,
|
|
},
|
|
{
|
|
name: "truncated stream",
|
|
chunk: terminalTestGeminiChunk("partial", "", false),
|
|
wantFinishReason: types.FinishReasonStop,
|
|
wantFinishOnFinalize: true,
|
|
wantEmptyFinishDelta: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
info := &convmeta.Values{
|
|
ChannelMetaAttached: true,
|
|
UpstreamModelName: "upstream-model",
|
|
}
|
|
state, err := NewResponseStreamState(
|
|
types.RelayFormatGemini,
|
|
types.RelayFormatOpenAI,
|
|
ResponseStreamOptions{
|
|
ID: "chatcmpl-fixed",
|
|
Created: 1700000000,
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
results, err := ConvertStreamResponseChunk(nil, info, state, tt.chunk)
|
|
require.NoError(t, err)
|
|
chunkFinishes := terminalTestFinishedChatChunks(t, results)
|
|
if tt.wantFinishOnFinalize {
|
|
assert.Empty(t, chunkFinishes)
|
|
} else {
|
|
require.Len(t, chunkFinishes, 1)
|
|
}
|
|
|
|
finalResults, err := FinalizeStreamResponse(nil, info, state)
|
|
require.NoError(t, err)
|
|
finalFinishes := terminalTestFinishedChatChunks(t, finalResults)
|
|
if tt.wantFinishOnFinalize {
|
|
require.Len(t, finalFinishes, 1)
|
|
} else {
|
|
assert.Empty(t, finalFinishes)
|
|
}
|
|
|
|
finishes := append(chunkFinishes, finalFinishes...)
|
|
require.Len(t, finishes, 1)
|
|
finish := finishes[0]
|
|
require.Len(t, finish.Choices, 1)
|
|
require.NotNil(t, finish.Choices[0].FinishReason)
|
|
assert.Equal(t, tt.wantFinishReason, *finish.Choices[0].FinishReason)
|
|
assert.Equal(t, "chatcmpl-fixed", finish.Id)
|
|
assert.Equal(t, int64(1700000000), finish.Created)
|
|
assert.Equal(t, "upstream-model", finish.Model)
|
|
require.NotNil(t, finish.Usage)
|
|
assert.Equal(t, 4, finish.Usage.PromptTokens)
|
|
assert.Equal(t, 2, finish.Usage.CompletionTokens)
|
|
assert.Equal(t, 6, finish.Usage.TotalTokens)
|
|
if tt.wantEmptyFinishDelta {
|
|
assert.Nil(t, finish.Choices[0].Delta.Content)
|
|
assert.Empty(t, finish.Choices[0].Delta.ToolCalls)
|
|
}
|
|
|
|
repeatedFinal, err := FinalizeStreamResponse(nil, info, state)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, repeatedFinal)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestClaudeTargetStatefulStreamTerminalTail(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
from types.RelayFormat
|
|
chunks []any
|
|
wantFinalizerTerminals bool
|
|
wantStopReason string
|
|
}{
|
|
{
|
|
name: "gemini to claude",
|
|
from: types.RelayFormatGemini,
|
|
chunks: []any{
|
|
terminalTestGeminiChunkWithoutUsage("Hello", ""),
|
|
terminalTestGeminiChunk(" world", "STOP", false),
|
|
},
|
|
wantStopReason: "end_turn",
|
|
},
|
|
{
|
|
name: "gemini tool call with split usage",
|
|
from: types.RelayFormatGemini,
|
|
chunks: []any{
|
|
terminalTestGeminiToolChunkWithoutUsage(),
|
|
terminalTestGeminiChunk("", "STOP", false),
|
|
},
|
|
wantStopReason: "tool_use",
|
|
},
|
|
{
|
|
name: "gemini non stop finish with split usage",
|
|
from: types.RelayFormatGemini,
|
|
chunks: []any{
|
|
terminalTestGeminiChunkWithoutUsage("partial", "MAX_TOKENS"),
|
|
terminalTestGeminiUsageOnlyChunk(),
|
|
},
|
|
wantStopReason: "max_tokens",
|
|
},
|
|
{
|
|
name: "responses to claude",
|
|
from: types.RelayFormatOpenAIResponses,
|
|
chunks: []any{
|
|
&dto.ResponsesStreamResponse{
|
|
Type: "response.output_text.delta",
|
|
Delta: "Hello",
|
|
},
|
|
&dto.ResponsesStreamResponse{
|
|
Type: "response.output_text.delta",
|
|
Delta: " world",
|
|
},
|
|
&dto.ResponsesStreamResponse{
|
|
Type: "response.completed",
|
|
Response: &dto.OpenAIResponsesResponse{
|
|
ID: "resp-fixed",
|
|
Object: "response",
|
|
Model: "upstream-model",
|
|
Status: []byte(`"completed"`),
|
|
Usage: &dto.Usage{
|
|
InputTokens: 4,
|
|
OutputTokens: 2,
|
|
TotalTokens: 6,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
wantFinalizerTerminals: true,
|
|
wantStopReason: "end_turn",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
info := &convmeta.Values{
|
|
ChannelMetaAttached: true,
|
|
UpstreamModelName: "upstream-model",
|
|
ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{
|
|
LastMessagesType: convmeta.LastMessageTypeNone,
|
|
},
|
|
}
|
|
state, err := NewResponseStreamState(
|
|
tt.from,
|
|
types.RelayFormatClaude,
|
|
ResponseStreamOptions{
|
|
ID: "stream-fixed",
|
|
Model: "upstream-model",
|
|
Created: 1700000000,
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
var results []ResponseResult
|
|
for _, chunk := range tt.chunks {
|
|
chunkResults, err := ConvertStreamResponseChunk(nil, info, state, chunk)
|
|
require.NoError(t, err)
|
|
results = append(results, chunkResults...)
|
|
}
|
|
|
|
finalResults, err := FinalizeStreamResponse(nil, info, state)
|
|
require.NoError(t, err)
|
|
if tt.wantFinalizerTerminals {
|
|
require.Len(t, finalResults, 3)
|
|
} else {
|
|
assert.Empty(t, finalResults)
|
|
}
|
|
results = append(results, finalResults...)
|
|
|
|
terminalTestAssertClaudeTail(t, results, tt.wantStopReason)
|
|
assert.True(t, info.ClaudeConvertInfo.Done)
|
|
|
|
repeatedFinal, err := FinalizeStreamResponse(nil, info, state)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, repeatedFinal)
|
|
})
|
|
}
|
|
|
|
t.Run("preserves preseeded usage", func(t *testing.T) {
|
|
info := &convmeta.Values{
|
|
ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{
|
|
LastMessagesType: convmeta.LastMessageTypeNone,
|
|
},
|
|
}
|
|
state, err := NewResponseStreamState(
|
|
types.RelayFormatOpenAIResponses,
|
|
types.RelayFormatClaude,
|
|
ResponseStreamOptions{Model: "upstream-model"},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
chunks := []*dto.ResponsesStreamResponse{
|
|
{
|
|
Type: "response.output_text.delta",
|
|
Delta: "Hello",
|
|
},
|
|
{
|
|
Type: "response.completed",
|
|
Response: &dto.OpenAIResponsesResponse{
|
|
ID: "resp-fixed",
|
|
Object: "response",
|
|
Model: "upstream-model",
|
|
Status: []byte(`"completed"`),
|
|
Usage: &dto.Usage{
|
|
InputTokens: 4,
|
|
OutputTokens: 2,
|
|
TotalTokens: 6,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
for _, chunk := range chunks {
|
|
_, err := ConvertStreamResponseChunk(nil, info, state, chunk)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
preseeded := &dto.Usage{
|
|
PromptTokens: 11,
|
|
CompletionTokens: 7,
|
|
TotalTokens: 18,
|
|
}
|
|
info.ClaudeConvertInfo.Usage = preseeded
|
|
finalResults, err := FinalizeStreamResponse(nil, info, state)
|
|
require.NoError(t, err)
|
|
require.Len(t, finalResults, 3)
|
|
assert.Same(t, preseeded, info.ClaudeConvertInfo.Usage)
|
|
|
|
messageDelta, ok := finalResults[1].Value.(*dto.ClaudeResponse)
|
|
require.True(t, ok)
|
|
assert.Equal(t, "message_delta", messageDelta.Type)
|
|
require.NotNil(t, messageDelta.Usage)
|
|
assert.Equal(t, 11, messageDelta.Usage.InputTokens)
|
|
assert.Equal(t, 7, messageDelta.Usage.OutputTokens)
|
|
})
|
|
}
|
|
|
|
func TestConvertStreamResponseKeepsStatelessCompatibility(t *testing.T) {
|
|
t.Run("gemini to openai", func(t *testing.T) {
|
|
info := &convmeta.Values{
|
|
ChannelMetaAttached: true,
|
|
UpstreamModelName: "upstream-model",
|
|
}
|
|
result, err := ConvertStreamResponse(
|
|
nil,
|
|
info,
|
|
types.RelayFormatOpenAI,
|
|
terminalTestGeminiChunk("Hello", "STOP", false),
|
|
)
|
|
require.NoError(t, err)
|
|
require.IsType(t, &dto.ChatCompletionsStreamResponse{}, result.Value)
|
|
|
|
response := result.Value.(*dto.ChatCompletionsStreamResponse)
|
|
require.Len(t, response.Choices, 1)
|
|
require.NotNil(t, response.Choices[0].Delta.Content)
|
|
assert.Equal(t, "Hello", *response.Choices[0].Delta.Content)
|
|
assert.Nil(t, response.Choices[0].FinishReason)
|
|
assert.Equal(t, "upstream-model", response.Model)
|
|
require.NotNil(t, response.Usage)
|
|
assert.Equal(t, 4, response.Usage.PromptTokens)
|
|
assert.Equal(t, 2, response.Usage.CompletionTokens)
|
|
assert.Equal(t, 6, response.Usage.TotalTokens)
|
|
})
|
|
|
|
t.Run("openai to claude", func(t *testing.T) {
|
|
info := &convmeta.Values{
|
|
SendResponseCount: 1,
|
|
ClaudeConvertInfo: &convmeta.ClaudeConvertInfo{
|
|
LastMessagesType: convmeta.LastMessageTypeNone,
|
|
},
|
|
}
|
|
result, err := ConvertStreamResponse(
|
|
nil,
|
|
info,
|
|
types.RelayFormatClaude,
|
|
&dto.ChatCompletionsStreamResponse{
|
|
Id: "chatcmpl-fixed",
|
|
Object: "chat.completion.chunk",
|
|
Created: 1700000000,
|
|
Model: "upstream-model",
|
|
Choices: []dto.ChatCompletionsStreamResponseChoice{
|
|
{
|
|
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{
|
|
Role: "assistant",
|
|
Content: terminalTestPtr("Hello"),
|
|
},
|
|
},
|
|
},
|
|
Usage: &dto.Usage{
|
|
PromptTokens: 4,
|
|
CompletionTokens: 2,
|
|
TotalTokens: 6,
|
|
},
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
require.IsType(t, []*dto.ClaudeResponse{}, result.Value)
|
|
|
|
responses := result.Value.([]*dto.ClaudeResponse)
|
|
require.Len(t, responses, 3)
|
|
assert.Equal(t, "message_start", responses[0].Type)
|
|
assert.Equal(t, "content_block_start", responses[1].Type)
|
|
assert.Equal(t, "content_block_delta", responses[2].Type)
|
|
require.NotNil(t, responses[2].Delta)
|
|
require.NotNil(t, responses[2].Delta.Text)
|
|
assert.Equal(t, "Hello", *responses[2].Delta.Text)
|
|
assert.False(t, info.ClaudeConvertInfo.Done)
|
|
assert.Equal(t, 6, result.Usage.TotalTokens)
|
|
})
|
|
}
|
|
|
|
func terminalTestGeminiChunk(text string, finishReason string, toolCall bool) *dto.GeminiChatResponse {
|
|
response := terminalTestGeminiChunkWithoutUsage(text, finishReason)
|
|
response.HasUsageMetadata = true
|
|
response.UsageMetadata = dto.GeminiUsageMetadata{
|
|
PromptTokenCount: 4,
|
|
CandidatesTokenCount: 2,
|
|
TotalTokenCount: 6,
|
|
}
|
|
if toolCall {
|
|
response.Candidates[0].Content.Parts = []dto.GeminiPart{
|
|
{
|
|
FunctionCall: &dto.FunctionCall{
|
|
FunctionName: "lookup",
|
|
Arguments: map[string]any{"q": "x"},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
return response
|
|
}
|
|
|
|
func terminalTestGeminiChunkWithoutUsage(text string, finishReason string) *dto.GeminiChatResponse {
|
|
candidate := dto.GeminiChatCandidate{
|
|
Content: dto.GeminiChatContent{
|
|
Role: "model",
|
|
Parts: []dto.GeminiPart{{Text: text}},
|
|
},
|
|
}
|
|
if finishReason != "" {
|
|
candidate.FinishReason = terminalTestPtr(finishReason)
|
|
}
|
|
return &dto.GeminiChatResponse{
|
|
Candidates: []dto.GeminiChatCandidate{candidate},
|
|
}
|
|
}
|
|
|
|
func terminalTestGeminiToolChunkWithoutUsage() *dto.GeminiChatResponse {
|
|
return &dto.GeminiChatResponse{
|
|
Candidates: []dto.GeminiChatCandidate{
|
|
{
|
|
Content: dto.GeminiChatContent{
|
|
Role: "model",
|
|
Parts: []dto.GeminiPart{
|
|
{
|
|
FunctionCall: &dto.FunctionCall{
|
|
FunctionName: "lookup",
|
|
Arguments: map[string]any{"q": "x"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func terminalTestGeminiUsageOnlyChunk() *dto.GeminiChatResponse {
|
|
return &dto.GeminiChatResponse{
|
|
HasUsageMetadata: true,
|
|
UsageMetadata: dto.GeminiUsageMetadata{
|
|
PromptTokenCount: 4,
|
|
CandidatesTokenCount: 2,
|
|
TotalTokenCount: 6,
|
|
},
|
|
}
|
|
}
|
|
|
|
func terminalTestFinishedChatChunks(t *testing.T, results []ResponseResult) []*dto.ChatCompletionsStreamResponse {
|
|
t.Helper()
|
|
finished := make([]*dto.ChatCompletionsStreamResponse, 0, 1)
|
|
for _, result := range results {
|
|
response, ok := result.Value.(*dto.ChatCompletionsStreamResponse)
|
|
require.True(t, ok, "unexpected stream result type %T", result.Value)
|
|
if response.IsFinished() {
|
|
finished = append(finished, response)
|
|
}
|
|
}
|
|
return finished
|
|
}
|
|
|
|
func terminalTestAssertClaudeTail(t *testing.T, results []ResponseResult, wantStopReason string) {
|
|
t.Helper()
|
|
responses := make([]*dto.ClaudeResponse, 0, len(results))
|
|
eventCounts := make(map[string]int)
|
|
for _, result := range results {
|
|
response, ok := result.Value.(*dto.ClaudeResponse)
|
|
require.True(t, ok, "unexpected stream result type %T", result.Value)
|
|
responses = append(responses, response)
|
|
eventCounts[response.Type]++
|
|
}
|
|
|
|
require.GreaterOrEqual(t, len(responses), 4)
|
|
assert.Equal(t, "message_start", responses[0].Type)
|
|
tail := responses[len(responses)-3:]
|
|
assert.Equal(t, "content_block_stop", tail[0].Type)
|
|
require.NotNil(t, tail[0].Index)
|
|
assert.Equal(t, 0, *tail[0].Index)
|
|
assert.Equal(t, "message_delta", tail[1].Type)
|
|
require.NotNil(t, tail[1].Delta)
|
|
require.NotNil(t, tail[1].Delta.StopReason)
|
|
assert.Equal(t, wantStopReason, *tail[1].Delta.StopReason)
|
|
require.NotNil(t, tail[1].Usage)
|
|
assert.Equal(t, 4, tail[1].Usage.InputTokens)
|
|
assert.Equal(t, 2, tail[1].Usage.OutputTokens)
|
|
assert.Equal(t, "message_stop", tail[2].Type)
|
|
assert.Equal(t, 1, eventCounts["content_block_stop"])
|
|
assert.Equal(t, 1, eventCounts["message_delta"])
|
|
assert.Equal(t, 1, eventCounts["message_stop"])
|
|
}
|
|
|
|
func terminalTestPtr[T any](value T) *T {
|
|
return &value
|
|
}
|