mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
fix(antigravity): preserve tool_result images, native Gemini functionCall/thoughtSignature, and composite /v1beta routing
This commit is contained in:
@@ -39,9 +39,11 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return
|
||||
}
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组(composite 分组允许路由到 antigravity)
|
||||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||||
if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
|
||||
effectivePlatform := effectiveAPIKeyPlatform(c, apiKey)
|
||||
if !hasForcePlatform && effectivePlatform != service.PlatformGemini &&
|
||||
!servesCompositeAntigravityGemini(apiKey, effectivePlatform) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
@@ -223,9 +225,11 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
googleError(c, http.StatusUnauthorized, "Invalid API key")
|
||||
return
|
||||
}
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由),否则要求 gemini 分组(composite 分组允许路由到 antigravity)
|
||||
forcePlatform, hasForcePlatform := middleware.GetForcePlatformFromContext(c)
|
||||
if !hasForcePlatform && effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
|
||||
effectivePlatform := effectiveAPIKeyPlatform(c, apiKey)
|
||||
if !hasForcePlatform && effectivePlatform != service.PlatformGemini &&
|
||||
!servesCompositeAntigravityGemini(apiKey, effectivePlatform) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
@@ -244,9 +248,13 @@ func (h *GatewayHandler) GeminiV1BetaGetModel(c *gin.Context) {
|
||||
if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" {
|
||||
modelName = strings.TrimSpace(resolvedModel)
|
||||
}
|
||||
if !hasForcePlatform && servesCompositeAntigravityGemini(apiKey, effectivePlatform) && !isAntigravityGeminiModel(modelName) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
|
||||
// 强制 antigravity 模式:返回 antigravity 模型信息
|
||||
if forcePlatform == service.PlatformAntigravity {
|
||||
// 强制 antigravity 模式或 composite 路由到 antigravity:返回 antigravity 模型信息
|
||||
if forcePlatform == service.PlatformAntigravity || effectivePlatform == service.PlatformAntigravity {
|
||||
c.JSON(http.StatusOK, antigravity.FallbackGeminiModel(modelName))
|
||||
return
|
||||
}
|
||||
@@ -299,9 +307,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则要求 gemini 分组
|
||||
// 检查平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则要求 gemini 分组(composite 分组允许路由到 antigravity)
|
||||
effectivePlatform := effectiveAPIKeyPlatform(c, apiKey)
|
||||
if !middleware.HasForcePlatform(c) {
|
||||
if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformGemini {
|
||||
if effectivePlatform != service.PlatformGemini &&
|
||||
!servesCompositeAntigravityGemini(apiKey, effectivePlatform) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
@@ -321,6 +331,10 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" {
|
||||
modelName = strings.TrimSpace(resolvedModel)
|
||||
}
|
||||
if !middleware.HasForcePlatform(c) && servesCompositeAntigravityGemini(apiKey, effectivePlatform) && !isAntigravityGeminiModel(modelName) {
|
||||
googleError(c, http.StatusBadRequest, "API key group platform is not gemini")
|
||||
return
|
||||
}
|
||||
|
||||
stream := action == "streamGenerateContent"
|
||||
reqLog = reqLog.With(zap.String("model", modelName), zap.String("action", action), zap.Bool("stream", stream))
|
||||
@@ -567,7 +581,11 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, service.PlatformGemini)
|
||||
diagPlatform := service.PlatformGemini
|
||||
if p := effectiveAPIKeyPlatform(c, apiKey); p == service.PlatformAntigravity {
|
||||
diagPlatform = service.PlatformAntigravity
|
||||
}
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, modelName, modelName, diagPlatform)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
@@ -1053,3 +1071,16 @@ func derefGroupID(groupID *int64) int64 {
|
||||
}
|
||||
return *groupID
|
||||
}
|
||||
|
||||
func servesCompositeAntigravityGemini(apiKey *service.APIKey, effectivePlatform string) bool {
|
||||
return apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite && effectivePlatform == service.PlatformAntigravity
|
||||
}
|
||||
|
||||
func isAntigravityGeminiModel(model string) bool {
|
||||
trimmed := strings.ToLower(strings.TrimSpace(strings.TrimPrefix(model, "models/")))
|
||||
if trimmed == "" || strings.HasPrefix(trimmed, "claude-") {
|
||||
return false
|
||||
}
|
||||
detected, ok := service.DetectModelPlatform(trimmed)
|
||||
return ok && detected == service.PlatformGemini
|
||||
}
|
||||
|
||||
@@ -216,3 +216,75 @@ func TestShouldFallbackGeminiModel_DelegatesScopeFallback(t *testing.T) {
|
||||
}
|
||||
require.True(t, shouldFallbackGeminiModel("gemini-future-model", res))
|
||||
}
|
||||
|
||||
func TestGeminiV1BetaGetModel_CompositeAntigravityRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("composite route to antigravity gemini model succeeds", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-4-argon", nil)
|
||||
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
||||
Matched: true,
|
||||
TargetPlatform: service.PlatformAntigravity,
|
||||
UpstreamModel: "gemini-4-argon-high",
|
||||
}))
|
||||
c.Request = req
|
||||
c.Params = gin.Params{{Key: "model", Value: "gemini-4-argon"}}
|
||||
groupID := int64(4)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 1,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
|
||||
(&GatewayHandler{}).GeminiV1BetaGetModel(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||
require.Contains(t, rec.Body.String(), "models/gemini-4-argon-high")
|
||||
})
|
||||
|
||||
t.Run("composite route to antigravity claude model is rejected on v1beta", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/claude-sonnet-5-5-high", nil)
|
||||
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
||||
Matched: true,
|
||||
TargetPlatform: service.PlatformAntigravity,
|
||||
UpstreamModel: "claude-sonnet-5-5-high",
|
||||
}))
|
||||
c.Request = req
|
||||
c.Params = gin.Params{{Key: "model", Value: "claude-sonnet-5-5-high"}}
|
||||
groupID := int64(4)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 1,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
|
||||
(&GatewayHandler{}).GeminiV1BetaGetModel(c)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
|
||||
})
|
||||
|
||||
t.Run("composite route to antigravity claude model is rejected on generateContent", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1beta/models/claude-sonnet-5-5-high:generateContent", nil)
|
||||
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
||||
Matched: true,
|
||||
TargetPlatform: service.PlatformAntigravity,
|
||||
UpstreamModel: "claude-sonnet-5-5-high",
|
||||
}))
|
||||
c.Request = req
|
||||
c.Params = gin.Params{{Key: "modelAction", Value: "/claude-sonnet-5-5-high:generateContent"}}
|
||||
groupID := int64(4)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 1,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1, Concurrency: 1})
|
||||
|
||||
(&GatewayHandler{}).GeminiV1BetaModels(c)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -588,6 +588,7 @@ func buildParts(content json.RawMessage, toolIDToName map[string]string, allowDu
|
||||
ID: block.ToolUseID,
|
||||
},
|
||||
})
|
||||
parts = append(parts, extractToolResultInlineParts(block.Content)...)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -638,6 +639,31 @@ func parseToolResultContent(content json.RawMessage, isError bool) string {
|
||||
return string(content)
|
||||
}
|
||||
|
||||
func extractToolResultInlineParts(content json.RawMessage) []GeminiPart {
|
||||
if len(content) == 0 {
|
||||
return nil
|
||||
}
|
||||
var innerBlocks []ContentBlock
|
||||
if err := json.Unmarshal(content, &innerBlocks); err != nil {
|
||||
return nil
|
||||
}
|
||||
var inlineParts []GeminiPart
|
||||
for _, item := range innerBlocks {
|
||||
if (item.Type == "image" || item.Type == "document") &&
|
||||
item.Source != nil &&
|
||||
item.Source.Type == "base64" &&
|
||||
strings.TrimSpace(item.Source.Data) != "" {
|
||||
inlineParts = append(inlineParts, GeminiPart{
|
||||
InlineData: &GeminiInlineData{
|
||||
MimeType: item.Source.MediaType,
|
||||
Data: item.Source.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
return inlineParts
|
||||
}
|
||||
|
||||
// buildGenerationConfig 构建 generationConfig
|
||||
const (
|
||||
defaultMaxOutputTokens = 64000
|
||||
|
||||
@@ -694,3 +694,64 @@ func TestToolConfigAlwaysPresent(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransformClaudeToGemini_ToolResultImagesReachUpstream(t *testing.T) {
|
||||
claudeReq := &ClaudeRequest{
|
||||
Model: "claude-sonnet-4-5",
|
||||
Messages: []ClaudeMessage{
|
||||
{
|
||||
Role: "assistant",
|
||||
Content: json.RawMessage(`[
|
||||
{"type":"tool_use","id":"toolu_img_1","name":"read_image","input":{"path":"colors.png"}}
|
||||
]`),
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: json.RawMessage(`[
|
||||
{
|
||||
"type":"tool_result",
|
||||
"tool_use_id":"toolu_img_1",
|
||||
"content":[
|
||||
{"type":"text","text":"loaded colors.png"},
|
||||
{"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}}
|
||||
]
|
||||
},
|
||||
{
|
||||
"type":"tool_result",
|
||||
"tool_use_id":"toolu_img_1",
|
||||
"content":[
|
||||
{"type":"image","source":{"type":"base64","media_type":"image/jpeg","data":"d29ybGQ="}}
|
||||
]
|
||||
}
|
||||
]`),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
body, err := TransformClaudeToGeminiWithOptions(claudeReq, "project-1", "claude-sonnet-4-5", DefaultTransformOptions())
|
||||
require.NoError(t, err)
|
||||
|
||||
var req V1InternalRequest
|
||||
require.NoError(t, json.Unmarshal(body, &req))
|
||||
require.Len(t, req.Request.Contents, 2)
|
||||
|
||||
userParts := req.Request.Contents[1].Parts
|
||||
require.Len(t, userParts, 4)
|
||||
|
||||
require.NotNil(t, userParts[0].FunctionResponse)
|
||||
require.Equal(t, "read_image", userParts[0].FunctionResponse.Name)
|
||||
require.Equal(t, "toolu_img_1", userParts[0].FunctionResponse.ID)
|
||||
require.Equal(t, "loaded colors.png", userParts[0].FunctionResponse.Response["result"])
|
||||
|
||||
require.NotNil(t, userParts[1].InlineData)
|
||||
require.Equal(t, "image/png", userParts[1].InlineData.MimeType)
|
||||
require.Equal(t, "aGVsbG8=", userParts[1].InlineData.Data)
|
||||
|
||||
require.NotNil(t, userParts[2].FunctionResponse)
|
||||
require.Equal(t, "read_image", userParts[2].FunctionResponse.Name)
|
||||
require.Equal(t, "Command executed successfully.", userParts[2].FunctionResponse.Response["result"])
|
||||
|
||||
require.NotNil(t, userParts[3].InlineData)
|
||||
require.Equal(t, "image/jpeg", userParts[3].InlineData.MimeType)
|
||||
require.Equal(t, "d29ybGQ=", userParts[3].InlineData.Data)
|
||||
}
|
||||
|
||||
@@ -141,6 +141,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
|
||||
injectedBody = reconciled
|
||||
}
|
||||
|
||||
injectedBody = ensureGeminiFunctionCallThoughtSignatures(injectedBody)
|
||||
|
||||
// 包装请求
|
||||
wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody)
|
||||
if err != nil {
|
||||
|
||||
@@ -2012,3 +2012,175 @@ func generateLargeUnwrapJSON(minSize int) []byte {
|
||||
b, _ := json.Marshal(outer)
|
||||
return b
|
||||
}
|
||||
|
||||
func TestAntigravityGatewayService_ForwardGemini_NonStreamPreservesFunctionCallAndThoughtSignature(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","description":"Get weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}]}]}`)
|
||||
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:generateContent", bytes.NewReader(body))
|
||||
|
||||
sseStream := strings.Join([]string{
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_fc_1","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_1"}}]}}],"usageMetadata":{"promptTokenCount":300,"candidatesTokenCount":16,"totalTokenCount":328,"thoughtsTokenCount":12},"modelVersion":"gemini-4-argon","responseId":"resp_1"}}`,
|
||||
"",
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":300,"candidatesTokenCount":22,"totalTokenCount":334,"thoughtsTokenCount":12},"modelVersion":"gemini-4-argon","responseId":"resp_1"}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
|
||||
upstream := &queuedHTTPUpstreamStub{
|
||||
responses: []*http.Response{{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(sseStream)),
|
||||
}},
|
||||
}
|
||||
svc := &AntigravityGatewayService{
|
||||
settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}),
|
||||
tokenProvider: &AntigravityTokenProvider{},
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
account := &Account{
|
||||
ID: 109,
|
||||
Name: "antigravity-fc-nonstream",
|
||||
Platform: PlatformAntigravity,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "test-access-token",
|
||||
"project_id": "test-project",
|
||||
"model_mapping": map[string]any{
|
||||
"gemini-4-argon": "gemini-4-argon",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
result, err := svc.ForwardGemini(context.Background(), c, account, "gemini-4-argon", "generateContent", false, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, http.StatusOK, writer.Code)
|
||||
|
||||
var respPayload map[string]any
|
||||
require.NoError(t, json.Unmarshal(writer.Body.Bytes(), &respPayload))
|
||||
candidates, ok := respPayload["candidates"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, candidates, 1)
|
||||
candidate, ok := candidates[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "STOP", candidate["finishReason"])
|
||||
|
||||
content, ok := candidate["content"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
parts, ok := content["parts"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, parts, 1, "trailing empty text part should not overwrite or pollute functionCall part")
|
||||
|
||||
part0, ok := parts[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "sig_fc_1", part0["thoughtSignature"])
|
||||
fc, ok := part0["functionCall"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "get_weather", fc["name"])
|
||||
require.Equal(t, "call_1", fc["id"])
|
||||
require.Equal(t, map[string]any{"city": "Tokyo"}, fc["args"])
|
||||
}
|
||||
|
||||
func TestAntigravityGatewayService_ForwardGemini_StreamAndFunctionResponseFollowUp(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// 1) Stream request with functionDeclarations returns unwrapped functionCall and thoughtSignature
|
||||
streamReqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}]}`)
|
||||
streamWriter := httptest.NewRecorder()
|
||||
streamCtx, _ := gin.CreateTestContext(streamWriter)
|
||||
streamCtx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:streamGenerateContent?alt=sse", bytes.NewReader(streamReqBody))
|
||||
|
||||
streamSSE := strings.Join([]string{
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_stream_1","functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5}}}`,
|
||||
"",
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":6}}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
|
||||
// 2) Non-stream follow-up with functionCall (missing thoughtSignature) + functionResponse (with inlineData image)
|
||||
followUpReqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"Call get_weather for Tokyo."}]},{"role":"model","parts":[{"functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}}]},{"role":"user","parts":[{"functionResponse":{"name":"get_weather","id":"call_s1","response":{"weather":"Sunny"},"parts":[{"inlineData":{"mimeType":"image/png","data":"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="}}]}}]}],"tools":[{"functionDeclarations":[{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}}}}]}]}`)
|
||||
followUpWriter := httptest.NewRecorder()
|
||||
followUpCtx, _ := gin.CreateTestContext(followUpWriter)
|
||||
followUpCtx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-4-argon:generateContent", bytes.NewReader(followUpReqBody))
|
||||
|
||||
followUpSSE := strings.Join([]string{
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"It is sunny in "}]}}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":4}}}`,
|
||||
"",
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"Tokyo."}]}}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":6}}}`,
|
||||
"",
|
||||
`data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"thoughtSignature":"sig_final","text":""}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":8}}}`,
|
||||
"",
|
||||
}, "\n")
|
||||
|
||||
upstream := &queuedHTTPUpstreamStub{
|
||||
responses: []*http.Response{
|
||||
{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(streamSSE)),
|
||||
},
|
||||
{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(followUpSSE)),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &AntigravityGatewayService{
|
||||
settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}),
|
||||
tokenProvider: &AntigravityTokenProvider{},
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
account := &Account{
|
||||
ID: 110,
|
||||
Name: "antigravity-fc-stream",
|
||||
Platform: PlatformAntigravity,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "test-access-token",
|
||||
"project_id": "test-project",
|
||||
"model_mapping": map[string]any{
|
||||
"gemini-4-argon": "gemini-4-argon",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
res1, err := svc.ForwardGemini(context.Background(), streamCtx, account, "gemini-4-argon", "streamGenerateContent", true, streamReqBody, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, res1)
|
||||
require.Contains(t, streamWriter.Body.String(), `"functionCall":{"name":"get_weather","args":{"city":"Tokyo"},"id":"call_s1"}`)
|
||||
require.Contains(t, streamWriter.Body.String(), `"thoughtSignature":"sig_stream_1"`)
|
||||
|
||||
res2, err := svc.ForwardGemini(context.Background(), followUpCtx, account, "gemini-4-argon", "generateContent", false, followUpReqBody, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, res2)
|
||||
require.Len(t, upstream.requestBodies, 2)
|
||||
|
||||
// Verify follow-up request had dummy thoughtSignature injected on functionCall and preserved functionResponse image part
|
||||
secondSent := string(upstream.requestBodies[1])
|
||||
require.Contains(t, secondSent, `"thoughtSignature":"skip_thought_signature_validator"`)
|
||||
require.Contains(t, secondSent, `"mimeType":"image/png"`)
|
||||
|
||||
var followUpResp map[string]any
|
||||
require.NoError(t, json.Unmarshal(followUpWriter.Body.Bytes(), &followUpResp))
|
||||
cands, ok := followUpResp["candidates"].([]any)
|
||||
require.True(t, ok)
|
||||
cand0, ok := cands[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
content, ok := cand0["content"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
parts, ok := content["parts"].([]any)
|
||||
require.True(t, ok)
|
||||
require.NotEmpty(t, parts)
|
||||
firstPart, ok := parts[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "It is sunny in Tokyo.", firstPart["text"])
|
||||
}
|
||||
|
||||
@@ -350,8 +350,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Cont
|
||||
var firstTokenMs *int
|
||||
var last map[string]any
|
||||
var lastWithParts map[string]any
|
||||
var collectedImageParts []map[string]any // 收集所有包含图片的 parts
|
||||
var collectedTextParts []string // 收集所有文本片段
|
||||
var collectedParts []map[string]any // 收集所有 parts(text、thinking、functionCall、inlineData 等)
|
||||
|
||||
type scanEvent struct {
|
||||
line string
|
||||
@@ -467,19 +466,10 @@ func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Cont
|
||||
}
|
||||
}
|
||||
|
||||
// 保留最后一个有 parts 的响应
|
||||
// 保留最后一个有 parts 的响应,并收集所有 parts
|
||||
if parts := extractGeminiParts(parsed); len(parts) > 0 {
|
||||
lastWithParts = parsed
|
||||
// 收集包含图片和文本的 parts
|
||||
for _, part := range parts {
|
||||
if inlineData, ok := part["inlineData"].(map[string]any); ok {
|
||||
collectedImageParts = append(collectedImageParts, part)
|
||||
_ = inlineData // 避免 unused 警告
|
||||
}
|
||||
if text, ok := part["text"].(string); ok && text != "" {
|
||||
collectedTextParts = append(collectedTextParts, text)
|
||||
}
|
||||
}
|
||||
collectedParts = append(collectedParts, parts...)
|
||||
}
|
||||
|
||||
case <-intervalCh:
|
||||
@@ -506,14 +496,9 @@ returnResponse:
|
||||
}
|
||||
}
|
||||
|
||||
// 如果收集到了图片 parts,需要合并到最终响应中
|
||||
if len(collectedImageParts) > 0 {
|
||||
finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts)
|
||||
}
|
||||
|
||||
// 如果收集到了文本,需要合并到最终响应中
|
||||
if len(collectedTextParts) > 0 {
|
||||
finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts)
|
||||
// 将收集的所有 parts 合并到最终响应中
|
||||
if len(collectedParts) > 0 {
|
||||
finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts)
|
||||
}
|
||||
|
||||
respBody, err := json.Marshal(finalResponse)
|
||||
@@ -576,21 +561,27 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map
|
||||
return response
|
||||
}
|
||||
|
||||
result, _, setParts := getOrCreateGeminiParts(response)
|
||||
result, existingParts, setParts := getOrCreateGeminiParts(response)
|
||||
|
||||
// 合并策略:
|
||||
// 1. 保持原始顺序
|
||||
// 2. 连续的普通 text parts 合并为一个
|
||||
// 2. 连续的普通 text parts 合并为一个,并保留关联的 thoughtSignature
|
||||
// 3. thinking、functionCall、inlineData 等保持原样
|
||||
var mergedParts []any
|
||||
var textBuffer strings.Builder
|
||||
var textSignature string
|
||||
|
||||
flushTextBuffer := func() {
|
||||
if textBuffer.Len() > 0 {
|
||||
mergedParts = append(mergedParts, map[string]any{
|
||||
if textBuffer.Len() > 0 || textSignature != "" {
|
||||
part := map[string]any{
|
||||
"text": textBuffer.String(),
|
||||
})
|
||||
}
|
||||
if textSignature != "" {
|
||||
part["thoughtSignature"] = textSignature
|
||||
}
|
||||
mergedParts = append(mergedParts, part)
|
||||
textBuffer.Reset()
|
||||
textSignature = ""
|
||||
}
|
||||
}
|
||||
|
||||
@@ -605,6 +596,9 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map
|
||||
} else {
|
||||
// 普通 text,累积到 buffer
|
||||
_, _ = textBuffer.WriteString(text)
|
||||
if sig, ok := part["thoughtSignature"].(string); ok && sig != "" {
|
||||
textSignature = sig
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 非 text part(functionCall、inlineData 等),先刷新 text buffer,然后保留原样
|
||||
@@ -616,76 +610,14 @@ func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map
|
||||
// 刷新剩余的 text
|
||||
flushTextBuffer()
|
||||
|
||||
if len(mergedParts) == 0 {
|
||||
mergedParts = existingParts
|
||||
}
|
||||
|
||||
setParts(mergedParts)
|
||||
return result
|
||||
}
|
||||
|
||||
// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中
|
||||
func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any {
|
||||
if len(imageParts) == 0 {
|
||||
return response
|
||||
}
|
||||
|
||||
result, existingParts, setParts := getOrCreateGeminiParts(response)
|
||||
|
||||
// 检查现有 parts 中是否已经有图片
|
||||
for _, p := range existingParts {
|
||||
if pm, ok := p.(map[string]any); ok {
|
||||
if _, hasInline := pm["inlineData"]; hasInline {
|
||||
return result // 已有图片,不重复添加
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 添加收集到的图片 parts
|
||||
for _, imgPart := range imageParts {
|
||||
existingParts = append(existingParts, imgPart)
|
||||
}
|
||||
setParts(existingParts)
|
||||
return result
|
||||
}
|
||||
|
||||
// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中
|
||||
func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any {
|
||||
if len(textParts) == 0 {
|
||||
return response
|
||||
}
|
||||
|
||||
mergedText := strings.Join(textParts, "")
|
||||
result, existingParts, setParts := getOrCreateGeminiParts(response)
|
||||
|
||||
// 查找并更新第一个 text part,或创建新的
|
||||
newParts := make([]any, 0, len(existingParts)+1)
|
||||
textUpdated := false
|
||||
|
||||
for _, p := range existingParts {
|
||||
pm, ok := p.(map[string]any)
|
||||
if !ok {
|
||||
newParts = append(newParts, p)
|
||||
continue
|
||||
}
|
||||
if _, hasText := pm["text"]; hasText && !textUpdated {
|
||||
// 用累积的文本替换
|
||||
newPart := make(map[string]any)
|
||||
for k, v := range pm {
|
||||
newPart[k] = v
|
||||
}
|
||||
newPart["text"] = mergedText
|
||||
newParts = append(newParts, newPart)
|
||||
textUpdated = true
|
||||
} else {
|
||||
newParts = append(newParts, pm)
|
||||
}
|
||||
}
|
||||
|
||||
if !textUpdated {
|
||||
newParts = append([]any{map[string]any{"text": mergedText}}, newParts...)
|
||||
}
|
||||
|
||||
setParts(newParts)
|
||||
return result
|
||||
}
|
||||
|
||||
// antigravityStatusClientClosed 是客户端在上游响应前断开时回写的状态码(499 client closed request)。
|
||||
const antigravityStatusClientClosed = 499
|
||||
|
||||
|
||||
Reference in New Issue
Block a user