mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-23 09:12:20 +08:00
feat(antigravity): bridge Claude WebSearch to native googleSearch
Add a native Antigravity WebSearch path for Claude typed WebSearch requests. Detect Claude Messages requests whose tools are only typed WebSearch tools (web_search_20250305 / web_search_20260209), and convert them into an Antigravity requestType=web_search payload instead of sending the request through the normal tool-calling path. Preserve the user's requested model. The native path is enabled only when that Antigravity model is known to support Google Search. Capability data fetched from Antigravity model info is used only as an enhancement to the local model registry, not as a replacement for the existing registry fallback behavior. Unsupported models keep the existing Antigravity request behavior and are not silently rerouted to another web-search-capable model. Translate Claude WebSearch request options to the verified Antigravity googleSearch shape: - max_uses -> googleSearch.enhancedContent.imageSearch.maxResultCount - allowed_domains -> googleSearch.includedDomains Leave blocked_domains and user_location unmapped because the Antigravity googleSearch request shape has no verified equivalent for them. This avoids sending speculative fields or pretending unsupported Claude WebSearch options are enforced upstream. Translate Antigravity web-search responses back into Claude-compatible output: server_tool_use blocks, web_search_tool_result blocks, cited text blocks, grounding URLs, and usage-compatible stream/non-stream responses. Cover the behavior with tests for request conversion, response conversion, grounding URL resolution, domain filter mapping, fetched capability hints, excluded-model handling, and unsupported-model behavior.
This commit is contained in:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -41,6 +41,7 @@ GEMINI.md
|
||||
.gemini/*
|
||||
.serena/*
|
||||
.agent/*
|
||||
.agents
|
||||
.agents/*
|
||||
.opencode/*
|
||||
.idea/*
|
||||
|
||||
@@ -85,6 +85,31 @@ func GetAntigravityModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Antigravity)
|
||||
}
|
||||
|
||||
// AntigravityWebSearchModelFor returns the Antigravity model that should run a
|
||||
// native web search request for modelID.
|
||||
func AntigravityWebSearchModelFor(modelID string) string {
|
||||
modelID = normalizeAntigravityCapabilityModelID(modelID)
|
||||
if modelID == "" {
|
||||
return ""
|
||||
}
|
||||
for _, model := range GetGlobalRegistry().GetAvailableModelsByProvider("antigravity") {
|
||||
if model == nil {
|
||||
continue
|
||||
}
|
||||
currentModelID := normalizeAntigravityCapabilityModelID(model.ID)
|
||||
if currentModelID == "" {
|
||||
continue
|
||||
}
|
||||
if currentModelID == modelID {
|
||||
if model.SupportsWebSearch {
|
||||
return currentModelID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// GetXAIModels returns the standard xAI Grok model definitions.
|
||||
func GetXAIModels() []*ModelInfo {
|
||||
return WithXAIBuiltins(cloneModelInfos(getModels().XAI))
|
||||
@@ -103,6 +128,14 @@ func WithXAIBuiltins(models []*ModelInfo) []*ModelInfo {
|
||||
return upsertModelInfos(models, xaiBuiltinImageModelInfo(), xaiBuiltinImageQualityModelInfo(), xaiBuiltinVideoModelInfo(), xaiBuiltinVideo15PreviewModelInfo())
|
||||
}
|
||||
|
||||
func normalizeAntigravityCapabilityModelID(modelID string) string {
|
||||
modelID = strings.ToLower(strings.TrimSpace(modelID))
|
||||
if open := strings.LastIndex(modelID, "("); open >= 0 && strings.HasSuffix(modelID, ")") {
|
||||
modelID = strings.TrimSpace(modelID[:open])
|
||||
}
|
||||
return modelID
|
||||
}
|
||||
|
||||
func codexBuiltinImageModelInfo() *ModelInfo {
|
||||
return &ModelInfo{
|
||||
ID: codexBuiltinImageModelID,
|
||||
|
||||
@@ -16,3 +16,35 @@ func TestWithXAIBuiltinsIncludesVideoPreviewModel(t *testing.T) {
|
||||
|
||||
t.Fatalf("expected xAI builtin model %s", xaiBuiltinVideo15PreviewModelID)
|
||||
}
|
||||
|
||||
func TestAntigravityWebSearchModelForRequiresRequestedModelCapability(t *testing.T) {
|
||||
registryRef := GetGlobalRegistry()
|
||||
registryRef.RegisterClient("test-antigravity-websearch-route", "antigravity", []*ModelInfo{
|
||||
{ID: "gemini-route-test"},
|
||||
{ID: "gemini-web-search-test", SupportsWebSearch: true},
|
||||
})
|
||||
registryRef.RegisterClient("test-gemini-websearch-route", "gemini", []*ModelInfo{
|
||||
{ID: "gemini-cross-provider-route"},
|
||||
{ID: "gemini-cross-provider-search", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
registryRef.UnregisterClient("test-antigravity-websearch-route")
|
||||
registryRef.UnregisterClient("test-gemini-websearch-route")
|
||||
})
|
||||
|
||||
if got := AntigravityWebSearchModelFor("gemini-route-test"); got != "" {
|
||||
t.Fatalf("route model without web search support should not get fallback model, got %q", got)
|
||||
}
|
||||
if got := AntigravityWebSearchModelFor("gemini-route-test(high)"); got != "" {
|
||||
t.Fatalf("suffix route model without web search support should not get fallback model, got %q", got)
|
||||
}
|
||||
if got := AntigravityWebSearchModelFor("gemini-web-search-test"); got != "gemini-web-search-test" {
|
||||
t.Fatalf("AntigravityWebSearchModelFor capable model = %q, want itself", got)
|
||||
}
|
||||
if got := AntigravityWebSearchModelFor("gemini-cross-provider-route"); got != "" {
|
||||
t.Fatalf("cross-provider model should not get Antigravity web search model, got %q", got)
|
||||
}
|
||||
if got := AntigravityWebSearchModelFor("unknown-model"); got != "" {
|
||||
t.Fatalf("unknown model should not get Antigravity web search model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,6 +54,9 @@ type ModelInfo struct {
|
||||
SupportedInputModalities []string `json:"supportedInputModalities,omitempty"`
|
||||
// SupportedOutputModalities lists supported output modalities (e.g., TEXT, IMAGE)
|
||||
SupportedOutputModalities []string `json:"supportedOutputModalities,omitempty"`
|
||||
// SupportsWebSearch indicates this Antigravity model is listed by
|
||||
// fetchAvailableModels.webSearchModelIds and can execute native googleSearch.
|
||||
SupportsWebSearch bool `json:"supports_web_search,omitempty"`
|
||||
|
||||
// Thinking holds provider-specific reasoning/thinking budget capabilities.
|
||||
// This is optional and currently used for Gemini thinking budget normalization.
|
||||
|
||||
@@ -271,6 +271,46 @@ func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []b
|
||||
return rawJSON, nil
|
||||
}
|
||||
|
||||
func hasAntigravityClaudeTypedWebSearchTool(payload []byte) bool {
|
||||
tools := gjson.GetBytes(payload, "tools")
|
||||
if !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
switch tool.Get("type").String() {
|
||||
case "web_search_20250305", "web_search_20260209":
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func hasAntigravityGoogleSearchTool(payload []byte) bool {
|
||||
tools := gjson.GetBytes(payload, "request.tools")
|
||||
if !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if tool.Get("googleSearch").Exists() {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func shouldResolveAntigravityWebSearchGroundingURLs(from sdktranslator.Format, originalRequestRawJSON, requestRawJSON []byte) bool {
|
||||
return from.String() == "claude" &&
|
||||
hasAntigravityClaudeTypedWebSearchTool(originalRequestRawJSON) &&
|
||||
hasAntigravityGoogleSearchTool(requestRawJSON)
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) resolveWebSearchGroundingURLs(ctx context.Context, auth *cliproxyauth.Auth, from sdktranslator.Format, originalRequestRawJSON, requestRawJSON, responseRawJSON []byte) []byte {
|
||||
if !shouldResolveAntigravityWebSearchGroundingURLs(from, originalRequestRawJSON, requestRawJSON) {
|
||||
return responseRawJSON
|
||||
}
|
||||
return helps.ResolveAntigravityGroundingURLs(ctx, e.cfg, auth, responseRawJSON)
|
||||
}
|
||||
|
||||
func countClaudeThinkingBlocks(rawJSON []byte) int {
|
||||
messages := gjson.GetBytes(rawJSON, "messages")
|
||||
if !messages.IsArray() {
|
||||
@@ -709,6 +749,7 @@ attemptLoop:
|
||||
if useCredits {
|
||||
clearAntigravityCreditsFailureState(auth)
|
||||
}
|
||||
bodyBytes = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, bodyBytes)
|
||||
reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes))
|
||||
var param any
|
||||
converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m)
|
||||
@@ -973,6 +1014,7 @@ attemptLoop:
|
||||
}
|
||||
resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())}
|
||||
|
||||
resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload)
|
||||
reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload))
|
||||
var param any
|
||||
converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m)
|
||||
@@ -1414,6 +1456,7 @@ attemptLoop:
|
||||
reporter.Publish(ctx, detail)
|
||||
}
|
||||
|
||||
payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload)
|
||||
chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m)
|
||||
for i := range chunks {
|
||||
select {
|
||||
@@ -2473,14 +2516,15 @@ func geminiToAntigravity(modelName string, payload []byte, projectID string) []b
|
||||
template, _ = sjson.SetBytes(template, "userAgent", "antigravity")
|
||||
|
||||
isImageModel := strings.Contains(modelName, "image")
|
||||
|
||||
var reqType string
|
||||
if isImageModel {
|
||||
reqType = "image_gen"
|
||||
} else {
|
||||
reqType = "agent"
|
||||
reqType := strings.TrimSpace(gjson.GetBytes(template, "requestType").String())
|
||||
if reqType == "" {
|
||||
if isImageModel {
|
||||
reqType = "image_gen"
|
||||
} else {
|
||||
reqType = "agent"
|
||||
}
|
||||
template, _ = sjson.SetBytes(template, "requestType", reqType)
|
||||
}
|
||||
template, _ = sjson.SetBytes(template, "requestType", reqType)
|
||||
|
||||
if projectID != "" {
|
||||
template, _ = sjson.SetBytes(template, "project", projectID)
|
||||
@@ -2490,7 +2534,7 @@ func geminiToAntigravity(modelName string, payload []byte, projectID string) []b
|
||||
|
||||
if isImageModel {
|
||||
template, _ = sjson.SetBytes(template, "requestId", generateImageGenRequestID())
|
||||
} else {
|
||||
} else if reqType != "web_search" {
|
||||
template, _ = sjson.SetBytes(template, "requestId", generateRequestID())
|
||||
template, _ = sjson.SetBytes(template, "request.sessionId", generateStableSessionID(payload))
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
)
|
||||
|
||||
func TestAntigravityBuildRequest_SanitizesGeminiToolSchema(t *testing.T) {
|
||||
@@ -110,6 +111,86 @@ func TestAntigravityBuildRequest_UsesAuthProjectID(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityBuildRequest_UsesRouteModelWhenPayloadContainsDifferentModel(t *testing.T) {
|
||||
body := buildRequestBodyFromRawPayload(t, "gemini-3-flash-agent", []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"request": {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "Perform a web search"}]
|
||||
}
|
||||
],
|
||||
"tools": [{"googleSearch": {}}]
|
||||
}
|
||||
}`))
|
||||
|
||||
if got, ok := body["model"].(string); !ok || got != "gemini-3-flash-agent" {
|
||||
t.Fatalf("request model should stay on route model, got=%v", body["model"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityBuildRequest_PreservesIndependentWebSearchRequestType(t *testing.T) {
|
||||
body := buildRequestBodyFromRawPayload(t, "gemini-3.1-flash-lite", []byte(`{
|
||||
"requestType": "web_search",
|
||||
"request": {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": "北京天气 2026-06-12"}]
|
||||
}
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"googleSearch": {
|
||||
"enhancedContent": {
|
||||
"imageSearch": {
|
||||
"maxResultCount": 5
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"candidateCount": 1
|
||||
}
|
||||
}
|
||||
}`))
|
||||
|
||||
if got, ok := body["requestType"].(string); !ok || got != "web_search" {
|
||||
t.Fatalf("requestType should stay web_search, got=%v", body["requestType"])
|
||||
}
|
||||
if _, ok := body["requestId"]; ok {
|
||||
t.Fatalf("web_search request should not add requestId: %v", body["requestId"])
|
||||
}
|
||||
request, ok := body["request"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("request missing or invalid: %v", body["request"])
|
||||
}
|
||||
if _, ok := request["sessionId"]; ok {
|
||||
t.Fatalf("web_search request should not add request.sessionId: %v", request["sessionId"])
|
||||
}
|
||||
if got, ok := body["project"].(string); !ok || got != "project-1" {
|
||||
t.Fatalf("project should come from auth metadata, got=%v", body["project"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldResolveAntigravityWebSearchGroundingURLsRequiresTypedWebSearchAndSearchRequest(t *testing.T) {
|
||||
original := []byte(`{"tools":[{"type":"web_search_20250305","name":"web_search"}]}`)
|
||||
translatedWithGoogleSearch := []byte(`{"requestType":"web_search","request":{"tools":[{"googleSearch":{}}]}}`)
|
||||
translatedWithoutGoogleSearch := []byte(`{"request":{"contents":[]}}`)
|
||||
|
||||
if !shouldResolveAntigravityWebSearchGroundingURLs(sdktranslator.FormatClaude, original, translatedWithGoogleSearch) {
|
||||
t.Fatal("expected typed Claude web search translated to web_search request to resolve grounding URLs")
|
||||
}
|
||||
if shouldResolveAntigravityWebSearchGroundingURLs(sdktranslator.FormatClaude, original, translatedWithoutGoogleSearch) {
|
||||
t.Fatal("expected request without googleSearch to skip grounding URL resolution")
|
||||
}
|
||||
if shouldResolveAntigravityWebSearchGroundingURLs(sdktranslator.FormatOpenAI, original, translatedWithGoogleSearch) {
|
||||
t.Fatal("expected non-Claude source format to skip grounding URL resolution")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityPrepareRequestAuth_FetchesMissingProjectID(t *testing.T) {
|
||||
executor := &AntigravityExecutor{}
|
||||
auth := &cliproxyauth.Auth{Metadata: map[string]any{
|
||||
|
||||
104
internal/runtime/executor/helps/antigravity_grounding_urls.go
Normal file
104
internal/runtime/executor/helps/antigravity_grounding_urls.go
Normal file
@@ -0,0 +1,104 @@
|
||||
package helps
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func isAntigravityVertexSearchRedirect(rawURL string) bool {
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return parsed.Scheme == "https" &&
|
||||
parsed.Host == "vertexaisearch.cloud.google.com" &&
|
||||
strings.HasPrefix(parsed.Path, "/grounding-api-redirect/")
|
||||
}
|
||||
|
||||
func resolveAntigravityGroundingURL(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, rawURL string) string {
|
||||
if !isAntigravityVertexSearchRedirect(rawURL) {
|
||||
return rawURL
|
||||
}
|
||||
client := NewProxyAwareHTTPClient(ctx, cfg, auth, 0)
|
||||
client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
req, errReq := http.NewRequestWithContext(ctx, http.MethodHead, rawURL, nil)
|
||||
if errReq != nil {
|
||||
log.WithError(errReq).Debug("antigravity grounding url: create redirect request failed")
|
||||
return rawURL
|
||||
}
|
||||
resp, errDo := client.Do(req)
|
||||
if errDo != nil {
|
||||
log.WithError(errDo).Debug("antigravity grounding url: resolve redirect failed")
|
||||
return rawURL
|
||||
}
|
||||
defer func() {
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.WithError(errClose).Debug("antigravity grounding url: close redirect response failed")
|
||||
}
|
||||
}()
|
||||
|
||||
if resp.StatusCode < http.StatusMultipleChoices || resp.StatusCode >= http.StatusBadRequest {
|
||||
return rawURL
|
||||
}
|
||||
location := strings.TrimSpace(resp.Header.Get("Location"))
|
||||
if location == "" {
|
||||
return rawURL
|
||||
}
|
||||
parsed, errParse := url.Parse(location)
|
||||
if errParse != nil || parsed.Scheme != "https" || parsed.Host == "" {
|
||||
return rawURL
|
||||
}
|
||||
return location
|
||||
}
|
||||
|
||||
// ResolveAntigravityGroundingURLs replaces Vertex Search redirect URLs in grounding chunks with their target URLs.
|
||||
func ResolveAntigravityGroundingURLs(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte) []byte {
|
||||
if len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
|
||||
basePath := "response.candidates.0.groundingMetadata.groundingChunks"
|
||||
chunks := gjson.GetBytes(payload, basePath)
|
||||
if !chunks.IsArray() {
|
||||
basePath = "candidates.0.groundingMetadata.groundingChunks"
|
||||
chunks = gjson.GetBytes(payload, basePath)
|
||||
}
|
||||
if !chunks.IsArray() {
|
||||
return payload
|
||||
}
|
||||
|
||||
output := payload
|
||||
resolved := map[string]string{}
|
||||
for i, chunk := range chunks.Array() {
|
||||
uri := strings.TrimSpace(chunk.Get("web.uri").String())
|
||||
if uri == "" {
|
||||
continue
|
||||
}
|
||||
resolvedURI, ok := resolved[uri]
|
||||
if !ok {
|
||||
resolvedURI = resolveAntigravityGroundingURL(ctx, cfg, auth, uri)
|
||||
resolved[uri] = resolvedURI
|
||||
}
|
||||
if resolvedURI == uri {
|
||||
continue
|
||||
}
|
||||
updated, errSet := sjson.SetBytes(output, fmt.Sprintf("%s.%d.web.uri", basePath, i), resolvedURI)
|
||||
if errSet != nil {
|
||||
log.WithError(errSet).Debug("antigravity grounding url: set resolved url failed")
|
||||
continue
|
||||
}
|
||||
output = updated
|
||||
}
|
||||
return output
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package helps
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type groundingURLRoundTripper func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f groundingURLRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func TestResolveAntigravityGroundingURLsResolvesVertexRedirects(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const redirectURL = "https://vertexaisearch.cloud.google.com/grounding-api-redirect/example-token"
|
||||
const resolvedURL = "https://example.com/weather"
|
||||
|
||||
var sawRedirectRequest bool
|
||||
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", groundingURLRoundTripper(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Method != http.MethodHead {
|
||||
t.Fatalf("method = %s, want HEAD", req.Method)
|
||||
}
|
||||
if req.URL.String() != redirectURL {
|
||||
t.Fatalf("url = %s, want %s", req.URL.String(), redirectURL)
|
||||
}
|
||||
sawRedirectRequest = true
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusFound,
|
||||
Header: http.Header{
|
||||
"Location": []string{resolvedURL},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader("")),
|
||||
}, nil
|
||||
}))
|
||||
|
||||
input := []byte(`{
|
||||
"response": {
|
||||
"candidates": [{
|
||||
"groundingMetadata": {
|
||||
"groundingChunks": [
|
||||
{"web": {"uri": "` + redirectURL + `", "title": "Weather"}},
|
||||
{"web": {"uri": "https://already.example/source", "title": "Existing"}}
|
||||
]
|
||||
}
|
||||
}]
|
||||
}
|
||||
}`)
|
||||
|
||||
output := ResolveAntigravityGroundingURLs(ctx, nil, nil, input)
|
||||
if !sawRedirectRequest {
|
||||
t.Fatal("expected resolver to request the vertex redirect")
|
||||
}
|
||||
if got := gjson.GetBytes(output, "response.candidates.0.groundingMetadata.groundingChunks.0.web.uri").String(); got != resolvedURL {
|
||||
t.Fatalf("resolved uri = %q, want %q; output=%s", got, resolvedURL, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "response.candidates.0.groundingMetadata.groundingChunks.1.web.uri").String(); got != "https://already.example/source" {
|
||||
t.Fatalf("non-vertex uri = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -256,6 +256,9 @@ func logDroppedAntigravityToolUseSignature(modelName string, messageIndex, conte
|
||||
func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _ bool) []byte {
|
||||
enableThoughtTranslate := true
|
||||
rawJSON := inputRawJSON
|
||||
if shouldBuildAntigravityWebSearchRequest(modelName, rawJSON) {
|
||||
return buildAntigravityWebSearchRequest(modelName, rawJSON)
|
||||
}
|
||||
|
||||
// system instruction
|
||||
var systemInstructionJSON []byte
|
||||
@@ -595,10 +598,13 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
allowedToolKeys := []string{"name", "description", "behavior", "parameters", "parametersJsonSchema", "response", "responseJsonSchema"}
|
||||
toolsResult := gjson.GetBytes(rawJSON, "tools")
|
||||
if toolsResult.IsArray() {
|
||||
toolsJSON = []byte(`[{"functionDeclarations":[]}]`)
|
||||
functionToolNode := []byte(`{"functionDeclarations":[]}`)
|
||||
toolsResults := toolsResult.Array()
|
||||
for i := 0; i < len(toolsResults); i++ {
|
||||
toolResult := toolsResults[i]
|
||||
if isClaudeTypedWebSearchToolType(toolResult.Get("type").String()) {
|
||||
continue
|
||||
}
|
||||
inputSchemaResult := toolResult.Get("input_schema")
|
||||
if inputSchemaResult.Exists() && inputSchemaResult.IsObject() {
|
||||
// Sanitize the input schema for Antigravity API compatibility
|
||||
@@ -612,10 +618,14 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
}
|
||||
tool, _ = sjson.DeleteBytes(tool, toolKey)
|
||||
}
|
||||
toolsJSON, _ = sjson.SetRawBytes(toolsJSON, "0.functionDeclarations.-1", tool)
|
||||
functionToolNode, _ = sjson.SetRawBytes(functionToolNode, "functionDeclarations.-1", tool)
|
||||
toolDeclCount++
|
||||
}
|
||||
}
|
||||
if toolDeclCount > 0 {
|
||||
toolsJSON = []byte(`[]`)
|
||||
toolsJSON, _ = sjson.SetRawBytes(toolsJSON, "-1", functionToolNode)
|
||||
}
|
||||
}
|
||||
|
||||
// Build output Gemini CLI request JSON
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/cache"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/sirupsen/logrus/hooks/test"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -180,6 +181,144 @@ func TestConvertClaudeRequestToAntigravity_ConvertsMessageSystemRoleToUserConten
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_MapsTypedWebSearchToIndependentSearchRequest(t *testing.T) {
|
||||
registry.GetGlobalRegistry().RegisterClient("test-antigravity-claude-websearch", "antigravity", []*registry.ModelInfo{
|
||||
{ID: "gemini-3.1-flash-lite", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("test-antigravity-claude-websearch") })
|
||||
|
||||
inputJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"messages": [{"role": "user", "content": "北京天气 2026-06-12"}],
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 8, "allowed_domains": ["www.baidu.com", "weather.com.cn"]}]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("gemini-3.1-flash-lite", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "requestType").String(); got != "web_search" {
|
||||
t.Fatalf("requestType = %q, want web_search: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.contents.0.parts.0.text").String(); got != "北京天气 2026-06-12" {
|
||||
t.Fatalf("search query = %q, want original user query: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.systemInstruction.parts.0.text").String(); got != antigravityWebSearchSystemInstruction {
|
||||
t.Fatalf("unexpected search system instruction: %q", got)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.0.googleSearch.enhancedContent.imageSearch.maxResultCount").Int(); got != 8 {
|
||||
t.Fatalf("image search maxResultCount = %d, want 8: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.0.googleSearch.includedDomains.0").String(); got != "www.baidu.com" {
|
||||
t.Fatalf("includedDomains.0 = %q, want www.baidu.com: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.0.googleSearch.includedDomains.1").String(); got != "weather.com.cn" {
|
||||
t.Fatalf("includedDomains.1 = %q, want weather.com.cn: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.generationConfig.candidateCount").Int(); got != 1 {
|
||||
t.Fatalf("candidateCount = %d, want 1: %s", got, output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_UsesDefaultWebSearchMaxResultCountWithoutMaxUses(t *testing.T) {
|
||||
registry.GetGlobalRegistry().RegisterClient("test-antigravity-claude-websearch-default-max", "antigravity", []*registry.ModelInfo{
|
||||
{ID: "gemini-3.1-flash-lite", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("test-antigravity-claude-websearch-default-max") })
|
||||
|
||||
inputJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"messages": [{"role": "user", "content": "北京天气 2026-06-12"}],
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("gemini-3.1-flash-lite", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "request.tools.0.googleSearch.enhancedContent.imageSearch.maxResultCount").Int(); got != 5 {
|
||||
t.Fatalf("image search maxResultCount = %d, want default 5: %s", got, output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_DoesNotMapTypedWebSearchWhenMixedWithCustomTools(t *testing.T) {
|
||||
registry.GetGlobalRegistry().RegisterClient("test-antigravity-claude-websearch-mixed", "antigravity", []*registry.ModelInfo{
|
||||
{ID: "gemini-3.1-flash-lite", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("test-antigravity-claude-websearch-mixed") })
|
||||
|
||||
inputJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"messages": [{"role": "user", "content": "Search current weather"}],
|
||||
"tools": [
|
||||
{"type": "web_search_20250305", "name": "web_search", "max_uses": 8},
|
||||
{"name": "lookup", "description": "Lookup local data", "input_schema": {"type": "object", "properties": {}}}
|
||||
]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("gemini-3.1-flash-lite", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "requestType").String(); got == "web_search" {
|
||||
t.Fatalf("mixed tools should not become independent web_search request: %s", output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.#(googleSearch)").Raw; got != "" {
|
||||
t.Fatalf("mixed tools should not inject native googleSearch into chat request: %s", output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, `request.tools.#.functionDeclarations.#(name=="lookup")`).Raw; got == "" {
|
||||
t.Fatalf("custom tool declaration should be preserved: %s", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_DoesNotMapTypedWebSearchForUnsupportedRouteModel(t *testing.T) {
|
||||
registry.GetGlobalRegistry().RegisterClient("test-antigravity-claude-websearch-route", "antigravity", []*registry.ModelInfo{
|
||||
{ID: "gemini-3.5-flash"},
|
||||
{ID: "gemini-3.1-flash-lite", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("test-antigravity-claude-websearch-route") })
|
||||
|
||||
inputJSON := []byte(`{
|
||||
"model": "gemini-3.5-flash",
|
||||
"messages": [{"role": "user", "content": "Perform a web search"}],
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("gemini-3.5-flash", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "model").String(); got != "gemini-3.5-flash" {
|
||||
t.Fatalf("web search request model = %q, want original route model: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.#(googleSearch)").Raw; got != "" {
|
||||
t.Fatalf("typed web_search should not become native googleSearch for unsupported route model: %s", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_DoesNotMapTypedWebSearchForFlashAgentWithoutCapability(t *testing.T) {
|
||||
registry.GetGlobalRegistry().RegisterClient("test-antigravity-claude-websearch-flash-agent", "antigravity", []*registry.ModelInfo{
|
||||
{ID: "gemini-3-flash-agent"},
|
||||
{ID: "gemini-3.1-flash-lite", SupportsWebSearch: true},
|
||||
})
|
||||
t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("test-antigravity-claude-websearch-flash-agent") })
|
||||
|
||||
inputJSON := []byte(`{
|
||||
"model": "gemini-3-flash-agent",
|
||||
"messages": [{"role": "user", "content": "Perform a web search"}],
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("gemini-3-flash-agent", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "model").String(); got != "gemini-3-flash-agent" {
|
||||
t.Fatalf("web search request model = %q, want original route model: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "request.tools.#(googleSearch)").Raw; got != "" {
|
||||
t.Fatalf("typed web_search should not become native googleSearch for flash-agent without capability: %s", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertClaudeRequestToAntigravity_DoesNotMapTypedWebSearchForOtherModels(t *testing.T) {
|
||||
inputJSON := []byte(`{
|
||||
"model": "claude-sonnet-4-6",
|
||||
"messages": [{"role": "user", "content": "Search current weather"}],
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}]
|
||||
}`)
|
||||
|
||||
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-6", inputJSON, true)
|
||||
if got := gjson.GetBytes(output, "request.tools.#(googleSearch)").Raw; got != "" {
|
||||
t.Fatalf("model without Antigravity web search capability should not get native googleSearch: %s", output)
|
||||
}
|
||||
}
|
||||
|
||||
func testNonAnthropicRawSignature(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
|
||||
@@ -69,6 +69,9 @@ type Params struct {
|
||||
HasSentFinalEvents bool // Indicates if final content/message events have been sent
|
||||
HasToolUse bool // Indicates if tool use was observed in the stream
|
||||
HasContent bool // Tracks whether any content (text, thinking, or tool use) has been output
|
||||
HasWebSearchTool bool
|
||||
WebSearchRequests int64
|
||||
WebSearchTextBuffer strings.Builder
|
||||
|
||||
// Signature caching support
|
||||
CurrentThinkingText strings.Builder // Accumulates thinking text for signature caching
|
||||
@@ -125,6 +128,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
appendEvent := func(event, payload string) {
|
||||
output = translatorcommon.AppendSSEEventString(output, event, payload, 3)
|
||||
}
|
||||
webSearchStreamMode := shouldTranslateWebSearchGrounding(originalRequestRawJSON, requestRawJSON)
|
||||
appendThinkingSignature := func(signature string) {
|
||||
if signature == "" || params.ResponseType != 2 {
|
||||
return
|
||||
@@ -150,7 +154,7 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
if promptTokenCount := gjson.GetBytes(rawJSON, "response.cpaUsageMetadata.promptTokenCount"); promptTokenCount.Exists() {
|
||||
messageStartTemplate, _ = sjson.SetBytes(messageStartTemplate, "message.usage.input_tokens", promptTokenCount.Int())
|
||||
}
|
||||
if candidatesTokenCount := gjson.GetBytes(rawJSON, "response.cpaUsageMetadata.candidatesTokenCount"); candidatesTokenCount.Exists() {
|
||||
if candidatesTokenCount := gjson.GetBytes(rawJSON, "response.cpaUsageMetadata.candidatesTokenCount"); candidatesTokenCount.Exists() && !webSearchStreamMode {
|
||||
messageStartTemplate, _ = sjson.SetBytes(messageStartTemplate, "message.usage.output_tokens", candidatesTokenCount.Int())
|
||||
}
|
||||
|
||||
@@ -166,10 +170,28 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
params.HasFirstResponse = true
|
||||
}
|
||||
|
||||
handledWebSearchGrounding := false
|
||||
if webSearchStreamMode && !params.HasWebSearchTool {
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
if groundingMetadata := antigravityGroundingMetadata(root); groundingMetadata.Exists() {
|
||||
toolUseID := newClaudeWebSearchToolUseID()
|
||||
textContent := params.WebSearchTextBuffer.String() + antigravityTextContent(root)
|
||||
params.WebSearchTextBuffer.Reset()
|
||||
params.ResponseIndex = appendClaudeWebSearchStreamBlocks(appendEvent, params.ResponseIndex, toolUseID, textContent, groundingMetadata)
|
||||
params.HasWebSearchTool = true
|
||||
params.WebSearchRequests = 1
|
||||
params.HasContent = true
|
||||
params.ResponseType = 0
|
||||
handledWebSearchGrounding = true
|
||||
}
|
||||
}
|
||||
|
||||
// Process the response parts array from the backend client
|
||||
// Each part can contain text content, thinking content, or function calls
|
||||
partsResult := gjson.GetBytes(rawJSON, "response.candidates.0.content.parts")
|
||||
if partsResult.IsArray() {
|
||||
if partsResult.IsArray() && webSearchStreamMode && !params.HasWebSearchTool && !handledWebSearchGrounding {
|
||||
appendWebSearchBufferedText(partsResult, ¶ms.WebSearchTextBuffer)
|
||||
} else if partsResult.IsArray() && !handledWebSearchGrounding {
|
||||
partResults := partsResult.Array()
|
||||
for i := 0; i < len(partResults); i++ {
|
||||
partResult := partResults[i]
|
||||
@@ -337,6 +359,10 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
}
|
||||
}
|
||||
|
||||
if webSearchStreamMode && !params.HasWebSearchTool && params.HasFinishReason && params.WebSearchTextBuffer.Len() > 0 {
|
||||
appendBufferedWebSearchTextBlock(params, appendEvent)
|
||||
}
|
||||
|
||||
if params.HasUsageMetadata && params.HasFinishReason {
|
||||
appendFinalEvents(params, &output, false)
|
||||
}
|
||||
@@ -344,6 +370,30 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
return [][]byte{output}
|
||||
}
|
||||
|
||||
func appendWebSearchBufferedText(partsResult gjson.Result, buffer *strings.Builder) {
|
||||
for _, partResult := range partsResult.Array() {
|
||||
if partResult.Get("thought").Bool() || partResult.Get("functionCall").Exists() {
|
||||
continue
|
||||
}
|
||||
if partTextResult := partResult.Get("text"); partTextResult.Exists() {
|
||||
buffer.WriteString(partTextResult.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func appendBufferedWebSearchTextBlock(params *Params, appendEvent func(string, string)) {
|
||||
text := params.WebSearchTextBuffer.String()
|
||||
params.WebSearchTextBuffer.Reset()
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
appendEvent("content_block_start", fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, params.ResponseIndex))
|
||||
data, _ := sjson.SetBytes([]byte(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, params.ResponseIndex)), "delta.text", text)
|
||||
appendEvent("content_block_delta", string(data))
|
||||
params.ResponseType = 1
|
||||
params.HasContent = true
|
||||
}
|
||||
|
||||
func appendFinalEvents(params *Params, output *[]byte, force bool) {
|
||||
if params.HasSentFinalEvents {
|
||||
return
|
||||
@@ -373,6 +423,9 @@ func appendFinalEvents(params *Params, output *[]byte, force bool) {
|
||||
}
|
||||
|
||||
delta := []byte(fmt.Sprintf(`{"type":"message_delta","delta":{"stop_reason":"%s","stop_sequence":null},"usage":{"input_tokens":%d,"output_tokens":%d}}`, stopReason, params.PromptTokenCount, usageOutputTokens))
|
||||
if params.WebSearchRequests > 0 {
|
||||
delta, _ = sjson.SetBytes(delta, "usage.server_tool_use.web_search_requests", params.WebSearchRequests)
|
||||
}
|
||||
// Add cache_read_input_tokens if cached tokens are present (indicates prompt caching is working)
|
||||
if params.CachedTokenCount > 0 {
|
||||
var err error
|
||||
@@ -443,6 +496,16 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
|
||||
}
|
||||
}
|
||||
|
||||
if shouldTranslateWebSearchGrounding(originalRequestRawJSON, requestRawJSON) {
|
||||
if groundingMetadata := antigravityGroundingMetadata(root); groundingMetadata.Exists() {
|
||||
toolUseID := newClaudeWebSearchToolUseID()
|
||||
responseJSON, _ = sjson.SetRawBytes(responseJSON, "content", buildClaudeWebSearchContent(toolUseID, antigravityTextContent(root), groundingMetadata))
|
||||
responseJSON, _ = sjson.SetBytes(responseJSON, "stop_reason", "end_turn")
|
||||
responseJSON, _ = sjson.SetBytes(responseJSON, "usage.server_tool_use.web_search_requests", 1)
|
||||
return responseJSON
|
||||
}
|
||||
}
|
||||
|
||||
contentArrayInitialized := false
|
||||
ensureContentArray := func() {
|
||||
if contentArrayInitialized {
|
||||
|
||||
@@ -3,6 +3,7 @@ package claude
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -14,6 +15,296 @@ import (
|
||||
// Signature Caching Tests
|
||||
// ============================================================================
|
||||
|
||||
func TestConvertAntigravityResponseToClaudeNonStream_WebSearchGrounding(t *testing.T) {
|
||||
requestJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
translatedRequestJSON := []byte(`{"model":"gemini-3.1-flash-lite","request":{"tools":[{"googleSearch":{}}]}}`)
|
||||
responseJSON := testAntigravityGroundingResponse()
|
||||
|
||||
output := ConvertAntigravityResponseToClaudeNonStream(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, responseJSON, nil)
|
||||
|
||||
if got := gjson.GetBytes(output, "content.0.type").String(); got != "server_tool_use" {
|
||||
t.Fatalf("first content block = %q, want server_tool_use: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "content.1.type").String(); got != "web_search_tool_result" {
|
||||
t.Fatalf("second content block = %q, want web_search_tool_result: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "usage.server_tool_use.web_search_requests").Int(); got != 1 {
|
||||
t.Fatalf("web_search_requests = %d, want 1: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "content.1.content.0.url").String(); got != "https://example.com/weather" {
|
||||
t.Fatalf("search result url = %q: %s", got, output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "content.2.citations.0.url").String(); got != "https://example.com/weather" {
|
||||
t.Fatalf("citation url = %q: %s", got, output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertAntigravityResponseToClaudeNonStream_WebSearchGroundingRequiresNativeGoogleSearch(t *testing.T) {
|
||||
requestJSON := []byte(`{
|
||||
"model": "gemini-3-flash-agent",
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
translatedRequestJSON := []byte(`{"model":"gemini-3-flash-agent","request":{"contents":[]}}`)
|
||||
responseJSON := testAntigravityGroundingResponse()
|
||||
|
||||
output := ConvertAntigravityResponseToClaudeNonStream(context.Background(), "gemini-3-flash-agent", requestJSON, translatedRequestJSON, responseJSON, nil)
|
||||
|
||||
if got := gjson.GetBytes(output, "content.0.type").String(); got == "server_tool_use" {
|
||||
t.Fatalf("non-native translated request should not synthesize server_tool_use: %s", output)
|
||||
}
|
||||
if got := gjson.GetBytes(output, "usage.server_tool_use.web_search_requests").Int(); got != 0 {
|
||||
t.Fatalf("web_search_requests = %d, want 0: %s", got, output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertAntigravityResponseToClaudeStream_WebSearchGrounding(t *testing.T) {
|
||||
requestJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
translatedRequestJSON := []byte(`{"model":"gemini-3.1-flash-lite","request":{"tools":[{"googleSearch":{}}]}}`)
|
||||
|
||||
var param any
|
||||
output := bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, testAntigravityGroundingResponse(), ¶m), nil)
|
||||
output = append(output, bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, []byte("[DONE]"), ¶m), nil)...)
|
||||
outputText := string(output)
|
||||
|
||||
for _, needle := range []string{
|
||||
`"type":"server_tool_use"`,
|
||||
`"type":"web_search_tool_result"`,
|
||||
`"web_search_requests":1`,
|
||||
`"type":"citations_delta"`,
|
||||
`event: message_stop`,
|
||||
} {
|
||||
if !strings.Contains(outputText, needle) {
|
||||
t.Fatalf("stream output missing %s:\n%s", needle, outputText)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertAntigravityResponseToClaudeStream_WebSearchBuffersTextUntilGrounding(t *testing.T) {
|
||||
requestJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
translatedRequestJSON := []byte(`{"model":"gemini-3.1-flash-lite","request":{"tools":[{"googleSearch":{}}]}}`)
|
||||
|
||||
var param any
|
||||
firstChunk := []byte(`{
|
||||
"response": {
|
||||
"modelVersion": "gemini-3.1-flash-lite",
|
||||
"responseId": "resp-web-search-stream",
|
||||
"candidates": [{
|
||||
"content": {
|
||||
"parts": [{"text": "Beijing weather "}]
|
||||
}
|
||||
}],
|
||||
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 2, "totalTokenCount": 12}
|
||||
}
|
||||
}`)
|
||||
finalChunk := []byte(`{
|
||||
"response": {
|
||||
"modelVersion": "gemini-3.1-flash-lite",
|
||||
"responseId": "resp-web-search-stream",
|
||||
"candidates": [{
|
||||
"content": {
|
||||
"parts": [{"text": "is clear today."}]
|
||||
},
|
||||
"groundingMetadata": {
|
||||
"webSearchQueries": ["Beijing weather"],
|
||||
"groundingChunks": [{"web": {"uri": "https://example.com/weather", "title": "Beijing Weather"}}],
|
||||
"groundingSupports": [{
|
||||
"segment": {"startIndex": 0, "endIndex": 31, "text": "Beijing weather is clear today."},
|
||||
"groundingChunkIndices": [0]
|
||||
}]
|
||||
},
|
||||
"finishReason": "STOP"
|
||||
}],
|
||||
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 6, "totalTokenCount": 16}
|
||||
}
|
||||
}`)
|
||||
|
||||
output := bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, firstChunk, ¶m), nil)
|
||||
output = append(output, bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, finalChunk, ¶m), nil)...)
|
||||
output = append(output, bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, []byte("[DONE]"), ¶m), nil)...)
|
||||
outputText := string(output)
|
||||
|
||||
textStart := strings.Index(outputText, `"content_block":{"type":"text"`)
|
||||
serverToolStart := strings.Index(outputText, `"content_block":{"type":"server_tool_use"`)
|
||||
if serverToolStart < 0 {
|
||||
t.Fatalf("stream output missing server_tool_use:\n%s", outputText)
|
||||
}
|
||||
if textStart >= 0 && textStart < serverToolStart {
|
||||
t.Fatalf("text block was emitted before server_tool_use:\n%s", outputText)
|
||||
}
|
||||
if strings.Contains(outputText, `"index":0,"content_block":{"type":"text"`) {
|
||||
t.Fatalf("index 0 must be reserved for server_tool_use:\n%s", outputText)
|
||||
}
|
||||
if !strings.Contains(outputText, `"index":0,"content_block":{"type":"server_tool_use"`) {
|
||||
t.Fatalf("server_tool_use must use index 0:\n%s", outputText)
|
||||
}
|
||||
if !strings.Contains(outputText, `"index":1,"content_block":{"type":"web_search_tool_result"`) {
|
||||
t.Fatalf("web_search_tool_result must use index 1:\n%s", outputText)
|
||||
}
|
||||
if !strings.Contains(outputText, `Beijing weather is clear today.`) {
|
||||
t.Fatalf("buffered text was not emitted after web search blocks:\n%s", outputText)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertAntigravityResponseToClaudeStream_WebSearchMessageStartOutputTokensZero(t *testing.T) {
|
||||
requestJSON := []byte(`{
|
||||
"model": "gemini-3.1-flash-lite",
|
||||
"tools": [{"type": "web_search_20250305", "name": "web_search"}]
|
||||
}`)
|
||||
translatedRequestJSON := []byte(`{"model":"gemini-3.1-flash-lite","request":{"tools":[{"googleSearch":{}}]}}`)
|
||||
responseJSON := []byte(`{
|
||||
"response": {
|
||||
"modelVersion": "gemini-3.1-flash-lite",
|
||||
"responseId": "resp-web-search-start",
|
||||
"candidates": [{
|
||||
"content": {"parts": [{"text": "Beijing weather"}]}
|
||||
}],
|
||||
"cpaUsageMetadata": {"promptTokenCount": 85, "candidatesTokenCount": 43}
|
||||
}
|
||||
}`)
|
||||
|
||||
var param any
|
||||
output := bytes.Join(ConvertAntigravityResponseToClaude(context.Background(), "gemini-3.1-flash-lite", requestJSON, translatedRequestJSON, responseJSON, ¶m), nil)
|
||||
messageStart := sseDataForEvent(t, string(output), "message_start")
|
||||
|
||||
if got := gjson.Get(messageStart, "message.usage.output_tokens").Int(); got != 0 {
|
||||
t.Fatalf("message_start output_tokens = %d, want 0: %s", got, messageStart)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearchResultsFromGrounding_DeduplicatesAndSkipsEmptyURLs(t *testing.T) {
|
||||
groundingMetadata := gjson.Parse(`{
|
||||
"groundingChunks": [
|
||||
{"web": {"uri": "https://example.com/a", "title": "A"}},
|
||||
{"web": {"uri": "https://example.com/b", "title": "B"}},
|
||||
{"web": {"uri": "https://example.com/a", "title": "A duplicate"}},
|
||||
{"web": {"uri": "", "title": "Empty"}}
|
||||
]
|
||||
}`)
|
||||
|
||||
results := webSearchResultsFromGrounding(groundingMetadata)
|
||||
|
||||
if got := gjson.GetBytes(results, "#").Int(); got != 2 {
|
||||
t.Fatalf("result count = %d, want 2: %s", got, string(results))
|
||||
}
|
||||
if got := gjson.GetBytes(results, "0.url").String(); got != "https://example.com/a" {
|
||||
t.Fatalf("first url = %q: %s", got, string(results))
|
||||
}
|
||||
if got := gjson.GetBytes(results, "1.url").String(); got != "https://example.com/b" {
|
||||
t.Fatalf("second url = %q: %s", got, string(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildWebSearchCitedTextBlocks_TrimsOverlappingGroundingSupports(t *testing.T) {
|
||||
first := "北京今天晴"
|
||||
second := "北京今天晴,气温19到31度"
|
||||
textContent := second + "。"
|
||||
|
||||
blocks := buildWebSearchCitedTextBlocks(textContent, []webSearchGroundingSupport{
|
||||
{
|
||||
StartIndex: 0,
|
||||
EndIndex: int64(len([]byte(first))),
|
||||
Text: first,
|
||||
ChunkURLs: []string{"https://example.com/weather"},
|
||||
ChunkTitle: "Weather",
|
||||
},
|
||||
{
|
||||
StartIndex: 0,
|
||||
EndIndex: int64(len([]byte(second))),
|
||||
Text: second,
|
||||
ChunkURLs: []string{"https://example.com/weather"},
|
||||
ChunkTitle: "Weather",
|
||||
},
|
||||
})
|
||||
|
||||
var got strings.Builder
|
||||
for _, block := range blocks {
|
||||
got.WriteString(block.Text)
|
||||
}
|
||||
if got.String() != textContent {
|
||||
t.Fatalf("joined text = %q, want %q", got.String(), textContent)
|
||||
}
|
||||
if len(blocks) < 2 || blocks[1].Text != ",气温19到31度" {
|
||||
t.Fatalf("overlap suffix block not trimmed correctly: %#v", blocks)
|
||||
}
|
||||
if gotCitation := blocks[1].Citations[0]["cited_text"]; gotCitation != blocks[1].Text {
|
||||
t.Fatalf("cited_text = %q, want emitted text %q", gotCitation, blocks[1].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func sseDataForEvent(t *testing.T, output string, eventName string) string {
|
||||
t.Helper()
|
||||
|
||||
currentEvent := ""
|
||||
for _, line := range strings.Split(output, "\n") {
|
||||
if strings.HasPrefix(line, "event: ") {
|
||||
currentEvent = strings.TrimPrefix(line, "event: ")
|
||||
continue
|
||||
}
|
||||
if currentEvent == eventName && strings.HasPrefix(line, "data: ") {
|
||||
return strings.TrimPrefix(line, "data: ")
|
||||
}
|
||||
}
|
||||
|
||||
t.Fatalf("event %q not found in:\n%s", eventName, output)
|
||||
return ""
|
||||
}
|
||||
|
||||
func testAntigravityGroundingResponse() []byte {
|
||||
resp := map[string]any{
|
||||
"response": map[string]any{
|
||||
"responseId": "resp-web-search",
|
||||
"modelVersion": "gemini-3.1-flash-lite",
|
||||
"candidates": []any{
|
||||
map[string]any{
|
||||
"content": map[string]any{
|
||||
"parts": []any{
|
||||
map[string]any{"text": "Beijing weather is clear today."},
|
||||
},
|
||||
},
|
||||
"groundingMetadata": map[string]any{
|
||||
"webSearchQueries": []any{"Beijing weather June 10 2026"},
|
||||
"groundingChunks": []any{
|
||||
map[string]any{
|
||||
"web": map[string]any{
|
||||
"uri": "https://example.com/weather",
|
||||
"title": "Beijing Weather",
|
||||
},
|
||||
},
|
||||
},
|
||||
"groundingSupports": []any{
|
||||
map[string]any{
|
||||
"segment": map[string]any{
|
||||
"startIndex": int64(0),
|
||||
"endIndex": int64(31),
|
||||
"text": "Beijing weather is clear today.",
|
||||
},
|
||||
"groundingChunkIndices": []any{0},
|
||||
},
|
||||
},
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
},
|
||||
},
|
||||
"usageMetadata": map[string]any{
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 6,
|
||||
"totalTokenCount": 16,
|
||||
},
|
||||
},
|
||||
}
|
||||
raw, _ := json.Marshal(resp)
|
||||
return raw
|
||||
}
|
||||
|
||||
func TestConvertAntigravityResponseToClaude_ParamsInitialized(t *testing.T) {
|
||||
cache.ClearSignatureCache("")
|
||||
|
||||
|
||||
502
internal/translator/antigravity/claude/web_search.go
Normal file
502
internal/translator/antigravity/claude/web_search.go
Normal file
@@ -0,0 +1,502 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
type webSearchGroundingSupport struct {
|
||||
StartIndex int64
|
||||
EndIndex int64
|
||||
Text string
|
||||
ChunkURLs []string
|
||||
ChunkTitle string
|
||||
}
|
||||
|
||||
type webSearchCitedTextBlock struct {
|
||||
Text string
|
||||
Citations []map[string]any
|
||||
}
|
||||
|
||||
const antigravityWebSearchSystemInstruction = "You are a search engine bot. You will be given a query from a user. Your task is to search the web for relevant information that will help the user. You MUST perform a web search. Do not respond or interact with the user, please respond as if they typed the query into a search bar."
|
||||
|
||||
func antigravitySupportsNativeGoogleSearch(model string) bool {
|
||||
return registry.AntigravityWebSearchModelFor(model) != ""
|
||||
}
|
||||
|
||||
func isClaudeTypedWebSearchToolType(toolType string) bool {
|
||||
return toolType == "web_search_20250305" || toolType == "web_search_20260209"
|
||||
}
|
||||
|
||||
func hasClaudeTypedWebSearchTool(payload []byte) bool {
|
||||
tools := gjson.GetBytes(payload, "tools")
|
||||
if !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if isClaudeTypedWebSearchToolType(tool.Get("type").String()) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func hasOnlyClaudeTypedWebSearchTools(payload []byte) bool {
|
||||
tools := gjson.GetBytes(payload, "tools")
|
||||
if !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
hasWebSearch := false
|
||||
for _, tool := range tools.Array() {
|
||||
if isClaudeTypedWebSearchToolType(tool.Get("type").String()) {
|
||||
hasWebSearch = true
|
||||
continue
|
||||
}
|
||||
return false
|
||||
}
|
||||
return hasWebSearch
|
||||
}
|
||||
|
||||
func allowsClaudeWebSearchToolChoice(payload []byte) bool {
|
||||
toolChoice := gjson.GetBytes(payload, "tool_choice")
|
||||
if !toolChoice.Exists() {
|
||||
return true
|
||||
}
|
||||
if toolChoice.Type == gjson.String {
|
||||
switch toolChoice.String() {
|
||||
case "", "auto", "any":
|
||||
return true
|
||||
case "none":
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
if !toolChoice.IsObject() {
|
||||
return false
|
||||
}
|
||||
switch toolChoice.Get("type").String() {
|
||||
case "", "auto", "any":
|
||||
return true
|
||||
case "tool":
|
||||
return toolChoice.Get("name").String() == "web_search"
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func shouldBuildAntigravityWebSearchRequest(model string, payload []byte) bool {
|
||||
return antigravitySupportsNativeGoogleSearch(model) &&
|
||||
hasOnlyClaudeTypedWebSearchTools(payload) &&
|
||||
allowsClaudeWebSearchToolChoice(payload)
|
||||
}
|
||||
|
||||
func buildAntigravityWebSearchRequest(model string, payload []byte) []byte {
|
||||
query := extractClaudeWebSearchQuery(payload)
|
||||
maxResultCount := extractClaudeWebSearchMaxUses(payload)
|
||||
includedDomains := extractClaudeWebSearchAllowedDomains(payload)
|
||||
out := []byte(`{"model":"","requestType":"web_search","request":{"contents":[{"role":"user","parts":[{"text":""}]}],"systemInstruction":{"role":"user","parts":[{"text":""}]},"tools":[{"googleSearch":{"enhancedContent":{"imageSearch":{"maxResultCount":5}}}}],"generationConfig":{"candidateCount":1}}}`)
|
||||
out, _ = sjson.SetBytes(out, "model", model)
|
||||
out, _ = sjson.SetBytes(out, "request.contents.0.parts.0.text", query)
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.parts.0.text", antigravityWebSearchSystemInstruction)
|
||||
out, _ = sjson.SetBytes(out, "request.tools.0.googleSearch.enhancedContent.imageSearch.maxResultCount", maxResultCount)
|
||||
if len(includedDomains) > 0 {
|
||||
if domainsJSON, err := json.Marshal(includedDomains); err == nil {
|
||||
out, _ = sjson.SetRawBytes(out, "request.tools.0.googleSearch.includedDomains", domainsJSON)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func extractClaudeWebSearchMaxUses(payload []byte) int64 {
|
||||
const defaultMaxResultCount int64 = 5
|
||||
|
||||
tools := gjson.GetBytes(payload, "tools")
|
||||
if !tools.IsArray() {
|
||||
return defaultMaxResultCount
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if !isClaudeTypedWebSearchToolType(tool.Get("type").String()) {
|
||||
continue
|
||||
}
|
||||
maxUses := tool.Get("max_uses").Int()
|
||||
if maxUses > 0 {
|
||||
return maxUses
|
||||
}
|
||||
}
|
||||
return defaultMaxResultCount
|
||||
}
|
||||
|
||||
func extractClaudeWebSearchAllowedDomains(payload []byte) []string {
|
||||
tools := gjson.GetBytes(payload, "tools")
|
||||
if !tools.IsArray() {
|
||||
return nil
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if !isClaudeTypedWebSearchToolType(tool.Get("type").String()) {
|
||||
continue
|
||||
}
|
||||
allowedDomains := tool.Get("allowed_domains")
|
||||
if !allowedDomains.IsArray() {
|
||||
return nil
|
||||
}
|
||||
domains := make([]string, 0, len(allowedDomains.Array()))
|
||||
for _, domain := range allowedDomains.Array() {
|
||||
if domain.Type != gjson.String {
|
||||
continue
|
||||
}
|
||||
if trimmed := strings.TrimSpace(domain.String()); trimmed != "" {
|
||||
domains = append(domains, trimmed)
|
||||
}
|
||||
}
|
||||
return domains
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func extractClaudeWebSearchQuery(payload []byte) string {
|
||||
messages := gjson.GetBytes(payload, "messages")
|
||||
if !messages.IsArray() {
|
||||
return ""
|
||||
}
|
||||
messageResults := messages.Array()
|
||||
for i := len(messageResults) - 1; i >= 0; i-- {
|
||||
message := messageResults[i]
|
||||
if role := message.Get("role").String(); role != "" && role != "user" {
|
||||
continue
|
||||
}
|
||||
if query := extractClaudeTextContent(message.Get("content")); query != "" {
|
||||
return query
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func extractClaudeTextContent(content gjson.Result) string {
|
||||
if content.Type == gjson.String {
|
||||
return strings.TrimSpace(content.String())
|
||||
}
|
||||
if !content.IsArray() {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, part := range content.Array() {
|
||||
if text := strings.TrimSpace(part.Get("text").String()); text != "" {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(text)
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(b.String())
|
||||
}
|
||||
|
||||
func hasAntigravityGoogleSearchTool(payload []byte) bool {
|
||||
tools := gjson.GetBytes(payload, "request.tools")
|
||||
if !tools.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, tool := range tools.Array() {
|
||||
if tool.Get("googleSearch").Exists() {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func shouldTranslateWebSearchGrounding(originalRequestRawJSON, requestRawJSON []byte) bool {
|
||||
return hasClaudeTypedWebSearchTool(originalRequestRawJSON) && hasAntigravityGoogleSearchTool(requestRawJSON)
|
||||
}
|
||||
|
||||
func antigravityGroundingMetadata(root gjson.Result) gjson.Result {
|
||||
groundingMetadata := root.Get("response.candidates.0.groundingMetadata")
|
||||
if groundingMetadata.Exists() {
|
||||
return groundingMetadata
|
||||
}
|
||||
return root.Get("candidates.0.groundingMetadata")
|
||||
}
|
||||
|
||||
func antigravityTextContent(root gjson.Result) string {
|
||||
var textBuilder strings.Builder
|
||||
parts := root.Get("response.candidates.0.content.parts")
|
||||
if !parts.IsArray() {
|
||||
parts = root.Get("candidates.0.content.parts")
|
||||
}
|
||||
if parts.IsArray() {
|
||||
for _, part := range parts.Array() {
|
||||
if text := part.Get("text"); text.Exists() {
|
||||
textBuilder.WriteString(text.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
return textBuilder.String()
|
||||
}
|
||||
|
||||
func antigravityUsageTokens(root gjson.Result) (int64, int64) {
|
||||
usage := root.Get("response.usageMetadata")
|
||||
if !usage.Exists() {
|
||||
usage = root.Get("usageMetadata")
|
||||
}
|
||||
inputTokens := usage.Get("promptTokenCount").Int()
|
||||
outputTokens := usage.Get("candidatesTokenCount").Int() + usage.Get("thoughtsTokenCount").Int()
|
||||
if outputTokens == 0 {
|
||||
totalTokens := usage.Get("totalTokenCount").Int()
|
||||
if totalTokens > 0 {
|
||||
outputTokens = totalTokens - inputTokens
|
||||
if outputTokens < 0 {
|
||||
outputTokens = 0
|
||||
}
|
||||
}
|
||||
}
|
||||
return inputTokens, outputTokens
|
||||
}
|
||||
|
||||
func webSearchQueryFromGrounding(groundingMetadata gjson.Result) string {
|
||||
if queries := groundingMetadata.Get("webSearchQueries"); queries.IsArray() && len(queries.Array()) > 0 {
|
||||
return queries.Array()[0].String()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func webSearchResultsFromGrounding(groundingMetadata gjson.Result) []byte {
|
||||
results := []byte(`[]`)
|
||||
groundingChunks := groundingMetadata.Get("groundingChunks")
|
||||
if !groundingChunks.IsArray() {
|
||||
return results
|
||||
}
|
||||
seenURLs := make(map[string]struct{})
|
||||
for _, chunk := range groundingChunks.Array() {
|
||||
web := chunk.Get("web")
|
||||
if !web.Exists() {
|
||||
continue
|
||||
}
|
||||
uri := strings.TrimSpace(web.Get("uri").String())
|
||||
if uri == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seenURLs[uri]; ok {
|
||||
continue
|
||||
}
|
||||
seenURLs[uri] = struct{}{}
|
||||
|
||||
result := []byte(`{"type":"web_search_result","page_age":null}`)
|
||||
if title := web.Get("title"); title.Exists() {
|
||||
result, _ = sjson.SetBytes(result, "title", title.String())
|
||||
}
|
||||
result, _ = sjson.SetBytes(result, "url", uri)
|
||||
results, _ = sjson.SetRawBytes(results, "-1", result)
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
func parseWebSearchGroundingSupports(groundingMetadata gjson.Result) []webSearchGroundingSupport {
|
||||
groundingChunks := groundingMetadata.Get("groundingChunks")
|
||||
if !groundingChunks.IsArray() {
|
||||
return nil
|
||||
}
|
||||
chunks := groundingChunks.Array()
|
||||
chunkData := make([]struct {
|
||||
URL string
|
||||
Title string
|
||||
}, len(chunks))
|
||||
for i, chunk := range chunks {
|
||||
web := chunk.Get("web")
|
||||
if web.Exists() {
|
||||
chunkData[i].URL = web.Get("uri").String()
|
||||
chunkData[i].Title = web.Get("title").String()
|
||||
}
|
||||
}
|
||||
|
||||
groundingSupports := groundingMetadata.Get("groundingSupports")
|
||||
if !groundingSupports.IsArray() {
|
||||
return nil
|
||||
}
|
||||
supports := make([]webSearchGroundingSupport, 0, len(groundingSupports.Array()))
|
||||
for _, support := range groundingSupports.Array() {
|
||||
segment := support.Get("segment")
|
||||
if !segment.Exists() {
|
||||
continue
|
||||
}
|
||||
parsed := webSearchGroundingSupport{
|
||||
StartIndex: segment.Get("startIndex").Int(),
|
||||
EndIndex: segment.Get("endIndex").Int(),
|
||||
Text: segment.Get("text").String(),
|
||||
}
|
||||
if chunkIndices := support.Get("groundingChunkIndices"); chunkIndices.IsArray() {
|
||||
for _, idx := range chunkIndices.Array() {
|
||||
chunkIndex := int(idx.Int())
|
||||
if chunkIndex < 0 || chunkIndex >= len(chunkData) {
|
||||
continue
|
||||
}
|
||||
parsed.ChunkURLs = append(parsed.ChunkURLs, chunkData[chunkIndex].URL)
|
||||
if parsed.ChunkTitle == "" {
|
||||
parsed.ChunkTitle = chunkData[chunkIndex].Title
|
||||
}
|
||||
}
|
||||
}
|
||||
supports = append(supports, parsed)
|
||||
}
|
||||
return supports
|
||||
}
|
||||
|
||||
func buildWebSearchCitedTextBlocks(textContent string, supports []webSearchGroundingSupport) []webSearchCitedTextBlock {
|
||||
if len(supports) == 0 {
|
||||
if textContent == "" {
|
||||
return nil
|
||||
}
|
||||
return []webSearchCitedTextBlock{{Text: textContent}}
|
||||
}
|
||||
|
||||
textBytes := []byte(textContent)
|
||||
blocks := make([]webSearchCitedTextBlock, 0, len(supports)+1)
|
||||
lastEnd := int64(0)
|
||||
for _, support := range supports {
|
||||
if support.EndIndex <= lastEnd {
|
||||
continue
|
||||
}
|
||||
if support.StartIndex > lastEnd {
|
||||
start := int(lastEnd)
|
||||
end := min(int(support.StartIndex), len(textBytes))
|
||||
if start < end {
|
||||
blocks = append(blocks, webSearchCitedTextBlock{Text: string(textBytes[start:end])})
|
||||
}
|
||||
}
|
||||
|
||||
citedStart := support.StartIndex
|
||||
if citedStart < lastEnd {
|
||||
citedStart = lastEnd
|
||||
}
|
||||
citedText := ""
|
||||
if citedStart < support.EndIndex {
|
||||
start := min(int(citedStart), len(textBytes))
|
||||
end := min(int(support.EndIndex), len(textBytes))
|
||||
if start < end {
|
||||
citedText = string(textBytes[start:end])
|
||||
}
|
||||
}
|
||||
if citedText != "" && len(support.ChunkURLs) > 0 {
|
||||
citation := map[string]any{
|
||||
"type": "web_search_result_location",
|
||||
"cited_text": citedText,
|
||||
"url": support.ChunkURLs[0],
|
||||
"title": support.ChunkTitle,
|
||||
}
|
||||
blocks = append(blocks, webSearchCitedTextBlock{
|
||||
Text: citedText,
|
||||
Citations: []map[string]any{citation},
|
||||
})
|
||||
}
|
||||
if support.EndIndex > lastEnd {
|
||||
lastEnd = support.EndIndex
|
||||
}
|
||||
}
|
||||
if int(lastEnd) < len(textBytes) {
|
||||
blocks = append(blocks, webSearchCitedTextBlock{Text: string(textBytes[lastEnd:])})
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
func buildClaudeWebSearchContent(toolUseID string, textContent string, groundingMetadata gjson.Result) []byte {
|
||||
content := []byte(`[]`)
|
||||
|
||||
serverToolUse := []byte(`{"type":"server_tool_use","id":"","name":"web_search","input":{}}`)
|
||||
serverToolUse, _ = sjson.SetBytes(serverToolUse, "id", toolUseID)
|
||||
if query := webSearchQueryFromGrounding(groundingMetadata); query != "" {
|
||||
serverToolUse, _ = sjson.SetBytes(serverToolUse, "input.query", query)
|
||||
}
|
||||
content, _ = sjson.SetRawBytes(content, "-1", serverToolUse)
|
||||
|
||||
webSearchToolResult := []byte(`{"type":"web_search_tool_result","tool_use_id":"","content":[]}`)
|
||||
webSearchToolResult, _ = sjson.SetBytes(webSearchToolResult, "tool_use_id", toolUseID)
|
||||
webSearchToolResult, _ = sjson.SetRawBytes(webSearchToolResult, "content", webSearchResultsFromGrounding(groundingMetadata))
|
||||
content, _ = sjson.SetRawBytes(content, "-1", webSearchToolResult)
|
||||
|
||||
for _, block := range buildWebSearchCitedTextBlocks(textContent, parseWebSearchGroundingSupports(groundingMetadata)) {
|
||||
if block.Text == "" {
|
||||
continue
|
||||
}
|
||||
textBlock := []byte(`{"type":"text","text":""}`)
|
||||
textBlock, _ = sjson.SetBytes(textBlock, "text", block.Text)
|
||||
if len(block.Citations) > 0 {
|
||||
citationsJSON, _ := json.Marshal(block.Citations)
|
||||
textBlock, _ = sjson.SetRawBytes(textBlock, "citations", citationsJSON)
|
||||
}
|
||||
content, _ = sjson.SetRawBytes(content, "-1", textBlock)
|
||||
}
|
||||
|
||||
return content
|
||||
}
|
||||
|
||||
func appendClaudeWebSearchStreamBlocks(appendEvent func(string, string), startIndex int, toolUseID string, textContent string, groundingMetadata gjson.Result) int {
|
||||
contentIndex := startIndex
|
||||
|
||||
serverToolUseStart := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"server_tool_use","id":"%s","name":"web_search","input":{}}}`,
|
||||
contentIndex, toolUseID)
|
||||
appendEvent("content_block_start", serverToolUseStart)
|
||||
if query := webSearchQueryFromGrounding(groundingMetadata); query != "" {
|
||||
queryJSON, _ := sjson.Set(`{}`, "query", query)
|
||||
inputDelta := fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"input_json_delta","partial_json":""}}`, contentIndex)
|
||||
inputDelta, _ = sjson.Set(inputDelta, "delta.partial_json", queryJSON)
|
||||
appendEvent("content_block_delta", inputDelta)
|
||||
}
|
||||
appendEvent("content_block_stop", fmt.Sprintf(`{"type":"content_block_stop","index":%d}`, contentIndex))
|
||||
contentIndex++
|
||||
|
||||
webSearchToolResultStart := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"web_search_tool_result","tool_use_id":"%s","content":[]}}`,
|
||||
contentIndex, toolUseID)
|
||||
webSearchToolResultStart, _ = sjson.SetRaw(webSearchToolResultStart, "content_block.content", string(webSearchResultsFromGrounding(groundingMetadata)))
|
||||
appendEvent("content_block_start", webSearchToolResultStart)
|
||||
appendEvent("content_block_stop", fmt.Sprintf(`{"type":"content_block_stop","index":%d}`, contentIndex))
|
||||
contentIndex++
|
||||
|
||||
for _, block := range buildWebSearchCitedTextBlocks(textContent, parseWebSearchGroundingSupports(groundingMetadata)) {
|
||||
if block.Text == "" {
|
||||
continue
|
||||
}
|
||||
textBlockStart := fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"text","text":""}}`, contentIndex)
|
||||
if len(block.Citations) > 0 {
|
||||
textBlockStart = fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"citations":[],"type":"text","text":""}}`, contentIndex)
|
||||
}
|
||||
appendEvent("content_block_start", textBlockStart)
|
||||
for _, citation := range block.Citations {
|
||||
citationJSON, _ := json.Marshal(citation)
|
||||
citationDelta := fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"citations_delta","citation":%s}}`, contentIndex, string(citationJSON))
|
||||
appendEvent("content_block_delta", citationDelta)
|
||||
}
|
||||
for _, chunk := range splitRunesForWebSearch(block.Text, 50) {
|
||||
textDelta := fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"text_delta","text":""}}`, contentIndex)
|
||||
textDelta, _ = sjson.Set(textDelta, "delta.text", chunk)
|
||||
appendEvent("content_block_delta", textDelta)
|
||||
}
|
||||
appendEvent("content_block_stop", fmt.Sprintf(`{"type":"content_block_stop","index":%d}`, contentIndex))
|
||||
contentIndex++
|
||||
}
|
||||
|
||||
return contentIndex
|
||||
}
|
||||
|
||||
func splitRunesForWebSearch(text string, chunkSize int) []string {
|
||||
if chunkSize <= 0 || text == "" {
|
||||
return nil
|
||||
}
|
||||
runes := []rune(text)
|
||||
chunks := make([]string, 0, (len(runes)+chunkSize-1)/chunkSize)
|
||||
for start := 0; start < len(runes); start += chunkSize {
|
||||
end := start + chunkSize
|
||||
if end > len(runes) {
|
||||
end = len(runes)
|
||||
}
|
||||
chunks = append(chunks, string(runes[start:end]))
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
func newClaudeWebSearchToolUseID() string {
|
||||
return fmt.Sprintf("srvtoolu_%d", time.Now().UnixNano())
|
||||
}
|
||||
150
sdk/cliproxy/antigravity_models.go
Normal file
150
sdk/cliproxy/antigravity_models.go
Normal file
@@ -0,0 +1,150 @@
|
||||
package cliproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/misc"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityModelBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
|
||||
antigravityModelBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
)
|
||||
|
||||
type antigravityFetchAvailableModelsResponse struct {
|
||||
WebSearchModelIDs []string `json:"webSearchModelIds"`
|
||||
}
|
||||
|
||||
type antigravityModelCapabilityHints struct {
|
||||
WebSearchModelIDs map[string]struct{}
|
||||
}
|
||||
|
||||
func (s *Service) fetchAntigravityModelCapabilityHintsForAuth(ctx context.Context, auth *coreauth.Auth) antigravityModelCapabilityHints {
|
||||
if auth == nil || auth.Metadata == nil {
|
||||
return antigravityModelCapabilityHints{}
|
||||
}
|
||||
accessToken, _ := auth.Metadata["access_token"].(string)
|
||||
accessToken = strings.TrimSpace(accessToken)
|
||||
if accessToken == "" {
|
||||
return antigravityModelCapabilityHints{}
|
||||
}
|
||||
|
||||
client := &http.Client{}
|
||||
if transport, _, errProxy := proxyutil.BuildHTTPTransport(s.antigravityModelFetchProxyURL(auth)); errProxy == nil && transport != nil {
|
||||
client.Transport = transport
|
||||
}
|
||||
|
||||
for _, baseURL := range antigravityModelBaseURLs(auth) {
|
||||
req, errReq := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(baseURL, "/")+antigravityModelsPath, strings.NewReader(`{}`))
|
||||
if errReq != nil {
|
||||
continue
|
||||
}
|
||||
req.Close = true
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
req.Header.Set("User-Agent", misc.AntigravityUserAgent())
|
||||
|
||||
resp, errDo := client.Do(req)
|
||||
if errDo != nil {
|
||||
continue
|
||||
}
|
||||
body, errRead := io.ReadAll(resp.Body)
|
||||
if errClose := resp.Body.Close(); errClose != nil {
|
||||
log.Debugf("antigravity model fetch: close response body: %v", errClose)
|
||||
}
|
||||
if errRead != nil {
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
continue
|
||||
}
|
||||
hints := parseAntigravityModelCapabilityHints(body)
|
||||
if len(hints.WebSearchModelIDs) > 0 {
|
||||
return hints
|
||||
}
|
||||
}
|
||||
return antigravityModelCapabilityHints{}
|
||||
}
|
||||
|
||||
func (s *Service) antigravityModelFetchProxyURL(auth *coreauth.Auth) string {
|
||||
if auth != nil {
|
||||
if proxyURL := strings.TrimSpace(auth.ProxyURL); proxyURL != "" {
|
||||
return proxyURL
|
||||
}
|
||||
}
|
||||
if s != nil && s.cfg != nil {
|
||||
return strings.TrimSpace(s.cfg.ProxyURL)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func antigravityModelBaseURLs(auth *coreauth.Auth) []string {
|
||||
if baseURL := resolveAntigravityModelBaseURL(auth); baseURL != "" {
|
||||
return []string{baseURL}
|
||||
}
|
||||
return []string{antigravityModelBaseURLDaily, antigravityModelBaseURLProd}
|
||||
}
|
||||
|
||||
func resolveAntigravityModelBaseURL(auth *coreauth.Auth) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
if auth.Attributes != nil {
|
||||
if value := strings.TrimSpace(auth.Attributes["base_url"]); value != "" {
|
||||
return strings.TrimRight(value, "/")
|
||||
}
|
||||
}
|
||||
if auth.Metadata != nil {
|
||||
if value, ok := auth.Metadata["base_url"].(string); ok {
|
||||
value = strings.TrimSpace(value)
|
||||
if value != "" {
|
||||
return strings.TrimRight(value, "/")
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func parseAntigravityModelCapabilityHints(body []byte) antigravityModelCapabilityHints {
|
||||
var parsed antigravityFetchAvailableModelsResponse
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return antigravityModelCapabilityHints{}
|
||||
}
|
||||
webSearchModels := make(map[string]struct{}, len(parsed.WebSearchModelIDs))
|
||||
for _, modelID := range parsed.WebSearchModelIDs {
|
||||
modelID = normalizeAntigravityFetchedModelID(modelID)
|
||||
if modelID != "" {
|
||||
webSearchModels[modelID] = struct{}{}
|
||||
}
|
||||
}
|
||||
return antigravityModelCapabilityHints{WebSearchModelIDs: webSearchModels}
|
||||
}
|
||||
|
||||
func applyAntigravityFetchedModelCapabilities(models []*ModelInfo, hints antigravityModelCapabilityHints) []*ModelInfo {
|
||||
if len(models) == 0 || len(hints.WebSearchModelIDs) == 0 {
|
||||
return models
|
||||
}
|
||||
|
||||
for _, model := range models {
|
||||
if model == nil {
|
||||
continue
|
||||
}
|
||||
modelID := normalizeAntigravityFetchedModelID(model.ID)
|
||||
if _, ok := hints.WebSearchModelIDs[modelID]; ok {
|
||||
model.SupportsWebSearch = true
|
||||
}
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
func normalizeAntigravityFetchedModelID(modelID string) string {
|
||||
return strings.ToLower(strings.TrimSpace(modelID))
|
||||
}
|
||||
@@ -1782,6 +1782,7 @@ func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) {
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "antigravity":
|
||||
models = registry.GetAntigravityModels()
|
||||
models = applyAntigravityFetchedModelCapabilities(models, s.fetchAntigravityModelCapabilityHintsForAuth(ctx, a))
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "claude":
|
||||
models = registry.GetClaudeModels()
|
||||
|
||||
@@ -2,6 +2,8 @@ package cliproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -133,3 +135,106 @@ func TestRegisterModelsForAuth_OpenAICompatibilityImageModelType(t *testing.T) {
|
||||
t.Fatal("expected chat model to keep default thinking support")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterModelsForAuth_AntigravityFetchesWebSearchCapability(t *testing.T) {
|
||||
var sawFetch bool
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != antigravityModelsPath {
|
||||
t.Fatalf("path = %q, want %s", r.URL.Path, antigravityModelsPath)
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer token" {
|
||||
t.Fatalf("Authorization = %q, want bearer token", got)
|
||||
}
|
||||
sawFetch = true
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{
|
||||
"models": {
|
||||
"gemini-3.1-flash-lite": {
|
||||
"displayName": "Gemini 3.1 Flash Lite",
|
||||
"maxTokens": 1,
|
||||
"maxOutputTokens": 2
|
||||
},
|
||||
"fetched-only-search-model": {
|
||||
"displayName": "Fetched Only Search Model"
|
||||
}
|
||||
},
|
||||
"webSearchModelIds": ["gemini-3.1-flash-lite", "fetched-only-search-model"]
|
||||
}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
service := &Service{cfg: &config.Config{}}
|
||||
auth := &coreauth.Auth{
|
||||
ID: "auth-antigravity-fetch-models",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"access_token": "token",
|
||||
},
|
||||
}
|
||||
|
||||
registry := internalregistry.GetGlobalRegistry()
|
||||
registry.UnregisterClient(auth.ID)
|
||||
t.Cleanup(func() {
|
||||
registry.UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
service.registerModelsForAuth(context.Background(), auth)
|
||||
if !sawFetch {
|
||||
t.Fatal("expected fetchAvailableModels request")
|
||||
}
|
||||
|
||||
models := registry.GetModelsForClient(auth.ID)
|
||||
staticModels := internalregistry.GetAntigravityModels()
|
||||
staticByID := make(map[string]*internalregistry.ModelInfo, len(staticModels))
|
||||
for _, model := range staticModels {
|
||||
if model != nil {
|
||||
staticByID[model.ID] = model
|
||||
}
|
||||
}
|
||||
|
||||
var webSearchModel, agentModel, staticOnlyModel, fetchedOnlyModel *internalregistry.ModelInfo
|
||||
for _, model := range models {
|
||||
if model == nil {
|
||||
continue
|
||||
}
|
||||
switch strings.TrimSpace(model.ID) {
|
||||
case "gemini-3.1-flash-lite":
|
||||
webSearchModel = model
|
||||
case "gemini-3-flash-agent":
|
||||
agentModel = model
|
||||
case "gpt-oss-120b-medium":
|
||||
staticOnlyModel = model
|
||||
case "fetched-only-search-model":
|
||||
fetchedOnlyModel = model
|
||||
}
|
||||
}
|
||||
if webSearchModel == nil {
|
||||
t.Fatal("expected gemini-3.1-flash-lite to be registered")
|
||||
}
|
||||
if !webSearchModel.SupportsWebSearch {
|
||||
t.Fatal("expected gemini-3.1-flash-lite to support web search")
|
||||
}
|
||||
staticWebSearchModel := staticByID["gemini-3.1-flash-lite"]
|
||||
if staticWebSearchModel == nil {
|
||||
t.Fatal("expected static gemini-3.1-flash-lite definition")
|
||||
}
|
||||
if webSearchModel.ContextLength != staticWebSearchModel.ContextLength || webSearchModel.MaxCompletionTokens != staticWebSearchModel.MaxCompletionTokens {
|
||||
t.Fatalf("static token limits should be preserved, got=%#v static=%#v", webSearchModel, staticWebSearchModel)
|
||||
}
|
||||
if agentModel == nil {
|
||||
t.Fatal("expected gemini-3-flash-agent to be registered")
|
||||
}
|
||||
if agentModel.SupportsWebSearch {
|
||||
t.Fatal("gemini-3-flash-agent should not support web search")
|
||||
}
|
||||
if staticOnlyModel == nil {
|
||||
t.Fatal("expected static-only Antigravity model to remain registered")
|
||||
}
|
||||
if fetchedOnlyModel != nil {
|
||||
t.Fatalf("fetched-only model should not be registered: %#v", fetchedOnlyModel)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user