langfuse

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

3
import { prisma } from "@/src/server/db";
4
import { makeAPICall, pruneDatabase } from "@/src/__tests__/test-utils";
5
import { v4 } from "uuid";
6

7
describe("/api/public/datasets and /api/public/dataset-items API Endpoints", () => {
8
  beforeEach(async () => await pruneDatabase());
9
  afterEach(async () => await pruneDatabase());
10

11
  it("should create and get a dataset", async () => {
12
    await makeAPICall("POST", "/api/public/datasets", {
13
      name: "dataset-name",
14
    });
15

16
    const dbDataset = await prisma.dataset.findMany({
17
      where: {
18
        name: "dataset-name",
19
      },
20
    });
21

22
    expect(dbDataset.length).toBeGreaterThan(0);
23

24
    const getDataset = await makeAPICall(
25
      "GET",
26
      `/api/public/datasets/dataset-name`,
27
    );
28

29
    expect(getDataset.status).toBe(200);
30
    expect(getDataset.body).toMatchObject({
31
      name: "dataset-name",
32
    });
33
  });
34

35
  it("should create and get a dataset item (via datasets and individually)", async () => {
36
    await makeAPICall("POST", "/api/public/datasets", {
37
      name: "dataset-name",
38
    });
39
    await makeAPICall("POST", "/api/public/dataset-items", {
40
      datasetName: "dataset-name",
41
      input: { key: "value" },
42
      expectedOutput: { key: "value" },
43
    });
44
    const dbDatasetItem = await prisma.datasetItem.findFirst({
45
      where: {
46
        dataset: {
47
          name: "dataset-name",
48
        },
49
      },
50
    });
51

52
    expect(dbDatasetItem).not.toBeNull();
53

54
    const getDataset = await makeAPICall(
55
      "GET",
56
      `/api/public/datasets/dataset-name`,
57
    );
58
    expect(getDataset.status).toBe(200);
59
    expect(getDataset.body).toMatchObject({
60
      name: "dataset-name",
61
      items: [
62
        {
63
          id: dbDatasetItem!.id,
64
          input: { key: "value" },
65
          expectedOutput: { key: "value" },
66
        },
67
      ],
68
    });
69

70
    const getDatasetItem = await makeAPICall(
71
      "GET",
72
      `/api/public/dataset-items/${dbDatasetItem!.id}`,
73
    );
74
    expect(getDatasetItem.status).toBe(200);
75
    expect(getDatasetItem.body).toMatchObject({
76
      id: dbDatasetItem!.id,
77
      input: { key: "value" },
78
      expectedOutput: { key: "value" },
79
    });
80
  });
81

82
  it("should upsert a dataset item", async () => {
83
    await makeAPICall("POST", "/api/public/datasets", {
84
      name: "dataset-name",
85
    });
86

87
    const item1 = await makeAPICall("POST", "/api/public/dataset-items", {
88
      id: "dataset-item-id",
89
      datasetName: "dataset-name",
90
      input: { key: "value" },
91
    });
92
    expect(item1.status).toBe(200);
93
    expect(item1.body).toMatchObject({
94
      id: "dataset-item-id",
95
    });
96

97
    const item2 = await makeAPICall("POST", "/api/public/dataset-items", {
98
      id: "dataset-item-id",
99
      datasetName: "dataset-name",
100
      input: { key: "value2" },
101
    });
102
    expect(item2.status).toBe(200);
103
    expect(item2.body).toMatchObject({
104
      id: "dataset-item-id",
105
      input: { key: "value2" },
106
    });
107

108
    const dbDatasetItem = await prisma.datasetItem.findFirst({
109
      where: { id: "dataset-item-id" },
110
    });
111
    expect(dbDatasetItem).not.toBeNull();
112
    expect(dbDatasetItem?.input).toMatchObject({ key: "value2" });
113
  });
114

115
  it("should create and get a dataset run", async () => {
116
    const dataset = await makeAPICall("POST", "/api/public/datasets", {
117
      name: "dataset-name",
118
    });
119
    expect(dataset.status).toBe(200);
120
    expect(dataset.body).toMatchObject({
121
      name: "dataset-name",
122
    });
123
    await makeAPICall("POST", "/api/public/dataset-items", {
124
      datasetName: "dataset-name",
125
      id: "dataset-item-id",
126
      input: { key: "value" },
127
      expectedOutput: { key: "value" },
128
    });
129
    const traceId = v4();
130
    const observationId = v4();
131
    const response = await makeAPICall("POST", "/api/public/ingestion", {
132
      batch: [
133
        {
134
          id: v4(),
135
          type: "trace-create",
136
          timestamp: new Date().toISOString(),
137
          body: {
138
            id: traceId,
139
            name: "trace-name",
140
            userId: "user-1",
141
            metadata: { key: "value" },
142
            release: "1.0.0",
143
            version: "2.0.0",
144
          },
145
        },
146
        {
147
          id: v4(),
148
          type: "observation-create",
149
          timestamp: new Date().toISOString(),
150
          body: {
151
            id: observationId,
152
            traceId: traceId,
153
            type: "GENERATION",
154
            name: "generation-name",
155
            startTime: "2021-01-01T00:00:00.000Z",
156
            endTime: "2021-01-01T00:00:00.000Z",
157
            modelParameters: { key: "value" },
158
            input: { key: "value" },
159
            metadata: { key: "value" },
160
            version: "2.0.0",
161
          },
162
        },
163
      ],
164
    });
165
    expect(response.status).toBe(207);
166

167
    const runItem = await makeAPICall("POST", "/api/public/dataset-run-items", {
168
      datasetItemId: "dataset-item-id",
169
      observationId: observationId,
170
      runName: "run-name",
171
    });
172
    const dbRun = await prisma.datasetRuns.findFirst({
173
      where: {
174
        name: "run-name",
175
      },
176
    });
177
    expect(dbRun).not.toBeNull();
178
    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
179
    expect(dbRun?.datasetId).toBe(dataset.body.id);
180
    expect(runItem.status).toBe(200);
181
    expect(runItem.body).toMatchObject({
182
      datasetItemId: "dataset-item-id",
183
      observationId: observationId,
184
      datasetRunId: dbRun?.id,
185
    });
186
  });
187
});
188