langfuse

1
import { type Trace, type Score } from "@prisma/client";
2
import { ObservationTree } from "./ObservationTree";
3
import { ObservationPreview } from "./ObservationPreview";
4
import { TracePreview } from "./TracePreview";
5

6
import Header from "@/src/components/layouts/header";
7
import { Badge } from "@/src/components/ui/badge";
8
import { TraceAggUsageBadge } from "@/src/components/token-usage-badge";
9
import { StringParam, useQueryParam } from "use-query-params";
10
import { PublishTraceSwitch } from "@/src/components/publish-object-switch";
11
import { DetailPageNav } from "@/src/features/navigate-detail-pages/DetailPageNav";
12
import { useRouter } from "next/router";
13
import { type ObservationReturnType } from "@/src/server/api/routers/traces";
14
import { api } from "@/src/utils/api";
15
import { StarTraceDetailsToggle } from "@/src/components/star-toggle";
16
import Link from "next/link";
17
import { NoAccessError } from "@/src/components/no-access";
18
import { TagTraceDetailsPopover } from "@/src/features/tag/components/TagTraceDetailsPopover";
19
import useLocalStorage from "@/src/components/useLocalStorage";
20
import { Toggle } from "@/src/components/ui/toggle";
21
import { Award, ChevronsDownUp, ChevronsUpDown } from "lucide-react";
22
import { usdFormatter } from "@/src/utils/numbers";
23
import Decimal from "decimal.js";
24
import { useCallback, useState } from "react";
25
import { DeleteButton } from "@/src/components/deleteButton";
26

27
export function Trace(props: {
28
  observations: Array<ObservationReturnType>;
29
  trace: Trace;
30
  scores: Score[];
31
  projectId: string;
32
}) {
33
  const [currentObservationId, setCurrentObservationId] = useQueryParam(
34
    "observation",
35
    StringParam,
36
  );
37
  const [metricsOnObservationTree, setMetricsOnObservationTree] =
38
    useLocalStorage("metricsOnObservationTree", true);
39
  const [scoresOnObservationTree, setScoresOnObservationTree] = useLocalStorage(
40
    "scoresOnObservationTree",
41
    true,
42
  );
43

44
  const [collapsedObservations, setCollapsedObservations] = useState<string[]>(
45
    [],
46
  );
47

48
  const toggleCollapsedObservation = useCallback(
49
    (id: string) => {
50
      if (collapsedObservations.includes(id)) {
51
        setCollapsedObservations(collapsedObservations.filter((i) => i !== id));
52
      } else {
53
        setCollapsedObservations([...collapsedObservations, id]);
54
      }
55
    },
56
    [collapsedObservations],
57
  );
58

59
  const collapseAll = useCallback(() => {
60
    // exclude all parents of the current observation
61
    let excludeParentObservations = new Set<string>();
62
    let newExcludeParentObservations = new Set<string>();
63
    do {
64
      excludeParentObservations = new Set<string>([
65
        ...excludeParentObservations,
66
        ...newExcludeParentObservations,
67
      ]);
68
      newExcludeParentObservations = new Set<string>(
69
        props.observations
70
          .filter(
71
            (o) =>
72
              o.parentObservationId !== null &&
73
              (o.id === currentObservationId ||
74
                excludeParentObservations.has(o.id)),
75
          )
76
          .map((o) => o.parentObservationId as string)
77
          .filter((id) => !excludeParentObservations.has(id)),
78
      );
79
    } while (newExcludeParentObservations.size > 0);
80

81
    setCollapsedObservations(
82
      props.observations
83
        .map((o) => o.id)
84
        .filter((id) => !excludeParentObservations.has(id)),
85
    );
86
  }, [props.observations, currentObservationId]);
87

88
  const expandAll = useCallback(() => {
89
    setCollapsedObservations([]);
90
  }, [setCollapsedObservations]);
91

92
  return (
93
    <div className="grid gap-4 md:h-full md:grid-cols-5 lg:grid-cols-6 xl:grid-cols-7">
94
      <div className="overflow-y-auto md:col-span-3 md:h-full lg:col-span-4 xl:col-span-5">
95
        {currentObservationId === undefined ||
96
        currentObservationId === "" ||
97
        currentObservationId === null ? (
98
          <TracePreview
99
            trace={props.trace}
100
            observations={props.observations}
101
            scores={props.scores}
102
          />
103
        ) : (
104
          <ObservationPreview
105
            observations={props.observations}
106
            scores={props.scores}
107
            projectId={props.projectId}
108
            currentObservationId={currentObservationId}
109
            traceId={props.trace.id}
110
          />
111
        )}
112
      </div>
113
      <div className="md:col-span-2 md:flex md:h-full md:flex-col md:overflow-hidden">
114
        <div className="mb-2 flex flex-shrink-0 flex-row justify-end gap-2">
115
          <Toggle
116
            pressed={scoresOnObservationTree}
117
            onPressedChange={(e) => {
118
              setScoresOnObservationTree(e);
119
            }}
120
            size="sm"
121
            title="Show scores"
122
          >
123
            <Award className="h-4 w-4" />
124
          </Toggle>
125
          <Toggle
126
            pressed={metricsOnObservationTree}
127
            onPressedChange={(e) => {
128
              setMetricsOnObservationTree(e);
129
            }}
130
            size="sm"
131
            title="Show metrics"
132
          >
133
            {metricsOnObservationTree ? (
134
              <ChevronsDownUp className="h-4 w-4" />
135
            ) : (
136
              <ChevronsUpDown className="h-4 w-4" />
137
            )}
138
          </Toggle>
139
        </div>
140

141
        <ObservationTree
142
          observations={props.observations}
143
          collapsedObservations={collapsedObservations}
144
          toggleCollapsedObservation={toggleCollapsedObservation}
145
          collapseAll={collapseAll}
146
          expandAll={expandAll}
147
          trace={props.trace}
148
          scores={props.scores}
149
          currentObservationId={currentObservationId ?? undefined}
150
          setCurrentObservationId={setCurrentObservationId}
151
          showMetrics={metricsOnObservationTree}
152
          showScores={scoresOnObservationTree}
153
          className="flex w-full flex-col overflow-y-auto"
154
        />
155
      </div>
156
    </div>
157
  );
158
}
159

