Add projects

backend/sdk/cliproxy/antigravity_models.go

@@ -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))
+}