langfuse

1
import { pruneDatabase } from "@/src/__tests__/test-utils";
2
import {
3
  createQuery,
4
  enrichAndCreateQuery,
5
  executeQuery,
6
} from "@/src/server/api/services/query-builder";
7
import { type aggregations } from "@/src/server/api/services/sqlInterface";
8
import { prisma } from "@/src/server/db";
9
import { type z } from "zod";
10

11
describe("Build valid SQL queries", () => {
12
  beforeEach(async () => await pruneDatabase());
13

14
  describe("should enrich mandatory filters", () => {
15
    [
16
      {
17
        table: "traces",
18
        values: ["project-id"],
19
        strings: [' FROM  traces t  WHERE  t."project_id" = ', ";"],
20
      } as const,
21
      {
22
        table: "traces_metrics",
23
        values: ["project-id"],
24
        strings: [' FROM traces_view t  WHERE  t."project_id" = ', ";"],
25
      } as const,
26
      {
27
        table: "traces_observations",
28
        values: ["project-id", "project-id"],
29
        strings: [
30
          ' FROM  traces t LEFT JOIN observations o ON t.id = o.trace_id  WHERE  t."project_id" = ',
31
          ' AND o."project_id" = ',
32
          ";",
33
        ],
34
      } as const,
35
      {
36
        table: "traces_observationsview",
37
        values: ["project-id", "project-id"],
38
        strings: [
39
          ' FROM  traces t LEFT JOIN observations_view o ON t.id = o.trace_id  WHERE  t."project_id" = ',
40
          ' AND o."project_id" = ',
41
          ";",
42
        ],
43
      } as const,
44
      {
45
        table: "observations",
46
        values: ["project-id"],
47
        strings: [' FROM  observations_view o  WHERE  o."project_id" = ', ";"],
48
      } as const,
49
      {
50
        table: "traces_scores",
51
        values: ["project-id"],
52
        strings: [
53
          ' FROM  traces t JOIN scores s ON t.id = s.trace_id  WHERE  t."project_id" = ',
54
          ";",
55
        ],
56
      } as const,
57
    ].forEach((prop) => {
58
      it(`should enrich mandatory filters ${prop.table}`, () => {
59
        const preparedQuery = enrichAndCreateQuery("project-id", {
60
          from: prop.table,
61
          select: [],
62
        });
63
        expect(preparedQuery.values).toEqual(prop.values);
64
        expect(preparedQuery.strings).toEqual(prop.strings);
65
      });
66
    });
67
  });
68

69
  describe("should build safe SQL", () => {
70
    it("should build a simple filter query", () => {
71
      const preparedQuery = createQuery({
72
        from: "traces",
73
        filter: [
74
          {
75
            type: "string" as const,
76
            column: "tracesProjectId",
77
            operator: "=" as const,
78
            value: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
79
          },
80
        ],
81

82
        select: [{ column: "traceId" }],
83
      });
84

85
      expect(preparedQuery.values).toEqual([
86
        "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
87
      ]);
88
    });
89

90
    it("should build a simple group by and filter query", () => {
91
      const preparedQuery = createQuery({
92
        from: "traces",
93
        filter: [
94
          {
95
            type: "string" as const,
96
            column: "tracesProjectId",
97
            operator: "=" as const,
98
            value: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
99
          },
100
        ],
101
        groupBy: [{ type: "string", column: "version" }],
102
        select: [{ column: "traceId" }],
103
      });
104

105
      expect(preparedQuery.values).toEqual([
106
        "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
107
      ]);
108
    });
109

110
    it("should build a time series group", () => {
111
      const preparedQuery = createQuery({
112
        from: "observations",
113
        filter: [
114
          {
115
            type: "datetime",
116
            column: "startTime",
117
            operator: ">=",
118
            value: new Date("2021-01-01T00:00:00.000Z"),
119
          },
120
          {
121
            type: "datetime",
122
            column: "startTime",
123
            operator: "<=",
124
            value: new Date("2021-01-04T00:00:00.000Z"),
125
          },
126
        ],
127
        groupBy: [
128
          { type: "datetime", column: "startTime", temporalUnit: "day" },
129
        ],
130
        select: [{ column: "completionTokens", agg: "SUM" }],
131
      });
132

133
      expect(preparedQuery.values).toEqual([
134
        new Date("2021-01-01T00:00:00.000Z"),
135
        new Date("2021-01-04T00:00:00.000Z"),
136
        new Date("2021-01-01T00:00:00.000Z"),
137
        new Date("2021-01-04T00:00:00.000Z"),
138
      ]);
139
    });
140

141
    it("should not filter an unknown column", () => {
142
      expect(() =>
143
        createQuery({
144
          from: "traces",
145
          filter: [
146
            { type: "string", column: "unknown", operator: "=", value: "" },
147
          ],
148
          select: [],
149
        }),
150
      ).toThrow("Invalid filter column: unknown");
151
    });
152

153
    it("should not select an unknown column", () => {
154
      expect(() =>
155
        createQuery({
156
          from: "traces",
157
          select: [{ column: "unknown" }],
158
        }),
159
      ).toThrow('Column "unknown" not found in table traces');
160
    });
161

162
    it("should not group by an unknown column", () => {
163
      expect(() =>
164
        createQuery({
165
          from: "traces",
166
          groupBy: [{ column: "unknown", type: "string" }],
167
          select: [],
168
        }),
169
      ).toThrow('Column "unknown" not found in table traces');
170
    });
171

172
    it("should not order by an unknown column", () => {
173
      expect(() =>
174
        createQuery({
175
          from: "traces",
176
          select: [],
177
          orderBy: [{ column: "unknown", direction: "ASC" }],
178
        }),
179
      ).toThrow('Column "unknown" not found in table traces');
180
    });
181
  });
182

183
  describe("should retrieve data", () => {
184
    it("should get a simple trace", async () => {
185
      await prisma.project.upsert({
186
        where: { id: "different-project-id" },
187
        create: {
188
          id: "different-project-id",
189
          name: "test-project",
190
        },
191
        update: {},
192
      });
193

194
      await prisma.trace.createMany({
195
        data: [
196
          {
197
            id: "trace-1",
198
            name: "trace-1",
199
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
200
            userId: "user-1",
201
            metadata: { key: "value" },
202
            release: "1.0.0",
203
            version: "2.0.0",
204
          },
205
          {
206
            id: "trace-2",
207
            name: "trace-1",
208
            projectId: "different-project-id",
209
            userId: "user-1",
210
            metadata: { key: "value" },
211
            release: "1.0.0",
212
            version: "2.0.0",
213
          },
214
        ],
215
      });
216

217
      const result = await executeQuery(
218
        prisma,
219
        "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
220
        {
221
          from: "traces",
222
          select: [{ column: "traceId" }],
223
        },
224
      );
225

226
      expect(result).toEqual([{ traceId: "trace-1" }]);
227
    });
228

229
    [
230
      {
231
        agg: "SUM",
232
        first: { sumCompletionTokens: 8, name: "trace-1" },
233
        second: { sumCompletionTokens: 4, name: "trace-2" },
234
      },
235
      {
236
        agg: "AVG",
237
        first: { avgCompletionTokens: 4, name: "trace-1" },
238
        second: { avgCompletionTokens: 4, name: "trace-2" },
239
      },
240
      {
241
        agg: "MIN",
242
        first: { minCompletionTokens: 3, name: "trace-1" },
243
        second: { minCompletionTokens: 4, name: "trace-2" },
244
      },
245
      {
246
        agg: "MAX",
247
        first: { maxCompletionTokens: 5, name: "trace-1" },
248
        second: { maxCompletionTokens: 4, name: "trace-2" },
249
      },
250
      {
251
        agg: "COUNT",
252
        first: { countCompletionTokens: 2, name: "trace-1" },
253
        second: { countCompletionTokens: 1, name: "trace-2" },
254
      },
255
    ].forEach((prop) => {
256
      it(`should group by name and aggregate ${prop.agg}`, async () => {
257
        await prisma.trace.create({
258
          data: {
259
            id: "trace-1",
260
            name: "trace-1",
261
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
262
          },
263
        });
264

265
        await prisma.observation.createMany({
266
          data: [
267
            {
268
              traceId: "trace-1",
269
              name: "trace-1",
270
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
271
              type: "GENERATION",
272
              completionTokens: 5,
273
            },
274
            {
275
              traceId: "trace-1",
276
              name: "trace-1",
277
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
278
              type: "GENERATION",
279
              completionTokens: 3,
280
            },
281
            {
282
              traceId: "trace-1",
283
              name: "trace-2",
284
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
285
              type: "GENERATION",
286
              completionTokens: 4,
287
            },
288
          ],
289
        });
290

291
        const result = await executeQuery(
292
          prisma,
293
          "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
294
          {
295
            from: "observations",
296
            groupBy: [{ type: "string", column: "name" }],
297
            select: [
298
              { column: "completionTokens", agg: prop.agg as "SUM" | "AVG" },
299
              { column: "name" },
300
            ],
301
          },
302
        );
303

304
        expect(result[0]!).toStrictEqual(prop.first);
305
        expect(result[1]!).toStrictEqual(prop.second);
306
      });
307
    });
308

309
    it("should  order by a column", async () => {
310
      await prisma.trace.create({
311
        data: {
312
          id: "trace-1",
313
          name: "trace-1",
314
          projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
315
        },
316
      });
317

318
      await prisma.observation.createMany({
319
        data: [
320
          {
321
            traceId: "trace-1",
322
            name: "trace-1",
323
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
324
            type: "GENERATION",
325
            completionTokens: 5,
326
            startTime: new Date("2021-01-01T00:00:00.000Z"),
327
          },
328
          {
329
            traceId: "trace-1",
330
            name: "trace-1",
331
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
332
            type: "GENERATION",
333
            completionTokens: 3,
334
            startTime: new Date("2021-01-01T00:00:00.000Z"),
335
          },
336
          {
337
            traceId: "trace-1",
338
            name: "trace-2",
339
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
340
            type: "GENERATION",
341
            completionTokens: 4,
342
            startTime: new Date("2021-01-02T00:00:00.000Z"),
343
          },
344
        ],
345
      });
346

347
      const result = await executeQuery(
348
        prisma,
349
        "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
350
        {
351
          from: "observations",
352
          filter: [
353
            {
354
              type: "datetime",
355
              column: "startTime",
356
              operator: ">=",
357
              value: new Date("2021-01-01T00:00:00.000Z"),
358
            },
359
            {
360
              type: "datetime",
361
              column: "startTime",
362
              operator: "<=",
363
              value: new Date("2021-01-04T00:00:00.000Z"),
364
            },
365
          ],
366

367
          select: [{ column: "completionTokens" }],
368
          orderBy: [{ column: "completionTokens", direction: "ASC" }],
369
        },
370
      );
371

372
      expect(result).toStrictEqual([
373
        { completionTokens: 3 },
374
        { completionTokens: 4 },
375
        { completionTokens: 5 },
376
      ]);
377
    });
378

379
    [{ agg: "SUM", one: 8, two: 4 }].forEach((prop) => {
380
      it(`should aggregate time series ${prop.agg}`, async () => {
381
        await prisma.trace.create({
382
          data: {
383
            id: "trace-1",
384
            name: "trace-1",
385
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
386
          },
387
        });
388

389
        await prisma.observation.createMany({
390
          data: [
391
            {
392
              traceId: "trace-1",
393
              name: "trace-1",
394
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
395
              type: "GENERATION",
396
              completionTokens: 5,
397
              startTime: new Date("2021-01-01T00:00:00.000Z"),
398
            },
399
            {
400
              traceId: "trace-1",
401
              name: "trace-1",
402
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
403
              type: "GENERATION",
404
              completionTokens: 3,
405
              startTime: new Date("2021-01-01T00:00:00.000Z"),
406
            },
407
            {
408
              traceId: "trace-1",
409
              name: "trace-2",
410
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
411
              type: "GENERATION",
412
              completionTokens: 4,
413
              startTime: new Date("2021-01-02T00:00:00.000Z"),
414
            },
415
          ],
416
        });
417

418
        const result = await executeQuery(
419
          prisma,
420
          "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
421
          {
422
            from: "observations",
423
            filter: [
424
              {
425
                type: "datetime",
426
                column: "startTime",
427
                operator: ">=",
428
                value: new Date("2021-01-01T00:00:00.000Z"),
429
              },
430
              {
431
                type: "datetime",
432
                column: "startTime",
433
                operator: "<=",
434
                value: new Date("2021-01-04T00:00:00.000Z"),
435
              },
436
            ],
437
            groupBy: [
438
              { type: "datetime", column: "startTime", temporalUnit: "day" },
439
            ],
440
            select: [
441
              { column: "completionTokens", agg: prop.agg as "SUM" | "AVG" },
442
            ],
443
          },
444
        );
445

446
        expect(result).toStrictEqual([
447
          {
448
            startTime: new Date("2021-01-01T00:00:00.000Z"),
449
            sumCompletionTokens: 8,
450
          },
451
          {
452
            startTime: new Date("2021-01-02T00:00:00.000Z"),
453
            sumCompletionTokens: 4,
454
          },
455
          {
456
            startTime: new Date("2021-01-03T00:00:00.000Z"),
457
            sumCompletionTokens: null,
458
          },
459
          {
460
            startTime: new Date("2021-01-04T00:00:00.000Z"),
461
            sumCompletionTokens: null,
462
          },
463
        ]);
464
      });
465
    });
466

467
    [
468
      {
469
        percentile: "50thPercentile",
470
        expectedOutcome: [
471
          {
472
            startTime: new Date("2021-01-01T00:00:00.000Z"),
473
            percentile50Duration: 8,
474
          },
475
          {
476
            startTime: new Date("2021-01-02T00:00:00.000Z"),
477
            percentile50Duration: 5,
478
          },
479
          {
480
            startTime: new Date("2021-01-03T00:00:00.000Z"),
481
            percentile50Duration: null,
482
          },
483
          {
484
            startTime: new Date("2021-01-04T00:00:00.000Z"),
485
            percentile50Duration: null,
486
          },
487
        ],
488
      },
489
      {
490
        percentile: "99thPercentile",
491
        expectedOutcome: [
492
          {
493
            startTime: new Date("2021-01-01T00:00:00.000Z"),
494
            percentile99Duration: 10,
495
          },
496
          {
497
            startTime: new Date("2021-01-02T00:00:00.000Z"),
498
            percentile99Duration: 5,
499
          },
500
          {
501
            startTime: new Date("2021-01-03T00:00:00.000Z"),
502
            percentile99Duration: null,
503
          },
504
          {
505
            startTime: new Date("2021-01-04T00:00:00.000Z"),
506
            percentile99Duration: null,
507
          },
508
        ],
509
      },
510
      {
511
        percentile: "90thPercentile",
512
        expectedOutcome: [
513
          {
514
            startTime: new Date("2021-01-01T00:00:00.000Z"),
515
            percentile90Duration: 10,
516
          },
517
          {
518
            startTime: new Date("2021-01-02T00:00:00.000Z"),
519
            percentile90Duration: 5,
520
          },
521
          {
522
            startTime: new Date("2021-01-03T00:00:00.000Z"),
523
            percentile90Duration: null,
524
          },
525
          {
526
            startTime: new Date("2021-01-04T00:00:00.000Z"),
527
            percentile90Duration: null,
528
          },
529
        ],
530
      },
531
    ].forEach((props) => {
532
      it(`should calculate right percentiles ${props.percentile}`, async () => {
533
        await prisma.trace.create({
534
          data: {
535
            id: "trace-1",
536
            name: "trace-1",
537
            projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
538
          },
539
        });
540

541
        await prisma.observation.createMany({
542
          data: [
543
            {
544
              traceId: "trace-1",
545
              name: "trace-1",
546
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
547
              type: "GENERATION",
548
              completionTokens: 5,
549
              startTime: new Date("2021-01-01T00:00:00.000Z"),
550
              endTime: new Date("2021-01-01T00:00:10.000Z"),
551
            },
552
            {
553
              traceId: "trace-1",
554
              name: "trace-1",
555
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
556
              type: "GENERATION",
557
              completionTokens: 3,
558
              startTime: new Date("2021-01-01T00:00:00.000Z"),
559
              endTime: new Date("2021-01-01T00:00:08.000Z"),
560
            },
561
            {
562
              traceId: "trace-1",
563
              name: "trace-1",
564
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
565
              type: "GENERATION",
566
              completionTokens: 3,
567
              startTime: new Date("2021-01-01T00:00:00.000Z"),
568
              endTime: new Date("2021-01-01T00:00:01.000Z"),
569
            },
570
            {
571
              traceId: "trace-1",
572
              name: "trace-2",
573
              projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
574
              type: "GENERATION",
575
              completionTokens: 4,
576
              startTime: new Date("2021-01-02T00:00:00.000Z"),
577
              endTime: new Date("2021-01-02T00:00:05.000Z"),
578
            },
579
          ],
580
        });
581

582
        const result = await executeQuery(
583
          prisma,
584
          "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
585
          {
586
            from: "observations",
587
            filter: [
588
              {
589
                type: "datetime",
590
                column: "startTime",
591
                operator: ">=",
592
                value: new Date("2021-01-01T00:00:00.000Z"),
593
              },
594
              {
595
                type: "datetime",
596
                column: "startTime",
597
                operator: "<=",
598
                value: new Date("2021-01-04T00:00:00.000Z"),
599
              },
600
            ],
601
            groupBy: [
602
              { type: "datetime", column: "startTime", temporalUnit: "day" },
603
            ],
604
            select: [
605
              {
606
                column: "duration",
607
                agg: props.percentile as z.infer<typeof aggregations>,
608
              },
609
            ],
610
          },
611
        );
612

613
        expect(result).toStrictEqual(props.expectedOutcome);
614
      });
615
    });
616
  });
617
});
618