160
export function TracePage({ traceId }: { traceId: string }) {
161
  const router = useRouter();
162
  const utils = api.useUtils();
163
  const trace = api.traces.byId.useQuery(
164
    { traceId },
165
    {
166
      retry(failureCount, error) {
167
        if (error.data?.code === "UNAUTHORIZED") return false;
168
        return failureCount < 3;
169
      },
170
    },
171
  );
172

173
  const traceFilterOptions = api.traces.filterOptions.useQuery(
174
    {
175
      projectId: trace.data?.projectId ?? "",
176
    },
177
    {
178
      trpc: {
179
        context: {
180
          skipBatch: true,
181
        },
182
      },
183
      enabled: !!trace.data?.projectId && trace.isSuccess,
184
    },
185
  );
186

187
  const filterOptionTags = traceFilterOptions.data?.tags ?? [];
188
  const allTags = filterOptionTags.map((t) => t.value);
189

190
  const totalCost = calculateDisplayTotalCost(trace.data?.observations ?? []);
191

192
  if (trace.error?.data?.code === "UNAUTHORIZED") return <NoAccessError />;
193
  if (!trace.data) return <div>loading...</div>;
194
  return (
195
    <div className="flex flex-col overflow-hidden 2xl:container md:h-[calc(100vh-2rem)]">
196
      <Header
197
        title="Trace Detail"
198
        breadcrumb={[
199
          {
200
            name: "Traces",
201
            href: `/project/${router.query.projectId as string}/traces`,
202
          },
203
          { name: traceId },
204
        ]}
205
        actionButtons={
206
          <>
207
            <StarTraceDetailsToggle
208
              traceId={trace.data.id}
209
              projectId={trace.data.projectId}
210
              value={trace.data.bookmarked}
211
            />
212
            <PublishTraceSwitch
213
              traceId={trace.data.id}
214
              projectId={trace.data.projectId}
215
              isPublic={trace.data.public}
216
            />
217
            <DetailPageNav
218
              currentId={traceId}
219
              path={(id) =>
220
                `/project/${router.query.projectId as string}/traces/${id}`
221
              }
222
              listKey="traces"
223
            />
224
            <DeleteButton
225
              itemId={traceId}
226
              projectId={trace.data.projectId}
227
              scope="traces:delete"
228
              invalidateFunc={() => void utils.traces.invalidate()}
229
              type="trace"
230
              redirectUrl={`/project/${router.query.projectId as string}/traces`}
231
            />
232
          </>
233
        }
234
      />
235
      <div className="flex flex-wrap gap-2">
236
        {trace.data.sessionId ? (
237
          <Link
238
            href={`/project/${
239
              router.query.projectId as string
240
            }/sessions/${encodeURIComponent(trace.data.sessionId)}`}
241
          >
242
            <Badge>Session: {trace.data.sessionId}</Badge>
243
          </Link>
244
        ) : null}
245
        {trace.data.userId ? (
246
          <Link
247
            href={`/project/${
248
              router.query.projectId as string
249
            }/users/${encodeURIComponent(trace.data.userId)}`}
250
          >
251
            <Badge>User ID: {trace.data.userId}</Badge>
252
          </Link>
253
        ) : null}
254
        <TraceAggUsageBadge observations={trace.data.observations} />
255
        {totalCost ? (
256
          <Badge variant="outline">
257
            Total cost: {usdFormatter(totalCost.toNumber())}
258
          </Badge>
259
        ) : undefined}
260
      </div>
261
      <div className="mt-5 rounded-lg border bg-card font-semibold text-card-foreground shadow-sm">
262
        <div className="flex flex-row items-center gap-3 p-2.5">
263
          Tags
264
          <TagTraceDetailsPopover
265
            tags={trace.data.tags}
266
            availableTags={allTags}
267
            traceId={trace.data.id}
268
            projectId={trace.data.projectId}
269
          />
270
        </div>
271
      </div>
272
      <div className="mt-5 flex-1 overflow-hidden border-t pt-5">
273
        <Trace
274
          key={trace.data.id}
275
          trace={trace.data}
276
          scores={trace.data.scores}
277
          projectId={trace.data.projectId}
278
          observations={trace.data.observations}
279
        />
280
      </div>
281
    </div>
282
  );
283
}
284

