feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -13,7 +13,9 @@ import (
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/tokenkit"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -494,8 +496,6 @@ func TestGeminiChatHandlerMissingUsageMetadataBuildsEstimatedBillingUsage(t *tes
|
||||
|
||||
func TestGeminiStreamHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
oldStreamingTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 300
|
||||
@@ -503,53 +503,67 @@ func TestGeminiStreamHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *
|
||||
constant.StreamingTimeout = oldStreamingTimeout
|
||||
})
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gemini-3-flash-preview",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: "gemini-3-flash-preview",
|
||||
},
|
||||
}
|
||||
info.SetEstimatePromptTokens(20)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
part dto.GeminiPart
|
||||
usageText string
|
||||
}{
|
||||
{"text", dto.GeminiPart{Text: "partial streamed answer before disconnect"}, "partial streamed answer before disconnect"},
|
||||
// An agent turn that only calls a tool still produced output tokens.
|
||||
{"function call", dto.GeminiPart{FunctionCall: &dto.FunctionCall{FunctionName: "read_file", Arguments: map[string]any{"path": "main.go"}}}, `read_file{"path":"main.go"}`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
// Simulates a client aborting the stream before the final chunk: text was
|
||||
// streamed but the last observed usageMetadata only carries prompt tokens.
|
||||
chunk := dto.GeminiChatResponse{
|
||||
Candidates: []dto.GeminiChatCandidate{
|
||||
{
|
||||
Content: dto.GeminiChatContent{
|
||||
Role: "model",
|
||||
Parts: []dto.GeminiPart{
|
||||
{Text: "partial streamed answer before disconnect"},
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gemini-3-flash-preview",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: "gemini-3-flash-preview",
|
||||
},
|
||||
}
|
||||
info.SetEstimatePromptTokens(20)
|
||||
|
||||
// Simulates a client aborting the stream before the final chunk: output was
|
||||
// streamed but the last observed usageMetadata only carries prompt tokens.
|
||||
chunk := dto.GeminiChatResponse{
|
||||
Candidates: []dto.GeminiChatCandidate{
|
||||
{
|
||||
Content: dto.GeminiChatContent{
|
||||
Role: "model",
|
||||
Parts: []dto.GeminiPart{tc.part},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
UsageMetadata: dto.GeminiUsageMetadata{
|
||||
PromptTokenCount: 151,
|
||||
TotalTokenCount: 151,
|
||||
},
|
||||
UsageMetadata: dto.GeminiUsageMetadata{
|
||||
PromptTokenCount: 151,
|
||||
TotalTokenCount: 151,
|
||||
},
|
||||
}
|
||||
|
||||
chunkData, err := common.Marshal(chunk)
|
||||
require.NoError(t, err)
|
||||
|
||||
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
|
||||
resp := &http.Response{
|
||||
Body: io.NopCloser(bytes.NewReader(streamBody)),
|
||||
}
|
||||
|
||||
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
|
||||
return true
|
||||
})
|
||||
require.Nil(t, newAPIError)
|
||||
require.NotNil(t, usage)
|
||||
require.Equal(t, 151, usage.PromptTokens)
|
||||
require.Equal(t, tokenkit.Estimate("gemini-3-flash-preview", tc.usageText), usage.CompletionTokens)
|
||||
require.Greater(t, usage.CompletionTokens, 0)
|
||||
require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens)
|
||||
require.NotNil(t, usage.BillingUsage)
|
||||
require.True(t, usage.BillingUsage.Estimated)
|
||||
require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata)
|
||||
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
|
||||
})
|
||||
}
|
||||
|
||||
chunkData, err := common.Marshal(chunk)
|
||||
require.NoError(t, err)
|
||||
|
||||
streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n")
|
||||
resp := &http.Response{
|
||||
Body: io.NopCloser(bytes.NewReader(streamBody)),
|
||||
}
|
||||
|
||||
usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool {
|
||||
return true
|
||||
})
|
||||
require.Nil(t, newAPIError)
|
||||
require.NotNil(t, usage)
|
||||
require.Equal(t, 151, usage.PromptTokens)
|
||||
require.Greater(t, usage.CompletionTokens, 0)
|
||||
require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens)
|
||||
require.NotNil(t, usage.BillingUsage)
|
||||
require.True(t, usage.BillingUsage.Estimated)
|
||||
require.NotNil(t, usage.BillingUsage.GeminiUsageMetadata)
|
||||
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
|
||||
}
|
||||
|
||||
func TestGeminiChatHandlerPromptOnlyUsageMetadataEstimatesCompletionTokens(t *testing.T) {
|
||||
@@ -639,3 +653,104 @@ func TestGeminiStreamHandlerEmptyUsageMetadataBuildsEstimatedBillingUsage(t *tes
|
||||
require.Equal(t, usage.CompletionTokens, usage.BillingUsage.GeminiUsageMetadata.CandidatesTokenCount)
|
||||
require.True(t, common.GetContextKeyBool(c, constant.ContextKeyLocalCountTokens))
|
||||
}
|
||||
|
||||
// Google Search grounding records the unique trimmed non-empty queries
|
||||
// (webSearchQueries and imageSearchQueries of every candidate, across every
|
||||
// stream frame) under google_search and one grounded prompt when any query
|
||||
// exists; the tool price table decides which of the two a model pays for.
|
||||
func TestGeminiHandlersCountGroundingQueries(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
oldStreamingTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 300
|
||||
t.Cleanup(func() {
|
||||
constant.StreamingTimeout = oldStreamingTimeout
|
||||
})
|
||||
const text = `"content":{"role":"model","parts":[{"text":"ok"}]}`
|
||||
grounded := func(queries int) map[string]int {
|
||||
return map[string]int{dto.BuildInToolGoogleSearch: queries, relaycommon.GoogleSearchGroundedPromptTool: 1}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
stream bool
|
||||
frames []string
|
||||
want map[string]int
|
||||
}{
|
||||
{
|
||||
name: "native non-stream",
|
||||
frames: []string{`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["Shanghai weather today","current weather in Shanghai"],"groundingChunks":[{"web":{"uri":"https://a.example"}}]}}]}`},
|
||||
want: grounded(2),
|
||||
},
|
||||
{
|
||||
name: "duplicates, whitespace and empty queries",
|
||||
frames: []string{`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["go release",""," go release "," ","tokyo"]}}]}`},
|
||||
want: grounded(2),
|
||||
},
|
||||
{
|
||||
name: "union across candidates",
|
||||
frames: []string{`{"candidates":[{"index":0,` + text + `,"groundingMetadata":{"webSearchQueries":["q1","q2"]}},{"index":1,` + text + `,"groundingMetadata":{"webSearchQueries":["q2","q3"]}}]}`},
|
||||
want: grounded(3),
|
||||
},
|
||||
{
|
||||
name: "image queries count like web queries",
|
||||
frames: []string{`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["cat breeds","orange cat"],"imageSearchQueries":["orange tabby cat photo","orange cat"]}}]}`},
|
||||
want: grounded(3),
|
||||
},
|
||||
{
|
||||
name: "only empty queries",
|
||||
frames: []string{`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":[""," "],"imageSearchQueries":[""]}}]}`},
|
||||
want: map[string]int{},
|
||||
},
|
||||
{
|
||||
name: "no grounding metadata",
|
||||
frames: []string{`{"candidates":[{` + text + `,"finishReason":"STOP"}]}`},
|
||||
want: map[string]int{},
|
||||
},
|
||||
{
|
||||
name: "stream union across frames",
|
||||
stream: true,
|
||||
frames: []string{
|
||||
`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["q1"]}}]}`,
|
||||
`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["q2","q1 "]}}]}`,
|
||||
`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["q3"]}}]}`,
|
||||
},
|
||||
want: grounded(3),
|
||||
},
|
||||
{
|
||||
name: "stream repeats are not added up",
|
||||
stream: true,
|
||||
frames: []string{
|
||||
`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["q1","q2"]}}]}`,
|
||||
`{"candidates":[{` + text + `,"groundingMetadata":{"webSearchQueries":["q1","q2"]}}]}`,
|
||||
},
|
||||
want: grounded(2),
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.7-flash:generateContent", nil)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gemini-3.7-flash",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.7-flash"},
|
||||
}
|
||||
var apiErr *types.NewAPIError
|
||||
if tt.stream {
|
||||
body := "data: " + strings.Join(tt.frames, "\ndata: ") + "\ndata: [DONE]\n"
|
||||
_, apiErr = geminiStreamHandler(c, info, &http.Response{Body: io.NopCloser(strings.NewReader(body))}, func(string, *dto.GeminiChatResponse) bool {
|
||||
return true
|
||||
})
|
||||
} else {
|
||||
_, apiErr = GeminiTextGenerationHandler(c, info, &http.Response{Body: io.NopCloser(strings.NewReader(tt.frames[0]))})
|
||||
}
|
||||
require.Nil(t, apiErr)
|
||||
counts := map[string]int{}
|
||||
if info.ResponsesUsageInfo != nil {
|
||||
for name, tool := range info.ResponsesUsageInfo.BuiltInTools {
|
||||
counts[name] = tool.CallCount
|
||||
}
|
||||
}
|
||||
assert.Equal(t, tt.want, counts)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user