Setting the file. One moment.
Resolve Source Model · LLM To Bedrock · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page 81
Creating Amazon Aurora Db Cluster With Instances
104
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Creating API Gateway Stage
— line 193
This file
Number 30.18
Position 18 of 35
Type Python
Size 8 KB
Lines 221 scripts/ resolve_source_model.py
Python · 221 lines · 8 KB
16 references/helpers/run-source-model-baseline/run-source-model-baseline.md.
17
18 Security: the key is read from the env file and sent ONLY as an auth header to
19 the provider's own official endpoint (api.openai.com / api.anthropic.com /
20 generativelanguage.googleapis.com) — never to any other host and never as a
21 URL query parameter. The key value is never printed.
22 """
23
24 from __future__ import annotations
25
26 import json
27 import os
28 import re
29 import sys
30 import urllib.request
31
32 SAFE_SUFFIX = re.compile( r " ^\d {4} - \d {2} - \d {2} $ | ^\d[\d . ] * $ " )
33
34
35 PROVIDER_KEYS = ( "OPENAI_API_KEY" , "ANTHROPIC_API_KEY" , "GEMINI_API_KEY" )
36
37
38 def load_env_file (path: str ) -> dict[ str , str ]:
39 """Parse the env file and return ITS pairs. The file is authoritative for
40 provider selection — an ambient OPENAI_API_KEY in the inherited process
41 environment must not select OpenAI when the migration's file carries a
42 GEMINI_API_KEY (the model ID would go to the wrong provider). Parsed pairs
43 are also exported so downstream reads use the file's values."""
44 pairs: dict[ str , str ] = {}
45 with open (path) as f:
46 for line in f:
47 line = line.strip()
48 if line and "=" in line:
49 k, v = line.split( "=" , 1 )
50 pairs[k] = v
51 os.environ[k] = v
52 return pairs
53
54
55 PROVIDER_TO_KEY = {
56 "openai" : "OPENAI_API_KEY" ,
57 "anthropic" : "ANTHROPIC_API_KEY" ,
58 "google" : "GEMINI_API_KEY" ,
59 "gemini" : "GEMINI_API_KEY" ,
60 }
61
62
63 def pick_provider (file_pairs: dict[ str , str ]) -> str | None :
64 """Select the provider key name. SOURCE_PROVIDER (the helper's declared
65 input) is authoritative when set — a file carrying several provider keys
66 must not fall back to tuple-order guessing. Without it, fall back to the
67 first recognized key IN THE FILE (never the ambient environment)."""
68 stated = os.environ.get( "SOURCE_PROVIDER" , "" ).lower()
69 if stated:
70 key = PROVIDER_TO_KEY .get(stated)
71 if key is None or key not in file_pairs:
72 return None
73 return key
74 return next ((k for k in PROVIDER_KEYS if k in file_pairs), None )
75
76
77 def _secret_variants (secrets) -> list :
78 """Every textual form a secret can take in error text. A raw CR/LF inside
79 a credential is rendered ESCAPED in exception messages (repr turns one
80 control char into backslash-r text), so literal replacement alone misses
81 it; the control-char-split fragments catch any remaining partial echo."""
82 import re as _re
83 out: list = []
84 for s in secrets:
85 if not s:
86 continue
87 out.append(s)
88 esc = s.encode( "unicode_escape" ).decode( "ascii" )
89 if esc != s:
90 out.append(esc)
91 for frag in _re.split( r " [ \x00 - \x1f ] + " , s):
92 if len (frag) >= 6 and frag not in out:
93 out.append(frag)
94 return sorted (out, key = len , reverse = True )
95
96
97 def redact (text: str , secrets) -> str :
98 """Strip secret values from failure text. A raised exception can embed a
99 header value verbatim (http.client rejects an invalid header with the full
100 'Bearer <key>' in the ValueError message), and the docstring promise is
101 that the key never reaches the output JSONL."""
102 for s in _secret_variants(secrets):
103 text = text.replace(s, "***" )
104 return text
105
106
107 def _get_json (url: str , headers: dict ) -> dict :
108 req = urllib.request.Request(url, headers = headers)
109 with urllib.request.urlopen(req, timeout = 30 ) as r: # nosec B310 — fixed https hosts
110 return json.loads(r.read())
111
112
113 def list_openai () -> list[ str ]:
114 # OpenAI's list endpoint returns the full catalog in one response (no cursor).
115 data = _get_json(
116 "https://api.openai.com/v1/models" ,
117 { "Authorization" : f "Bearer { os.environ[ 'OPENAI_API_KEY' ] } " },
118 )
119 return [m[ "id" ] for m in data[ "data" ]]
120
121
122 def list_anthropic () -> list[ str ]:
123 # Paginated: default page size is 20, so a valid model can sit past page 1
124 # and must not be reported not_found. Cursor: has_more + last_id → after_id.
125 headers = {
126 "x-api-key" : os.environ[ "ANTHROPIC_API_KEY" ],
127 "anthropic-version" : "2023-06-01" ,
128 }
129 ids: list[ str ] = []
130 url = "https://api.anthropic.com/v1/models?limit=1000"
131 while True :
132 data = _get_json(url, headers)
133 ids += [m[ "id" ] for m in data[ "data" ]]
134 if not data.get( "has_more" ):
135 return ids
136 url = f "https://api.anthropic.com/v1/models?limit=1000&after_id= { data[ 'last_id' ] } "
137
138
139 def list_gemini () -> list[ str ]:
140 # Key travels as a header, not a query parameter — URLs end up in logs.
141 # Paginated via nextPageToken.
142 headers = { "x-goog-api-key" : os.environ[ "GEMINI_API_KEY" ]}
143 ids: list[ str ] = []
144 url = "https://generativelanguage.googleapis.com/v1beta/models?pageSize=1000"
145 while True :
146 data = _get_json(url, headers)
147 # Gemini returns names like "models/gemini-1.5-pro"; strip the prefix.
148 ids += [m[ "name" ].split( "/" , 1 )[ - 1 ] for m in data.get( "models" , [])]
149 token = data.get( "nextPageToken" )
150 if not token:
151 return ids
152 url = ( "https://generativelanguage.googleapis.com/v1beta/models"
153 f "?pageSize=1000&pageToken= { token } " )
154
155
156 def safe_variant (catalog_id: str , plan_id: str ) -> bool :
157 if catalog_id == plan_id:
158 return True
159 suffix = catalog_id[ len (plan_id) + 1 :] # strip "PLAN_ID-"
160 return bool ( SAFE_SUFFIX .match(suffix))
161
162
163 def _lcp (a: str , b: str ) -> int :
164 n = min ( len (a), len (b))
165 i = 0
166 while i < n and a[i] == b[i]:
167 i += 1
168 return i
169
170
171 def resolve (catalog: list[ str ], plan_id: str ) -> dict :
172 """Pure resolution decision — no network, unit-tested."""
173 if plan_id in catalog:
174 return { "status" : "exact" , "resolved_id" : plan_id}
175
176 prefix_hits = [m for m in catalog if m == plan_id or m.startswith(plan_id + "-" )]
177
178 # A bare prefix match is NOT enough to auto-resolve: "gpt-4o-mini",
179 # "claude-3-5-sonnet-latest" start with a plausible plan ID but are
180 # different model lines / non-deterministic aliases. Only a date or pure
181 # version suffix is the same line.
182 safe_hits = sorted ((m for m in prefix_hits if safe_variant(m, plan_id)), key = len )
183 if safe_hits:
184 return { "status" : "prefix" , "resolved_id" : safe_hits[ 0 ], "all_hits" : safe_hits}
185
186 if prefix_hits:
187 return { "status" : "not_found" , "candidates" : prefix_hits[: 5 ], "ambiguous_prefix" : True }
188
189 ranked = sorted (catalog, key =lambda m: - _lcp(m, plan_id))[: 5 ]
190 return { "status" : "not_found" , "candidates" : ranked}
191
192
193 def main () -> int :
194 if len (sys.argv) != 2 :
195 print ( "usage: PLAN_MODEL_ID=<id> resolve_source_model.py <env-file>" , file = sys.stderr)
196 return 2
197 file_pairs = load_env_file(sys.argv[ 1 ])
198 plan_id = os.environ[ "PLAN_MODEL_ID" ]
199
200 provider = pick_provider(file_pairs)
201 listers = { "OPENAI_API_KEY" : list_openai, "ANTHROPIC_API_KEY" : list_anthropic,
202 "GEMINI_API_KEY" : list_gemini}
203 if provider not in listers:
204 print (json.dumps({ "status" : "no_key" }))
205 return 2
206
207 secrets = [v for k, v in file_pairs.items() if k in PROVIDER_KEYS ]
208 try :
209 catalog = listers[provider]()
210 except Exception as e: # noqa: BLE001 — an unhandled traceback would put the
211 # exception text (which can embed the auth header) on stderr unredacted
212 print (json.dumps({ "status" : "error" ,
213 "detail" : redact( f " { type (e). __name__ } : { e } " , secrets)}))
214 return 3
215
216 print (json.dumps(resolve(catalog, plan_id)))
217 return 0
218
219
220 if __name__ == "__main__" :
221 sys.exit(main())