@@ -77,6 +77,7 @@ def __init__(
7777 model_name : str ,
7878 provider_name : str ,
7979 context : Context ,
80+ graph_key : Optional [str ] = None ,
8081 ):
8182 """
8283 Initialize an AI Config tracker.
@@ -88,6 +89,8 @@ def __init__(
8889 :param model_name: Name of the model used.
8990 :param provider_name: Name of the provider used.
9091 :param context: Context for evaluation.
92+ :param graph_key: When set, include ``graphKey`` in all event payloads
93+ (e.g. config-level metrics inside a graph).
9194 """
9295 self ._ld_client = ld_client
9396 self ._variation_key = variation_key
@@ -96,13 +99,13 @@ def __init__(
9699 self ._model_name = model_name
97100 self ._provider_name = provider_name
98101 self ._context = context
102+ self ._graph_key = graph_key
99103 self ._summary = LDAIMetricSummary ()
100104
101- def __get_track_data (self , graph_key : Optional [ str ] = None ) -> dict :
105+ def __get_track_data (self ) -> dict :
102106 """
103107 Get tracking data for events.
104108
105- :param graph_key: When set, include ``graphKey`` in the payload.
106109 :return: Dictionary containing variation and config keys.
107110 """
108111 data = {
@@ -112,82 +115,72 @@ def __get_track_data(self, graph_key: Optional[str] = None) -> dict:
112115 "modelName" : self ._model_name ,
113116 "providerName" : self ._provider_name ,
114117 }
115- if graph_key is not None :
116- data ['graphKey' ] = graph_key
118+ if self . _graph_key is not None :
119+ data ['graphKey' ] = self . _graph_key
117120 return data
118121
119- def track_duration (self , duration : int , * , graph_key : Optional [ str ] = None ) -> None :
122+ def track_duration (self , duration : int ) -> None :
120123 """
121124 Manually track the duration of an AI operation.
122125
123126 :param duration: Duration in milliseconds.
124- :param graph_key: When set, include ``graphKey`` in the event payload
125- (e.g. config-level metrics inside a graph).
126127 """
127128 self ._summary ._duration = duration
128129 self ._ld_client .track (
129- "$ld:ai:duration:total" , self ._context , self .__get_track_data (graph_key ), duration
130+ "$ld:ai:duration:total" , self ._context , self .__get_track_data (), duration
130131 )
131132
132- def track_time_to_first_token (
133- self , time_to_first_token : int , * , graph_key : Optional [str ] = None
134- ) -> None :
133+ def track_time_to_first_token (self , time_to_first_token : int ) -> None :
135134 """
136135 Manually track the time to first token of an AI operation.
137136
138137 :param time_to_first_token: Time to first token in milliseconds.
139- :param graph_key: When set, include ``graphKey`` in the event payload.
140138 """
141139 self ._summary ._time_to_first_token = time_to_first_token
142140 self ._ld_client .track (
143141 "$ld:ai:tokens:ttf" ,
144142 self ._context ,
145- self .__get_track_data (graph_key ),
143+ self .__get_track_data (),
146144 time_to_first_token ,
147145 )
148146
149- def track_duration_of (self , func , * , graph_key : Optional [ str ] = None ):
147+ def track_duration_of (self , func ):
150148 """
151149 Automatically track the duration of an AI operation.
152150
153151 An exception occurring during the execution of the function will still
154152 track the duration. The exception will be re-thrown.
155153
156154 :param func: Function to track (synchronous only).
157- :param graph_key: When set, passed through to :meth:`track_duration`.
158155 :return: Result of the tracked function.
159156 """
160157 start_ns = time .perf_counter_ns ()
161158 try :
162159 result = func ()
163160 finally :
164161 duration = (time .perf_counter_ns () - start_ns ) // 1_000_000 # duration in milliseconds
165- self .track_duration (duration , graph_key = graph_key )
162+ self .track_duration (duration )
166163
167164 return result
168165
169166 def _track_from_metrics_extractor (
170167 self ,
171168 result : Any ,
172169 metrics_extractor : Callable [[Any ], Any ],
173- * ,
174- graph_key : Optional [str ] = None ,
175170 ) -> Any :
176171 metrics = metrics_extractor (result )
177172 if metrics .success :
178- self .track_success (graph_key = graph_key )
173+ self .track_success ()
179174 else :
180- self .track_error (graph_key = graph_key )
175+ self .track_error ()
181176 if metrics .usage :
182- self .track_tokens (metrics .usage , graph_key = graph_key )
177+ self .track_tokens (metrics .usage )
183178 return result
184179
185180 def track_metrics_of (
186181 self ,
187182 func : Callable [[], Any ],
188183 metrics_extractor : Callable [[Any ], Any ],
189- * ,
190- graph_key : Optional [str ] = None ,
191184 ) -> Any :
192185 """
193186 Track metrics for a synchronous AI operation.
@@ -203,33 +196,29 @@ def track_metrics_of(
203196
204197 :param func: Synchronous callable that runs the operation
205198 :param metrics_extractor: Function that extracts LDAIMetrics from the operation result
206- :param graph_key: When set, include ``graphKey`` on emitted config-level events.
207199 :return: The result of the operation
208200 """
209201 start_ns = time .perf_counter_ns ()
210202 try :
211203 result = func ()
212204 except Exception as err :
213205 duration = (time .perf_counter_ns () - start_ns ) // 1_000_000
214- self .track_duration (duration , graph_key = graph_key )
215- self .track_error (graph_key = graph_key )
206+ self .track_duration (duration )
207+ self .track_error ()
216208 raise err
217209
218210 duration = (time .perf_counter_ns () - start_ns ) // 1_000_000
219- self .track_duration (duration , graph_key = graph_key )
220- return self ._track_from_metrics_extractor (result , metrics_extractor , graph_key = graph_key )
211+ self .track_duration (duration )
212+ return self ._track_from_metrics_extractor (result , metrics_extractor )
221213
222- async def track_metrics_of_async (
223- self , func , metrics_extractor , * , graph_key : Optional [str ] = None
224- ):
214+ async def track_metrics_of_async (self , func , metrics_extractor ):
225215 """
226216 Track metrics for an async AI operation (``func`` is awaited).
227217
228218 Same event semantics as :meth:`track_metrics_of`.
229219
230220 :param func: Async callable or zero-arg callable that returns an awaitable when called
231221 :param metrics_extractor: Function that extracts LDAIMetrics from the operation result
232- :param graph_key: When set, include ``graphKey`` on emitted config-level events.
233222 :return: The result of the operation
234223 """
235224 start_ns = time .perf_counter_ns ()
@@ -238,20 +227,19 @@ async def track_metrics_of_async(
238227 result = await func ()
239228 except Exception as err :
240229 duration = (time .perf_counter_ns () - start_ns ) // 1_000_000
241- self .track_duration (duration , graph_key = graph_key )
242- self .track_error (graph_key = graph_key )
230+ self .track_duration (duration )
231+ self .track_error ()
243232 raise err
244233
245234 duration = (time .perf_counter_ns () - start_ns ) // 1_000_000
246- self .track_duration (duration , graph_key = graph_key )
247- return self ._track_from_metrics_extractor (result , metrics_extractor , graph_key = graph_key )
235+ self .track_duration (duration )
236+ return self ._track_from_metrics_extractor (result , metrics_extractor )
248237
249- def track_eval_scores (self , scores : Dict [str , Any ], * , graph_key : Optional [ str ] = None ) -> None :
238+ def track_eval_scores (self , scores : Dict [str , Any ]) -> None :
250239 """
251240 Track evaluation scores for multiple metrics.
252241
253242 :param scores: Dictionary mapping metric keys to their evaluation scores (EvalScore objects)
254- :param graph_key: When set, include ``graphKey`` in the event payload.
255243 """
256244 from ldai .providers .types import EvalScore
257245
@@ -261,23 +249,22 @@ def track_eval_scores(self, scores: Dict[str, Any], *, graph_key: Optional[str]
261249 self ._ld_client .track (
262250 metric_key ,
263251 self ._context ,
264- self .__get_track_data (graph_key = graph_key ),
252+ self .__get_track_data (),
265253 eval_score .score
266254 )
267255
268- def track_judge_response (self , judge_response : Any , * , graph_key : Optional [ str ] = None ) -> None :
256+ def track_judge_response (self , judge_response : Any ) -> None :
269257 """
270258 Track a judge response, including evaluation scores with judge config key.
271259
272260 :param judge_response: JudgeResponse object containing evals and success status
273- :param graph_key: When set, include ``graphKey`` in the event payload.
274261 """
275262 from ldai .providers .types import EvalScore , JudgeResponse
276263
277264 if isinstance (judge_response , JudgeResponse ):
278265 # Track evaluation scores with judge config key included in metadata
279266 if judge_response .evals :
280- track_data = self .__get_track_data (graph_key = graph_key )
267+ track_data = self .__get_track_data ()
281268 if judge_response .judge_config_key :
282269 track_data = {** track_data , 'judgeConfigKey' : judge_response .judge_config_key }
283270
@@ -290,49 +277,44 @@ def track_judge_response(self, judge_response: Any, *, graph_key: Optional[str]
290277 eval_score .score
291278 )
292279
293- def track_feedback (self , feedback : Dict [str , FeedbackKind ], * , graph_key : Optional [ str ] = None ) -> None :
280+ def track_feedback (self , feedback : Dict [str , FeedbackKind ]) -> None :
294281 """
295282 Track user feedback for an AI operation.
296283
297284 :param feedback: Dictionary containing feedback kind.
298- :param graph_key: When set, include ``graphKey`` in the event payload.
299285 """
300286 self ._summary ._feedback = feedback
301287 if feedback ["kind" ] == FeedbackKind .Positive :
302288 self ._ld_client .track (
303289 "$ld:ai:feedback:user:positive" ,
304290 self ._context ,
305- self .__get_track_data (graph_key = graph_key ),
291+ self .__get_track_data (),
306292 1 ,
307293 )
308294 elif feedback ["kind" ] == FeedbackKind .Negative :
309295 self ._ld_client .track (
310296 "$ld:ai:feedback:user:negative" ,
311297 self ._context ,
312- self .__get_track_data (graph_key = graph_key ),
298+ self .__get_track_data (),
313299 1 ,
314300 )
315301
316- def track_success (self , * , graph_key : Optional [ str ] = None ) -> None :
302+ def track_success (self ) -> None :
317303 """
318304 Track a successful AI generation.
319-
320- :param graph_key: When set, include ``graphKey`` in the event payload.
321305 """
322306 self ._summary ._success = True
323307 self ._ld_client .track (
324- "$ld:ai:generation:success" , self ._context , self .__get_track_data (graph_key = graph_key ), 1
308+ "$ld:ai:generation:success" , self ._context , self .__get_track_data (), 1
325309 )
326310
327- def track_error (self , * , graph_key : Optional [ str ] = None ) -> None :
311+ def track_error (self ) -> None :
328312 """
329313 Track an unsuccessful AI generation attempt.
330-
331- :param graph_key: When set, include ``graphKey`` in the event payload.
332314 """
333315 self ._summary ._success = False
334316 self ._ld_client .track (
335- "$ld:ai:generation:error" , self ._context , self .__get_track_data (graph_key = graph_key ), 1
317+ "$ld:ai:generation:error" , self ._context , self .__get_track_data (), 1
336318 )
337319
338320 def track_openai_metrics (self , func ):
@@ -390,15 +372,14 @@ def track_bedrock_converse_metrics(self, res: dict) -> dict:
390372 self .track_tokens (_bedrock_to_token_usage (res ["usage" ]))
391373 return res
392374
393- def track_tokens (self , tokens : TokenUsage , * , graph_key : Optional [ str ] = None ) -> None :
375+ def track_tokens (self , tokens : TokenUsage ) -> None :
394376 """
395377 Track token usage metrics.
396378
397379 :param tokens: Token usage data from either custom, OpenAI, or Bedrock sources.
398- :param graph_key: When set, include ``graphKey`` in the event payload.
399380 """
400381 self ._summary ._usage = tokens
401- td = self .__get_track_data (graph_key = graph_key )
382+ td = self .__get_track_data ()
402383 if tokens .total > 0 :
403384 self ._ld_client .track (
404385 "$ld:ai:tokens:total" ,
@@ -421,32 +402,28 @@ def track_tokens(self, tokens: TokenUsage, *, graph_key: Optional[str] = None) -
421402 tokens .output ,
422403 )
423404
424- def track_tool_call (self , tool_key : str , * , graph_key : Optional [ str ] = None ) -> None :
405+ def track_tool_call (self , tool_key : str ) -> None :
425406 """
426407 Track a tool invocation for this configuration (standalone or within a graph).
427408
428409 :param tool_key: Identifier of the tool that was invoked.
429- :param graph_key: When set, include ``graphKey`` in the event payload.
430410 """
431- track_data = {** self .__get_track_data (graph_key = graph_key ), "toolKey" : tool_key }
411+ track_data = {** self .__get_track_data (), "toolKey" : tool_key }
432412 self ._ld_client .track (
433413 "$ld:ai:tool_call" ,
434414 self ._context ,
435415 track_data ,
436416 1 ,
437417 )
438418
439- def track_tool_calls (
440- self , tool_keys : Iterable [str ], * , graph_key : Optional [str ] = None
441- ) -> None :
419+ def track_tool_calls (self , tool_keys : Iterable [str ]) -> None :
442420 """
443421 Track multiple tool invocations for this configuration.
444422
445423 :param tool_keys: Tool identifiers (e.g. from a model response).
446- :param graph_key: When set, include ``graphKey`` on each event.
447424 """
448425 for tool_key in tool_keys :
449- self .track_tool_call (tool_key , graph_key = graph_key )
426+ self .track_tool_call (tool_key )
450427
451428 def get_summary (self ) -> LDAIMetricSummary :
452429 """
0 commit comments