align bedrock with spec

Author: jsonbailey

lib/server/ai/config_tracker.rb

@@ -74,7 +74,7 @@ def track_duration(duration)
         # @yield The block to measure
         # @return The result of the block
         #
-        def track_duration_of
+        def track_duration_of(&block)
           start_time = Time.now
           yield
         ensure
@@ -207,21 +207,17 @@ def track_openai_metrics(&block)
         # Track AWS Bedrock conversation operations.
         # This method tracks the duration, token usage, and success/error status.
         #
-        # @param res [Hash] Response hash from Bedrock.
+        # @yield The block to track.
         # @return [Hash] The original response hash.
         #
-        def track_bedrock_converse_metrics(res)
-          status_code = res.dig(:'$metadata', :httpStatusCode) || 0
-          if status_code == 200
-            track_success
-          elsif status_code >= 400
-            track_error
-          end
-
-          track_duration(res[:metrics][:latency_ms]) if res.dig(:metrics, :latency_ms)
-          track_tokens(bedrock_to_token_usage(res[:usage])) if res[:usage]
-
-          res
+        def track_bedrock_converse_metrics(&block)
+          result = track_duration_of(&block)
+          track_success
+          track_tokens(bedrock_to_token_usage(result[:usage])) if result[:usage]
+          result
+        rescue StandardError
+          track_error
+          raise
         end
 
         private def flag_data

spec/server/ai/config_tracker_spec.rb

@@ -122,19 +122,17 @@
   end
 
   describe '#track_bedrock_metrics' do
-    it 'tracks duration and tokens' do
-      # TODO: Verify the $metadata field in bedrock result. I don't see anything like this in the docs.
-      bedrock_result = {
-        '$metadata': { httpStatusCode: 200 },
+    let(:bedrock_result) do
+      {
         usage: {
           total_tokens: 300,
           input_tokens: 200,
           output_tokens: 100,
         },
-        metrics: {
-          latency_ms: 50,
-        },
       }
+    end
+
+    it 'tracks duration and tokens' do
       expect(ld_client).to receive(:track).with(
         '$ld:ai:generation',
         context,
@@ -172,30 +170,19 @@
         100
       )
 
-      result = tracker.track_bedrock_converse_metrics(bedrock_result)
+      result = tracker.track_bedrock_converse_metrics { bedrock_result }
       expect(result).to eq(bedrock_result)
       expect(tracker.summary).to be_a(LaunchDarkly::Server::AI::MetricSummary)
       expect(tracker.summary.usage).to be_a(LaunchDarkly::Server::AI::TokenUsage)
       expect(tracker.summary.usage.total).to eq(300)
       expect(tracker.summary.usage.input).to eq(200)
       expect(tracker.summary.usage.output).to eq(100)
-      expect(tracker.summary.duration).to eq(50)
+      expect(tracker.summary.duration).to be_a(Integer)
+      expect(tracker.summary.duration).to be >= 0
       expect(tracker.summary.success).to be true
     end
 
-    it 'tracks duration and tokens with error' do
-      # Verify the $metadata field in bedrock result. I don't see anything like this in the docs.
-      bedrock_result = {
-        '$metadata': { httpStatusCode: 400 },
-        usage: {
-          total_tokens: 300,
-          input_tokens: 200,
-          output_tokens: 100,
-        },
-        metrics: {
-          latency_ms: 50,
-        },
-      }
+    it 'tracks error for failed operation' do
       expect(ld_client).to receive(:track).with(
         '$ld:ai:generation',
         context,
@@ -214,33 +201,11 @@
         tracker_flag_data,
         kind_of(Integer)
       )
-      expect(ld_client).to receive(:track).with(
-        '$ld:ai:tokens:total',
-        context,
-        tracker_flag_data,
-        300
-      )
-      expect(ld_client).to receive(:track).with(
-        '$ld:ai:tokens:input',
-        context,
-        tracker_flag_data,
-        200
-      )
-      expect(ld_client).to receive(:track).with(
-        '$ld:ai:tokens:output',
-        context,
-        tracker_flag_data,
-        100
-      )
 
-      result = tracker.track_bedrock_converse_metrics(bedrock_result)
-      expect(result).to eq(bedrock_result)
-      expect(tracker.summary).to be_a(LaunchDarkly::Server::AI::MetricSummary)
-      expect(tracker.summary.usage).to be_a(LaunchDarkly::Server::AI::TokenUsage)
-      expect(tracker.summary.usage.total).to eq(300)
-      expect(tracker.summary.usage.input).to eq(200)
-      expect(tracker.summary.usage.output).to eq(100)
-      expect(tracker.summary.duration).to eq(50)
+      expect { tracker.track_bedrock_converse_metrics { raise 'test error' } }.to raise_error('test error')
+      expect(tracker.summary.usage).to be_nil
+      expect(tracker.summary.duration).to be_a(Integer)
+      expect(tracker.summary.duration).to be >= 0
       expect(tracker.summary.success).to be false
     end
   end
@@ -299,6 +264,8 @@
       expect(tracker.summary.usage.total).to eq(300)
       expect(tracker.summary.usage.input).to eq(200)
       expect(tracker.summary.usage.output).to eq(100)
+      expect(tracker.summary.duration).to be_a(Integer)
+      expect(tracker.summary.duration).to be >= 0
       expect(tracker.summary.success).to be true
     end
 
@@ -324,6 +291,9 @@
 
       expect { tracker.track_openai_metrics { raise 'test error' } }.to raise_error('test error')
       expect(tracker.summary.usage).to be_nil
+      expect(tracker.summary.duration).to be_a(Integer)
+      expect(tracker.summary.duration).to be >= 0
+      expect(tracker.summary.success).to be false
     end
   end