Improve trace viewer for large trace logs

Author: edbeeching

trackio/frontend/src/lib/api.js

@@ -105,6 +105,130 @@ export async function getLogsBatch(project, runs) {
   return await callApi("/get_logs_batch", payload);
 }
 
+function sqlString(value) {
+  return `'${String(value ?? "").replaceAll("'", "''")}'`;
+}
+
+function scalarLogsFromQueryRows(rows) {
+  return (rows || []).map((row) => {
+    const metrics = (() => {
+      try {
+        return JSON.parse(row.metrics || "{}");
+      } catch {
+        return {};
+      }
+    })();
+    return {
+      ...metrics,
+      timestamp: row.timestamp,
+      step: row.step,
+    };
+  });
+}
+
+function scalarOnlyLogs(logs) {
+  return (logs || []).map((row) => {
+    const out = {
+      timestamp: row.timestamp,
+      step: row.step,
+    };
+    for (const [key, value] of Object.entries(row)) {
+      if (key === "timestamp" || key === "step") continue;
+      if (typeof value === "number" && Number.isFinite(value)) out[key] = value;
+    }
+    return out;
+  });
+}
+
+async function queryProject(project, query) {
+  return await callApi("/query_project", { project, query });
+}
+
+export async function getTraceStepCounts(project, run) {
+  if (await isStaticMode()) return [];
+
+  const normalized = normalizeRun(run);
+  const where = normalized.run_id
+    ? `m.run_id = ${sqlString(normalized.run_id)}`
+    : `m.run_name = ${sqlString(normalized.run)}`;
+  const result = await queryProject(
+    project,
+    `
+      SELECT
+        m.step AS step,
+        SUM(
+          CASE
+            WHEN json_type(j.value) = 'array' THEN json_array_length(j.value)
+            ELSE 1
+          END
+        ) AS trace_count
+      FROM metrics m, json_each(CAST(m.metrics AS TEXT)) j
+      WHERE ${where}
+        AND m.step IS NOT NULL
+        AND (
+          (
+            json_type(j.value) = 'array'
+            AND (
+              json_extract(j.value, '$[0]._type') = 'trackio.trace'
+              OR json_extract(j.value, '$[0].messages') IS NOT NULL
+            )
+          )
+          OR (
+            json_type(j.value) = 'object'
+            AND (
+              json_extract(j.value, '$._type') = 'trackio.trace'
+              OR json_extract(j.value, '$.messages') IS NOT NULL
+            )
+          )
+        )
+      GROUP BY m.step
+      ORDER BY m.step
+    `,
+  );
+  return (result.rows || []).map((row) => ({
+    step: Number(row.step),
+    count: Number(row.trace_count || 0),
+  })).filter((row) => Number.isFinite(row.step) && row.count > 0);
+}
+
+export async function getScalarLogsBatch(project, runs) {
+  if (await isStaticMode()) {
+    const out = [];
+    for (const run of runs) {
+      const logs = await staticApi.getLogs(project, run);
+      out.push({ ...normalizeRun(run), logs: scalarOnlyLogs(logs) });
+    }
+    return out;
+  }
+
+  const out = [];
+  for (const run of runs) {
+    const normalized = normalizeRun(run);
+    const where = normalized.run_id
+      ? `m.run_id = ${sqlString(normalized.run_id)}`
+      : `m.run_name = ${sqlString(normalized.run)}`;
+    const result = await queryProject(
+      project,
+      `
+        SELECT
+          m.timestamp,
+          m.step,
+          json_group_object(j.key, j.value) AS metrics
+        FROM metrics m, json_each(CAST(m.metrics AS TEXT)) j
+        WHERE ${where}
+          AND json_type(j.value) IN ('integer', 'real')
+        GROUP BY m.id, m.timestamp, m.step
+        ORDER BY m.timestamp
+      `,
+    );
+    out.push({
+      ...normalized,
+      logs: scalarLogsFromQueryRows(result.rows),
+    });
+  }
+  return out;
+}
+
 export async function getTraces(project, run, options = {}) {
   const params = { project, ...normalizeRun(run), ...options };
   if (await isStaticMode()) return staticApi.getTraces(project, run, options);

trackio/frontend/src/pages/Metrics.svelte

@@ -5,7 +5,7 @@
   import BarPlot from "../components/BarPlot.svelte";
   import Accordion from "../components/Accordion.svelte";
   import LoadingTrackio from "../components/LoadingTrackio.svelte";
-  import { getLogsBatch } from "../lib/api.js";
+  import { getScalarLogsBatch } from "../lib/api.js";
   import {
     getMetricsPollIntervalMs,
     isRateLimitCooldownActive,
@@ -168,7 +168,7 @@
     const results = [];
     for (let i = 0; i < runs.length; i += MAX_BATCH_RUNS) {
       const chunk = runs.slice(i, i + MAX_BATCH_RUNS);
-      const batch = await getLogsBatch(project, chunk);
+      const batch = await getScalarLogsBatch(project, chunk);
       results.push(...batch);
     }
     return results;