Skip to content

Commit ef2147d

Browse files
knfreemLDknfreemLD
committed
Added provider and model to ai tracker
1 parent df8cc20 commit ef2147d

File tree

3 files changed

+61
-47
lines changed

3 files changed

+61
-47
lines changed

‎ldai/client.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,8 @@ def __evaluate(
402402
variation.get('_ldMeta', {}).get('variationKey', ''),
403403
key,
404404
int(variation.get('_ldMeta', {}).get('version', 1)),
405+
model.name if model else '',
406+
provider_config.name if provider_config else '',
405407
context,
406408
)
407409

‎ldai/testing/test_tracker.py

Lines changed: 51 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def client(td: TestData) -> LDClient:
4242

4343
def test_summary_starts_empty(client: LDClient):
4444
context = Context.create("user-key")
45-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 1, context)
45+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 1, "fakeModel", "fakeProvider", context)
4646

4747
assert tracker.get_summary().duration is None
4848
assert tracker.get_summary().feedback is None
@@ -52,13 +52,13 @@ def test_summary_starts_empty(client: LDClient):
5252

5353
def test_tracks_duration(client: LDClient):
5454
context = Context.create("user-key")
55-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
55+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
5656
tracker.track_duration(100)
5757

5858
client.track.assert_called_with( # type: ignore
5959
"$ld:ai:duration:total",
6060
context,
61-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
61+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
6262
100,
6363
)
6464

@@ -67,7 +67,7 @@ def test_tracks_duration(client: LDClient):
6767

6868
def test_tracks_duration_of(client: LDClient):
6969
context = Context.create("user-key")
70-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
70+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
7171
tracker.track_duration_of(lambda: sleep(0.01))
7272

7373
calls = client.track.mock_calls # type: ignore
@@ -79,19 +79,21 @@ def test_tracks_duration_of(client: LDClient):
7979
"variationKey": "variation-key",
8080
"configKey": "config-key",
8181
"version": 3,
82+
"modelName": "fakeModel",
83+
"providerName": "fakeProvider",
8284
}
8385
assert calls[0].args[3] == pytest.approx(10, rel=10)
8486

8587

8688
def test_tracks_time_to_first_token(client: LDClient):
8789
context = Context.create("user-key")
88-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
90+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
8991
tracker.track_time_to_first_token(100)
9092

9193
client.track.assert_called_with( # type: ignore
9294
"$ld:ai:tokens:ttf",
9395
context,
94-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
96+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
9597
100,
9698
)
9799

@@ -100,7 +102,7 @@ def test_tracks_time_to_first_token(client: LDClient):
100102

101103
def test_tracks_duration_of_with_exception(client: LDClient):
102104
context = Context.create("user-key")
103-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
105+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
104106

105107
def sleep_and_throw():
106108
sleep(0.01)
@@ -121,13 +123,15 @@ def sleep_and_throw():
121123
"variationKey": "variation-key",
122124
"configKey": "config-key",
123125
"version": 3,
126+
"modelName": "fakeModel",
127+
"providerName": "fakeProvider",
124128
}
125129
assert calls[0].args[3] == pytest.approx(10, rel=10)
126130

127131

128132
def test_tracks_token_usage(client: LDClient):
129133
context = Context.create("user-key")
130-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
134+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
131135

132136
tokens = TokenUsage(300, 200, 100)
133137
tracker.track_tokens(tokens)
@@ -136,19 +140,19 @@ def test_tracks_token_usage(client: LDClient):
136140
call(
137141
"$ld:ai:tokens:total",
138142
context,
139-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
143+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
140144
300,
141145
),
142146
call(
143147
"$ld:ai:tokens:input",
144148
context,
145-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
149+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
146150
200,
147151
),
148152
call(
149153
"$ld:ai:tokens:output",
150154
context,
151-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
155+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
152156
100,
153157
),
154158
]
@@ -160,7 +164,7 @@ def test_tracks_token_usage(client: LDClient):
160164

161165
def test_tracks_bedrock_metrics(client: LDClient):
162166
context = Context.create("user-key")
163-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
167+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
164168

165169
bedrock_result = {
166170
"ResponseMetadata": {"HTTPStatusCode": 200},
@@ -179,37 +183,37 @@ def test_tracks_bedrock_metrics(client: LDClient):
179183
call(
180184
"$ld:ai:generation",
181185
context,
182-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
186+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
183187
1,
184188
),
185189
call(
186190
"$ld:ai:generation:success",
187191
context,
188-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
192+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
189193
1,
190194
),
191195
call(
192196
"$ld:ai:duration:total",
193197
context,
194-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
198+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
195199
50,
196200
),
197201
call(
198202
"$ld:ai:tokens:total",
199203
context,
200-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
204+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
201205
330,
202206
),
203207
call(
204208
"$ld:ai:tokens:input",
205209
context,
206-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
210+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
207211
220,
208212
),
209213
call(
210214
"$ld:ai:tokens:output",
211215
context,
212-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
216+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
213217
110,
214218
),
215219
]
@@ -223,7 +227,7 @@ def test_tracks_bedrock_metrics(client: LDClient):
223227

