langfuse

1
/** @jest-environment node */
2

3
import { v4 as uuidv4 } from "uuid";
4

5
import { makeAPICall, pruneDatabase } from "@/src/__tests__/test-utils";
6
import { ModelUsageUnit } from "@/src/constants";
7
import { prisma } from "@/src/server/db";
8

9
describe("/api/public/generations API Endpoint", () => {
10
  beforeEach(async () => await pruneDatabase());
11

12
  [
13
    {
14
      usage: {
15
        input: 100,
16
        output: 200,
17
        total: 100,
18
        unit: ModelUsageUnit.Characters,
19
      },
20
      expectedUnit: ModelUsageUnit.Characters,
21
      expectedPromptTokens: 100,
22
      expectedCompletionTokens: 200,
23
      expectedTotalTokens: 100,
24
    },
25
    {
26
      usage: {
27
        total: 100,
28
        unit: ModelUsageUnit.Characters,
29
      },
30
      expectedUnit: ModelUsageUnit.Characters,
31
      expectedPromptTokens: 0,
32
      expectedCompletionTokens: 0,
33
      expectedTotalTokens: 100,
34
    },
35
    {
36
      usage: {
37
        total: 100,
38
      },
39
      expectedUnit: null,
40
      expectedPromptTokens: 0,
41
      expectedCompletionTokens: 0,
42
      expectedTotalTokens: 100,
43
    },
44
    {
45
      usage: {
46
        promptTokens: 100,
47
        completionTokens: 200,
48
        totalTokens: 100,
49
      },
50
      expectedPromptTokens: 100,
51
      expectedCompletionTokens: 200,
52
      expectedTotalTokens: 100,
53
      expectedUnit: ModelUsageUnit.Tokens,
54
    },
55
    {
56
      usage: {
57
        totalTokens: 100,
58
      },
59
      expectedPromptTokens: 0,
60
      expectedCompletionTokens: 0,
61
      expectedTotalTokens: 100,
62
      expectedUnit: ModelUsageUnit.Tokens,
63
    },
64
    {
65
      usage: undefined,
66
      expectedPromptTokens: 0,
67
      expectedCompletionTokens: 0,
68
      expectedTotalTokens: 0,
69
      expectedUnit: null,
70
    },
71
    {
72
      usage: null,
73
      expectedPromptTokens: 0,
74
      expectedCompletionTokens: 0,
75
      expectedTotalTokens: 0,
76
      expectedUnit: null,
77
    },
78
    {
79
      usage: {},
80
      expectedPromptTokens: 0,
81
      expectedCompletionTokens: 0,
82
      expectedTotalTokens: 0,
83
      expectedUnit: null,
84
    },
85
  ].forEach((testConfig) => {
86
    it(`should create generation after trace 1 ${JSON.stringify(
87
      testConfig,
88
    )}`, async () => {
89
      await pruneDatabase();
90

91
      const traceId = uuidv4();
92

93
      await makeAPICall("POST", "/api/public/traces", {
94
        id: traceId,
95
        name: "trace-name",
96
        userId: "user-1",
97
        projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
98
        metadata: { key: "value" },
99
        release: "1.0.0",
100
        version: "2.0.0",
101
      });
102

103
      const dbTrace = await prisma.trace.findMany({
104
        where: {
105
          id: traceId,
106
        },
107
      });
108

109
      expect(dbTrace.length).toBeGreaterThan(0);
110
      expect(dbTrace[0]?.id).toBe(traceId);
111

112
      const generationId = uuidv4();
113
      const createGeneration = await makeAPICall(
114
        "POST",
115
        "/api/public/generations",
116
        {
117
          id: generationId,
118
          traceId: traceId,
119
          name: "generation-name",
120
          startTime: "2021-01-01T00:00:00.000Z",
121
          endTime: "2021-01-01T00:00:00.000Z",
122
          model: "model-name",
123
          modelParameters: { key: "value" },
124
          prompt: { key: "value" },
125
          metadata: { key: "value" },
126
          version: "2.0.0",
127
          usage: testConfig.usage,
128
        },
129
      );
130

131
      expect(createGeneration.status).toBe(200);
132
      const dbGeneration = await prisma.observation.findUnique({
133
        where: {
134
          id: generationId,
135
        },
136
      });
137

138
      expect(dbGeneration?.id).toBe(generationId);
139
      expect(dbGeneration?.traceId).toBe(traceId);
140
      expect(dbGeneration?.name).toBe("generation-name");
141
      expect(dbGeneration?.startTime).toEqual(
142
        new Date("2021-01-01T00:00:00.000Z"),
143
      );
144
      expect(dbGeneration?.endTime).toEqual(
145
        new Date("2021-01-01T00:00:00.000Z"),
146
      );
147
      expect(dbGeneration?.model).toBe("model-name");
148
      expect(dbGeneration?.modelParameters).toEqual({ key: "value" });
149
      expect(dbGeneration?.input).toEqual({ key: "value" });
150
      expect(dbGeneration?.metadata).toEqual({ key: "value" });
151
      expect(dbGeneration?.version).toBe("2.0.0");
152
      expect(dbGeneration?.unit).toBe(testConfig.expectedUnit);
153
      expect(dbGeneration?.promptTokens).toBe(testConfig.expectedPromptTokens);
154
      expect(dbGeneration?.completionTokens).toBe(
155
        testConfig.expectedCompletionTokens,
156
      );
157
      expect(dbGeneration?.totalTokens).toBe(testConfig.expectedTotalTokens);
158
    });
159
  });
160

161
  it("should create generation before trace", async () => {
162
    await pruneDatabase();
163

164
    const traceId = uuidv4();
165

166
    const generationId = uuidv4();
167
    const createGeneration = await makeAPICall(
168
      "POST",
169
      "/api/public/generations",
170
      {
171
        id: generationId,
172
        traceId: traceId,
173
        name: "generation-name",
174
        startTime: "2021-01-01T00:00:00.000Z",
175
        endTime: "2021-01-01T00:00:00.000Z",
176
        model: "model-name",
177
        modelParameters: { key: "value" },
178
        prompt: { key: "value" },
179
        metadata: { key: "value" },
180
        version: "2.0.0",
181
      },
182
    );
183

184
    expect(createGeneration.status).toBe(200);
185
    const dbGeneration = await prisma.observation.findUnique({
186
      where: {
187
        id: generationId,
188
      },
189
    });
190

191
    expect(dbGeneration?.id).toBe(generationId);
192
    expect(dbGeneration?.traceId).toBe(traceId);
193
    expect(dbGeneration?.name).toBe("generation-name");
194
    expect(dbGeneration?.startTime).toEqual(
195
      new Date("2021-01-01T00:00:00.000Z"),
196
    );
197
    expect(dbGeneration?.endTime).toEqual(new Date("2021-01-01T00:00:00.000Z"));
198
    expect(dbGeneration?.model).toBe("model-name");
199
    expect(dbGeneration?.modelParameters).toEqual({ key: "value" });
200
    expect(dbGeneration?.input).toEqual({ key: "value" });
201
    expect(dbGeneration?.metadata).toEqual({ key: "value" });
202
    expect(dbGeneration?.version).toBe("2.0.0");
203

204
    await makeAPICall("POST", "/api/public/traces", {
205
      id: traceId,
206
      name: "trace-name",
207
      userId: "user-1",
208
      projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
209
      metadata: { key: "value" },
210
      release: "1.0.0",
211
      version: "2.0.0",
212
    });
213

214
    const dbTrace = await prisma.trace.findMany({
215
      where: {
216
        id: traceId,
217
      },
218
    });
219

220
    expect(dbTrace.length).toBeGreaterThan(0);
221
    expect(dbTrace[0]?.id).toBe(traceId);
222
  });
223

224
  it("should create generation after trace ignoring externalId", async () => {
225
    const traceId = uuidv4();
226

227
    const response = await makeAPICall("POST", "/api/public/traces", {
228
      externalId: traceId,
229
      name: "trace-name",
230
      userId: "user-1",
231
      projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
232
      metadata: { key: "value" },
233
      release: "1.0.0",
234
      version: "2.0.0",
235
    });
236

237
    expect(response.status).toBe(200);
238

239
    const dbTrace = await prisma.trace.findMany({
240
      where: {
241
        name: "trace-name",
242
      },
243
    });
244

245
    expect(dbTrace.length).toBeGreaterThan(0);
246
    expect(dbTrace[0]?.externalId).toBeNull();
247
    expect(dbTrace[0]?.id).not.toBe(traceId);
248

249
    const generationId = uuidv4();
250
    const createGeneration = await makeAPICall(
251
      "POST",
252
      "/api/public/generations",
253
      {
254
        id: generationId,
255
        traceIdType: "EXTERNAL",
256
        traceId: traceId,
257
        name: "generation-name",
258
        startTime: "2021-01-01T00:00:00.000Z",
259
        endTime: "2021-01-01T00:00:00.000Z",
260
        model: "model-name",
261
        modelParameters: { key: "value" },
262
        prompt: { key: "value" },
263
        metadata: { key: "value" },
264
        version: "2.0.0",
265
      },
266
    );
267

268
    expect(createGeneration.status).toBe(200);
269
    const dbGeneration = await prisma.observation.findUnique({
270
      where: {
271
        id: generationId,
272
      },
273
    });
274

275
    expect(dbGeneration?.id).toBe(generationId);
276
    expect(dbGeneration?.traceId).toBe(traceId);
277
    expect(dbGeneration?.name).toBe("generation-name");
278
    expect(dbGeneration?.startTime).toEqual(
279
      new Date("2021-01-01T00:00:00.000Z"),
280
    );
281
    expect(dbGeneration?.endTime).toEqual(new Date("2021-01-01T00:00:00.000Z"));
282
    expect(dbGeneration?.model).toBe("model-name");
283
    expect(dbGeneration?.modelParameters).toEqual({ key: "value" });
284
    expect(dbGeneration?.input).toEqual({ key: "value" });
285
    expect(dbGeneration?.metadata).toEqual({ key: "value" });
286
    expect(dbGeneration?.version).toBe("2.0.0");
287
  });
288

289
  it("should create trace when creating generation without existing trace", async () => {
290
    const generationName = uuidv4();
291

292
    const generationId = uuidv4();
293
    const createGeneration = await makeAPICall(
294
      "POST",
295
      "/api/public/generations",
296
      {
297
        id: generationId,
298
        name: generationName,
299
        startTime: "2021-01-01T00:00:00.000Z",
300
        endTime: "2021-01-01T00:00:00.000Z",
301
        model: "model-name",
302
        modelParameters: { key: "value" },
303
        prompt: { key: "value" },
304
        metadata: { key: "value" },
305
        version: "2.0.0",
306
      },
307
    );
308

309
    const dbTrace = await prisma.trace.findMany({
310
      where: {
311
        name: generationName,
312
      },
313
    });
314

315
    expect(dbTrace.length).toBe(1);
316
    expect(dbTrace[0]?.name).toBe(generationName);
317

318
    expect(createGeneration.status).toBe(200);
319
    const dbGeneration = await prisma.observation.findUnique({
320
      where: {
321
        id: generationId,
322
      },
323
    });
324

325
    expect(dbGeneration?.id).toBe(generationId);
326
    expect(dbGeneration?.traceId).toBe(dbTrace[0]?.id);
327
    expect(dbGeneration?.name).toBe(generationName);
328
    expect(dbGeneration?.startTime).toEqual(
329
      new Date("2021-01-01T00:00:00.000Z"),
330
    );
331
    expect(dbGeneration?.endTime).toEqual(new Date("2021-01-01T00:00:00.000Z"));
332
    expect(dbGeneration?.model).toBe("model-name");
333
    expect(dbGeneration?.modelParameters).toEqual({ key: "value" });
334
    expect(dbGeneration?.input).toEqual({ key: "value" });
335
    expect(dbGeneration?.metadata).toEqual({ key: "value" });
336
    expect(dbGeneration?.version).toBe("2.0.0");
337
  });
338

339
  it("should create nested generations", async () => {
340
    const generationName = uuidv4();
341

342
    const generationId = uuidv4();
343
    const createGeneration = await makeAPICall(
344
      "POST",
345
      "/api/public/generations",
346
      {
347
        id: generationId,
348
        name: generationName,
349
        startTime: "2021-01-01T00:00:00.000Z",
350
        endTime: "2021-01-01T00:00:00.000Z",
351
        model: "model-name",
352
        modelParameters: { key: "value" },
353
        prompt: { key: "value" },
354
        metadata: { key: "value" },
355
        version: "2.0.0",
356
      },
357
    );
358

359
    expect(createGeneration.status).toBe(200);
360
    const dbGeneration = await prisma.observation.findUnique({
361
      where: {
362
        id: generationId,
363
      },
364
    });
365

366
    expect(dbGeneration?.id).toBe(generationId);
367

368
    const generationId2 = uuidv4();
369
    const generationName2 = uuidv4();
370

371
    const createGeneration2 = await makeAPICall(
372
      "POST",
373
      "/api/public/generations",
374
      {
375
        id: generationId2,
376
        name: generationName2,
377
        startTime: "2021-01-01T00:00:00.000Z",
378
        endTime: "2021-01-01T00:00:00.000Z",
379
        model: "model-name",
380
        modelParameters: { key: "value" },
381
        prompt: { key: "value" },
382
        metadata: { key: "value" },
383
        version: "2.0.0",
384
        parentObservationId: generationId,
385
      },
386
    );
387
    expect(createGeneration2.status).toBe(200);
388

389
    const dbGeneration2 = await prisma.observation.findUnique({
390
      where: {
391
        id: generationId2,
392
      },
393
    });
394

395
    expect(dbGeneration2?.id).toBe(generationId2);
396
    expect(dbGeneration2?.parentObservationId).toBe(generationId);
397
  });
398

399
  it("should not create trace when creating generation without existing trace with externalId", async () => {
400
    const generationName = uuidv4();
401

402
    const generationId = uuidv4();
403
    const externalTraceId = uuidv4();
404
    const createGeneration = await makeAPICall(
405
      "POST",
406
      "/api/public/generations",
407
      {
408
        id: generationId,
409
        traceIdType: "EXTERNAL",
410
        traceId: externalTraceId,
411
        name: generationName,
412
        startTime: "2021-01-01T00:00:00.000Z",
413
        endTime: "2021-01-01T00:00:00.000Z",
414
        model: "model-name",
415
        modelParameters: { key: "value" },
416
        prompt: { key: "value" },
417
        metadata: { key: "value" },
418
        version: "2.0.0",
419
      },
420
    );
421

422
    expect(createGeneration.status).toBe(200);
423

424
    const dbGeneration = await prisma.observation.findFirstOrThrow({
425
      where: {
426
        name: generationName,
427
      },
428
    });
429
    expect(dbGeneration.id).toBe(generationId);
430
    expect(dbGeneration.traceId).toBe(externalTraceId);
431

432
    const dbTraces = await prisma.trace.findMany();
433
    expect(dbTraces.length).toBe(0);
434
  });
435

436
  it("should create trace when creating generation without existing trace without traceId", async () => {
437
    const generationName = uuidv4();
438

439
    const generationId = uuidv4();
440
    const createGeneration = await makeAPICall(
441
      "POST",
442
      "/api/public/generations",
443
      {
444
        id: generationId,
445
        name: generationName,
446
        startTime: "2021-01-01T00:00:00.000Z",
447
        endTime: "2021-01-01T00:00:00.000Z",
448
        model: "model-name",
449
        modelParameters: { key: "value" },
450
        prompt: { key: "value" },
451
        metadata: { key: "value" },
452
        version: "2.0.0",
453
      },
454
    );
455

456
    const dbGeneration = await prisma.observation.findFirstOrThrow({
457
      where: {
458
        name: generationName,
459
      },
460
    });
461

462
    const dbTrace = await prisma.trace.findMany({
463
      where: {
464
        id: dbGeneration.traceId!,
465
      },
466
    });
467

468
    expect(dbTrace.length).toBe(1);
469
    expect(dbTrace[0]?.name).toBe(generationName);
470

471
    expect(createGeneration.status).toBe(200);
472

473
    expect(dbGeneration.id).toBe(generationId);
474
    expect(dbGeneration.traceId).toBe(dbTrace[0]?.id);
475
    expect(dbGeneration.name).toBe(generationName);
476
    expect(dbGeneration.startTime).toEqual(
477
      new Date("2021-01-01T00:00:00.000Z"),
478
    );
479
    expect(dbGeneration.endTime).toEqual(new Date("2021-01-01T00:00:00.000Z"));
480
    expect(dbGeneration.model).toBe("model-name");
481
    expect(dbGeneration.modelParameters).toEqual({ key: "value" });
482
    expect(dbGeneration.input).toEqual({ key: "value" });
483
    expect(dbGeneration.metadata).toEqual({ key: "value" });
484
    expect(dbGeneration.version).toBe("2.0.0");
485
  });
486

487
  it("should update generation", async () => {
488
    const generationName = uuidv4();
489

490
    const generationId = uuidv4();
491
    const createGeneration = await makeAPICall(
492
      "POST",
493
      "/api/public/generations",
494
      {
495
        id: generationId,
496
        name: generationName,
497
        startTime: "2021-01-01T00:00:00.000Z",
498
        endTime: "2021-01-01T00:00:00.000Z",
499
        model: "model-name",
500
        modelParameters: { key: "value" },
501
        prompt: { key: "value" },
502
        metadata: { key: "value" },
503
        version: "2.0.0",
504
      },
505
    );
506

507
    expect(createGeneration.status).toBe(200);
508

509
    const updateGeneration = await makeAPICall(
510
      "PATCH",
511
      "/api/public/generations",
512
      {
513
        generationId: generationId,
514
        completion: "this is a great gpt response",
515
      },
516
    );
517
    expect(updateGeneration.status).toBe(200);
518

519
    const dbGeneration = await prisma.observation.findUnique({
520
      where: {
521
        id: generationId,
522
      },
523
    });
524

525
    expect(dbGeneration?.id).toBe(generationId);
526
    expect(dbGeneration?.name).toBe(generationName);
527
    expect(dbGeneration?.startTime).toEqual(
528
      new Date("2021-01-01T00:00:00.000Z"),
529
    );
530
    expect(dbGeneration?.endTime).toEqual(new Date("2021-01-01T00:00:00.000Z"));
531
    expect(dbGeneration?.model).toBe("model-name");
532
    expect(dbGeneration?.modelParameters).toEqual({ key: "value" });
533
    expect(dbGeneration?.input).toEqual({ key: "value" });
534
    expect(dbGeneration?.output).toEqual("this is a great gpt response");
535
    expect(dbGeneration?.metadata).toEqual({ key: "value" });
536
    expect(dbGeneration?.version).toBe("2.0.0");
537
  });
538

539
  it("should accept objects as i/o", async () => {
540
    const generationName = uuidv4();
541

542
    const generationId = uuidv4();
543
    const createGeneration = await makeAPICall(
544
      "POST",
545
      "/api/public/generations",
546
      {
547
        id: generationId,
548
        name: generationName,
549
        prompt: { key: "value" },
550
        completion: [
551
          {
552
            foo: "bar",
553
          },
554
        ],
555
        metadata: [
556
          {
557
            tags: ["example tag", "second tag"],
558
          },
559
        ],
560
      },
561
    );
562

563
    expect(createGeneration.status).toBe(200);
564

565
    const dbGeneration = await prisma.observation.findUnique({
566
      where: {
567
        id: generationId,
568
      },
569
    });
570

571
    expect(dbGeneration?.id).toBe(generationId);
572
    expect(dbGeneration?.name).toBe(generationName);
573
    expect(dbGeneration?.input).toEqual({ key: "value" });
574
    expect(dbGeneration?.output).toEqual([
575
      {
576
        foo: "bar",
577
      },
578
    ]);
579
    expect(dbGeneration?.metadata).toEqual([
580
      {
581
        tags: ["example tag", "second tag"],
582
      },
583
    ]);
584
  });
585

586
  it("should not succeed update if generation does not exist", async () => {
587
    const generationId = uuidv4();
588

589
    const updateGeneration = await makeAPICall(
590
      "PATCH",
591
      "/api/public/generations",
592
      {
593
        generationId: generationId,
594
        completion: "this is a great gpt response",
595
      },
596
    );
597
    expect(updateGeneration.status).toBe(404);
598

599
    const dbGeneration = await prisma.observation.findUnique({
600
      where: {
601
        id: generationId,
602
      },
603
    });
604

605
    expect(dbGeneration).toBeNull();
606
  });
607
});
608