Setting the file. One moment.
Source Baseline · 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
scripts/ source_baseline.py
Python · 277 lines · 10 KB
Security: the key is read from the env file into the process environment only
17 and sent ONLY as an auth header to its own provider's official endpoint
18 (api.openai.com / api.anthropic.com / generativelanguage.googleapis.com) —
19 never to any other host, never as a URL query parameter, never to stdout or
20 the output JSONL.
21 """
22
23 from __future__ import annotations
24
25 import json
26 import os
27 import sys
28 import urllib.error
29 import urllib.request
30
31 # Requests intentionally use provider defaults for temperature/top_p — the
32 # golden dataset doesn't record per-request sampling params, and the same
33 # defaults-only shape is used for all three providers so the comparison is
34 # apples-to-apples. The 4096-token cap matches the Bedrock eval side.
35 MAX_TOKENS = 4096
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 must not select OpenAI when
41 the migration's file carries a GEMINI_API_KEY. Parsed pairs are also
42 exported so the request builders read the file's values."""
43 pairs: dict[ str , str ] = {}
44 with open (path) as f:
45 for line in f:
46 line = line.strip()
47 if line and "=" in line:
48 k, v = line.split( "=" , 1 )
49 pairs[k] = v
50 os.environ[k] = v
51 return pairs
52
53
54 def build_openai_request (model: str , system: str , user_text: str ) -> tuple[ str , dict , dict ]:
55 body = {
56 "model" : model,
57 "messages" : ([{ "role" : "system" , "content" : system}] if system else [])
58 + [{ "role" : "user" , "content" : user_text}],
59 # gpt-5.x rejects max_tokens (HTTP 400 unsupported_parameter) and
60 # requires max_completion_tokens. The newer name is accepted by all
61 # current models, so it is sent unconditionally.
62 "max_completion_tokens" : MAX_TOKENS ,
63 }
64 headers = {
65 "Authorization" : f "Bearer { os.environ[ 'OPENAI_API_KEY' ] } " ,
66 "Content-Type" : "application/json" ,
67 }
68 return "https://api.openai.com/v1/chat/completions" , headers, body
69
70
71 def build_anthropic_request (model: str , system: str , user_text: str ) -> tuple[ str , dict , dict ]:
72 body = {
73 "model" : model,
74 "max_tokens" : MAX_TOKENS ,
75 "messages" : [{ "role" : "user" , "content" : user_text}],
76 }
77 if system:
78 body[ "system" ] = system
79 headers = {
80 "x-api-key" : os.environ[ "ANTHROPIC_API_KEY" ],
81 "anthropic-version" : "2023-06-01" ,
82 "Content-Type" : "application/json" ,
83 }
84 return "https://api.anthropic.com/v1/messages" , headers, body
85
86
87 def build_gemini_request (model: str , system: str , user_text: str ) -> tuple[ str , dict , dict ]:
88 body = {
89 "contents" : [{ "parts" : [{ "text" : user_text}]}],
90 "generationConfig" : { "maxOutputTokens" : MAX_TOKENS },
91 }
92 if system:
93 # systemInstruction mirrors how the customer's app passes system
94 # prompts — concatenating into the user turn would change behavior.
95 body[ "systemInstruction" ] = { "parts" : [{ "text" : system}]}
96 # Key travels as a header, not a query parameter — URLs end up in logs.
97 headers = {
98 "x-goog-api-key" : os.environ[ "GEMINI_API_KEY" ],
99 "Content-Type" : "application/json" ,
100 }
101 url = f "https://generativelanguage.googleapis.com/v1beta/models/ { model } :generateContent"
102 return url, headers, body
103
104
105 PROVIDER_TO_KEY = {
106 "openai" : "OPENAI_API_KEY" ,
107 "anthropic" : "ANTHROPIC_API_KEY" ,
108 "google" : "GEMINI_API_KEY" ,
109 "gemini" : "GEMINI_API_KEY" ,
110 }
111
112
113 def pick_provider_key (file_pairs: dict ) -> str | None :
114 """Select the provider key name. SOURCE_PROVIDER (the helper's declared
115 input) is authoritative when set — a file carrying several provider keys
116 must not fall back to dict-order guessing. Without it, fall back to the
117 first recognized key IN THE FILE (never the ambient environment)."""
118 stated = os.environ.get( "SOURCE_PROVIDER" , "" ).lower()
119 if stated:
120 key = PROVIDER_TO_KEY .get(stated)
121 if key is None or key not in file_pairs:
122 return None
123 return key
124 return next ((k for k in PROVIDERS if k in file_pairs), None )
125
126
127 def _secret_variants (secrets) -> list :
128 """Every textual form a secret can take in error text. A raw CR/LF inside
129 a credential is rendered ESCAPED in exception messages (repr turns one
130 control char into backslash-r text), so literal replacement alone misses
131 it; the control-char-split fragments catch any remaining partial echo."""
132 import re as _re
133 out: list = []
134 for s in secrets:
135 if not s:
136 continue
137 out.append(s)
138 esc = s.encode( "unicode_escape" ).decode( "ascii" )
139 if esc != s:
140 out.append(esc)
141 for frag in _re.split( r " [ \x00 - \x1f ] + " , s):
142 if len (frag) >= 6 and frag not in out:
143 out.append(frag)
144 return sorted (out, key = len , reverse = True )
145
146
147 def redact (text: str , secrets) -> str :
148 """Strip secret values from failure text. A raised exception can embed a
149 header value verbatim (http.client rejects an invalid header with the full
150 'Bearer <key>' in the ValueError message), and the docstring promise is
151 that the key never reaches the output JSONL."""
152 for s in _secret_variants(secrets):
153 text = text.replace(s, "***" )
154 return text
155
156
157 PROVIDERS = {
158 # env key → (request builder, response-text extractor)
159 "OPENAI_API_KEY" : (
160 build_openai_request,
161 lambda d: d[ "choices" ][ 0 ][ "message" ][ "content" ],
162 ),
163 "ANTHROPIC_API_KEY" : (
164 build_anthropic_request,
165 lambda d: d[ "content" ][ 0 ][ "text" ],
166 ),
167 "GEMINI_API_KEY" : (
168 build_gemini_request,
169 lambda d: d[ "candidates" ][ 0 ][ "content" ][ "parts" ][ 0 ][ "text" ],
170 ),
171 }
172
173
174 ERROR_BODY_CAP = 2048
175 AUTH_CODES = ( 401 , 403 )
176
177
178 def error_detail (e: urllib.error.HTTPError, secrets = ()) -> str :
179 """Bounded, credential-free extract of the provider's error body.
180
181 The evaluator contract needs the HTTP 400 body's message/param to tell a
182 request-shape bug from a quota or auth problem — reason alone is usually
183 just "Bad Request" and the body is irrecoverable after this process exits.
184
185 Auth errors (401/403) return no detail at all: their classification uses
186 the status code alone, and an auth endpoint can echo the submitted
187 credential in its body. Every remaining path is redacted against the known
188 key values, including the non-JSON fallback.
189 """
190 if e.code in AUTH_CODES :
191 return ""
192 try :
193 raw = e.read( ERROR_BODY_CAP ).decode( "utf-8" , "replace" )
194 except Exception :
195 return ""
196 try :
197 err = json.loads(raw).get( "error" , {})
198 if isinstance (err, dict ):
199 msg = err.get( "message" , "" )
200 param = err.get( "param" )
201 detail = f " { msg } (param: { param } )" if param else msg
202 return redact(detail, secrets)
203 except (json.JSONDecodeError, AttributeError ):
204 pass
205 return redact(raw[: 200 ], secrets)
206
207
208 def _send (url: str , headers: dict , body: dict ) -> dict :
209 req = urllib.request.Request(url, data = json.dumps(body).encode( "utf-8" ), headers = headers)
210 with urllib.request.urlopen(req, timeout = 60 ) as resp: # nosec B310 — fixed https hosts
211 return json.loads(resp.read())
212
213
214 def main () -> int :
215 if len (sys.argv) != 2 :
216 print ( __doc__ .strip().splitlines()[ 2 ].strip(), file = sys.stderr)
217 return 2
218 file_pairs = load_env_file(sys.argv[ 1 ])
219
220 try :
221 model = os.environ[ "SOURCE_MODEL_ID" ]
222 dataset_path = os.environ[ "GOLDEN_DATASET_PATH" ]
223 output_path = os.environ[ "OUTPUT_PATH" ]
224 except KeyError as e:
225 print ( f "FAIL: required env var { e } not set" , file = sys.stderr)
226 return 2
227
228 provider = pick_provider_key(file_pairs)
229 if provider is None :
230 print ( "FAIL: no usable provider key in the env file (check SOURCE_PROVIDER)" , file = sys.stderr)
231 return 2
232 build, extract = PROVIDERS [provider]
233 secrets = [v for k, v in file_pairs.items() if k in PROVIDERS ]
234
235 with open (dataset_path) as f:
236 prompts = [json.loads(line) for line in f if line.strip()]
237
238 results = []
239 done_live = set ()
240 if os.path.exists(output_path):
241 with open (output_path) as f:
242 for line in f:
243 if line.strip():
244 row = json.loads(line)
245 if row.get( "status" ) == "live" :
246 results.append(row)
247 done_live.add(row[ "id" ])
248 prompts = [p for p in prompts if p[ "id" ] not in done_live]
249 if done_live:
250 print ( f "RESUME: { len (done_live) } live baselines kept, { len (prompts) } to fetch" )
251
252 for p in prompts:
253 try :
254 url, headers, body = build(model, p.get( "system_prompt" ) or "" , p[ "user_prompt" ])
255 out = extract(_send(url, headers, body))
256 results.append({ "id" : p[ "id" ], "source_response" : out, "status" : "live" })
257 except urllib.error.HTTPError as e:
258 detail = error_detail(e, secrets)
259 status = f "http_ { e.code } : { e.reason } " + ( f " — { detail } " if detail else "" )
260 results.append({ "id" : p[ "id" ], "source_response" : "" ,
261 "status" : redact(status, secrets)})
262 except Exception as e: # noqa: BLE001 — every row must land in the JSONL
263 results.append({ "id" : p[ "id" ], "source_response" : "" ,
264 "status" : redact( f "error: { type (e). __name__ } : { e } " , secrets)})
265
266 os.makedirs(os.path.dirname(output_path), exist_ok = True )
267 with open (output_path, "w" ) as f:
268 for r in results:
269 f.write(json.dumps(r) + " \n " )
270
271 ok = sum ( 1 for r in results if r[ "status" ] == "live" )
272 print ( f "live source baselines: { ok } / { len (results) } " )
273 return 0
274
275
276 if __name__ == "__main__" :
277 sys.exit(main())