mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-12 15:21:09 +00:00
feat: support Responses to Chat (#5787)
* fix(openai): harden Chat-to-Responses compatibility Add a shared Responses-to-Chat stream state machine and use it from the OpenAI relay path. Preserve assistant text alongside tool calls, bind tool argument deltas by output_index, map incomplete finish reasons, support reasoning/custom tool events, and buffer upstream SSE for non-stream Chat clients. Add deterministic service tests and relay SSE tests for the conversion path. Related to #5745. * refactor: rename openaicompat to relayconvert for improved clarity * feat(gemini): support responses request conversion * feat: add responses to chat conversion support * fix: harden responses chat conversion edge cases
This commit is contained in:
@@ -12,6 +12,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/service/relayconvert"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/QuantumNous/new-api/setting/reasoning"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
@@ -238,8 +239,17 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
|
||||
// TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
request, err := preprocessGeminiOpenAIResponsesRequest(request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
chatRequest, err := relayconvert.ResponsesRequestToChatCompletionsRequest(&request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return a.ConvertOpenAIRequest(c, info, chatRequest)
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
|
||||
@@ -247,6 +257,13 @@ 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 == constant.RelayModeResponses {
|
||||
if info.IsStream {
|
||||
return GeminiResponsesStreamHandler(c, info, resp)
|
||||
}
|
||||
return GeminiResponsesHandler(c, info, resp)
|
||||
}
|
||||
|
||||
if info.RelayMode == constant.RelayModeGemini {
|
||||
if strings.Contains(info.RequestURLPath, ":embedContent") ||
|
||||
strings.Contains(info.RequestURLPath, ":batchEmbedContents") {
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
)
|
||||
|
||||
const (
|
||||
geminiResponsesInputTypeCustomToolCall = "custom_tool_call"
|
||||
geminiResponsesInputTypeCustomToolCallOutput = "custom_tool_call_output"
|
||||
geminiResponsesInputTypeFunctionCallOutput = "function_call_output"
|
||||
)
|
||||
|
||||
func preprocessGeminiOpenAIResponsesRequest(request dto.OpenAIResponsesRequest) (dto.OpenAIResponsesRequest, error) {
|
||||
tools, err := filterGeminiResponsesTools(request.Tools)
|
||||
if err != nil {
|
||||
return request, err
|
||||
}
|
||||
request.Tools = tools
|
||||
|
||||
input, err := filterGeminiResponsesInput(request.Input)
|
||||
if err != nil {
|
||||
return request, err
|
||||
}
|
||||
request.Input = input
|
||||
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func filterGeminiResponsesTools(raw []byte) ([]byte, error) {
|
||||
if !geminiRawJSONPresent(raw) || common.GetJsonType(raw) != "array" {
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
var tools []map[string]any
|
||||
if err := common.Unmarshal(raw, &tools); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filtered := make([]map[string]any, 0, len(tools))
|
||||
for _, tool := range tools {
|
||||
if strings.TrimSpace(common.Interface2String(tool["type"])) != "function" {
|
||||
// TODO: Support Responses custom/freeform tools when Gemini has a safe equivalent representation.
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, tool)
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
return common.Marshal(filtered)
|
||||
}
|
||||
|
||||
func filterGeminiResponsesInput(raw []byte) ([]byte, error) {
|
||||
if !geminiRawJSONPresent(raw) || common.GetJsonType(raw) != "array" {
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
var items []map[string]any
|
||||
if err := common.Unmarshal(raw, &items); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
skippedCustomCallIDs := make(map[string]struct{})
|
||||
for _, item := range items {
|
||||
if strings.TrimSpace(common.Interface2String(item["type"])) != geminiResponsesInputTypeCustomToolCall {
|
||||
continue
|
||||
}
|
||||
if callID := strings.TrimSpace(common.Interface2String(item["call_id"])); callID != "" {
|
||||
skippedCustomCallIDs[callID] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
filtered := make([]map[string]any, 0, len(items))
|
||||
for _, item := range items {
|
||||
itemType := strings.TrimSpace(common.Interface2String(item["type"]))
|
||||
switch itemType {
|
||||
case geminiResponsesInputTypeCustomToolCall, geminiResponsesInputTypeCustomToolCallOutput:
|
||||
// TODO: Support Responses custom/freeform tool calls once Gemini can preserve their semantics.
|
||||
continue
|
||||
case geminiResponsesInputTypeFunctionCallOutput:
|
||||
if _, ok := skippedCustomCallIDs[strings.TrimSpace(common.Interface2String(item["call_id"]))]; ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
|
||||
return common.Marshal(filtered)
|
||||
}
|
||||
|
||||
func geminiRawJSONPresent(raw []byte) bool {
|
||||
if len(raw) == 0 {
|
||||
return false
|
||||
}
|
||||
return common.GetJsonType(raw) != "null"
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestConvertOpenAIResponsesRequestToGeminiInstructionsAndInput(t *testing.T) {
|
||||
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
|
||||
Model: "gemini-test",
|
||||
Instructions: mustGeminiRawMessage(t, "system rules"),
|
||||
Input: mustGeminiRawMessage(t, "hello"),
|
||||
})
|
||||
|
||||
require.NotNil(t, got.SystemInstructions)
|
||||
require.Len(t, got.SystemInstructions.Parts, 1)
|
||||
assert.Equal(t, "system rules", got.SystemInstructions.Parts[0].Text)
|
||||
require.Len(t, got.Contents, 1)
|
||||
assert.Equal(t, "user", got.Contents[0].Role)
|
||||
require.Len(t, got.Contents[0].Parts, 1)
|
||||
assert.Equal(t, "hello", got.Contents[0].Parts[0].Text)
|
||||
}
|
||||
|
||||
func TestConvertOpenAIResponsesRequestToGeminiFunctionToolAndChoice(t *testing.T) {
|
||||
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
|
||||
Model: "gemini-test",
|
||||
Input: mustGeminiRawMessage(t, "lookup weather"),
|
||||
Tools: mustGeminiRawMessage(t, []map[string]any{
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Lookup data",
|
||||
"parameters": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"q": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{"type": "custom", "name": "freeform"},
|
||||
}),
|
||||
ToolChoice: mustGeminiRawMessage(t, map[string]any{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
}),
|
||||
})
|
||||
|
||||
tools := got.GetTools()
|
||||
require.Len(t, tools, 1)
|
||||
assert.Equal(t, "lookup", gjson.GetBytes(got.Tools, "0.functionDeclarations.0.name").String())
|
||||
assert.Equal(t, "Lookup data", gjson.GetBytes(got.Tools, "0.functionDeclarations.0.description").String())
|
||||
require.NotNil(t, got.ToolConfig)
|
||||
require.NotNil(t, got.ToolConfig.FunctionCallingConfig)
|
||||
assert.Equal(t, dto.FunctionCallingConfigMode("ANY"), got.ToolConfig.FunctionCallingConfig.Mode)
|
||||
assert.Equal(t, []string{"lookup"}, got.ToolConfig.FunctionCallingConfig.AllowedFunctionNames)
|
||||
}
|
||||
|
||||
func TestConvertOpenAIResponsesRequestToGeminiFunctionCallConversation(t *testing.T) {
|
||||
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
|
||||
Model: "gemini-test",
|
||||
Input: mustGeminiRawMessage(t, []map[string]any{
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": []map[string]any{
|
||||
{"type": "output_text", "text": "I will call."},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_1",
|
||||
"name": "lookup",
|
||||
"arguments": map[string]any{"q": "x"},
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": map[string]any{"ok": true},
|
||||
},
|
||||
}),
|
||||
Tools: mustGeminiRawMessage(t, []map[string]any{
|
||||
{"type": "function", "name": "lookup", "parameters": map[string]any{"type": "object"}},
|
||||
}),
|
||||
})
|
||||
|
||||
require.Len(t, got.Contents, 2)
|
||||
assert.Equal(t, "model", got.Contents[0].Role)
|
||||
require.Len(t, got.Contents[0].Parts, 2)
|
||||
require.NotNil(t, got.Contents[0].Parts[0].FunctionCall)
|
||||
assert.Equal(t, "lookup", got.Contents[0].Parts[0].FunctionCall.FunctionName)
|
||||
assert.Equal(t, map[string]interface{}{"q": "x"}, got.Contents[0].Parts[0].FunctionCall.Arguments)
|
||||
assert.Equal(t, "I will call.", got.Contents[0].Parts[1].Text)
|
||||
|
||||
assert.Equal(t, "user", got.Contents[1].Role)
|
||||
require.Len(t, got.Contents[1].Parts, 1)
|
||||
require.NotNil(t, got.Contents[1].Parts[0].FunctionResponse)
|
||||
assert.Equal(t, "lookup", got.Contents[1].Parts[0].FunctionResponse.Name)
|
||||
assert.Equal(t, map[string]interface{}{"ok": true}, got.Contents[1].Parts[0].FunctionResponse.Response)
|
||||
}
|
||||
|
||||
func TestConvertOpenAIResponsesRequestToGeminiSkipsCustomToolCalls(t *testing.T) {
|
||||
got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{
|
||||
Model: "gemini-test",
|
||||
Input: mustGeminiRawMessage(t, []map[string]any{
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": []map[string]any{
|
||||
{"type": "output_text", "text": "before custom"},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "custom_tool_call",
|
||||
"call_id": "call_custom",
|
||||
"name": "apply_patch",
|
||||
"input": "patch body",
|
||||
},
|
||||
{
|
||||
"type": "custom_tool_call_output",
|
||||
"call_id": "call_custom",
|
||||
"output": "ok",
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_custom",
|
||||
"output": "legacy custom output",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "next turn",
|
||||
},
|
||||
}),
|
||||
Tools: mustGeminiRawMessage(t, []map[string]any{
|
||||
{"type": "custom", "name": "apply_patch"},
|
||||
{"type": "unknown", "name": "unknown"},
|
||||
}),
|
||||
})
|
||||
|
||||
assert.Empty(t, got.GetTools())
|
||||
require.Len(t, got.Contents, 2)
|
||||
assert.Equal(t, "model", got.Contents[0].Role)
|
||||
require.Len(t, got.Contents[0].Parts, 1)
|
||||
assert.Equal(t, "before custom", got.Contents[0].Parts[0].Text)
|
||||
assert.Nil(t, got.Contents[0].Parts[0].FunctionCall)
|
||||
|
||||
assert.Equal(t, "user", got.Contents[1].Role)
|
||||
require.Len(t, got.Contents[1].Parts, 1)
|
||||
assert.Equal(t, "next turn", got.Contents[1].Parts[0].Text)
|
||||
assert.Nil(t, got.Contents[1].Parts[0].FunctionResponse)
|
||||
}
|
||||
|
||||
func mustConvertResponsesToGemini(t *testing.T, req dto.OpenAIResponsesRequest) *dto.GeminiChatRequest {
|
||||
t.Helper()
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: req.Model,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: req.Model,
|
||||
},
|
||||
}
|
||||
got, err := (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, req)
|
||||
require.NoError(t, err)
|
||||
geminiReq, ok := got.(*dto.GeminiChatRequest)
|
||||
require.True(t, ok)
|
||||
return geminiReq
|
||||
}
|
||||
|
||||
func mustGeminiRawMessage(t *testing.T, value any) []byte {
|
||||
t.Helper()
|
||||
raw, err := common.Marshal(value)
|
||||
require.NoError(t, err)
|
||||
return raw
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/service/relayconvert"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
|
||||
responseBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
logger.LogDebug(c, "Gemini responses response body: %s", responseBody)
|
||||
|
||||
var geminiResponse dto.GeminiChatResponse
|
||||
if err := common.Unmarshal(responseBody, &geminiResponse); err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
return &usage, types.NewOpenAIError(
|
||||
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
|
||||
types.ErrorCodePromptBlocked,
|
||||
http.StatusBadRequest,
|
||||
)
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "gemini_empty_candidates")
|
||||
return &usage, types.NewOpenAIError(
|
||||
errors.New("empty response from Gemini API"),
|
||||
types.ErrorCodeEmptyResponse,
|
||||
http.StatusInternalServerError,
|
||||
)
|
||||
}
|
||||
|
||||
chatResp := responseGeminiChat2OpenAI(c, &geminiResponse)
|
||||
chatResp.Model = info.UpstreamModelName
|
||||
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
|
||||
chatResp.Usage = usage
|
||||
|
||||
responsesResp, responsesUsage, err := service.ChatCompletionsResponseToResponsesResponse(chatResp, helper.GetResponseID(c))
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
if responsesUsage == nil || responsesUsage.TotalTokens == 0 {
|
||||
responsesResp.Usage = relayconvert.UsageFromChatUsage(&usage)
|
||||
}
|
||||
|
||||
responseBody, err = common.Marshal(responsesResp)
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
|
||||
}
|
||||
service.IOCopyBytesGracefully(c, resp, responseBody)
|
||||
return &usage, nil
|
||||
}
|
||||
|
||||
func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
responseID := helper.GetResponseID(c)
|
||||
created := common.GetTimestamp()
|
||||
state := relayconvert.NewChatToResponsesStreamState(responseID, info.UpstreamModelName)
|
||||
state.Created = created
|
||||
finishReason := constant.FinishReasonStop
|
||||
toolCallIndexByChoice := make(map[int]map[string]int)
|
||||
nextToolCallIndexByChoice := make(map[int]int)
|
||||
var streamErr *types.NewAPIError
|
||||
|
||||
sendEvent := func(event relayconvert.ChatToResponsesStreamEvent) bool {
|
||||
data, err := common.Marshal(event.Payload)
|
||||
if err != nil {
|
||||
streamErr = types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError)
|
||||
return false
|
||||
}
|
||||
helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data))
|
||||
return true
|
||||
}
|
||||
sendChunk := func(chunk *dto.ChatCompletionsStreamResponse) bool {
|
||||
events, err := relayconvert.ChatCompletionsStreamChunkToResponsesEvents(chunk, state)
|
||||
if err != nil {
|
||||
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError)
|
||||
return false
|
||||
}
|
||||
for _, event := range events {
|
||||
if !sendEvent(event) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool {
|
||||
response, isStop := streamResponseGeminiChat2OpenAI(geminiResponse)
|
||||
response.Id = responseID
|
||||
response.Created = created
|
||||
response.Model = info.UpstreamModelName
|
||||
|
||||
if response.IsToolCall() {
|
||||
finishReason = constant.FinishReasonToolCalls
|
||||
}
|
||||
for choiceIdx := range response.Choices {
|
||||
choiceKey := response.Choices[choiceIdx].Index
|
||||
for toolIdx := range response.Choices[choiceIdx].Delta.ToolCalls {
|
||||
tool := &response.Choices[choiceIdx].Delta.ToolCalls[toolIdx]
|
||||
if tool.ID == "" {
|
||||
continue
|
||||
}
|
||||
indexByID := toolCallIndexByChoice[choiceKey]
|
||||
if indexByID == nil {
|
||||
indexByID = make(map[string]int)
|
||||
toolCallIndexByChoice[choiceKey] = indexByID
|
||||
}
|
||||
if idx, ok := indexByID[tool.ID]; ok {
|
||||
tool.SetIndex(idx)
|
||||
continue
|
||||
}
|
||||
idx := nextToolCallIndexByChoice[choiceKey]
|
||||
nextToolCallIndexByChoice[choiceKey] = idx + 1
|
||||
indexByID[tool.ID] = idx
|
||||
tool.SetIndex(idx)
|
||||
}
|
||||
}
|
||||
|
||||
if !sendChunk(response) {
|
||||
return false
|
||||
}
|
||||
if isStop {
|
||||
return sendChunk(helper.GenerateStopResponse(responseID, created, info.UpstreamModelName, finishReason))
|
||||
}
|
||||
return true
|
||||
})
|
||||
if err != nil {
|
||||
return usage, err
|
||||
}
|
||||
if streamErr != nil {
|
||||
return nil, streamErr
|
||||
}
|
||||
|
||||
if usage != nil {
|
||||
state.Usage = relayconvert.UsageFromChatUsage(usage)
|
||||
}
|
||||
for _, event := range relayconvert.FinalizeChatCompletionsStreamToResponses(state) {
|
||||
if !sendEvent(event) {
|
||||
return nil, streamErr
|
||||
}
|
||||
}
|
||||
return usage, nil
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"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"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGeminiResponsesHandlerReturnsOpenAIResponsesJSON(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "gemini-responses-test")
|
||||
|
||||
info := newGeminiResponsesRelayInfo(false)
|
||||
payload := dto.GeminiChatResponse{
|
||||
Candidates: []dto.GeminiChatCandidate{
|
||||
{
|
||||
Content: dto.GeminiChatContent{
|
||||
Role: "model",
|
||||
Parts: []dto.GeminiPart{
|
||||
{Text: "hello"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
UsageMetadata: dto.GeminiUsageMetadata{
|
||||
PromptTokenCount: 2,
|
||||
CandidatesTokenCount: 3,
|
||||
TotalTokenCount: 5,
|
||||
},
|
||||
}
|
||||
body, err := common.Marshal(payload)
|
||||
require.NoError(t, err)
|
||||
|
||||
usage, newAPIError := GeminiResponsesHandler(c, info, &http.Response{
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
})
|
||||
require.Nil(t, newAPIError)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 2, usage.PromptTokens)
|
||||
assert.Equal(t, 3, usage.CompletionTokens)
|
||||
|
||||
got := recorder.Body.String()
|
||||
assert.Contains(t, got, `"object":"response"`)
|
||||
assert.Contains(t, got, `"status":"completed"`)
|
||||
assert.Contains(t, got, `"type":"output_text"`)
|
||||
assert.Contains(t, got, `"text":"hello"`)
|
||||
assert.Contains(t, got, `"input_tokens":2`)
|
||||
assert.Contains(t, got, `"output_tokens":3`)
|
||||
assert.NotContains(t, got, `"choices"`)
|
||||
assert.NotContains(t, got, `"candidates"`)
|
||||
}
|
||||
|
||||
func TestGeminiResponsesHandlerClosesBodyOnReadError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "gemini-responses-read-error-test")
|
||||
|
||||
body := &failingReadCloser{}
|
||||
usage, newAPIError := GeminiResponsesHandler(c, newGeminiResponsesRelayInfo(false), &http.Response{Body: body})
|
||||
|
||||
require.Nil(t, usage)
|
||||
require.NotNil(t, newAPIError)
|
||||
assert.True(t, body.closed)
|
||||
}
|
||||
|
||||
func TestGeminiResponsesStreamHandlerReturnsOpenAIResponsesSSE(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "gemini-responses-stream-test")
|
||||
|
||||
oldStreamingTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 300
|
||||
t.Cleanup(func() { constant.StreamingTimeout = oldStreamingTimeout })
|
||||
|
||||
info := newGeminiResponsesRelayInfo(true)
|
||||
first := dto.GeminiChatResponse{
|
||||
Candidates: []dto.GeminiChatCandidate{
|
||||
{
|
||||
Content: dto.GeminiChatContent{
|
||||
Role: "model",
|
||||
Parts: []dto.GeminiPart{
|
||||
{Text: "hello"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
UsageMetadata: dto.GeminiUsageMetadata{
|
||||
PromptTokenCount: 2,
|
||||
CandidatesTokenCount: 3,
|
||||
TotalTokenCount: 5,
|
||||
},
|
||||
}
|
||||
stop := "STOP"
|
||||
final := dto.GeminiChatResponse{
|
||||
Candidates: []dto.GeminiChatCandidate{
|
||||
{
|
||||
FinishReason: &stop,
|
||||
Content: dto.GeminiChatContent{
|
||||
Role: "model",
|
||||
Parts: []dto.GeminiPart{{Text: ""}},
|
||||
},
|
||||
},
|
||||
},
|
||||
UsageMetadata: dto.GeminiUsageMetadata{
|
||||
PromptTokenCount: 2,
|
||||
CandidatesTokenCount: 3,
|
||||
TotalTokenCount: 5,
|
||||
},
|
||||
}
|
||||
firstData, err := common.Marshal(first)
|
||||
require.NoError(t, err)
|
||||
finalData, err := common.Marshal(final)
|
||||
require.NoError(t, err)
|
||||
streamBody := strings.Join([]string{
|
||||
"data: " + string(firstData),
|
||||
"",
|
||||
"data: " + string(finalData),
|
||||
"",
|
||||
"data: [DONE]",
|
||||
"",
|
||||
}, "\n")
|
||||
|
||||
usage, newAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{
|
||||
Body: io.NopCloser(strings.NewReader(streamBody)),
|
||||
})
|
||||
require.Nil(t, newAPIError)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 5, usage.TotalTokens)
|
||||
|
||||
got := recorder.Body.String()
|
||||
assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type"))
|
||||
assert.Contains(t, got, `event: response.created`)
|
||||
assert.Contains(t, got, `event: response.output_text.delta`)
|
||||
assert.Contains(t, got, `"delta":"hello"`)
|
||||
assert.Contains(t, got, `event: response.completed`)
|
||||
assert.Contains(t, got, `"input_tokens":2`)
|
||||
assert.Contains(t, got, `"output_tokens":3`)
|
||||
assert.NotContains(t, got, `"choices"`)
|
||||
assert.NotContains(t, got, `"candidates"`)
|
||||
requireOrderedGeminiResponsesSubstrings(t, got,
|
||||
`event: response.created`,
|
||||
`event: response.output_item.added`,
|
||||
`event: response.output_text.delta`,
|
||||
`event: response.output_text.done`,
|
||||
`event: response.completed`,
|
||||
)
|
||||
}
|
||||
|
||||
func newGeminiResponsesRelayInfo(isStream bool) *relaycommon.RelayInfo {
|
||||
return &relaycommon.RelayInfo{
|
||||
IsStream: isStream,
|
||||
RelayMode: relayconstant.RelayModeResponses,
|
||||
RelayFormat: types.RelayFormatOpenAIResponses,
|
||||
RequestURLPath: "/v1/responses",
|
||||
DisablePing: true,
|
||||
OriginModelName: "gemini-test",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: "gemini-test",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type failingReadCloser struct {
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (r *failingReadCloser) Read([]byte) (int, error) {
|
||||
return 0, errors.New("read failed")
|
||||
}
|
||||
|
||||
func (r *failingReadCloser) Close() error {
|
||||
r.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func requireOrderedGeminiResponsesSubstrings(t *testing.T, s string, parts ...string) {
|
||||
t.Helper()
|
||||
offset := 0
|
||||
for _, part := range parts {
|
||||
idx := strings.Index(s[offset:], part)
|
||||
require.NotEqualf(t, -1, idx, "missing %q after byte offset %d", part, offset)
|
||||
offset += idx + len(part)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user