feat(tracing): promote session_id/user_id to first-class stream columns

src/lib/tracing/tracer.ts

@@ -13,6 +13,7 @@ import {
   stepFactory,
 } from './steps';
 import Openlayer, { type ClientOptions } from '../../index';
+import type { DataStreamParams } from '../../resources/inference-pipelines/data';
 import { OfflineBuffer } from './offlineBuffer';
 
 let currentTrace: Trace | null = null;
@@ -212,16 +213,7 @@ export function processAndUploadTrace(trace: Trace, openlayerInferencePipelineId
     console.debug('Uploading trace to Openlayer...');
 
     // Lifted to a const so the failure handler can buffer it for later replay.
-    const config = {
-      outputColumnName: 'output',
-      inputVariableNames: inputVariableNames,
-      groundTruthColumnName: 'groundTruth',
-      latencyColumnName: 'latency',
-      costColumnName: 'cost',
-      timestampColumnName: 'inferenceTimestamp',
-      inferenceIdColumnName: 'inferenceId',
-      numOfTokenColumnName: 'tokens',
-    };
+    const config = buildStreamConfig(processedTraceData, inputVariableNames);
 
     return openlayerClient.inferencePipelines.data
       .stream(inferencePipelineId, { config, rows: [processedTraceData] })
@@ -718,9 +710,51 @@ export function postProcessTrace(traceObj: Trace): { traceData: any; inputVariab
     Object.assign(traceData, input_variables);
   }
 
+  // Surface session_id / user_id from the root step's metadata as top-level
+  // columns so the stream config can map them to a first-class session/user
+  // (mirrors the Python SDK). An explicit input variable of the same name wins.
+  const rootMetadata = (rootStep!.metadata ?? {}) as Record<string, any>;
+  for (const key of ['session_id', 'user_id'] as const) {
+    if (rootMetadata[key] != null && !(key in traceData)) {
+      (traceData as Record<string, any>)[key] = rootMetadata[key];
+    }
+  }
+
   return { traceData, inputVariableNames };
 }
 
+/**
+ * Builds the inference-pipeline stream config for a processed trace row.
+ *
+ * The session/user column names are only included when the row actually carries
+ * a `session_id` / `user_id`, so the Openlayer platform records a first-class
+ * session/user for the trace. This mirrors the Python SDK's `post_process_trace`.
+ */
+export function buildStreamConfig(
+  traceData: Record<string, any>,
+  inputVariableNames: string[],
+): DataStreamParams.LlmData {
+  const config: DataStreamParams.LlmData = {
+    outputColumnName: 'output',
+    inputVariableNames: inputVariableNames,
+    groundTruthColumnName: 'groundTruth',
+    latencyColumnName: 'latency',
+    costColumnName: 'cost',
+    timestampColumnName: 'inferenceTimestamp',
+    inferenceIdColumnName: 'inferenceId',
+    numOfTokenColumnName: 'tokens',
+  };
+
+  if (traceData['session_id'] != null) {
+    config.sessionIdColumnName = 'session_id';
+  }
+  if (traceData['user_id'] != null) {
+    config.userIdColumnName = 'user_id';
+  }
+
+  return config;
+}
+
 export default trace;
 
 // ----------------------------------------------------------------------------

tests/tracer-session-config.test.ts

@@ -0,0 +1,83 @@
+import { buildStreamConfig, postProcessTrace } from '../src/lib/tracing/tracer';
+import { Trace } from '../src/lib/tracing/traces';
+import { stepFactory, StepType } from '../src/lib/tracing/steps';
+
+describe('buildStreamConfig', () => {
+  it('always includes the fixed column config', () => {
+    const config = buildStreamConfig({}, ['question']);
+    expect(config).toMatchObject({
+      outputColumnName: 'output',
+      inputVariableNames: ['question'],
+      groundTruthColumnName: 'groundTruth',
+      latencyColumnName: 'latency',
+      costColumnName: 'cost',
+      timestampColumnName: 'inferenceTimestamp',
+      inferenceIdColumnName: 'inferenceId',
+      numOfTokenColumnName: 'tokens',
+    });
+  });
+
+  it('omits the session/user column names when the row carries neither', () => {
+    const config = buildStreamConfig({ output: 'hi' }, []);
+    expect(config.sessionIdColumnName).toBeUndefined();
+    expect(config.userIdColumnName).toBeUndefined();
+  });
+
+  it('adds sessionIdColumnName only when the row carries a session_id', () => {
+    const config = buildStreamConfig({ session_id: 'thread-1' }, []);
+    expect(config.sessionIdColumnName).toBe('session_id');
+    expect(config.userIdColumnName).toBeUndefined();
+  });
+
+  it('adds userIdColumnName only when the row carries a user_id', () => {
+    const config = buildStreamConfig({ user_id: 'u-1' }, []);
+    expect(config.userIdColumnName).toBe('user_id');
+    expect(config.sessionIdColumnName).toBeUndefined();
+  });
+
+  it('adds both column names when the row carries both', () => {
+    const config = buildStreamConfig({ session_id: 's', user_id: 'u' }, []);
+    expect(config.sessionIdColumnName).toBe('session_id');
+    expect(config.userIdColumnName).toBe('user_id');
+  });
+});
+
+describe('postProcessTrace session/user promotion', () => {
+  function traceWithRootMetadata(metadata: Record<string, any>): Trace {
+    const trace = new Trace();
+    const root = stepFactory(StepType.CHAT_COMPLETION, 'root', { question: 'hi' }, 'answer', metadata);
+    trace.addStep(root);
+    return trace;
+  }
+
+  it('promotes session_id and user_id from root metadata to top-level columns', () => {
+    const { traceData } = postProcessTrace(
+      traceWithRootMetadata({ session_id: 'thread-42', user_id: 'user-7' }),
+    );
+
+    expect(traceData.session_id).toBe('thread-42');
+    expect(traceData.user_id).toBe('user-7');
+
+    // ...and they flow through to the stream config the platform reads.
+    const config = buildStreamConfig(traceData, []);
+    expect(config.sessionIdColumnName).toBe('session_id');
+    expect(config.userIdColumnName).toBe('user_id');
+  });
+
+  it('does not add session_id / user_id when the root metadata has none', () => {
+    const { traceData } = postProcessTrace(traceWithRootMetadata({ foo: 'bar' }));
+    expect('session_id' in traceData).toBe(false);
+    expect('user_id' in traceData).toBe(false);
+  });
+
+  it('lets an input variable named session_id win over the metadata promotion', () => {
+    const trace = new Trace();
+    const root = stepFactory(StepType.CHAT_COMPLETION, 'root', { session_id: 'from-input' }, 'answer', {
+      session_id: 'from-metadata',
+    });
+    trace.addStep(root);
+
+    const { traceData } = postProcessTrace(trace);
+    expect(traceData.session_id).toBe('from-input');
+  });
+});