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)
14package main
15
16import (
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
36const (
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
43func init() {
44 logging.SetupBaseLogger()
45 log.SetLevel(log.InfoLevel)
46}
47
48// modelOutput wraps the fetched model list with fetch metadata.
49type modelOutput struct {
50 Models []modelEntry `json:"models"`
51}
52
53// modelEntry contains only the fields we want to keep for static model definitions.
54type 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
66func 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
162func 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
261func 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}