Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
4.4 kB · 150 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151package 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))}