重构前

2026-08-23 21:35:45 +08:00
parent ada6383e49
commit dbb3e316dd
2004 changed files with 440326 additions and 0 deletions
@@ -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"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"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)
+}