From 1b9fab1ed4d75b9ef7e3c614bcec64fdae583bce Mon Sep 17 00:00:00 2001 From: Jonathan Haas Date: Fri, 29 May 2026 08:24:48 -0700 Subject: [PATCH] Fix ingest compatibility metadata gaps --- app/src/api/query-api.ts | 2 + app/src/api/saved-runs.ts | 1 + app/src/components/FlameTimeline.tsx | 4 + app/src/components/RunDetail.tsx | 10 +- app/src/components/SpanTree.tsx | 1 + app/src/pages/SavedPage.tsx | 4 +- app/src/pages/SearchPage.tsx | 2 + app/src/utils/types.ts | 1 + app/tests-e2e/helpers.ts | 5 +- drizzle/0002_volatile_mordo.sql | 1 + drizzle/meta/0002_snapshot.json | 852 +++++++++++++++++++++++++++ drizzle/meta/_journal.json | 7 + src/db.ts | 26 +- src/db/migration-assets.ts | 12 + src/db/schema.ts | 1 + src/parse.ts | 5 +- src/server.ts | 6 +- src/spans/adapters/traceloop.ts | 5 +- tests/parse.test.ts | 71 +++ 19 files changed, 992 insertions(+), 24 deletions(-) create mode 100644 drizzle/0002_volatile_mordo.sql create mode 100644 drizzle/meta/0002_snapshot.json create mode 100644 tests/parse.test.ts diff --git a/app/src/api/query-api.ts b/app/src/api/query-api.ts index bf1cf37..dca512e 100644 --- a/app/src/api/query-api.ts +++ b/app/src/api/query-api.ts @@ -44,6 +44,7 @@ const traceSpanSchema = z.object({ output: z.string().nullable(), input_tokens: z.number().nullable(), output_tokens: z.number().nullable(), + total_tokens: z.number().nullable().optional(), model: z.string().nullable(), provider: z.string().nullable(), attributes: z.record(z.string(), z.union([z.string(), z.number()])), @@ -176,6 +177,7 @@ export function mapTraceToSpans(traces: TraceSpan[], eventId: string): Span[] { provider: t.provider, input_tokens: t.input_tokens, output_tokens: t.output_tokens, + total_tokens: t.total_tokens ?? null, attributes: Object.keys(t.attributes).length > 0 ? JSON.stringify(t.attributes) : null, }; }); diff --git a/app/src/api/saved-runs.ts b/app/src/api/saved-runs.ts index dcd6b35..ea32dd0 100644 --- a/app/src/api/saved-runs.ts +++ b/app/src/api/saved-runs.ts @@ -79,6 +79,7 @@ const spanCacheSchema: z.ZodType = z.object({ provider: z.string().nullable(), input_tokens: z.number().nullable(), output_tokens: z.number().nullable(), + total_tokens: z.number().nullable().optional(), attributes: z.string().nullable(), normalized: z.custom().optional(), }); diff --git a/app/src/components/FlameTimeline.tsx b/app/src/components/FlameTimeline.tsx index 42fe1a4..7729a23 100644 --- a/app/src/components/FlameTimeline.tsx +++ b/app/src/components/FlameTimeline.tsx @@ -94,6 +94,7 @@ function SpanTooltip({ const type = spanTypeInfo(span); const inTok = span.input_tokens ?? 0; const outTok = span.output_tokens ?? 0; + const totalTok = span.total_tokens ?? 0; const inputRaw = span.input_payload?.trim() ?? ""; const outputRaw = span.output_payload?.trim() ?? ""; @@ -132,6 +133,9 @@ function SpanTooltip({ {(inTok > 0 || outTok > 0) && ( {inTok.toLocaleString()} in / {outTok.toLocaleString()} out )} + {inTok === 0 && outTok === 0 && totalTok > 0 && ( + {totalTok.toLocaleString()} total + )}
{inputRaw ? : null} diff --git a/app/src/components/RunDetail.tsx b/app/src/components/RunDetail.tsx index c912a4a..f0deb3d 100644 --- a/app/src/components/RunDetail.tsx +++ b/app/src/components/RunDetail.tsx @@ -221,7 +221,7 @@ function Badge({ label, copyValue }: { label: string; copyValue?: string }) { } function StatsLine({ stats, model, spans, active, startedAt }: { - stats: { spans: number; tools: number; llms: number; errors: number; dur: number; agents?: number; inTokens?: number; outTokens?: number }; + stats: { spans: number; tools: number; llms: number; errors: number; dur: number; agents?: number; inTokens?: number; outTokens?: number; totalTokens?: number }; model?: string | null; spans?: Span[]; active?: boolean; @@ -231,6 +231,8 @@ function StatsLine({ stats, model, spans, active, startedAt }: { const costRef = useRef(null); const inTok = stats.inTokens ?? 0; const outTok = stats.outTokens ?? 0; + const totalTok = stats.totalTokens ?? 0; + const splitTok = inTok + outTok; const [now, setNow] = useState(Date.now()); useEffect(() => { @@ -261,7 +263,8 @@ function StatsLine({ stats, model, spans, active, startedAt }: { {stats.errors > 0 && spans && <>} {stats.errors > 0 && !spans && <> error{stats.errors !== 1 ? "s" : ""}} {durMin > 0 ? <>m s : <>s} - {(inTok > 0 || outTok > 0) && <> in / out} + {(inTok > 0 || outTok > 0) && <> in / out{totalTok > splitTok && <> / total}} + {splitTok === 0 && totalTok > 0 && <> total} {(() => { const totalCost = breakdown.reduce((sum, b) => sum + (b.breakdown?.totalCost ?? 0), 0); const cost = totalCost > 0 ? fmtCost(totalCost) : null; @@ -449,7 +452,7 @@ function ViewHeader({ title: string; model?: string | null; active: boolean; - stats: { spans: number; tools: number; llms: number; errors: number; dur: number; agents?: number; inTokens?: number; outTokens?: number }; + stats: { spans: number; tools: number; llms: number; errors: number; dur: number; agents?: number; inTokens?: number; outTokens?: number; totalTokens?: number }; allSpans?: Span[]; startedAt?: number; anthropicModels?: string[]; @@ -1344,6 +1347,7 @@ export function RunDetail({ runId, routeBase, initialData, isReplay, source, onF agents: subAgents.length, inTokens: [...getTokensByModel(spans).values()].reduce((s, v) => s + v.inTok, 0), outTokens: [...getTokensByModel(spans).values()].reduce((s, v) => s + v.outTok, 0), + totalTokens: spans.reduce((sum, s) => sum + (s.total_tokens ?? 0), 0), }} allSpans={spans} run={run} diff --git a/app/src/components/SpanTree.tsx b/app/src/components/SpanTree.tsx index efd7659..2440cf9 100644 --- a/app/src/components/SpanTree.tsx +++ b/app/src/components/SpanTree.tsx @@ -153,6 +153,7 @@ function SpanDetail({ span }: { span: Span }) { {span.model && <>
model
{span.model}
} {span.input_tokens != null && <>
input tokens
{span.input_tokens.toLocaleString()}
} {span.output_tokens != null && <>
output tokens
{span.output_tokens.toLocaleString()}
} + {span.total_tokens != null && <>
total tokens
{span.total_tokens.toLocaleString()}
}
status
{span.status}
start
diff --git a/app/src/pages/SavedPage.tsx b/app/src/pages/SavedPage.tsx index abe3ec2..2a2800b 100644 --- a/app/src/pages/SavedPage.tsx +++ b/app/src/pages/SavedPage.tsx @@ -1015,7 +1015,7 @@ interface TraceSpan { span_name: string; span_type: string; status: string; start_time_ns: number; end_time_ns: number; duration_ns: number; input: string | null; output: string | null; - input_tokens: number | null; output_tokens: number | null; + input_tokens: number | null; output_tokens: number | null; total_tokens?: number | null; model: string | null; provider: string | null; attributes: Record; } @@ -1048,7 +1048,7 @@ function mapTraceToSpans(traces: TraceSpan[], eventId: string): Span[] { input_payload: inputPayload, output_payload: outputPayload, start_time_ms: t.start_time_ns / 1e6, end_time_ms: t.end_time_ns / 1e6, duration_ms: t.duration_ns / 1e6, model: t.model, provider: t.provider, - input_tokens: t.input_tokens, output_tokens: t.output_tokens, + input_tokens: t.input_tokens, output_tokens: t.output_tokens, total_tokens: t.total_tokens ?? null, attributes: Object.keys(t.attributes).length > 0 ? JSON.stringify(t.attributes) : null, }; }); diff --git a/app/src/pages/SearchPage.tsx b/app/src/pages/SearchPage.tsx index 8a34621..108033d 100644 --- a/app/src/pages/SearchPage.tsx +++ b/app/src/pages/SearchPage.tsx @@ -47,6 +47,7 @@ interface TraceSpan { output: string | null; input_tokens: number | null; output_tokens: number | null; + total_tokens?: number | null; model: string | null; provider: string | null; attributes: Record; @@ -134,6 +135,7 @@ function mapTraceToSpans(traces: TraceSpan[], eventId: string): Span[] { provider: t.provider, input_tokens: t.input_tokens, output_tokens: t.output_tokens, + total_tokens: t.total_tokens ?? null, attributes: Object.keys(t.attributes).length > 0 ? JSON.stringify(t.attributes) : null, }; }); } diff --git a/app/src/utils/types.ts b/app/src/utils/types.ts index 3d3dad6..6cc47c8 100644 --- a/app/src/utils/types.ts +++ b/app/src/utils/types.ts @@ -100,6 +100,7 @@ export interface Span { provider: string | null; input_tokens: number | null; output_tokens: number | null; + total_tokens: number | null; attributes: string | null; /** * SDK-agnostic typed view of this span's content. Populated by the server. diff --git a/app/tests-e2e/helpers.ts b/app/tests-e2e/helpers.ts index 55ad627..c7bf038 100644 --- a/app/tests-e2e/helpers.ts +++ b/app/tests-e2e/helpers.ts @@ -379,7 +379,7 @@ export type DbSpanRow = { async function fetchSpansViaApi( workshopUrl: string, runId: string, -): Promise<{ id: string; name: string; span_type: string | null; status: string | null; input_preview: string; output_preview: string; model: string | null; tokens: { in: number; out: number } }[]> { +): Promise<{ id: string; name: string; span_type: string | null; status: string | null; input_preview: string; output_preview: string; model: string | null; tokens: { in: number; out: number; total: number } }[]> { const res = await fetch( `${workshopUrl}/api/runs/${encodeURIComponent(runId)}/spans?limit=500&payload_preview_chars=400`, ); @@ -395,7 +395,7 @@ async function fetchSpansViaApi( input_preview: string; output_preview: string; model: string | null; - tokens: { in: number; out: number }; + tokens: { in: number; out: number; total: number }; }[] >; } @@ -431,6 +431,7 @@ export type LlmSpanRow = { model: string | null; input_tokens: number | null; output_tokens: number | null; + total_tokens: number | null; }; /** diff --git a/drizzle/0002_volatile_mordo.sql b/drizzle/0002_volatile_mordo.sql new file mode 100644 index 0000000..ac2ed65 --- /dev/null +++ b/drizzle/0002_volatile_mordo.sql @@ -0,0 +1 @@ +ALTER TABLE `spans` ADD `total_tokens` integer; \ No newline at end of file diff --git a/drizzle/meta/0002_snapshot.json b/drizzle/meta/0002_snapshot.json new file mode 100644 index 0000000..fcf6bd3 --- /dev/null +++ b/drizzle/meta/0002_snapshot.json @@ -0,0 +1,852 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "1bda1ea9-fab0-4c6f-982f-f5580bca56b3", + "prevId": "04b30dfa-cb94-4c6f-aec5-263daaf3eb00", + "tables": { + "annotations": { + "name": "annotations", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "span_id": { + "name": "span_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "kind": { + "name": "kind", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "note": { + "name": "note", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "source": { + "name": "source", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": { + "idx_annotations_run": { + "name": "idx_annotations_run", + "columns": [ + "run_id" + ], + "isUnique": false + }, + "idx_annotations_span": { + "name": "idx_annotations_span", + "columns": [ + "span_id" + ], + "isUnique": false, + "where": "\"annotations\".\"span_id\" IS NOT NULL" + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "live_events": { + "name": "live_events", + "columns": { + "id": { + "name": "id", + "type": "integer", + "primaryKey": true, + "notNull": true, + "autoincrement": true + }, + "trace_id": { + "name": "trace_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "span_id": { + "name": "span_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "type": { + "name": "type", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "timestamp": { + "name": "timestamp", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "metadata": { + "name": "metadata", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "idx_live_trace": { + "name": "idx_live_trace", + "columns": [ + "trace_id", + "timestamp" + ], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "messages": { + "name": "messages", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "session_id": { + "name": "session_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "content": { + "name": "content", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "state_updated_at": { + "name": "state_updated_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": { + "idx_messages_session_created": { + "name": "idx_messages_session_created", + "columns": [ + "session_id", + "\"created_at\" asc" + ], + "isUnique": false + }, + "idx_messages_state": { + "name": "idx_messages_state", + "columns": [ + "state", + "state_updated_at" + ], + "isUnique": false + } + }, + "foreignKeys": { + "messages_session_id_sessions_id_fk": { + "name": "messages_session_id_sessions_id_fk", + "tableFrom": "messages", + "tableTo": "sessions", + "columnsFrom": [ + "session_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "runs": { + "name": "runs", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "event_id": { + "name": "event_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_name": { + "name": "event_name", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "convo_id": { + "name": "convo_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "last_updated_at": { + "name": "last_updated_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "metadata": { + "name": "metadata", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "idx_runs_last_updated": { + "name": "idx_runs_last_updated", + "columns": [ + "last_updated_at" + ], + "isUnique": false + }, + "idx_runs_event_id": { + "name": "idx_runs_event_id", + "columns": [ + "event_id" + ], + "isUnique": false, + "where": "\"runs\".\"event_id\" IS NOT NULL" + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "saved_events": { + "name": "saved_events", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "event_name": { + "name": "event_name", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "convo_id": { + "name": "convo_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "timestamp": { + "name": "timestamp", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "user_input": { + "name": "user_input", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "assistant_output": { + "name": "assistant_output", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "signals": { + "name": "signals", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "properties": { + "name": "properties", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "saved_at": { + "name": "saved_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "summary": { + "name": "summary", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "source": { + "name": "source", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "folder": { + "name": "folder", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "idx_saved_events_saved_at": { + "name": "idx_saved_events_saved_at", + "columns": [ + "\"saved_at\" desc" + ], + "isUnique": false + }, + "idx_saved_events_folder": { + "name": "idx_saved_events_folder", + "columns": [ + "folder" + ], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "saved_folders": { + "name": "saved_folders", + "columns": { + "name": { + "name": "name", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "color": { + "name": "color", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "saved_run_cache": { + "name": "saved_run_cache", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "data": { + "name": "data", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "sessions": { + "name": "sessions", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "active": { + "name": "active", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "created_at": { + "name": "created_at", + "type": "integer", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "deactivated_at": { + "name": "deactivated_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "idx_sessions_active_created": { + "name": "idx_sessions_active_created", + "columns": [ + "active", + "\"created_at\" desc" + ], + "isUnique": false + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "spans": { + "name": "spans", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "parent_span_id": { + "name": "parent_span_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true, + "autoincrement": false + }, + "span_type": { + "name": "span_type", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false, + "default": "'UNSET'" + }, + "input_payload": { + "name": "input_payload", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "output_payload": { + "name": "output_payload", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "start_time_ms": { + "name": "start_time_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "end_time_ms": { + "name": "end_time_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "duration_ms": { + "name": "duration_ms", + "type": "real", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "provider": { + "name": "provider", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "input_tokens": { + "name": "input_tokens", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "output_tokens": { + "name": "output_tokens", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "total_tokens": { + "name": "total_tokens", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "attributes": { + "name": "attributes", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": { + "idx_spans_run_id": { + "name": "idx_spans_run_id", + "columns": [ + "run_id" + ], + "isUnique": false + }, + "idx_spans_parent": { + "name": "idx_spans_parent", + "columns": [ + "parent_span_id" + ], + "isUnique": false + } + }, + "foreignKeys": { + "spans_run_id_runs_id_fk": { + "name": "spans_run_id_runs_id_fk", + "tableFrom": "spans", + "tableTo": "runs", + "columnsFrom": [ + "run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": { + "runs_with_hints": { + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_id": { + "name": "event_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "event_name": { + "name": "event_name", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "convo_id": { + "name": "convo_id", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "started_at": { + "name": "started_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "last_updated_at": { + "name": "last_updated_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "metadata": { + "name": "metadata", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "finished": { + "name": "finished", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "span_count": { + "name": "span_count", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "live_event_count": { + "name": "live_event_count", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "payload_total_chars": { + "name": "payload_total_chars", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "name": "runs_with_hints", + "isExisting": false, + "definition": "\n SELECT\n r.*,\n (SELECT s.model FROM spans s WHERE s.run_id = r.id AND s.model IS NOT NULL LIMIT 1) AS model,\n (SELECT CASE WHEN COUNT(*) > 0\n AND COUNT(*) = COUNT(CASE WHEN s.status IN ('OK','ERROR') THEN 1 END)\n THEN 1 ELSE 0 END\n FROM spans s WHERE s.run_id = r.id AND s.parent_span_id IS NULL) AS finished,\n (SELECT COUNT(*) FROM spans s WHERE s.run_id = r.id) AS span_count,\n (SELECT COUNT(*) FROM live_events e WHERE e.trace_id = r.id) AS live_event_count,\n (SELECT COALESCE(SUM(LENGTH(COALESCE(s.input_payload, '')) + LENGTH(COALESCE(s.output_payload, ''))), 0)\n FROM spans s WHERE s.run_id = r.id) AS payload_total_chars\n FROM runs r\n" + } + }, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": { + "idx_messages_session_created": { + "columns": { + "\"created_at\" asc": { + "isExpression": true + } + } + }, + "idx_saved_events_saved_at": { + "columns": { + "\"saved_at\" desc": { + "isExpression": true + } + } + }, + "idx_sessions_active_created": { + "columns": { + "\"created_at\" desc": { + "isExpression": true + } + } + } + } + } +} \ No newline at end of file diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index 7dce1c6..a45a388 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -15,6 +15,13 @@ "when": 1778475541625, "tag": "0001_cynical_betty_brant", "breakpoints": true + }, + { + "idx": 2, + "version": "6", + "when": 1780068068711, + "tag": "0002_volatile_mordo", + "breakpoints": true } ] } \ No newline at end of file diff --git a/src/db.ts b/src/db.ts index c5df704..3b49f77 100644 --- a/src/db.ts +++ b/src/db.ts @@ -321,7 +321,7 @@ export function deleteRun(runId: string) { }); } -export function insertSpan(span: { id: string; run_id: string; parent_span_id?: string; name: string; span_type?: string; status?: string; input_payload?: string; output_payload?: string; start_time_ms: number; end_time_ms: number; duration_ms: number; model?: string; provider?: string; input_tokens?: number; output_tokens?: number; attributes?: string }) { +export function insertSpan(span: { id: string; run_id: string; parent_span_id?: string; name: string; span_type?: string; status?: string; input_payload?: string; output_payload?: string; start_time_ms: number; end_time_ms: number; duration_ms: number; model?: string; provider?: string; input_tokens?: number; output_tokens?: number; total_tokens?: number; attributes?: string }) { const row = { id: span.id, run_id: span.run_id, @@ -338,6 +338,7 @@ export function insertSpan(span: { id: string; run_id: string; parent_span_id?: provider: span.provider ?? null, input_tokens: span.input_tokens ?? null, output_tokens: span.output_tokens ?? null, + total_tokens: span.total_tokens ?? null, attributes: span.attributes ?? null, }; getDrizzleDb() @@ -481,6 +482,7 @@ export interface SpanMetaRow { duration_ms: number | null; input_tokens: number | null; output_tokens: number | null; + total_tokens: number | null; model: string | null; provider: string | null; attributes: string | null; @@ -504,6 +506,7 @@ export function getSpanMeta(spanId: string): SpanMetaRow | null { duration_ms: schema.spans.duration_ms, input_tokens: schema.spans.input_tokens, output_tokens: schema.spans.output_tokens, + total_tokens: schema.spans.total_tokens, model: schema.spans.model, provider: schema.spans.provider, attributes: schema.spans.attributes, @@ -811,7 +814,7 @@ export interface OutlineSpan { span_type: string | null; status: string; duration_ms: number; - tokens: { in: number; out: number }; + tokens: { in: number; out: number; total: number }; model: string | null; input_preview: string; output_preview: string; @@ -889,6 +892,7 @@ export function getRunOutline(runId: string, payloadPreviewChars = 80): RunOutli (SELECT COUNT(*) FROM live_events e WHERE e.trace_id = r.id) AS live_event_count, (SELECT COALESCE(SUM(input_tokens), 0) FROM spans s WHERE s.run_id = r.id) AS total_input_tokens, (SELECT COALESCE(SUM(output_tokens), 0) FROM spans s WHERE s.run_id = r.id) AS total_output_tokens, + (SELECT COALESCE(SUM(total_tokens), 0) FROM spans s WHERE s.run_id = r.id) AS total_tokens, (SELECT COALESCE(SUM(LENGTH(COALESCE(s.input_payload, '')) + LENGTH(COALESCE(s.output_payload, ''))), 0) FROM spans s WHERE s.run_id = r.id) AS payload_total_chars, (SELECT s.model FROM spans s WHERE s.run_id = r.id AND s.model IS NOT NULL LIMIT 1) AS model, @@ -917,7 +921,7 @@ export function getRunOutline(runId: string, payloadPreviewChars = 80): RunOutli const rawSpans = rawQuery(` SELECT id, parent_span_id, name, span_type, status, start_time_ms, end_time_ms, duration_ms, - input_tokens, output_tokens, model, + input_tokens, output_tokens, total_tokens, model, SUBSTR(COALESCE(input_payload, ''), 1, ?) AS input_head, LENGTH(COALESCE(input_payload, '')) AS input_chars, SUBSTR(COALESCE(output_payload, ''), 1, ?) AS output_head, @@ -954,7 +958,7 @@ export function getRunOutline(runId: string, payloadPreviewChars = 80): RunOutli span_type: s.span_type, status: s.status ?? "UNSET", duration_ms: s.duration_ms ?? 0, - tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0 }, + tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0, total: s.total_tokens ?? 0 }, model: s.model, input_preview: s.input_chars > cap ? s.input_head.slice(0, cap) + "…" : s.input_head, output_preview: s.output_chars > cap ? s.output_head.slice(0, cap) + "…" : s.output_head, @@ -1054,7 +1058,7 @@ export function getRunOutline(runId: string, payloadPreviewChars = 80): RunOutli })); // detectSubAgents only reads id / parent_span_id / name / span_type / status / - // start_time_ms / end_time_ms / duration_ms / model / input_tokens / output_tokens — + // start_time_ms / end_time_ms / duration_ms / model / input_tokens / output_tokens / total_tokens — // all already projected into rawSpans. Reuse it instead of issuing a SELECT * that // would pull the fat input_payload / output_payload columns we worked to avoid. const sub_agents: RunOutline["sub_agents"] = []; @@ -1106,7 +1110,7 @@ export interface SpanSkeleton { status: string; start_time_ms: number; duration_ms: number; - tokens: { in: number; out: number }; + tokens: { in: number; out: number; total: number }; model: string | null; input_chars: number; output_chars: number; @@ -1139,7 +1143,7 @@ export function listSpansFiltered(runId: string, opts: ListSpansOpts = {}): Span } if (typeof f.min_duration_ms === "number") { where.push("duration_ms >= ?"); params.push(f.min_duration_ms); } if (typeof f.min_tokens === "number") { - where.push("(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)) >= ?"); + where.push("MAX(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0), COALESCE(total_tokens, 0)) >= ?"); params.push(f.min_tokens); } @@ -1152,13 +1156,13 @@ export function listSpansFiltered(runId: string, opts: ListSpansOpts = {}): Span start_asc: "start_time_ms ASC", start_desc: "start_time_ms DESC", duration_desc: "duration_ms DESC", - tokens_desc: "(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)) DESC", + tokens_desc: "MAX(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0), COALESCE(total_tokens, 0)) DESC", } as const) as Record)[opts.sort ?? "start_asc"] ?? "start_time_ms ASC"; const sql = ` SELECT id, parent_span_id, name, span_type, status, start_time_ms, end_time_ms, duration_ms, - input_tokens, output_tokens, model, + input_tokens, output_tokens, total_tokens, model, SUBSTR(COALESCE(input_payload, ''), 1, ?) AS input_head, LENGTH(COALESCE(input_payload, '')) AS input_chars, SUBSTR(COALESCE(output_payload, ''), 1, ?) AS output_head, @@ -1185,7 +1189,7 @@ export function listSpansFiltered(runId: string, opts: ListSpansOpts = {}): Span status: s.status ?? "UNSET", start_time_ms: s.start_time_ms, duration_ms: s.duration_ms ?? 0, - tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0 }, + tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0, total: s.total_tokens ?? 0 }, model: s.model, input_chars: s.input_chars, output_chars: s.output_chars, @@ -1359,7 +1363,7 @@ function spanRowToContextSkeleton(s: any) { status: s.status ?? "UNSET", start_time_ms: s.start_time_ms, duration_ms: s.duration_ms ?? 0, - tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0 }, + tokens: { in: s.input_tokens ?? 0, out: s.output_tokens ?? 0, total: s.total_tokens ?? 0 }, model: s.model, }; } diff --git a/src/db/migration-assets.ts b/src/db/migration-assets.ts index 5c6799f..0414b43 100644 --- a/src/db/migration-assets.ts +++ b/src/db/migration-assets.ts @@ -4,6 +4,7 @@ import migration0000Path from "../../drizzle/0000_massive_winter_soldier.sql" with { type: "file" }; import migration0001Path from "../../drizzle/0001_cynical_betty_brant.sql" with { type: "file" }; +import migration0002Path from "../../drizzle/0002_volatile_mordo.sql" with { type: "file" }; export const embeddedMigrationJournal = { "version": "7", @@ -22,6 +23,13 @@ export const embeddedMigrationJournal = { "when": 1778475541625, "tag": "0001_cynical_betty_brant", "breakpoints": true + }, + { + "idx": 2, + "version": "6", + "when": 1780068068711, + "tag": "0002_volatile_mordo", + "breakpoints": true } ] } as const; @@ -35,4 +43,8 @@ export const embeddedMigrationFiles = [ relativePath: "0001_cynical_betty_brant.sql", sourcePath: migration0001Path, }, + { + relativePath: "0002_volatile_mordo.sql", + sourcePath: migration0002Path, + }, ] as const; diff --git a/src/db/schema.ts b/src/db/schema.ts index 3d8b37f..7221695 100644 --- a/src/db/schema.ts +++ b/src/db/schema.ts @@ -38,6 +38,7 @@ export const spans = sqliteTable( provider: text("provider"), input_tokens: integer("input_tokens"), output_tokens: integer("output_tokens"), + total_tokens: integer("total_tokens"), attributes: text("attributes"), }, (table) => [ diff --git a/src/parse.ts b/src/parse.ts index f70ca83..3f75518 100644 --- a/src/parse.ts +++ b/src/parse.ts @@ -6,7 +6,7 @@ export interface ParsedSpan { traceId: string; spanId: string; parentSpanId?: string; name: string; spanType: string; status: string; inputPayload?: string; outputPayload?: string; startTimeMs: number; endTimeMs: number; durationMs: number; model?: string; provider?: string; - inputTokens?: number; outputTokens?: number; attributes: Record; + inputTokens?: number; outputTokens?: number; totalTokens?: number; attributes: Record; eventId?: string; eventName?: string; userId?: string; convoId?: string; replayRunId?: string; /** @@ -179,6 +179,7 @@ export function parseOtlpRequest(body: any): ParsedSpan[] { const provider = first(attrs, "ai.model.provider", "gen_ai.system", "gen_ai.provider.name", "llm.system") as string | undefined; const inputTokens = first(attrs, "ai.usage.inputTokens", "ai.usage.promptTokens", "ai.usage.prompt_tokens", "gen_ai.usage.input_tokens") as number | undefined; const outputTokens = first(attrs, "ai.usage.outputTokens", "ai.usage.completionTokens", "ai.usage.completion_tokens", "gen_ai.usage.output_tokens") as number | undefined; + const totalTokens = first(attrs, "ai.usage.totalTokens", "ai.usage.total_tokens", "gen_ai.usage.total_tokens", "llm.usage.total_tokens") as number | undefined; const eventId = first(attrs, "ai.telemetry.metadata.raindrop.eventId", "ai.telemetry.metadata.traceloop.association.properties.event_id", "traceloop.association.properties.event_id", "traceloop.association.properties.traceloop.association.properties.event_id") as string | undefined; const eventName = first(attrs, "ai.telemetry.metadata.raindrop.eventName", "ai.telemetry.metadata.traceloop.association.properties.event_name", "traceloop.association.properties.event_name", "traceloop.association.properties.traceloop.association.properties.event_name") as string | undefined; @@ -201,6 +202,7 @@ export function parseOtlpRequest(body: any): ParsedSpan[] { // is "put it in properties" — works across every SDK wrapper. let replayRunId = first( attrs, + "raindrop.replayRunId", "ai.telemetry.metadata.raindrop.replayRunId", "traceloop.association.properties.replayRunId", ) as string | undefined; @@ -231,6 +233,7 @@ export function parseOtlpRequest(body: any): ParsedSpan[] { model, provider, inputTokens: typeof inputTokens === "number" ? inputTokens : undefined, outputTokens: typeof outputTokens === "number" ? outputTokens : undefined, + totalTokens: typeof totalTokens === "number" ? totalTokens : undefined, attributes: allAttrs, eventId, eventName, userId, convoId, replayRunId, normalized: match.normalized, diff --git a/src/server.ts b/src/server.ts index 42bada6..88a3fe9 100644 --- a/src/server.ts +++ b/src/server.ts @@ -582,7 +582,7 @@ export async function createServer(port: number) { input_payload: s.inputPayload, output_payload: s.outputPayload, start_time_ms: s.startTimeMs, end_time_ms: s.endTimeMs, duration_ms: s.durationMs, model: s.model, provider: s.provider, - input_tokens: s.inputTokens, output_tokens: s.outputTokens, + input_tokens: s.inputTokens, output_tokens: s.outputTokens, total_tokens: s.totalTokens, attributes: JSON.stringify(s.attributes), }); } @@ -1693,7 +1693,9 @@ export async function createServer(port: number) { input_payload: s.input_payload ?? undefined, output_payload: s.output_payload ?? undefined, start_time_ms: s.start_time_ms, end_time_ms: s.end_time_ms, duration_ms: s.duration_ms, model: s.model ?? undefined, provider: s.provider ?? undefined, - input_tokens: s.input_tokens ?? undefined, output_tokens: s.output_tokens ?? undefined, + input_tokens: s.input_tokens ?? undefined, + output_tokens: s.output_tokens ?? undefined, + total_tokens: s.total_tokens ?? undefined, attributes: s.attributes ?? undefined, }); } diff --git a/src/spans/adapters/traceloop.ts b/src/spans/adapters/traceloop.ts index cf6091d..a898849 100644 --- a/src/spans/adapters/traceloop.ts +++ b/src/spans/adapters/traceloop.ts @@ -165,7 +165,7 @@ function normalizeLegacyIndexedGenAiMessages(attrs: Record ({ role: attrs[`gen_ai.prompt.${index}.role`], - content: attrs[`gen_ai.prompt.${index}.content`], + content: attrs[`gen_ai.prompt.${index}.content`] ?? attrs[`gen_ai.prompt.${index}.user`], }))) : undefined, outputPayload, @@ -350,4 +350,3 @@ export const traceloopToolAdapter: SpanAdapter = { }; }, }; - diff --git a/tests/parse.test.ts b/tests/parse.test.ts new file mode 100644 index 0000000..35c7696 --- /dev/null +++ b/tests/parse.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, test } from "bun:test"; +import { parseOtlpRequest } from "../src/parse"; + +function attr(key: string, value: string | number | boolean) { + if (typeof value === "number") return { key, value: { intValue: value } }; + if (typeof value === "boolean") return { key, value: { boolValue: value } }; + return { key, value: { stringValue: value } }; +} + +function otlpRequest(attributes: ReturnType[]) { + return { + resourceSpans: [ + { + scopeSpans: [ + { + spans: [ + { + traceId: "0123456789abcdef0123456789abcdef", + spanId: "0123456789abcdef", + name: "cohere.rerank", + startTimeUnixNano: "0", + endTimeUnixNano: "1000000", + attributes, + status: { code: 1 }, + }, + ], + }, + ], + }, + ], + }; +} + +describe("parseOtlpRequest", () => { + test("reads top-level replayRunId metadata", () => { + const spans = parseOtlpRequest(otlpRequest([ + attr("raindrop.replayRunId", "replay_123"), + ])); + + expect(spans).toHaveLength(1); + expect(spans[0].replayRunId).toBe("replay_123"); + }); + + test("normalizes legacy Cohere gen_ai.prompt user attributes", () => { + const spans = parseOtlpRequest(otlpRequest([ + attr("llm.request.type", "rerank"), + attr("gen_ai.prompt.0.role", "user"), + attr("gen_ai.prompt.0.user", "Which result is best?"), + attr("gen_ai.completion.0.content", "Document 2"), + ])); + + expect(spans).toHaveLength(1); + expect(spans[0].inputPayload).toContain("Which result is best?"); + expect(spans[0].normalized.kind).toBe("llm"); + if (spans[0].normalized.kind === "llm") { + expect(spans[0].normalized.messages[0]?.content).toBe("Which result is best?"); + } + }); + + test("reads total-only token usage without split input or output tokens", () => { + const spans = parseOtlpRequest(otlpRequest([ + attr("llm.request.type", "rerank"), + attr("llm.usage.total_tokens", 7), + ])); + + expect(spans).toHaveLength(1); + expect(spans[0].inputTokens).toBeUndefined(); + expect(spans[0].outputTokens).toBeUndefined(); + expect(spans[0].totalTokens).toBe(7); + }); +});