fix: bedrock traces do not have correct latency in TypeScript client

Author: vikasnair

src/lib/integrations/bedrockAgentTracer.ts

@@ -199,6 +199,8 @@ function createTracedCompletion(
           rawOutput: JSON.stringify(rawOutputChunks, null, 2),
           metadata: metadata,
           provider: 'Bedrock',
+          startTime: startTime,
+          endTime: endTime,
         };
 
         addChatCompletionStepToTrace(traceStepData);

src/lib/integrations/langchainCallback.ts

@@ -136,6 +136,8 @@ export class OpenlayerHandler extends BaseCallbackHandler {
       modelParameters: this.modelParameters,
       metadata: this.metadata,
       provider: this.provider ?? '',
+      startTime: this.startTime,
+      endTime: this.endTime,
     });
   }
 }

src/lib/integrations/openAiTracer.ts

@@ -69,6 +69,8 @@ export function traceOpenAI(openai: OpenAI): OpenAI {
               completionTokens: completionTokens,
               promptTokens: 0,
               tokens: completionTokens,
+              startTime: startTime,
+              endTime: endTime,
             };
             addChatCompletionStepToTrace(traceData);
           }
@@ -101,6 +103,8 @@ export function traceOpenAI(openai: OpenAI): OpenAI {
           rawOutput: JSON.stringify(response, null, 2),
           metadata: {},
           provider: 'OpenAI',
+          startTime: startTime,
+          endTime: endTime,
         };
         addChatCompletionStepToTrace(traceData);
         return response;

src/lib/tracing/steps.ts

@@ -39,14 +39,15 @@ export class Step {
   latency: number | null = null;
   steps: Step[] = [];
 
-  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
+  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}, startTime?: number | null, endTime?: number | null) {
     this.name = name;
     this.id = uuidv4();
     this.inputs = inputs;
     this.output = output;
     this.metadata = metadata;
 
-    this.startTime = Date.now();
+    this.startTime = startTime ?? Date.now();
+    this.endTime = endTime ?? null;
   }
 
   addNestedStep(nestedStep: Step): void {
@@ -80,8 +81,8 @@ export class Step {
 }
 
 export class UserCallStep extends Step {
-  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
-    super(name, inputs, output, metadata);
+  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}, startTime?: number | null, endTime?: number | null) {
+    super(name, inputs, output, metadata, startTime, endTime);
     this.stepType = StepType.USER_CALL;
   }
 }
@@ -96,8 +97,8 @@ export class ChatCompletionStep extends Step {
   modelParameters: Record<string, any> | null = null;
   rawOutput: string | null = null;
 
-  constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}) {
-    super(name, inputs, output, metadata);
+    constructor(name: string, inputs: any = null, output: any = null, metadata: Record<string, any> = {}, startTime?: number | null, endTime?: number | null) {
+    super(name, inputs, output, metadata, startTime, endTime);
     this.stepType = StepType.CHAT_COMPLETION;
   }
 

src/lib/tracing/tracer.ts

@@ -31,13 +31,15 @@ function createStep(
   inputs?: any,
   output?: any,
   metadata: Record<string, any> | null = null,
+  startTime?: number | null,
+  endTime?: number | null,
 ): [Step, () => void] {
   metadata = metadata || {};
   let newStep: Step;
   if (stepType === StepType.CHAT_COMPLETION) {
-    newStep = new ChatCompletionStep(name, inputs, output, metadata);
+    newStep = new ChatCompletionStep(name, inputs, output, metadata, startTime, endTime);
   } else {
-    newStep = new UserCallStep(name, inputs, output, metadata);
+    newStep = new UserCallStep(name, inputs, output, metadata, startTime, endTime);
   }
 
   const parentStep = getCurrentStep();
@@ -167,6 +169,8 @@ export function addChatCompletionStepToTrace({
   rawOutput = null,
   metadata = {},
   provider = 'OpenAI',
+  startTime = null,
+  endTime = null,
 }: {
   name: string;
   inputs: any;
@@ -180,8 +184,10 @@ export function addChatCompletionStepToTrace({
   rawOutput?: string | null;
   metadata?: Record<string, any>;
   provider?: string;
+  startTime?: number | null;
+  endTime?: number | null;
 }) {
-  const [step, endStep] = createStep(name, StepType.CHAT_COMPLETION, inputs, output, metadata);
+  const [step, endStep] = createStep(name, StepType.CHAT_COMPLETION, inputs, output, metadata, startTime, endTime);
 
   if (step.stepType === StepType.CHAT_COMPLETION) {
     (step as ChatCompletionStep).provider = provider;