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
import { type ObservationView } from "@prisma/client";
9

10
describe("/api/public/observations API Endpoint", () => {
11
  beforeEach(async () => await pruneDatabase());
12
  afterEach(async () => await pruneDatabase());
13

14
  it("should fetch all observations", async () => {
15
    await pruneDatabase();
16

17
    const traceId = uuidv4();
18

19
    await prisma.trace.create({
20
      data: {
21
        id: traceId,
22
        name: "trace-name",
23
        userId: "user-1",
24
        projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
25
        metadata: { key: "value" },
26
        release: "1.0.0",
27
        version: "2.0.0",
28
      },
29
    });
30

31
    const model = await prisma.model.create({
32
      data: {
33
        id: "model-1",
34
        modelName: "gpt-3.5-turbo",
35
        inputPrice: "0.0000010",
36
        outputPrice: "0.0000020",
37
        totalPrice: "0.1",
38
        matchPattern: "(.*)(gpt-)(35|3.5)(-turbo)?(.*)",
39
        projectId: null,
40
        unit: ModelUsageUnit.Tokens,
41
      },
42
    });
43

44
    await prisma.observation.create({
45
      data: {
46
        id: uuidv4(),
47
        traceId: traceId,
48
        name: "generation-name",
49
        startTime: new Date("2021-01-01T00:00:00.000Z"),
50
        endTime: new Date("2021-01-01T00:00:00.000Z"),
51
        model: "gpt-3.5-turbo",
52
        modelParameters: { key: "value" },
53
        input: { key: "input" },
54
        output: { key: "output" },
55
        promptTokens: 10,
56
        completionTokens: 20,
57
        totalTokens: 30,
58
        version: "2.0.0",
59
        type: "GENERATION",
60
        project: {
61
          connect: { id: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a" },
62
        },
63
        internalModel: "gpt-3.5-turbo",
64
        unit: ModelUsageUnit.Tokens,
65
      },
66
    });
67

68
    const fetchedObservations = await makeAPICall(
69
      "GET",
70
      "/api/public/observations",
71
      undefined,
72
    );
73

74
    expect(fetchedObservations.status).toBe(200);
75

76
    if (!isObservationList(fetchedObservations.body)) {
77
      throw new Error(
78
        "Expected body to be an array of observations" +
79
          JSON.stringify(fetchedObservations.body),
80
      );
81
    }
82

83
    expect(fetchedObservations.body.data.length).toBe(1);
84
    expect(fetchedObservations.body.data[0]?.traceId).toBe(traceId);
85
    expect(fetchedObservations.body.data[0]?.input).toEqual({ key: "input" });
86
    expect(fetchedObservations.body.data[0]?.output).toEqual({ key: "output" });
87
    expect(fetchedObservations.body.data[0]?.model).toEqual("gpt-3.5-turbo");
88
    expect(fetchedObservations.body.data[0]?.modelId).toEqual(model.id);
89
    expect(
90
      fetchedObservations.body.data[0]?.calculatedInputCost,
91
    ).toBeGreaterThan(0);
92
    expect(
93
      fetchedObservations.body.data[0]?.calculatedOutputCost,
94
    ).toBeGreaterThan(0);
95
    expect(
96
      fetchedObservations.body.data[0]?.calculatedTotalCost,
97
    ).toBeGreaterThan(0);
98
  });
99
  it("should fetch all observations, filtered by generations", async () => {
100
    await pruneDatabase();
101

102
    const traceId = uuidv4();
103

104
    await prisma.trace.create({
105
      data: {
106
        id: traceId,
107
        name: "trace-name",
108
        userId: "user-1",
109
        projectId: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a",
110
        metadata: { key: "value" },
111
        release: "1.0.0",
112
        version: "2.0.0",
113
      },
114
    });
115

116
    await prisma.model.create({
117
      data: {
118
        id: "model-1",
119
        modelName: "gpt-3.5-turbo",
120
        inputPrice: "0.0000010",
121
        outputPrice: "0.0000020",
122
        totalPrice: "0.1",
123
        matchPattern: "(.*)(gpt-)(35|3.5)(-turbo)?(.*)",
124
        projectId: null,
125
        unit: ModelUsageUnit.Tokens,
126
      },
127
    });
128

129
    await prisma.observation.create({
130
      data: {
131
        id: uuidv4(),
132
        traceId: traceId,
133
        name: "generation-name",
134
        startTime: new Date("2021-01-01T00:00:00.000Z"),
135
        endTime: new Date("2021-01-01T00:00:00.000Z"),
136
        model: "gpt-3.5-turbo",
137
        internalModel: "gpt-3.5-turbo",
138
        modelParameters: { key: "value" },
139
        input: { key: "input" },
140
        output: { key: "output" },
141
        promptTokens: 10,
142
        completionTokens: 20,
143
        totalTokens: 30,
144
        version: "2.0.0",
145
        unit: ModelUsageUnit.Tokens,
146
        type: "GENERATION",
147
        project: {
148
          connect: { id: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a" },
149
        },
150
      },
151
    });
152

153
    await prisma.observation.create({
154
      data: {
155
        id: uuidv4(),
156
        traceId: traceId,
157
        name: "generation-name",
158
        startTime: new Date("2021-01-01T00:00:00.000Z"),
159
        endTime: new Date("2021-01-01T00:00:00.000Z"),
160
        modelParameters: { key: "value" },
161
        input: { key: "input" },
162
        output: { key: "output" },
163
        version: "2.0.0",
164
        type: "SPAN",
165
        project: {
166
          connect: { id: "7a88fb47-b4e2-43b8-a06c-a5ce950dc53a" },
167
        },
168
      },
169
    });
170

171
    const fetchedObservations = await makeAPICall(
172
      "GET",
173
      "/api/public/observations?type=GENERATION",
174
      undefined,
175
    );
176

177
    console.log(fetchedObservations.body);
178

179
    expect(fetchedObservations.status).toBe(200);
180

181
    if (!isObservationList(fetchedObservations.body)) {
182
      throw new Error("Expected body to be an array of observations");
183
    }
184

185
    expect(fetchedObservations.body.data.length).toBe(1);
186
    expect(fetchedObservations.body.data[0]?.traceId).toBe(traceId);
187
    expect(fetchedObservations.body.data[0]?.input).toEqual({ key: "input" });
188
    expect(fetchedObservations.body.data[0]?.output).toEqual({ key: "output" });
189
    expect(fetchedObservations.body.data[0]?.type).toEqual("GENERATION");
190
  });
191
});
192

193
const isObservationList = (val: unknown): val is ObservationResponse => {
194
  return (
195
    typeof val === "object" &&
196
    val !== null &&
197
    "data" in val &&
198
    Array.isArray(val.data) &&
199
    val.data.every(
200
      (element) =>
201
        typeof element === "object" &&
202
        element !== null &&
203
        "id" in element &&
204
        "traceId" in element &&
205
        "name" in element &&
206
        "startTime" in element &&
207
        "endTime" in element &&
208
        "model" in element &&
209
        "input" in element &&
210
        "output" in element &&
211
        "metadata" in element &&
212
        "version" in element &&
213
        "modelId" in element &&
214
        "inputPrice" in element &&
215
        "outputPrice" in element &&
216
        "totalPrice" in element &&
217
        "calculatedInputCost" in element &&
218
        "calculatedOutputCost" in element &&
219
        "calculatedTotalCost" in element,
220
    )
221
  );
222
};
223

224
type ObservationResponse = {
225
  data: ObservationView[];
226
};
227