Refactor perf load observation handling

Author: marcusschiesser

packages/perf/src/load.ts

@@ -3,6 +3,7 @@ import {
 	propagateAttributes,
 	startActiveObservation,
 } from "@langfuse/tracing";
+import type { LangfuseSpanAttributes } from "@langfuse/tracing";
 
 type ObservationType = "span" | "generation" | "tool";
 
@@ -91,10 +92,10 @@ async function generateChildObservation(
 ): Promise<void> {
 	const type = observationType(traceIndex, childIndex);
 	const name = `${type} step ${childIndex}`;
-	await startActiveObservation(
-		name,
-		async (span) => {
-			if (type === "generation") {
+	if (type === "generation") {
+		await startActiveObservation(
+			name,
+			async (span) => {
 				const inputTokens = 120 + ((traceIndex + childIndex) % 700);
 				const outputTokens = 40 + ((traceIndex * 3 + childIndex) % 260);
 				span.update({
@@ -120,32 +121,41 @@ async function generateChildObservation(
 						total: roundCost(inputTokens * 0.0000015 + outputTokens * 0.000006),
 					},
 				});
-				return;
-			}
+			},
+			{ asType: "generation" },
+		);
+		return;
+	}
 
-			span.update({
-				input: {
-					step: childIndex,
-					payload: sizedText("step input", traceIndex, childIndex),
-				},
-				output: {
-					ok: true,
-					items: Array.from({ length: 1 + (childIndex % 4) }, (_, offset) => ({
-						id: `item-${traceIndex}-${childIndex}-${offset}`,
-						score: Number(((traceIndex + offset) % 100) / 100).toFixed(2),
-					})),
-				},
-				metadata: {
-					component: type,
-					retry_count: traceIndex % 5 === 0 ? 1 : 0,
-					latency_bucket: ["fast", "normal", "slow"][
-						(traceIndex + childIndex) % 3
-					],
-				},
-			});
-		},
-		{ asType: type },
-	);
+	const updateSpan = async (span: {
+		update(attributes: LangfuseSpanAttributes): unknown;
+	}) => {
+		span.update({
+			input: {
+				step: childIndex,
+				payload: sizedText("step input", traceIndex, childIndex),
+			},
+			output: {
+				ok: true,
+				items: Array.from({ length: 1 + (childIndex % 4) }, (_, offset) => ({
+					id: `item-${traceIndex}-${childIndex}-${offset}`,
+					score: Number(((traceIndex + offset) % 100) / 100).toFixed(2),
+				})),
+			},
+			metadata: {
+				component: type,
+				retry_count: traceIndex % 5 === 0 ? 1 : 0,
+				latency_bucket: ["fast", "normal", "slow"][
+					(traceIndex + childIndex) % 3
+				],
+			},
+		});
+	};
+	if (type === "tool") {
+		await startActiveObservation(name, updateSpan, { asType: "tool" });
+		return;
+	}
+	await startActiveObservation(name, updateSpan);
 }
 
 function observationType(