Spaces:
Paused
feat(antigravity): enable token counting via API with resilient routing
Browse filesIntroduces the capability to count tokens for Antigravity-backed requests. This implementation leverages the `countTokens` endpoint of the Antigravity API, replacing the prior unsupported stub.
Key aspects of this update include:
- **API Integration**: Direct integration with the Antigravity `countTokens` API, including necessary request payload translation and authentication.
- **Resilient Infrastructure**: A fallback mechanism has been established, allowing the system to attempt connections across multiple Antigravity base URLs to ensure request success even in the event of temporary service interruptions.
- **Model Aliasing**: Added mappings for `gemini-3-flash` and `gemini-3-flash-preview` to ensure compatibility with the latest model variants.
- **Robust Error Handling**: Comprehensive error handling and logging are in place to manage failures during API interactions.
|
@@ -32,15 +32,16 @@ import (
|
|
| 32 |
const (
|
| 33 |
antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
| 34 |
// antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
|
| 35 |
-
antigravityBaseURLProd
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
|
|
|
| 44 |
)
|
| 45 |
|
| 46 |
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
@@ -646,9 +647,131 @@ func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Au
|
|
| 646 |
return updated, nil
|
| 647 |
}
|
| 648 |
|
| 649 |
-
// CountTokens counts tokens for the given request
|
| 650 |
-
func (e *AntigravityExecutor) CountTokens(context.Context, *cliproxyauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
| 651 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 652 |
}
|
| 653 |
|
| 654 |
// FetchAntigravityModels retrieves available models using the supplied auth.
|
|
@@ -1114,6 +1237,8 @@ func modelName2Alias(modelName string) string {
|
|
| 1114 |
return "gemini-3-pro-image-preview"
|
| 1115 |
case "gemini-3-pro-high":
|
| 1116 |
return "gemini-3-pro-preview"
|
|
|
|
|
|
|
| 1117 |
case "claude-sonnet-4-5":
|
| 1118 |
return "gemini-claude-sonnet-4-5"
|
| 1119 |
case "claude-sonnet-4-5-thinking":
|
|
@@ -1135,6 +1260,8 @@ func alias2ModelName(modelName string) string {
|
|
| 1135 |
return "gemini-3-pro-image"
|
| 1136 |
case "gemini-3-pro-preview":
|
| 1137 |
return "gemini-3-pro-high"
|
|
|
|
|
|
|
| 1138 |
case "gemini-claude-sonnet-4-5":
|
| 1139 |
return "claude-sonnet-4-5"
|
| 1140 |
case "gemini-claude-sonnet-4-5-thinking":
|
|
|
|
| 32 |
const (
|
| 33 |
antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
| 34 |
// antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
|
| 35 |
+
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
| 36 |
+
antigravityCountTokensPath = "/v1internal:countTokens"
|
| 37 |
+
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
| 38 |
+
antigravityGeneratePath = "/v1internal:generateContent"
|
| 39 |
+
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
| 40 |
+
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
| 41 |
+
antigravityClientSecret = "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"
|
| 42 |
+
defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
|
| 43 |
+
antigravityAuthType = "antigravity"
|
| 44 |
+
refreshSkew = 3000 * time.Second
|
| 45 |
)
|
| 46 |
|
| 47 |
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
|
|
|
| 647 |
return updated, nil
|
| 648 |
}
|
| 649 |
|
| 650 |
+
// CountTokens counts tokens for the given request using the Antigravity API.
|
| 651 |
+
func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
| 652 |
+
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
|
| 653 |
+
if errToken != nil {
|
| 654 |
+
return cliproxyexecutor.Response{}, errToken
|
| 655 |
+
}
|
| 656 |
+
if updatedAuth != nil {
|
| 657 |
+
auth = updatedAuth
|
| 658 |
+
}
|
| 659 |
+
if strings.TrimSpace(token) == "" {
|
| 660 |
+
return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
|
| 661 |
+
}
|
| 662 |
+
|
| 663 |
+
from := opts.SourceFormat
|
| 664 |
+
to := sdktranslator.FromString("antigravity")
|
| 665 |
+
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
| 666 |
+
|
| 667 |
+
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
| 668 |
+
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
| 669 |
+
|
| 670 |
+
var authID, authLabel, authType, authValue string
|
| 671 |
+
if auth != nil {
|
| 672 |
+
authID = auth.ID
|
| 673 |
+
authLabel = auth.Label
|
| 674 |
+
authType, authValue = auth.AccountInfo()
|
| 675 |
+
}
|
| 676 |
+
|
| 677 |
+
var lastStatus int
|
| 678 |
+
var lastBody []byte
|
| 679 |
+
var lastErr error
|
| 680 |
+
|
| 681 |
+
for idx, baseURL := range baseURLs {
|
| 682 |
+
payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
| 683 |
+
payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
|
| 684 |
+
payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, payload)
|
| 685 |
+
payload = normalizeAntigravityThinking(req.Model, payload)
|
| 686 |
+
payload = deleteJSONField(payload, "project")
|
| 687 |
+
payload = deleteJSONField(payload, "model")
|
| 688 |
+
payload = deleteJSONField(payload, "request.safetySettings")
|
| 689 |
+
|
| 690 |
+
base := strings.TrimSuffix(baseURL, "/")
|
| 691 |
+
if base == "" {
|
| 692 |
+
base = buildBaseURL(auth)
|
| 693 |
+
}
|
| 694 |
+
|
| 695 |
+
var requestURL strings.Builder
|
| 696 |
+
requestURL.WriteString(base)
|
| 697 |
+
requestURL.WriteString(antigravityCountTokensPath)
|
| 698 |
+
if opts.Alt != "" {
|
| 699 |
+
requestURL.WriteString("?$alt=")
|
| 700 |
+
requestURL.WriteString(url.QueryEscape(opts.Alt))
|
| 701 |
+
}
|
| 702 |
+
|
| 703 |
+
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload))
|
| 704 |
+
if errReq != nil {
|
| 705 |
+
return cliproxyexecutor.Response{}, errReq
|
| 706 |
+
}
|
| 707 |
+
httpReq.Header.Set("Content-Type", "application/json")
|
| 708 |
+
httpReq.Header.Set("Authorization", "Bearer "+token)
|
| 709 |
+
httpReq.Header.Set("User-Agent", resolveUserAgent(auth))
|
| 710 |
+
httpReq.Header.Set("Accept", "application/json")
|
| 711 |
+
if host := resolveHost(base); host != "" {
|
| 712 |
+
httpReq.Host = host
|
| 713 |
+
}
|
| 714 |
+
|
| 715 |
+
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
| 716 |
+
URL: requestURL.String(),
|
| 717 |
+
Method: http.MethodPost,
|
| 718 |
+
Headers: httpReq.Header.Clone(),
|
| 719 |
+
Body: payload,
|
| 720 |
+
Provider: e.Identifier(),
|
| 721 |
+
AuthID: authID,
|
| 722 |
+
AuthLabel: authLabel,
|
| 723 |
+
AuthType: authType,
|
| 724 |
+
AuthValue: authValue,
|
| 725 |
+
})
|
| 726 |
+
|
| 727 |
+
httpResp, errDo := httpClient.Do(httpReq)
|
| 728 |
+
if errDo != nil {
|
| 729 |
+
recordAPIResponseError(ctx, e.cfg, errDo)
|
| 730 |
+
lastStatus = 0
|
| 731 |
+
lastBody = nil
|
| 732 |
+
lastErr = errDo
|
| 733 |
+
if idx+1 < len(baseURLs) {
|
| 734 |
+
log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
| 735 |
+
continue
|
| 736 |
+
}
|
| 737 |
+
return cliproxyexecutor.Response{}, errDo
|
| 738 |
+
}
|
| 739 |
+
|
| 740 |
+
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
| 741 |
+
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
| 742 |
+
if errClose := httpResp.Body.Close(); errClose != nil {
|
| 743 |
+
log.Errorf("antigravity executor: close response body error: %v", errClose)
|
| 744 |
+
}
|
| 745 |
+
if errRead != nil {
|
| 746 |
+
recordAPIResponseError(ctx, e.cfg, errRead)
|
| 747 |
+
return cliproxyexecutor.Response{}, errRead
|
| 748 |
+
}
|
| 749 |
+
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
| 750 |
+
|
| 751 |
+
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
| 752 |
+
count := gjson.GetBytes(bodyBytes, "totalTokens").Int()
|
| 753 |
+
translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, bodyBytes)
|
| 754 |
+
return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
|
| 755 |
+
}
|
| 756 |
+
|
| 757 |
+
lastStatus = httpResp.StatusCode
|
| 758 |
+
lastBody = append([]byte(nil), bodyBytes...)
|
| 759 |
+
lastErr = nil
|
| 760 |
+
if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) {
|
| 761 |
+
log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
| 762 |
+
continue
|
| 763 |
+
}
|
| 764 |
+
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}
|
| 765 |
+
}
|
| 766 |
+
|
| 767 |
+
switch {
|
| 768 |
+
case lastStatus != 0:
|
| 769 |
+
return cliproxyexecutor.Response{}, statusErr{code: lastStatus, msg: string(lastBody)}
|
| 770 |
+
case lastErr != nil:
|
| 771 |
+
return cliproxyexecutor.Response{}, lastErr
|
| 772 |
+
default:
|
| 773 |
+
return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"}
|
| 774 |
+
}
|
| 775 |
}
|
| 776 |
|
| 777 |
// FetchAntigravityModels retrieves available models using the supplied auth.
|
|
|
|
| 1237 |
return "gemini-3-pro-image-preview"
|
| 1238 |
case "gemini-3-pro-high":
|
| 1239 |
return "gemini-3-pro-preview"
|
| 1240 |
+
case "gemini-3-flash":
|
| 1241 |
+
return "gemini-3-flash-preview"
|
| 1242 |
case "claude-sonnet-4-5":
|
| 1243 |
return "gemini-claude-sonnet-4-5"
|
| 1244 |
case "claude-sonnet-4-5-thinking":
|
|
|
|
| 1260 |
return "gemini-3-pro-image"
|
| 1261 |
case "gemini-3-pro-preview":
|
| 1262 |
return "gemini-3-pro-high"
|
| 1263 |
+
case "gemini-3-flash-preview":
|
| 1264 |
+
return "gemini-3-flash"
|
| 1265 |
case "gemini-claude-sonnet-4-5":
|
| 1266 |
return "claude-sonnet-4-5"
|
| 1267 |
case "gemini-claude-sonnet-4-5-thinking":
|