langfuse

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

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

7
interface GetTracesAPIResponse {
8
  data: Array<{
9
    id: string;
10
    [key: string]: unknown;
11
  }>;
12
}
13

14
describe("/api/public/traces API Endpoint", () => {
15
  beforeEach(async () => await pruneDatabase());
16
  afterEach(async () => await pruneDatabase());
17

18
  it("should create", async () => {
19
    await pruneDatabase();
20

21
    await makeAPICall("POST", "/api/public/traces", {
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
    const dbTrace = await prisma.trace.findMany({
31
      where: {
32
        name: "trace-name",
33
      },
34
    });
35

36
    expect(dbTrace.length).toBeGreaterThan(0);
37
    expect(dbTrace[0]?.name).toBe("trace-name");
38
    expect(dbTrace[0]?.release).toBe("1.0.0");
39
    expect(dbTrace[0]?.externalId).toBeNull();
40
    expect(dbTrace[0]?.version).toBe("2.0.0");
41
    expect(dbTrace[0]?.projectId).toBe("7a88fb47-b4e2-43b8-a06c-a5ce950dc53a");
42
  });
43

44
  it("should upsert second trace", async () => {
45
    await pruneDatabase();
46

47
    await makeAPICall("POST", "/api/public/traces", {
48
      id: "trace-id",
49
      name: "trace-name",
50
      userId: "user-1",
51
      metadata: { key: "value" },
52
      release: "1.0.0",
53
      version: "2.0.0",
54
      public: true,
55
    });
56

57
    const dbTrace1 = await prisma.trace.findFirst({
58
      where: {
59
        id: "trace-id",
60
      },
61
    });
62

63
    expect(dbTrace1).not.toBeNull();
64
    expect(dbTrace1).toMatchObject({
65
      name: "trace-name",
66
      release: "1.0.0",
67
      externalId: null,
68
      version: "2.0.0",
69
      public: true,
70
      userId: "user-1",
71
    });
72

73
    await makeAPICall("POST", "/api/public/traces", {
74
      id: "trace-id",
75
      metadata: { key: "value" },
76
      timestamp: "2021-01-01T00:00:00.000Z",
77
      release: "1.0.0",
78
      version: "5.0.0",
79
      public: false,
80
    });
81

82
    const dbTrace2 = await prisma.trace.findFirst({
83
      where: {
84
        name: "trace-name",
85
      },
86
    });
87

88
    expect(dbTrace2).not.toBeNull();
89
    expect(dbTrace2).toMatchObject({
90
      name: "trace-name",
91
      release: "1.0.0",
92
      externalId: null,
93
      version: "5.0.0",
94
      public: false,
95
      userId: "user-1",
96
      timestamp: new Date("2021-01-01T00:00:00.000Z"),
97
    });
98
  });
99

100
  it("should use tags correctly on POST and GET", async () => {
101
    await pruneDatabase();
102

103
    await makeAPICall("POST", "/api/public/traces", {
104
      id: "trace-1",
105
      tags: ["tag-1", "tag-2", "tag-3"],
106
    });
107

108
    await makeAPICall("POST", "/api/public/traces", {
109
      id: "trace-2",
110
      tags: ["tag-1"],
111
    });
112

113
    await makeAPICall("POST", "/api/public/traces", {
114
      id: "trace-3",
115
      tags: ["tag-2", "tag-3"],
116
    });
117

118
    // multiple tags
119
    const traces = await makeAPICall<GetTracesAPIResponse>(
120
      "GET",
121
      "/api/public/traces?tags=tag-2&tags=tag-3",
122
    );
123
    const traceIds = traces.body.data.map((t) => t.id);
124
    // check for equality ok as ordered by timestamp
125
    expect(traceIds).toEqual(["trace-3", "trace-1"]);
126

127
    // single tag
128
    const traces2 = await makeAPICall<GetTracesAPIResponse>(
129
      "GET",
130
      "/api/public/traces?tags=tag-1",
131
    );
132
    const traceIds2 = traces2.body.data.map((t) => t.id);
133
    // check for equality ok as ordered by timestamp
134
    expect(traceIds2).toEqual(["trace-2", "trace-1"]);
135

136
    // wrong tag
137
    const traces3 = await makeAPICall<GetTracesAPIResponse>(
138
      "GET",
139
      "/api/public/traces?tags=tag-10",
140
    );
141
    const traceIds3 = traces3.body.data.map((t) => t.id);
142
    // check for equality ok as ordered by timestamp
143
    expect(traceIds3).toEqual([]);
144

145
    // no tag
146
    const traces4 = await makeAPICall<GetTracesAPIResponse>(
147
      "GET",
148
      "/api/public/traces?tags=",
149
    );
150
    const traceIds4 = traces4.body.data.map((t) => t.id);
151
    // check for equality ok as ordered by timestamp
152
    expect(traceIds4).toEqual(["trace-3", "trace-2", "trace-1"]);
153
  });
154

155
  it("should handle metrics correctly on GET traces and GET trace", async () => {
156
    await pruneDatabase();
157

158
    // Create a trace with some observations that have costs and latencies
159
    const traceId = uuidv4();
160
    await makeAPICall("POST", "/api/public/traces", {
161
      id: traceId,
162
      name: "trace-with-costs",
163
      userId: "user-costs",
164
      projectId: "project-costs",
165
      metadata: { key: "value" },
166
      release: "1.0.0",
167
      version: "2.0.0",
168
    });
169
    console.log(traceId);
170

171
    // Simulate observations with costs and latencies
172
    await makeAPICall("POST", "/api/public/generations", {
173
      traceId: traceId,
174
      usage: { totalCost: 10.5 },
175
      startTime: "2021-01-01T00:00:00.000Z",
176
      endTime: "2021-01-01T00:10:00.000Z",
177
    });
178
    await makeAPICall("POST", "/api/public/generations", {
179
      traceId: traceId,
180
      usage: { totalCost: 5.25 },
181
      startTime: "2021-01-01T00:10:00.000Z",
182
      endTime: "2021-01-01T00:20:00.000Z",
183
    });
184

185
    // GET traces
186
    // Retrieve the trace with totalCost and latency
187
    const traces = await makeAPICall<GetTracesAPIResponse>(
188
      "GET",
189
      `/api/public/traces`,
190
    );
191
    const traceData = traces.body.data[0];
192
    if (!traceData) throw new Error("traceData is undefined");
193

194
    // Check if the totalCost and latency are calculated correctly
195
    expect(traceData.totalCost).toBeCloseTo(15.75); // Sum of costs
196
    expect(traceData.latency).toBeCloseTo(1200); // Difference in seconds between min startTime and max endTime
197
    expect(traceData.id).toBe(traceId);
198
    expect(traceData.htmlPath).toContain(`/traces/${traceId}`);
199
    expect(traceData.htmlPath).toContain(`/project/`); // do not know the projectId
200

201
    // GET trace
202
    // Retrieve the trace with total
203
    const trace = await makeAPICall<{
204
      id: string;
205
      totalCost: number;
206
      htmlPath: string;
207
    }>("GET", `/api/public/traces/${traceId}`);
208
    console.log(trace.body);
209
    expect(trace.body.totalCost).toBeCloseTo(15.75);
210
    expect(trace.body.id).toBe(traceId);
211
    expect(trace.body.id).toBe(traceId);
212
    expect(trace.body.htmlPath).toContain(`/traces/${traceId}`);
213
    expect(trace.body.htmlPath).toContain(`/project/`); // do not know the projectId
214
  });
215
});
216