langfuse

1
import { type NestedObservation } from "@/src/utils/types";
2
import { cn } from "@/src/utils/tailwind";
3
import { type Trace, type Score, $Enums } from "@prisma/client";
4
import { GroupedScoreBadges } from "@/src/components/grouped-score-badge";
5
import { Fragment } from "react";
6
import { type ObservationReturnType } from "@/src/server/api/routers/traces";
7
import { LevelColors } from "@/src/components/level-colors";
8
import { formatIntervalSeconds } from "@/src/utils/dates";
9
import { MinusCircle, MinusIcon, PlusCircleIcon, PlusIcon } from "lucide-react";
10
import { Toggle } from "@/src/components/ui/toggle";
11
import { Button } from "@/src/components/ui/button";
12

13
export const ObservationTree = (props: {
14
  observations: ObservationReturnType[];
15
  collapsedObservations: string[];
16
  toggleCollapsedObservation: (id: string) => void;
17
  collapseAll: () => void;
18
  expandAll: () => void;
19
  trace: Trace;
20
  scores: Score[];
21
  currentObservationId: string | undefined;
22
  setCurrentObservationId: (id: string | undefined) => void;
23
  showMetrics: boolean;
24
  showScores: boolean;
25
  className?: string;
26
}) => {
27
  const nestedObservations = nestObservations(props.observations);
28
  return (
29
    <div className={props.className}>
30
      <ObservationTreeTraceNode
31
        expandAll={props.expandAll}
32
        collapseAll={props.collapseAll}
33
        trace={props.trace}
34
        scores={props.scores}
35
        currentObservationId={props.currentObservationId}
36
        setCurrentObservationId={props.setCurrentObservationId}
37
        showMetrics={props.showMetrics}
38
        showScores={props.showScores}
39
      />
40
      <ObservationTreeNode
41
        observations={nestedObservations}
42
        collapsedObservations={props.collapsedObservations}
43
        toggleCollapsedObservation={props.toggleCollapsedObservation}
44
        scores={props.scores}
45
        indentationLevel={1}
46
        currentObservationId={props.currentObservationId}
47
        setCurrentObservationId={props.setCurrentObservationId}
48
        showMetrics={props.showMetrics}
49
        showScores={props.showScores}
50
      />
51
    </div>
52
  );
53
};
54

55
const ObservationTreeTraceNode = (props: {
56
  trace: Trace & { latency?: number };
57
  expandAll: () => void;
58
  collapseAll: () => void;
59
  scores: Score[];
60
  currentObservationId: string | undefined;
61
  setCurrentObservationId: (id: string | undefined) => void;
62
  showMetrics?: boolean;
63
  showScores?: boolean;
64
}) => (
65
  <div
66
    className={cn(
67
      "group mb-0.5 flex cursor-pointer flex-col gap-1 rounded-sm p-1",
68
      props.currentObservationId === undefined ||
69
        props.currentObservationId === ""
70
        ? "bg-gray-100"
71
        : "hover:bg-gray-50",
72
    )}
73
    onClick={() => props.setCurrentObservationId(undefined)}
74
  >
75
    <div className="flex gap-2">
76
      <span className={cn("rounded-sm bg-gray-200 p-1 text-xs")}>TRACE</span>
77
      <span className="flex-1 break-all text-sm">{props.trace.name}</span>
78
      <Button
79
        onClick={(ev) => (ev.stopPropagation(), props.expandAll())}
80
        size="xs"
81
        variant="ghost"
82
        title="Expand all"
83
      >
84
        <PlusCircleIcon className="h-4 w-4" />
85
      </Button>
86
      <Button
87
        onClick={(ev) => (ev.stopPropagation(), props.collapseAll())}
88
        size="xs"
89
        variant="ghost"
90
        title="Collapse all"
91
      >
92
        <MinusCircle className="h-4 w-4" />
93
      </Button>
94
    </div>
95

96
    {props.showMetrics && props.trace.latency ? (
97
      <div className="flex gap-2">
98
        <span className="text-xs text-gray-500">
99
          {formatIntervalSeconds(props.trace.latency)}
100
        </span>
101
      </div>
102
    ) : null}
103
    {props.showScores && props.scores.find((s) => s.observationId === null) ? (
104
      <div className="flex flex-wrap gap-1">
105
        <GroupedScoreBadges
106
          scores={props.scores.filter((s) => s.observationId === null)}
107
        />
108
      </div>
109
    ) : null}
110
  </div>
111
);
112

