mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-14 00:01:53 +00:00
feat: configurable tool pricing, Sub2API channel, and alpha search billing
Add admin-configurable tool-call prices with cross-provider surcharge settlement, Sub2API channel support, /v1/alpha/search relay, and usage-log surcharge UI.
This commit is contained in:
@@ -114,6 +114,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
data = patchClaudeMessageDeltaUsageData(data, buildMessageDeltaPatchUsage(&claudeResponse, claudeInfo))
|
||||
}
|
||||
}
|
||||
countClaudeStreamBillableTools(c, info, &claudeResponse)
|
||||
helper.ClaudeChunkData(c, claudeResponse, data)
|
||||
} else if info.RelayFormat == types.RelayFormatOpenAI {
|
||||
response := StreamResponseClaude2OpenAI(&claudeResponse)
|
||||
@@ -122,6 +123,8 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
return nil
|
||||
}
|
||||
|
||||
countClaudeStreamBillableTools(c, info, &claudeResponse)
|
||||
|
||||
err = helper.ObjectData(c, response)
|
||||
if err != nil {
|
||||
logger.LogError(c, "send_stream_response_failed: "+err.Error())
|
||||
@@ -130,6 +133,23 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
return nil
|
||||
}
|
||||
|
||||
func countClaudeStreamBillableTools(c *gin.Context, info *relaycommon.RelayInfo, claudeResponse *dto.ClaudeResponse) {
|
||||
if claudeResponse == nil {
|
||||
return
|
||||
}
|
||||
if claudeResponse.Type == "content_block_start" &&
|
||||
claudeResponse.ContentBlock != nil &&
|
||||
claudeResponse.ContentBlock.Type == "tool_use" {
|
||||
info.CountBillableToolCall(dto.BuildInCallToolUse, claudeResponse.ContentBlock.Name)
|
||||
}
|
||||
if claudeResponse.Type == "message_delta" &&
|
||||
claudeResponse.Usage != nil &&
|
||||
claudeResponse.Usage.ServerToolUse != nil &&
|
||||
claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 {
|
||||
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
|
||||
}
|
||||
}
|
||||
|
||||
func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) {
|
||||
if claudeInfo.Usage.PromptTokens == 0 {
|
||||
//上游出错
|
||||
@@ -235,6 +255,12 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
|
||||
}
|
||||
|
||||
for _, block := range claudeResponse.Content {
|
||||
if block.Type == "tool_use" {
|
||||
info.CountBillableToolCall(dto.BuildInCallToolUse, block.Name)
|
||||
}
|
||||
}
|
||||
|
||||
service.IOCopyBytesGracefully(c, httpResp, responseData)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestHandleClaudeResponseDataCountsToolUse(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
operation_setting.SetToolPriceForTest("lookup_fn", 3.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("lookup_fn")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "claude-3-7-sonnet",
|
||||
RelayFormat: types.RelayFormatClaude,
|
||||
}
|
||||
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}}
|
||||
|
||||
data := []byte(`{
|
||||
"type":"message",
|
||||
"content":[
|
||||
{"type":"text","text":"hi"},
|
||||
{"type":"tool_use","id":"tu1","name":"lookup_fn","input":{}},
|
||||
{"type":"server_tool_use","id":"stu1","name":"web_search","input":{}}
|
||||
],
|
||||
"usage":{"input_tokens":1,"output_tokens":1}
|
||||
}`)
|
||||
|
||||
err := HandleClaudeResponseData(c, info, claudeInfo, nil, data)
|
||||
require.Nil(t, err)
|
||||
require.NotNil(t, info.ResponsesUsageInfo)
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "lookup_fn")
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["lookup_fn"].CallCount)
|
||||
assert.NotContains(t, info.ResponsesUsageInfo.BuiltInTools, "web_search")
|
||||
}
|
||||
|
||||
func TestCountClaudeStreamBillableToolsSetsWebSearchRequests(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "claude-3-7-sonnet"}
|
||||
|
||||
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
|
||||
Type: "message_delta",
|
||||
Usage: &dto.ClaudeUsage{
|
||||
ServerToolUse: &dto.ClaudeServerToolUse{WebSearchRequests: 3},
|
||||
},
|
||||
})
|
||||
assert.Equal(t, 3, c.GetInt("claude_web_search_requests"))
|
||||
|
||||
operation_setting.SetToolPriceForTest("stream_fn", 2.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("stream_fn")
|
||||
})
|
||||
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &dto.ClaudeMediaMessage{
|
||||
Type: "tool_use",
|
||||
Name: "stream_fn",
|
||||
},
|
||||
})
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "stream_fn")
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["stream_fn"].CallCount)
|
||||
}
|
||||
@@ -112,18 +112,20 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
if info.RelayMode != relayconstant.RelayModeResponses && info.RelayMode != relayconstant.RelayModeResponsesCompact {
|
||||
switch info.RelayMode {
|
||||
case relayconstant.RelayModeAlphaSearch:
|
||||
// Alpha search responses are handled by relay.AlphaSearchHelper.
|
||||
return nil, types.NewError(errors.New("codex channel: alpha search response should be handled by AlphaSearchHelper"), types.ErrorCodeInvalidRequest)
|
||||
case relayconstant.RelayModeResponsesCompact:
|
||||
return openai.OaiResponsesCompactionHandler(c, resp)
|
||||
case relayconstant.RelayModeResponses:
|
||||
if info.IsStream {
|
||||
return openai.OaiResponsesStreamHandler(c, info, resp)
|
||||
}
|
||||
return openai.OaiResponsesHandler(c, info, resp)
|
||||
default:
|
||||
return nil, types.NewError(errors.New("codex channel: endpoint not supported"), types.ErrorCodeInvalidRequest)
|
||||
}
|
||||
|
||||
if info.RelayMode == relayconstant.RelayModeResponsesCompact {
|
||||
return openai.OaiResponsesCompactionHandler(c, resp)
|
||||
}
|
||||
|
||||
if info.IsStream {
|
||||
return openai.OaiResponsesStreamHandler(c, info, resp)
|
||||
}
|
||||
return openai.OaiResponsesHandler(c, info, resp)
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
@@ -135,12 +137,16 @@ func (a *Adaptor) GetChannelName() string {
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
if info.RelayMode != relayconstant.RelayModeResponses && info.RelayMode != relayconstant.RelayModeResponsesCompact {
|
||||
return "", errors.New("codex channel: only /v1/responses and /v1/responses/compact are supported")
|
||||
}
|
||||
path := "/backend-api/codex/responses"
|
||||
if info.RelayMode == relayconstant.RelayModeResponsesCompact {
|
||||
var path string
|
||||
switch info.RelayMode {
|
||||
case relayconstant.RelayModeResponses:
|
||||
path = "/backend-api/codex/responses"
|
||||
case relayconstant.RelayModeResponsesCompact:
|
||||
path = "/backend-api/codex/responses/compact"
|
||||
case relayconstant.RelayModeAlphaSearch:
|
||||
path = "/backend-api/codex/alpha/search"
|
||||
default:
|
||||
return "", errors.New("codex channel: only /v1/responses, /v1/responses/compact and /v1/alpha/search are supported")
|
||||
}
|
||||
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, path, info.ChannelType), nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
package codex
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetRequestURLAlphaSearch(t *testing.T) {
|
||||
adaptor := &Adaptor{}
|
||||
info := &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeCodex,
|
||||
ChannelBaseUrl: "https://chatgpt.com",
|
||||
},
|
||||
RelayMode: relayconstant.RelayModeAlphaSearch,
|
||||
}
|
||||
|
||||
url, err := adaptor.GetRequestURL(info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "https://chatgpt.com/backend-api/codex/alpha/search", url)
|
||||
}
|
||||
@@ -76,6 +76,18 @@ func geminiResponseUsageText(response *dto.GeminiChatResponse) string {
|
||||
return text.String()
|
||||
}
|
||||
|
||||
func markGeminiGoogleSearchCall(c *gin.Context, response *dto.GeminiChatResponse) {
|
||||
if c == nil || response == nil {
|
||||
return
|
||||
}
|
||||
for _, candidate := range response.Candidates {
|
||||
if candidate.GroundingMetadata != nil && len(candidate.GroundingMetadata.WebSearchQueries) > 0 {
|
||||
c.Set("gemini_google_search_call", true)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func buildUsageFromGeminiResponse(c *gin.Context, info *relaycommon.RelayInfo, response *dto.GeminiChatResponse) dto.Usage {
|
||||
metadata := response.GetUsageMetadata()
|
||||
if dto.HasGeminiUsageMetadataTokens(metadata) {
|
||||
@@ -149,6 +161,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
}
|
||||
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
|
||||
// 统计图片数量
|
||||
for _, candidate := range geminiResponse.Candidates {
|
||||
for _, part := range candidate.Content.Parts {
|
||||
@@ -308,6 +322,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
|
||||
|
||||
|
||||
@@ -31,6 +31,7 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h
|
||||
if err := common.Unmarshal(responseBody, &geminiResponse); err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
|
||||
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
|
||||
@@ -119,6 +119,8 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
|
||||
var usage = &dto.Usage{}
|
||||
var lastStreamData string
|
||||
var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型
|
||||
seenStreamToolCalls := make(map[string]struct{})
|
||||
var streamFunctionCallNames []string
|
||||
|
||||
// 检查是否为音频模型
|
||||
isAudioModel := strings.Contains(strings.ToLower(model), "audio")
|
||||
@@ -137,6 +139,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
|
||||
}
|
||||
|
||||
lastStreamData = data
|
||||
collectStreamFunctionCallNames(data, seenStreamToolCalls, &streamFunctionCallNames)
|
||||
if err := processTokenData(info.RelayMode, data, &responseTextBuilder, &toolCount); err != nil {
|
||||
logger.LogError(c, "error processing stream token data: "+err.Error())
|
||||
sr.Error(err)
|
||||
@@ -182,11 +185,40 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
|
||||
|
||||
applyUsagePostProcessing(info, usage, common.StringToByteSlice(lastStreamData))
|
||||
|
||||
for _, name := range streamFunctionCallNames {
|
||||
info.CountBillableToolCall(dto.BuildInCallFunctionCall, name)
|
||||
}
|
||||
|
||||
HandleFinalResponse(c, info, lastStreamData, responseId, createAt, model, systemFingerprint, usage, containStreamUsage)
|
||||
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
func collectStreamFunctionCallNames(data string, seen map[string]struct{}, names *[]string) {
|
||||
var streamResponse dto.ChatCompletionsStreamResponse
|
||||
if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil {
|
||||
return
|
||||
}
|
||||
for _, choice := range streamResponse.Choices {
|
||||
for i, tc := range choice.Delta.ToolCalls {
|
||||
name := tc.Function.Name
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
toolIdx := i
|
||||
if tc.Index != nil {
|
||||
toolIdx = *tc.Index
|
||||
}
|
||||
key := fmt.Sprintf("%d-%d", choice.Index, toolIdx)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
*names = append(*names, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
|
||||
@@ -228,6 +260,12 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
|
||||
}
|
||||
}
|
||||
|
||||
for _, choice := range simpleResponse.Choices {
|
||||
for _, tc := range choice.Message.ParseToolCalls() {
|
||||
info.CountBillableToolCall(dto.BuildInCallFunctionCall, tc.Function.Name)
|
||||
}
|
||||
}
|
||||
|
||||
forceFormat := false
|
||||
if info.ChannelSetting.ForceFormat {
|
||||
forceFormat = true
|
||||
|
||||
@@ -34,12 +34,6 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
if responsesResponse.HasImageGenerationCall() {
|
||||
c.Set("image_generation_call", true)
|
||||
c.Set("image_generation_call_quality", responsesResponse.GetQuality())
|
||||
c.Set("image_generation_call_size", responsesResponse.GetSize())
|
||||
}
|
||||
|
||||
// 写入新的 response body
|
||||
service.IOCopyBytesGracefully(c, resp, responseBody)
|
||||
|
||||
@@ -54,18 +48,27 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
usage.PromptTokensDetails.CacheWriteTokens = responsesResponse.Usage.InputTokensDetails.CacheWriteTokens
|
||||
}
|
||||
}
|
||||
if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil {
|
||||
return &usage, nil
|
||||
}
|
||||
// 解析 Tools 用量
|
||||
for _, tool := range responsesResponse.Tools {
|
||||
buildToolinfo, ok := info.ResponsesUsageInfo.BuiltInTools[common.Interface2String(tool["type"])]
|
||||
if !ok || buildToolinfo == nil {
|
||||
logger.LogError(c, fmt.Sprintf("BuiltInTools not found for tool type: %v", tool["type"]))
|
||||
continue
|
||||
// Count actual tool invocations from Output (not tool declarations).
|
||||
for _, output := range responsesResponse.Output {
|
||||
switch output.Type {
|
||||
case dto.BuildInCallWebSearchCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallWebSearchCall, "")
|
||||
case dto.BuildInCallFileSearchCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFileSearchCall, "")
|
||||
case dto.BuildInCallFunctionCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFunctionCall, output.Name)
|
||||
}
|
||||
buildToolinfo.CallCount++
|
||||
}
|
||||
|
||||
imageCounter := &relaycommon.ImageGenerationCallCounter{}
|
||||
if !relaycommon.IsNonBillableResponsesStatus(responsesResponse.Status) {
|
||||
for i := range responsesResponse.Output {
|
||||
idx := i
|
||||
imageCounter.Observe(&responsesResponse.Output[i], &idx)
|
||||
}
|
||||
}
|
||||
imageCounter.Commit(info)
|
||||
|
||||
return &usage, nil
|
||||
}
|
||||
|
||||
@@ -79,6 +82,8 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
|
||||
var usage = &dto.Usage{}
|
||||
var responseTextBuilder strings.Builder
|
||||
imageCounter := &relaycommon.ImageGenerationCallCounter{}
|
||||
imageCommitted := false
|
||||
|
||||
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
|
||||
|
||||
@@ -91,7 +96,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
}
|
||||
sendResponsesStreamData(c, streamResponse, data)
|
||||
switch streamResponse.Type {
|
||||
case "response.completed":
|
||||
case "response.completed", "response.done":
|
||||
if streamResponse.Response != nil {
|
||||
if streamResponse.Response.Usage != nil {
|
||||
if streamResponse.Response.Usage.InputTokens != 0 {
|
||||
@@ -108,24 +113,45 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
usage.PromptTokensDetails.CacheWriteTokens = streamResponse.Response.Usage.InputTokensDetails.CacheWriteTokens
|
||||
}
|
||||
}
|
||||
if streamResponse.Response.HasImageGenerationCall() {
|
||||
c.Set("image_generation_call", true)
|
||||
c.Set("image_generation_call_quality", streamResponse.Response.GetQuality())
|
||||
c.Set("image_generation_call_size", streamResponse.Response.GetSize())
|
||||
if !imageCommitted {
|
||||
if relaycommon.IsNonBillableResponsesStatus(streamResponse.Response.Status) {
|
||||
imageCounter.Reset()
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
} else {
|
||||
for i := range streamResponse.Response.Output {
|
||||
idx := i
|
||||
imageCounter.Observe(&streamResponse.Response.Output[i], &idx)
|
||||
}
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
}
|
||||
} else if !imageCommitted {
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
case "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
|
||||
if !imageCommitted {
|
||||
imageCounter.Reset()
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
case "response.output_text.delta":
|
||||
// 处理输出文本
|
||||
responseTextBuilder.WriteString(streamResponse.Delta)
|
||||
case dto.ResponsesOutputTypeItemDone:
|
||||
// 函数调用处理
|
||||
if streamResponse.Item != nil {
|
||||
switch streamResponse.Item.Type {
|
||||
case dto.BuildInCallWebSearchCall:
|
||||
if info != nil && info.ResponsesUsageInfo != nil && info.ResponsesUsageInfo.BuiltInTools != nil {
|
||||
if webSearchTool, exists := info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool != nil {
|
||||
webSearchTool.CallCount++
|
||||
}
|
||||
info.CountBillableToolCall(dto.BuildInCallWebSearchCall, "")
|
||||
case dto.BuildInCallFileSearchCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFileSearchCall, "")
|
||||
case dto.BuildInCallFunctionCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFunctionCall, streamResponse.Item.Name)
|
||||
case dto.ResponsesOutputTypeImageGenerationCall:
|
||||
if !imageCommitted {
|
||||
imageCounter.Observe(streamResponse.Item, streamResponse.OutputIndex)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOaiResponsesHandlerCountsOutputCallsNotDeclarations(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
operation_setting.SetToolPriceForTest("priced_fn", 5.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("priced_fn")
|
||||
})
|
||||
|
||||
body, err := common.Marshal(dto.OpenAIResponsesResponse{
|
||||
Tools: []map[string]any{
|
||||
{"type": "web_search_preview"},
|
||||
{"type": "file_search"},
|
||||
},
|
||||
Output: []dto.ResponsesOutput{
|
||||
{Type: dto.BuildInCallWebSearchCall},
|
||||
{Type: dto.BuildInCallWebSearchCall},
|
||||
{Type: dto.BuildInCallFunctionCall, Name: "priced_fn"},
|
||||
{Type: dto.BuildInCallFunctionCall, Name: "unpriced_fn"},
|
||||
},
|
||||
Usage: &dto.Usage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gpt-5.1",
|
||||
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
||||
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
||||
dto.BuildInToolWebSearchPreview: {ToolName: dto.BuildInToolWebSearchPreview, CallCount: 0},
|
||||
dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch, CallCount: 0},
|
||||
},
|
||||
},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
}
|
||||
|
||||
usage, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 2, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview].CallCount)
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount)
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "priced_fn")
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["priced_fn"].CallCount)
|
||||
assert.NotContains(t, info.ResponsesUsageInfo.BuiltInTools, "unpriced_fn")
|
||||
}
|
||||
|
||||
func TestOaiResponsesHandlerDeclaredToolsWithoutOutputCountZero(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
body, err := common.Marshal(dto.OpenAIResponsesResponse{
|
||||
Tools: []map[string]any{
|
||||
{"type": "web_search_preview"},
|
||||
{"type": "file_search"},
|
||||
},
|
||||
Output: []dto.ResponsesOutput{
|
||||
{Type: "message", Role: "assistant"},
|
||||
},
|
||||
Usage: &dto.Usage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gpt-5.1",
|
||||
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
||||
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
||||
dto.BuildInToolWebSearchPreview: {ToolName: dto.BuildInToolWebSearchPreview, CallCount: 0},
|
||||
dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch, CallCount: 0},
|
||||
},
|
||||
},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview].CallCount)
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount)
|
||||
}
|
||||
|
||||
func TestOaiResponsesHandlerCountsCompletedImageGenerationOutputs(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
body, err := common.Marshal(dto.OpenAIResponsesResponse{
|
||||
Status: []byte(`"completed"`),
|
||||
Output: []dto.ResponsesOutput{
|
||||
{
|
||||
Type: dto.ResponsesOutputTypeImageGenerationCall,
|
||||
ID: "img_1",
|
||||
Status: "completed",
|
||||
Result: "base64-a",
|
||||
},
|
||||
{
|
||||
Type: dto.ResponsesOutputTypeImageGenerationCall,
|
||||
ID: "img_2",
|
||||
Status: "completed",
|
||||
Result: "base64-b",
|
||||
},
|
||||
{
|
||||
Type: dto.ResponsesOutputTypeImageGenerationCall,
|
||||
ID: "img_empty",
|
||||
Status: "completed",
|
||||
Result: "",
|
||||
},
|
||||
},
|
||||
Usage: &dto.Usage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "gpt-5.1"}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, dto.BuildInToolImageGeneration)
|
||||
assert.Equal(t, 2, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
assert.False(t, c.GetBool("image_generation_call"))
|
||||
}
|
||||
|
||||
func TestOaiResponsesHandlerIncompleteStatusCommitsZeroImageGeneration(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
body, err := common.Marshal(dto.OpenAIResponsesResponse{
|
||||
Status: []byte(`"incomplete"`),
|
||||
Output: []dto.ResponsesOutput{
|
||||
{
|
||||
Type: dto.ResponsesOutputTypeImageGenerationCall,
|
||||
ID: "img_1",
|
||||
Status: "completed",
|
||||
Result: "base64-a",
|
||||
},
|
||||
},
|
||||
Usage: &dto.Usage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gpt-5.1",
|
||||
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
||||
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
||||
dto.BuildInToolImageGeneration: {ToolName: dto.BuildInToolImageGeneration, CallCount: 0},
|
||||
},
|
||||
},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
}
|
||||
|
||||
func runResponsesImageBillingStream(t *testing.T, events ...string) *relaycommon.RelayInfo {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
oldTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() {
|
||||
constant.StreamingTimeout = oldTimeout
|
||||
})
|
||||
|
||||
var body strings.Builder
|
||||
for _, event := range events {
|
||||
body.WriteString("data: ")
|
||||
body.WriteString(event)
|
||||
body.WriteString("\n\n")
|
||||
}
|
||||
body.WriteString("data: [DONE]\n\n")
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "responses-image-billing-test")
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gpt-5.1",
|
||||
DisablePing: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: "gpt-5.1",
|
||||
},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body.String())),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesStreamHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, info.ResponsesUsageInfo)
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, dto.BuildInToolImageGeneration)
|
||||
return info
|
||||
}
|
||||
|
||||
func TestOaiResponsesStreamHandlerDeduplicatesCompletedImageOutput(t *testing.T) {
|
||||
item := `{"type":"image_generation_call","id":"img_1","call_id":"call_1","status":"completed","result":"base64-a"}`
|
||||
info := runResponsesImageBillingStream(
|
||||
t,
|
||||
`{"type":"response.output_item.done","output_index":0,"item":`+item+`}`,
|
||||
`{"type":"response.completed","response":{"status":"completed","output":[`+item+`],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`,
|
||||
)
|
||||
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
}
|
||||
|
||||
func TestOaiResponsesStreamHandlerDiscardsImageOutputOnIncomplete(t *testing.T) {
|
||||
info := runResponsesImageBillingStream(
|
||||
t,
|
||||
`{"type":"response.output_item.done","output_index":0,"item":{"type":"image_generation_call","id":"img_1","status":"completed","result":"base64-a"}}`,
|
||||
`{"type":"response.incomplete","response":{"status":"incomplete"}}`,
|
||||
)
|
||||
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
}
|
||||
|
||||
func TestOaiResponsesStreamHandlerDoesNotCountPartialImageEvent(t *testing.T) {
|
||||
info := runResponsesImageBillingStream(
|
||||
t,
|
||||
`{"type":"response.image_generation_call.partial_image","output_index":0,"partial_image_b64":"partial-bytes"}`,
|
||||
`{"type":"response.completed","response":{"status":"completed","output":[]}}`,
|
||||
)
|
||||
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCollectStreamFunctionCallNamesDedupesSameIndex(t *testing.T) {
|
||||
seen := make(map[string]struct{})
|
||||
var names []string
|
||||
|
||||
chunks := []string{
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"get_weather","arguments":""}}]}}]}`,
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"q\":"}}]}}]}`,
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"x\"}"}}]}}]}`,
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"c2","type":"function","function":{"name":"get_time","arguments":""}}]}}]}`,
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{}"}}]}}]}`,
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
collectStreamFunctionCallNames(chunk, seen, &names)
|
||||
}
|
||||
|
||||
require.Len(t, names, 2)
|
||||
assert.Equal(t, []string{"get_weather", "get_time"}, names)
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package sub2api
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
"github.com/QuantumNous/new-api/relay/channel/gemini"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Adaptor struct {
|
||||
openaiAdaptor openai.Adaptor
|
||||
claudeAdaptor claude.Adaptor
|
||||
geminiAdaptor gemini.Adaptor
|
||||
}
|
||||
|
||||
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
|
||||
a.openaiAdaptor.Init(info)
|
||||
a.claudeAdaptor.Init(info)
|
||||
a.geminiAdaptor.Init(info)
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
if info.RelayMode == relayconstant.RelayModeAlphaSearch {
|
||||
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, "/v1/alpha/search", info.ChannelType), nil
|
||||
}
|
||||
return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, info.RequestURLPath, info.ChannelType), nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
|
||||
channel.SetupApiRequestHeader(info, c, req)
|
||||
req.Set("Authorization", "Bearer "+info.ApiKey)
|
||||
|
||||
switch info.RelayFormat {
|
||||
case types.RelayFormatClaude:
|
||||
req.Set("x-api-key", info.ApiKey)
|
||||
if req.Get("anthropic-version") == "" {
|
||||
anthropicVersion := c.Request.Header.Get("anthropic-version")
|
||||
if anthropicVersion == "" {
|
||||
anthropicVersion = "2023-06-01"
|
||||
}
|
||||
req.Set("anthropic-version", anthropicVersion)
|
||||
}
|
||||
case types.RelayFormatGemini:
|
||||
req.Set("x-goog-api-key", info.ApiKey)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
|
||||
if request == nil {
|
||||
return nil, errors.New("request is nil")
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
|
||||
if request == nil {
|
||||
return nil, errors.New("request is nil")
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
|
||||
if request == nil {
|
||||
return nil, errors.New("request is nil")
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
|
||||
return nil, errors.New("endpoint not supported")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
|
||||
return nil, errors.New("endpoint not supported")
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
|
||||
return channel.DoApiRequest(a, c, info, requestBody)
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
switch info.RelayFormat {
|
||||
case types.RelayFormatClaude:
|
||||
return a.claudeAdaptor.DoResponse(c, resp, info)
|
||||
case types.RelayFormatGemini:
|
||||
return a.geminiAdaptor.DoResponse(c, resp, info)
|
||||
default:
|
||||
return a.openaiAdaptor.DoResponse(c, resp, info)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
return ModelList
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package sub2api
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetRequestURLAlphaSearch(t *testing.T) {
|
||||
adaptor := &Adaptor{}
|
||||
info := &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeSub2API,
|
||||
ChannelBaseUrl: "https://sub2api.example",
|
||||
},
|
||||
RequestURLPath: "/v1/alpha/search",
|
||||
RelayMode: relayconstant.RelayModeAlphaSearch,
|
||||
}
|
||||
|
||||
url, err := adaptor.GetRequestURL(info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "https://sub2api.example/v1/alpha/search", url)
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package sub2api
|
||||
|
||||
const ChannelName = "sub2api"
|
||||
|
||||
// ModelList is empty because models are fetched dynamically from upstream /v1/models.
|
||||
var ModelList = []string{}
|
||||
Reference in New Issue
Block a user