query light-weight trace info on span insert

lmnr-ai · Feb 25, 2025 · 84d28c6 · 84d28c6
1 parent 8171a61
commit 84d28c6
Show file tree

Hide file tree

Showing 5 changed files with 184 additions and 135 deletions.
diff --git a/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts b/frontend/app/api/projects/[projectId]/traces/[traceId]/route.ts
@@ -1,10 +1,8 @@
-import { and, asc, eq, inArray, sql } from 'drizzle-orm';
+import { and, eq } from 'drizzle-orm';
 import { NextRequest, NextResponse } from 'next/server';
 
-import { searchSpans } from '@/lib/clickhouse/spans';
-import { TimeRange } from '@/lib/clickhouse/utils';
 import { db } from '@/lib/db/drizzle';
-import { events, spans, traces } from '@/lib/db/migrations/schema';
+import { traces } from '@/lib/db/migrations/schema';
 
 export async function GET(
   req: NextRequest,
@@ -13,65 +11,14 @@ export async function GET(
   const params = await props.params;
   const projectId = params.projectId;
   const traceId = params.traceId;
-  const searchQuery = req.nextUrl.searchParams.get("search");
 
-  let searchSpanIds = null;
-  if (searchQuery) {
-    const timeRange = { pastHours: 'all' } as TimeRange;
-    const searchResult = await searchSpans(projectId, searchQuery, timeRange);
-    searchSpanIds = Array.from(searchResult.spanIds);
-  }
-
-  const traceQuery = db.query.traces.findFirst({
+  const trace = await db.query.traces.findFirst({
     where: and(eq(traces.id, traceId), eq(traces.projectId, projectId)),
   });
 
-  const spanEventsQuery = db.$with('span_events').as(
-    db.select({
-      spanId: events.spanId,
-      projectId: events.projectId,
-      events: sql`jsonb_agg(jsonb_build_object(
-        'id', events.id,
-        'spanId', events.span_id,
-        'timestamp', events.timestamp,
-        'name', events.name,
-        'attributes', events.attributes
-      ))`.as('events')
-    })
-      .from(events)
-      .groupBy(events.spanId, events.projectId)
-  );
-
-  const spansQuery = db.with(spanEventsQuery).select({
-    // inputs and outputs are ignored on purpose
-    spanId: spans.spanId,
-    startTime: spans.startTime,
-    endTime: spans.endTime,
-    traceId: spans.traceId,
-    parentSpanId: spans.parentSpanId,
-    name: spans.name,
-    attributes: spans.attributes,
-    spanType: spans.spanType,
-    events: sql`COALESCE(${spanEventsQuery.events}, '[]'::jsonb)`.as('events'),
-  })
-    .from(spans)
-    .leftJoin(spanEventsQuery,
-      and(
-        eq(spans.spanId, spanEventsQuery.spanId),
-        eq(spans.projectId, spanEventsQuery.projectId)
-      )
-    )
-    .where(
-      and(
-        eq(spans.traceId, traceId),
-        eq(spans.projectId, projectId),
-        ...(searchSpanIds ? [inArray(spans.spanId, searchSpanIds)] : [])
-      )
-    )
-    .orderBy(asc(spans.startTime));
-
-  const [trace, spanItems] = await Promise.all([traceQuery, spansQuery]);
+  if (!trace) {
+    return NextResponse.json({ error: 'Trace not found' }, { status: 404 });
+  }
 
-  return NextResponse.json({ ...trace, spans: spanItems });
+  return NextResponse.json(trace);
 }
-
diff --git a/frontend/app/api/projects/[projectId]/traces/[traceId]/spans/route.ts b/frontend/app/api/projects/[projectId]/traces/[traceId]/spans/route.ts
@@ -0,0 +1,71 @@
+import { and, asc, eq, inArray, sql } from 'drizzle-orm';
+import { NextRequest, NextResponse } from 'next/server';
+
+import { searchSpans } from '@/lib/clickhouse/spans';
+import { TimeRange } from '@/lib/clickhouse/utils';
+import { db } from '@/lib/db/drizzle';
+import { events, spans } from '@/lib/db/migrations/schema';
+
+export async function GET(
+  req: NextRequest,
+  props: { params: Promise<{ projectId: string; traceId: string }> }
+): Promise<Response> {
+  const params = await props.params;
+  const projectId = params.projectId;
+  const traceId = params.traceId;
+  const searchQuery = req.nextUrl.searchParams.get("search");
+
+  let searchSpanIds = null;
+  if (searchQuery) {
+    const timeRange = { pastHours: 'all' } as TimeRange;
+    const searchResult = await searchSpans(projectId, searchQuery, timeRange);
+    searchSpanIds = Array.from(searchResult.spanIds);
+  }
+
+  const spanEventsQuery = db.$with('span_events').as(
+    db.select({
+      spanId: events.spanId,
+      projectId: events.projectId,
+      events: sql`jsonb_agg(jsonb_build_object(
+        'id', events.id,
+        'spanId', events.span_id,
+        'timestamp', events.timestamp,
+        'name', events.name,
+        'attributes', events.attributes
+      ))`.as('events')
+    })
+      .from(events)
+      .groupBy(events.spanId, events.projectId)
+  );
+
+  const spanItems = await db.with(spanEventsQuery).select({
+    // inputs and outputs are ignored on purpose
+    spanId: spans.spanId,
+    startTime: spans.startTime,
+    endTime: spans.endTime,
+    traceId: spans.traceId,
+    parentSpanId: spans.parentSpanId,
+    name: spans.name,
+    attributes: spans.attributes,
+    spanType: spans.spanType,
+    events: sql`COALESCE(${spanEventsQuery.events}, '[]'::jsonb)`.as('events'),
+  })
+    .from(spans)
+    .leftJoin(spanEventsQuery,
+      and(
+        eq(spans.spanId, spanEventsQuery.spanId),
+        eq(spans.projectId, spanEventsQuery.projectId)
+      )
+    )
+    .where(
+      and(
+        eq(spans.traceId, traceId),
+        eq(spans.projectId, projectId),
+        ...(searchSpanIds ? [inArray(spans.spanId, searchSpanIds)] : [])
+      )
+    )
+    .orderBy(asc(spans.startTime));
+
+
+  return NextResponse.json(spanItems);
+}
diff --git a/frontend/components/traces/spans-table.tsx b/frontend/components/traces/spans-table.tsx
@@ -163,8 +163,8 @@ export default function SpansTable({ onRowClick }: SpansTableProps) {
     inputPreview: row.input_preview,
     outputPreview: row.output_preview,
     events: [],
-    inputUrl: null,
-    outputUrl: null,
+    inputUrl: row.input_url,
+    outputUrl: row.output_url,
     model: row.attributes['gen_ai.response.model'] ?? row.attributes['gen_ai.request.model'] ?? null,
   });