113
const ObservationTreeNode = (props: {
114
  observations: NestedObservation[];
115
  collapsedObservations: string[];
116
  toggleCollapsedObservation: (id: string) => void;
117
  scores: Score[];
118
  indentationLevel: number;
119
  currentObservationId: string | undefined;
120
  setCurrentObservationId: (id: string | undefined) => void;
121
  showMetrics?: boolean;
122
  showScores?: boolean;
123
}) => (
124
  <>
125
    {props.observations
126
      .sort((a, b) => a.startTime.getTime() - b.startTime.getTime())
127
      .map((observation) => {
128
        const collapsed = props.collapsedObservations.includes(observation.id);
129

130
        return (
131
          <Fragment key={observation.id}>
132
            <div className="flex">
133
              {Array.from({ length: props.indentationLevel }, (_, i) => (
134
                <div className="mx-2 border-r" key={i} />
135
              ))}
136
              <div
137
                className={cn(
138
                  "group my-0.5 flex flex-1 cursor-pointer flex-col gap-1 rounded-sm p-1",
139
                  props.currentObservationId === observation.id
140
                    ? "bg-gray-100"
141
                    : "hover:bg-gray-50",
142
                )}
143
                onClick={() => props.setCurrentObservationId(observation.id)}
144
              >
145
                <div className="flex gap-2">
146
                  <ColorCodedObservationType
147
                    observationType={observation.type}
148
                  />
149
                  <span className="flex-1 break-all text-sm">
150
                    {observation.name}
151
                  </span>
152
                  {observation.children.length === 0 ? null : (
153
                    <Toggle
154
                      onClick={(ev) => (
155
                        ev.stopPropagation(),
156
                        props.toggleCollapsedObservation(observation.id)
157
                      )}
158
                      variant="default"
159
                      pressed={collapsed}
160
                      size="xs"
161
                      className="w-7"
162
                      title={
163
                        collapsed ? "Expand children" : "Collapse children"
164
                      }
165
                    >
166
                      {collapsed ? (
167
                        <PlusIcon className="h-4 w-4" />
168
                      ) : (
169
                        <MinusIcon className="h-4 w-4" />
170
                      )}
171
                    </Toggle>
172
                  )}
173
                </div>
174
                {props.showMetrics &&
175
                  (observation.promptTokens ||
176
                    observation.completionTokens ||
177
                    observation.totalTokens ||
178
                    observation.endTime) && (
179
                    <div className="flex gap-2">
180
                      {observation.endTime ? (
181
                        <span className="text-xs text-gray-500">
182
                          {formatIntervalSeconds(
183
                            (observation.endTime.getTime() -
184
                              observation.startTime.getTime()) /
185
                              1000,
186
                          )}
187
                        </span>
188
                      ) : null}
189
                      {observation.promptTokens ||
190
                      observation.completionTokens ||
191
                      observation.totalTokens ? (
192
                        <span className="text-xs text-gray-500">
193
                          {observation.promptTokens} →{" "}
194
                          {observation.completionTokens} (∑{" "}
195
                          {observation.totalTokens})
196
                        </span>
197
                      ) : null}
198
                    </div>
199
                  )}
200
                {observation.level !== "DEFAULT" ? (
201
                  <div className="flex">
202
                    <span
203
                      className={cn(
204
                        "rounded-sm p-0.5 text-xs",
205
                        LevelColors[observation.level].bg,
206
                        LevelColors[observation.level].text,
207
                      )}
208
                    >
209
                      {observation.level}
210
                    </span>
211
                  </div>
212
                ) : null}
213
                {props.showScores &&
214
                props.scores.find((s) => s.observationId === observation.id) ? (
215
                  <div className="flex flex-wrap gap-1">
216
                    <GroupedScoreBadges
217
                      scores={props.scores.filter(
218
                        (s) => s.observationId === observation.id,
219
                      )}
220
                    />
221
                  </div>
222
                ) : null}
223
              </div>
224
            </div>
225
            {!collapsed && (
226
              <ObservationTreeNode
227
                observations={observation.children}
228
                collapsedObservations={props.collapsedObservations}
229
                toggleCollapsedObservation={props.toggleCollapsedObservation}
230
                scores={props.scores}
231
                indentationLevel={props.indentationLevel + 1}
232
                currentObservationId={props.currentObservationId}
233
                setCurrentObservationId={props.setCurrentObservationId}
234
                showMetrics={props.showMetrics}
235
                showScores={props.showScores}
236
              />
237
            )}
238
          </Fragment>
239
        );
240
      })}
241
  </>
242
);
243

244
const ColorCodedObservationType = (props: {
245
  observationType: $Enums.ObservationType;
246
}) => {
247
  const colors: Record<$Enums.ObservationType, string> = {
248
    [$Enums.ObservationType.SPAN]: "bg-blue-100",
249
    [$Enums.ObservationType.GENERATION]: "bg-orange-100",
250
    [$Enums.ObservationType.EVENT]: "bg-green-100",
251
  };
252

253
  return (
254
    <span
255
      className={cn(
256
        "self-start rounded-sm p-1 text-xs",
257
        colors[props.observationType],
258
      )}
259
    >
260
      {props.observationType}
261
    </span>
262
  );
263
};
264

265
export function nestObservations(
266
  list: ObservationReturnType[],
267
): NestedObservation[] {
268
  if (list.length === 0) return [];
269

270
  // Step 1: Create a map where the keys are object IDs, and the values are
271
  // the corresponding objects with an added 'children' property.
272
  const map = new Map<string, NestedObservation>();
273
  for (const obj of list) {
274
    map.set(obj.id, { ...obj, children: [] });
275
  }
276

277
  // Step 2: Create another map for the roots of all trees.
278
  const roots = new Map<string, NestedObservation>();
279

280
  // Step 3: Populate the 'children' arrays and root map.
281
  for (const obj of map.values()) {
282
    if (obj.parentObservationId) {
283
      const parent = map.get(obj.parentObservationId);
284
      if (parent) {
285
        parent.children.push(obj);
286
      }
287
    } else {
288
      roots.set(obj.id, obj);
289
    }
290
  }
291

292
  // TODO sum token amounts per level
293

294
  // Step 4: Return the roots.
295
  return Array.from(roots.values());
296
}
297