Setting the file. One moment.
Test Evaluate Traces · AWS Observability · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page 69
Creating Amazon Aurora Db Cluster With Instances
92
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Debugging Lambda Timeouts
— line 191
This file
Number 37.42
Position 42 of 67
Type Python
Size 12 KB
Lines 316 scripts/cloudwatch-omni/ test_evaluate_traces.py
Python · 316 lines · 12 KB
importlib.util
15 import io
16 import json
17 import os
18 import sys
19 import types
20
21 HERE = os.path.dirname(os.path.abspath( __file__ ))
22 _N = 0
23
24
25 class _ClientError ( Exception ):
26 def __init__ (self, response = None , op = None ):
27 self .response = response or {}
28 super (). __init__ ( str (response))
29
30
31 class _BotoCoreError ( Exception ):
32 pass
33
34
35 def _load (make_session):
36 """Import a fresh copy of evaluate_traces with boto3/botocore stubbed."""
37 global _N
38 _N += 1
39 be = types.ModuleType( "botocore.exceptions" )
40 be.ClientError = _ClientError
41 be.BotoCoreError = _BotoCoreError
42 bc = types.ModuleType( "botocore" )
43 bc.exceptions = be
44 b3 = types.ModuleType( "boto3" )
45 b3.Session = lambda ** _k: make_session()
46 sys.modules.update({ "boto3" : b3, "botocore" : bc, "botocore.exceptions" : be})
47
48 spec = importlib.util.spec_from_file_location(
49 "evaluate_traces_under_test_ %d " % _N, os.path.join( HERE , "evaluate_traces.py" )
50 )
51 mod = importlib.util.module_from_spec(spec)
52 spec.loader.exec_module(mod)
53 mod.time.sleep = lambda _: None # no-op poll delay so tests run instantly
54 return mod
55
56
57 def _make_omni_rows (docs):
58 """Convert span dicts to Omni SQL column-alias row format."""
59 rows = []
60 for d in docs:
61 attrs = d.get( "attributes" ) or {}
62 row = {
63 "traceId" : d.get( "traceId" ),
64 "spanId" : d.get( "spanId" ),
65 "parentSpanId" : d.get( "parentSpanId" ),
66 "name" : d.get( "name" , "span" ),
67 "kind" : d.get( "kind" , "INTERNAL" ),
68 "startTimeUnixNano" : "1000000000" ,
69 "endTimeUnixNano" : "2000000000" ,
70 "svc_name" : ((d.get( "resource" ) or {}).get( "attributes" ) or {}).get( "service.name" , "" ),
71 "scope_name" : "" ,
72 "status_code" : "" ,
73 # span-kind attributes
74 "oi_span_kind" : attrs.get( "openinference.span.kind" , "" ),
75 "aws_genai_span_kind" : attrs.get( "aws.genai.span_kind" , "" ),
76 "gen_ai_op_name" : attrs.get( "gen_ai.operation.name" , "" ),
77 # input/output
78 "input_value" : attrs.get( "input.value" , "" ),
79 "output_value" : attrs.get( "output.value" , "" ),
80 "gen_ai_input_value" : "" ,
81 "gen_ai_output_value" : "" ,
82 "gen_ai_input_messages" : "" ,
83 "gen_ai_output_messages" : "" ,
84 "llm_input_messages" : "" ,
85 "llm_output_messages" : "" ,
86 # tool
87 "tool_name" : attrs.get( "tool.name" , "" ),
88 "gen_ai_tool_name" : "" ,
89 "gen_ai_tool_call_id" : "" ,
90 # session
91 "session_id_attr" : attrs.get( "session.id" , "" ),
92 # LLM metadata
93 "gen_ai_req_model" : "" ,
94 "gen_ai_resp_model" : "" ,
95 "gen_ai_system" : "" ,
96 # token counts
97 "usage_input_tokens" : "" ,
98 "usage_output_tokens" : "" ,
99 "prompt_tokens" : "" ,
100 "completion_tokens" : "" ,
101 }
102 rows.append(row)
103 return rows
104
105
106 def _make_session (counters, * , tool_spans = 1 , session_id = None , eval_mode = "ok" , raise_cls = None ):
107 """Build a fake boto3 Session. `counters` records span-fetch and evaluate-call counts."""
108
109 def _span (span_id, is_tool):
110 attrs = {}
111 if session_id:
112 attrs[ "session.id" ] = session_id
113 if is_tool:
114 attrs[ "gen_ai.operation.name" ] = "execute_tool"
115 return { "traceId" : "t" , "spanId" : span_id, "kind" : "INTERNAL" ,
116 "attributes" : attrs, "resource" : { "attributes" : { "service.name" : "svc" }}}
117
118 docs = [_span( "0" * 16 , False )] + [_span( " %016x " % (i + 1 ), True ) for i in range (tool_spans)]
119
120 class FakeOmni :
121 """Fake cloudwatchomni client for Omni SQL span queries."""
122 _phase = {} # queryId -> "polled"
123
124 def start_telemetry_query_session (self, ** k):
125 return { "sessionId" : "sess-1" }
126
127 def start_telemetry_query (self, ** k):
128 qid = "q- %d " % len (FakeOmni._phase)
129 FakeOmni._phase[qid] = { "sql" : k.get( "queryString" , "" ), "polled" : False }
130 return { "queryId" : qid, "sessionId" : k.get( "sessionId" )}
131
132 def get_telemetry_query_results (self, queryId, maxResults = 1 , nextToken = None ):
133 if raise_cls is not None :
134 raise raise_cls(
135 { "Error" : { "Code" : "AccessDeniedException" , "Message" : "no cloudwatchomni access" }},
136 "StartTelemetryQuery" ,
137 )
138 state = FakeOmni._phase.get(queryId, {})
139 # Phase 1 poll (maxResults=1): just confirm Complete
140 if maxResults == 1 :
141 FakeOmni._phase[queryId][ "polled" ] = True
142 return { "status" : "Complete" , "rows" : []}
143 # Phase 2 fetch
144 sql = state.get( "sql" , "" )
145 if "session_id" in sql and "session.id" in sql and "IS NOT NULL" in sql:
146 # session resolution query
147 rows = [{ "session_id" : session_id}] if session_id else []
148 return { "status" : "Complete" , "rows" : rows}
149 # span fetch query
150 counters[ "fetch" ] = counters.get( "fetch" , 0 ) + 1
151 return { "status" : "Complete" , "rows" : _make_omni_rows(docs)}
152
153 def stop_telemetry_query_session (self, ** k):
154 return {}
155
156 class FakeLogs :
157 """Fake logs client — writeback path only (put_log_events)."""
158 def put_log_events (self, ** k):
159 return {}
160
161 def create_log_group (self, ** k):
162 return {}
163
164 def create_log_stream (self, ** k):
165 return {}
166
167 def put_retention_policy (self, ** k):
168 return {}
169
170 class FakeAC :
171 def evaluate (self, ** k):
172 counters[ "eval" ] = counters.get( "eval" , 0 ) + 1
173 if eval_mode == "error" :
174 return { "evaluationResults" : [{ "evaluatorId" : k[ "evaluatorId" ], "value" : None ,
175 "errorCode" : "AgentSpanMappingException" ,
176 "errorMessage" : "boom" }]}
177 return { "evaluationResults" : [{ "evaluatorId" : k[ "evaluatorId" ], "value" : 0.5 ,
178 "label" : "ok" , "explanation" : "why" }]}
179
180 class FakeSession :
181 def client (self, name):
182 if name == "cloudwatchomni" :
183 return FakeOmni()
184 if name == "logs" :
185 return FakeLogs()
186 return FakeAC()
187
188 return FakeSession()
189
190
191 def _run (mod, ** kw):
192 ns = mod.argparse.Namespace(
193 trace_id = None , trace_ids = None , session_id = None , evaluator_id = None , evaluator_ids = None ,
194 level = "TRACE" , region = "us-east-1" , window_days = 30 , no_writeback = True ,
195 retention_days = 30 , kms_key_id = None , include_explanations = False )
196 for k, v in kw.items():
197 setattr (ns, k, v)
198 out = io.StringIO()
199 try :
200 with contextlib.redirect_stdout(out):
201 mod._run(ns)
202 except SystemExit :
203 pass
204 return json.loads(out.getvalue())
205
206
207 RESULTS = []
208
209
210 def check (name, cond):
211 RESULTS .append((name, bool (cond)))
212 print (( "PASS " if cond else "FAIL " ) + name)
213
214
215 def test_single_trace_back_compat ():
216 c = {}
217 mod = _load( lambda : _make_session(c))
218 r = _run(mod, trace_id = "aa" )
219 check( "single-trace: one pair scored, one fetch" , r[ "pairsScored" ] == 1 and c[ "fetch" ] == 1 )
220 check( "single-trace: row tagged with traceId + evaluatorId" ,
221 { "traceId" , "evaluatorId" , "value" } <= set (r[ "results" ][ 0 ]))
222
223
224 def test_matrix_fetch_once ():
225 c = {}
226 mod = _load( lambda : _make_session(c))
227 r = _run(mod, trace_ids = "aa,bb" , evaluator_ids = "E1,E2" )
228 check( "matrix 2x2: 4 pairs scored" , r[ "pairsScored" ] == 4 )
229 check( "matrix 2x2: fetch-once (2 fetches, not 4)" , c[ "fetch" ] == 2 )
230
231
232 def test_ground_truth_gating ():
233 mod = _load( lambda : _make_session({}))
234 r = _run(mod, trace_id = "aa" , evaluator_ids = "Builtin.Helpfulness,Builtin.TrajectoryExactOrderMatch" )
235 check( "mixed GT: drops trajectory matcher, scores the rest" ,
236 r.get( "evaluators" ) == [ "Builtin.Helpfulness" ])
237 r2 = _run(mod, trace_id = "aa" , evaluator_ids = "Builtin.TrajectoryExactOrderMatch" )
238 check( "all-GT: refuses the run" , "error" in r2)
239
240
241 def test_caps ():
242 mod = _load( lambda : _make_session({}))
243 r = _run(mod, trace_ids = "," .join( " %02x " % i for i in range ( 60 )))
244 check( "trace cap: 60 -> 50 with a note" ,
245 r[ "tracesScored" ] == 50 and any ( "capped to 50 traces" in n for n in r.get( "notes" , [])))
246 r2 = _run(mod, trace_ids = "aa,bb,cc" , evaluator_ids = "," .join( "E %d " % i for i in range ( 12 )))
247 check( "evaluator cap: 12 -> 10 with a note" ,
248 len (r2[ "evaluators" ]) == 10 and any ( "capped to 10 evaluators" in n for n in r2.get( "notes" , [])))
249 r3 = _run(mod, trace_ids = "," .join( " %02x " % i for i in range ( 40 )), evaluator_ids = "E1,E2,E3" )
250 check( "pair cap: 40x3 -> 33 traces (pair bound bites below the 50-trace bound)" ,
251 r3[ "tracesScored" ] == 33 and r3[ "pairsScored" ] == 99
252 and any ( "100-pair" in n for n in r3.get( "notes" , [])))
253
254
255 def test_evaluate_call_budget ():
256 c = {}
257 mod = _load( lambda : _make_session(c, tool_spans = 25 )) # 25 tool spans -> 3 chunks/pair
258 mod. MAX_EVALUATE_CALLS = 5
259 r = _run(mod, trace_ids = "aa,bb,cc" , evaluator_ids = "E1,E2" , level = "TOOL_CALL" )
260 check( "TOOL_CALL: total evaluate calls bounded by budget" , c[ "eval" ] <= 5 )
261 check( "TOOL_CALL: budget/truncation noted" ,
262 any ( "budget" in n for n in r.get( "notes" , [])))
263
264
265 def test_pairs_scored_excludes_errors ():
266 mod = _load( lambda : _make_session({}, eval_mode = "error" ))
267 r = _run(mod, trace_id = "aa" )
268 check( "error-only result: not counted in pairsScored" , r[ "pairsScored" ] == 0 )
269 check( "error-only result: still surfaced in results" ,
270 any (x.get( "errorCode" ) for x in r[ "results" ]))
271
272
273 def test_traces_scored_honesty_under_budget ():
274 c = {}
275 mod = _load( lambda : _make_session(c, tool_spans = 25 ))
276 mod. MAX_EVALUATE_CALLS = 3 # only the first trace's chunks fit
277 r = _run(mod, trace_ids = "aa,bb,cc" , level = "TOOL_CALL" )
278 check( "tracesScored < tracesFetched when budget truncates" ,
279 r[ "tracesScored" ] < r[ "tracesFetched" ])
280 check( "tracesScored == distinct traces with a real score" ,
281 r[ "tracesScored" ] == len ({
282 x[ "traceId" ] for x in r[ "results" ]
283 if x.get( "value" ) is not None and not x.get( "errorCode" )
284 }))
285
286
287 def test_session_dedup ():
288 c = {}
289 mod = _load( lambda : _make_session(c, session_id = "S1" ))
290 r = _run(mod, trace_ids = "aa,bb" , evaluator_ids = "E1" , level = "SESSION" )
291 check( "same-session: fetched once" , c[ "fetch" ] == 1 )
292 check( "same-session: SESSION scored once (no duplicate)" ,
293 len ([x for x in r[ "results" ] if x.get( "value" ) is not None ]) == 1 )
294
295
296 def test_real_error_propagated ():
297 mod = _load( lambda : _make_session({}, raise_cls = _ClientError))
298 r = _run(mod, trace_id = "aa" )
299 blob = json.dumps(r)
300 check( "real Omni error surfaced (not the generic window message)" ,
301 "AccessDenied" in blob or "cloudwatchomni" in blob)
302
303
304 def main ():
305 for fn in [test_single_trace_back_compat, test_matrix_fetch_once, test_ground_truth_gating,
306 test_caps, test_evaluate_call_budget, test_pairs_scored_excludes_errors,
307 test_traces_scored_honesty_under_budget, test_session_dedup,
308 test_real_error_propagated]:
309 fn()
310 passed = sum ( 1 for _, ok in RESULTS if ok)
311 print ( " \n%d / %d passed" % (passed, len ( RESULTS )))
312 sys.exit( 0 if passed == len ( RESULTS ) else 1 )
313
314
315 if __name__ == "__main__" :
316 main()