285
export const calculateDisplayTotalCost = (
286
  observations: ObservationReturnType[],
287
) => {
288
  return observations.reduce(
289
    (prev: Decimal | undefined, curr: ObservationReturnType) => {
290
      // if we don't have any calculated costs, we can't do anything
291
      if (
292
        !curr.calculatedTotalCost &&
293
        !curr.calculatedInputCost &&
294
        !curr.calculatedOutputCost
295
      )
296
        return prev;
297

298
      // if we have either input or output cost, but not total cost, we can use that
299
      if (
300
        !curr.calculatedTotalCost &&
301
        (curr.calculatedInputCost || curr.calculatedOutputCost)
302
      ) {
303
        return prev
304
          ? prev.plus(
305
              curr.calculatedInputCost ??
306
                new Decimal(0).plus(
307
                  curr.calculatedOutputCost ?? new Decimal(0),
308
                ),
309
            )
310
          : curr.calculatedInputCost ?? curr.calculatedOutputCost ?? undefined;
311
      }
312

313
      if (!curr.calculatedTotalCost) return prev;
314

315
      // if we have total cost, we can use that
316
      return prev
317
        ? prev.plus(curr.calculatedTotalCost)
318
        : curr.calculatedTotalCost;
319
    },
320
    undefined,
321
  );
322
};
323