224228
def test_tracks_bedrock_metrics_with_error(client: LDClient):
225229
context = Context.create("user-key")
226-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
230+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
227231

228232
bedrock_result = {
229233
"ResponseMetadata": {"HTTPStatusCode": 500},
@@ -242,37 +246,37 @@ def test_tracks_bedrock_metrics_with_error(client: LDClient):
242246
call(
243247
"$ld:ai:generation",
244248
context,
245-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
249+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
246250
1,
247251
),
248252
call(
249253
"$ld:ai:generation:error",
250254
context,
251-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
255+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
252256
1,
253257
),
254258
call(
255259
"$ld:ai:duration:total",
256260
context,
257-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
261+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
258262
50,
259263
),
260264
call(
261265
"$ld:ai:tokens:total",
262266
context,
263-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
267+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
264268
330,
265269
),
266270
call(
267271
"$ld:ai:tokens:input",
268272
context,
269-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
273+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
270274
220,
271275
),
272276
call(
273277
"$ld:ai:tokens:output",
274278
context,
275-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
279+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
276280
110,
277281
),
278282
]
@@ -286,7 +290,7 @@ def test_tracks_bedrock_metrics_with_error(client: LDClient):
286290

287291
def test_tracks_openai_metrics(client: LDClient):
288292
context = Context.create("user-key")
289-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
293+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
290294

291295
class Result:
292296
def __init__(self):
@@ -306,31 +310,31 @@ def to_dict(self):
306310
call(
307311
"$ld:ai:generation",
308312
context,
309-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
313+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
310314
1,
311315
),
312316
call(
313317
"$ld:ai:generation:success",
314318
context,
315-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
319+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
316320
1,
317321
),
318322
call(
319323
"$ld:ai:tokens:total",
320324
context,
321-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
325+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
322326
330,
323327
),
324328
call(
325329
"$ld:ai:tokens:input",
326330
context,
327-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
331+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
328332
220,
329333
),
330334
call(
331335
"$ld:ai:tokens:output",
332336
context,
333-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
337+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
334338
110,
335339
),
336340
]
@@ -342,7 +346,7 @@ def to_dict(self):
342346

343347
def test_tracks_openai_metrics_with_exception(client: LDClient):
344348
context = Context.create("user-key")
345-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
349+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
346350

347351
def raise_exception():
348352
raise ValueError("Something went wrong")
@@ -357,13 +361,13 @@ def raise_exception():
357361
call(
358362
"$ld:ai:generation",
359363
context,
360-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
364+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
361365
1,
362366
),
363367
call(
364368
"$ld:ai:generation:error",
365369
context,
366-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
370+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
367371
1,
368372
),
369373
]
@@ -382,35 +386,35 @@ def raise_exception():
382386
)
383387
def test_tracks_feedback(client: LDClient, kind: FeedbackKind, label: str):
384388
context = Context.create("user-key")
385-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
389+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
386390

387391
tracker.track_feedback({"kind": kind})
388392

389393
client.track.assert_called_with( # type: ignore
390394
f"$ld:ai:feedback:user:{label}",
391395
context,
392-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
396+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
393397
1,
394398
)
395399
assert tracker.get_summary().feedback == {"kind": kind}
396400

397401

398402
def test_tracks_success(client: LDClient):
399403
context = Context.create("user-key")
400-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
404+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
401405
tracker.track_success()
402406

403407
calls = [
404408
call(
405409
"$ld:ai:generation",
406410
context,
407-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
411+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
408412
1,
409413
),
410414
call(
411415
"$ld:ai:generation:success",
412416
context,
413-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
417+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
414418
1,
415419
),
416420
]
@@ -422,20 +426,20 @@ def test_tracks_success(client: LDClient):
422426

423427
def test_tracks_error(client: LDClient):
424428
context = Context.create("user-key")
425-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
429+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
426430
tracker.track_error()
427431

428432
calls = [
429433
call(
430434
"$ld:ai:generation",
431435
context,
432-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
436+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
433437
1,
434438
),
435439
call(
436440
"$ld:ai:generation:error",
437441
context,
438-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
442+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
439443
1,
440444
),
441445
]
@@ -447,33 +451,33 @@ def test_tracks_error(client: LDClient):
447451

448452
def test_error_overwrites_success(client: LDClient):
449453
context = Context.create("user-key")
450-
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context)
454+
tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
451455
tracker.track_success()
452456
tracker.track_error()
453457

454458
calls = [
455459
call(
456460
"$ld:ai:generation",
457461
context,
458-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
462+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
459463
1,
460464
),
461465
call(
462466
"$ld:ai:generation:success",
463467
context,
464-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
468+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
465469
1,
466470
),
467471
call(
468472
"$ld:ai:generation",
469473
context,
470-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
474+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
471475
1,
472476
),
473477
call(
474478
"$ld:ai:generation:error",
475479
context,
476-
{"variationKey": "variation-key", "configKey": "config-key", "version": 3},
480+
{"variationKey": "variation-key", "configKey": "config-key", "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"},
477481
1,
478482
),
479483
]

0 commit comments

Comments
 (0)