History

feat(fetch_antigravity_models): add command to fetch and save Antigravity model list

hkfires committed 8 months ago

Commit ec24baf1 parent dea3e741 file · +275 −0Browse files

Unified · Split

cmd/fetch_antigravity_models/main.go+275−0
@@ -0,0 +1,275 @@
1+// Command fetch_antigravity_models connects to the Antigravity API using the
2+// stored auth credentials and saves the dynamically fetched model list to a
3+// JSON file for inspection or offline use.
4+//
5+// Usage:
6+//
7+// go run ./cmd/fetch_antigravity_models [flags]
8+//
9+// Flags:
10+//
11+// --auths-dir <path> Directory containing auth JSON files (default: "auths")
12+// --output <path> Output JSON file path (default: "antigravity_models.json")
13+// --pretty Pretty-print the output JSON (default: true)
14+package main
15+
16+import (
17+ "context"
18+ "encoding/json"
19+ "flag"
20+ "fmt"
21+ "io"
22+ "net/http"
23+ "os"
24+ "path/filepath"
25+ "strings"
26+ "time"
27+
28+ "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
29+ sdkauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
30+ coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
31+ "github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
32+ log "github.com/sirupsen/logrus"
33+ "github.com/tidwall/gjson"
34+)
35+
36+const (
37+ antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
38+ antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
39+ antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
40+ antigravityModelsPath = "/v1internal:fetchAvailableModels"
41+)
42+
43+func init() {
44+ logging.SetupBaseLogger()
45+ log.SetLevel(log.InfoLevel)
46+}
47+
48+// modelOutput wraps the fetched model list with fetch metadata.
49+type modelOutput struct {
50+ Models []modelEntry `json:"models"`
51+}
52+
53+// modelEntry contains only the fields we want to keep for static model definitions.
54+type modelEntry struct {
55+ ID string `json:"id"`
56+ Object string `json:"object"`
57+ OwnedBy string `json:"owned_by"`
58+ Type string `json:"type"`
59+ DisplayName string `json:"display_name"`
60+ Name string `json:"name"`
61+ Description string `json:"description"`
62+ ContextLength int `json:"context_length,omitempty"`
63+ MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
64+}
65+
66+func main() {
67+ var authsDir string
68+ var outputPath string
69+ var pretty bool
70+
71+ flag.StringVar(&authsDir, "auths-dir", "auths", "Directory containing auth JSON files")
72+ flag.StringVar(&outputPath, "output", "antigravity_models.json", "Output JSON file path")
73+ flag.BoolVar(&pretty, "pretty", true, "Pretty-print the output JSON")
74+ flag.Parse()
75+
76+ // Resolve relative paths against the working directory.
77+ wd, err := os.Getwd()
78+ if err != nil {
79+ fmt.Fprintf(os.Stderr, "error: cannot get working directory: %v\n", err)
80+ os.Exit(1)
81+ }
82+ if !filepath.IsAbs(authsDir) {
83+ authsDir = filepath.Join(wd, authsDir)
84+ }
85+ if !filepath.IsAbs(outputPath) {
86+ outputPath = filepath.Join(wd, outputPath)
87+ }
88+
89+ fmt.Printf("Scanning auth files in: %s\n", authsDir)
90+
91+ // Load all auth records from the directory.
92+ fileStore := sdkauth.NewFileTokenStore()
93+ fileStore.SetBaseDir(authsDir)
94+
95+ ctx := context.Background()
96+ auths, err := fileStore.List(ctx)
97+ if err != nil {
98+ fmt.Fprintf(os.Stderr, "error: failed to list auth files: %v\n", err)
99+ os.Exit(1)
100+ }
101+ if len(auths) == 0 {
102+ fmt.Fprintf(os.Stderr, "error: no auth files found in %s\n", authsDir)
103+ os.Exit(1)
104+ }
105+
106+ // Find the first enabled antigravity auth.
107+ var chosen *coreauth.Auth
108+ for _, a := range auths {
109+ if a == nil || a.Disabled {
110+ continue
111+ }
112+ if strings.EqualFold(strings.TrimSpace(a.Provider), "antigravity") {
113+ chosen = a
114+ break
115+ }
116+ }
117+ if chosen == nil {
118+ fmt.Fprintf(os.Stderr, "error: no enabled antigravity auth found in %s\n", authsDir)
119+ os.Exit(1)
120+ }
121+
122+ fmt.Printf("Using auth: id=%s label=%s\n", chosen.ID, chosen.Label)
123+
124+ // Fetch models from the upstream Antigravity API.
125+ fmt.Println("Fetching Antigravity model list from upstream...")
126+
127+ fetchCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
128+ defer cancel()
129+
130+ models := fetchModels(fetchCtx, chosen)
131+ if len(models) == 0 {
132+ fmt.Fprintln(os.Stderr, "warning: no models returned (API may be unavailable or token expired)")
133+ } else {
134+ fmt.Printf("Fetched %d models.\n", len(models))
135+ }
136+
137+ // Build the output payload.
138+ out := modelOutput{
139+ Models: models,
140+ }
141+
142+ // Marshal to JSON.
143+ var raw []byte
144+ if pretty {
145+ raw, err = json.MarshalIndent(out, "", " ")
146+ } else {
147+ raw, err = json.Marshal(out)
148+ }
149+ if err != nil {
150+ fmt.Fprintf(os.Stderr, "error: failed to marshal JSON: %v\n", err)
151+ os.Exit(1)
152+ }
153+
154+ if err = os.WriteFile(outputPath, raw, 0o644); err != nil {
155+ fmt.Fprintf(os.Stderr, "error: failed to write output file %s: %v\n", outputPath, err)
156+ os.Exit(1)
157+ }
158+
159+ fmt.Printf("Model list saved to: %s\n", outputPath)
160+}
161+
162+func fetchModels(ctx context.Context, auth *coreauth.Auth) []modelEntry {
163+ accessToken := metaStringValue(auth.Metadata, "access_token")
164+ if accessToken == "" {
165+ fmt.Fprintln(os.Stderr, "error: no access token found in auth")
166+ return nil
167+ }
168+
169+ baseURLs := []string{antigravityBaseURLProd, antigravityBaseURLDaily, antigravitySandboxBaseURLDaily}
170+
171+ for _, baseURL := range baseURLs {
172+ modelsURL := baseURL + antigravityModelsPath
173+
174+ var payload []byte
175+ if auth != nil && auth.Metadata != nil {
176+ if pid, ok := auth.Metadata["project_id"].(string); ok && strings.TrimSpace(pid) != "" {
177+ payload = []byte(fmt.Sprintf(`{"project": "%s"}`, strings.TrimSpace(pid)))
178+ }
179+ }
180+ if len(payload) == 0 {
181+ payload = []byte(`{}`)
182+ }
183+
184+ httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, modelsURL, strings.NewReader(string(payload)))
185+ if errReq != nil {
186+ continue
187+ }
188+ httpReq.Close = true
189+ httpReq.Header.Set("Content-Type", "application/json")
190+ httpReq.Header.Set("Authorization", "Bearer "+accessToken)
191+ httpReq.Header.Set("User-Agent", "antigravity/1.19.6 darwin/arm64")
192+
193+ httpClient := &http.Client{Timeout: 30 * time.Second}
194+ if transport, _, errProxy := proxyutil.BuildHTTPTransport(auth.ProxyURL); errProxy == nil && transport != nil {
195+ httpClient.Transport = transport
196+ }
197+ httpResp, errDo := httpClient.Do(httpReq)
198+ if errDo != nil {
199+ continue
200+ }
201+
202+ bodyBytes, errRead := io.ReadAll(httpResp.Body)
203+ httpResp.Body.Close()
204+ if errRead != nil {
205+ continue
206+ }
207+
208+ if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
209+ continue
210+ }
211+
212+ result := gjson.GetBytes(bodyBytes, "models")
213+ if !result.Exists() {
214+ continue
215+ }
216+
217+ var models []modelEntry
218+
219+ for originalName, modelData := range result.Map() {
220+ modelID := strings.TrimSpace(originalName)
221+ if modelID == "" {
222+ continue
223+ }
224+ // Skip internal/experimental models
225+ switch modelID {
226+ case "chat_20706", "chat_23310", "tab_flash_lite_preview", "tab_jump_flash_lite_preview", "gemini-2.5-flash-thinking", "gemini-2.5-pro":
227+ continue
228+ }
229+
230+ displayName := modelData.Get("displayName").String()
231+ if displayName == "" {
232+ displayName = modelID
233+ }
234+
235+ entry := modelEntry{
236+ ID: modelID,
237+ Object: "model",
238+ OwnedBy: "antigravity",
239+ Type: "antigravity",
240+ DisplayName: displayName,
241+ Name: modelID,
242+ Description: displayName,
243+ }
244+
245+ if maxTok := modelData.Get("maxTokens").Int(); maxTok > 0 {
246+ entry.ContextLength = int(maxTok)
247+ }
248+ if maxOut := modelData.Get("maxOutputTokens").Int(); maxOut > 0 {
249+ entry.MaxCompletionTokens = int(maxOut)
250+ }
251+
252+ models = append(models, entry)
253+ }
254+
255+ return models
256+ }
257+
258+ return nil
259+}
260+
261+func metaStringValue(m map[string]interface{}, key string) string {
262+ if m == nil {
263+ return ""
264+ }
265+ v, ok := m[key]
266+ if !ok {
267+ return ""
268+ }
269+ switch val := v.(type) {
270+ case string:
271+ return val
272+ default:
273+ return ""
274+ }
275+}