feat!: Remove bedrock-specific tracker method, use track_metrics_of method with custom extractor (#191)

Author: jsonbailey

packages/sdk/server-ai/src/ldai/tracker.py

@@ -500,35 +500,13 @@ def track_error(self) -> None:
             "$ld:ai:generation:error", self._context, self.__get_track_data(), 1
         )
 
-    def track_bedrock_converse_metrics(self, res: dict) -> dict:
-        """
-        Track AWS Bedrock conversation operations.
-
-
-        This function will track the duration of the operation, the token
-        usage, and the success or error status.
-
-        :param res: Response dictionary from Bedrock.
-        :return: The original response dictionary.
-        """
-        status_code = res.get("ResponseMetadata", {}).get("HTTPStatusCode", 0)
-        if status_code == 200:
-            self.track_success()
-        elif status_code >= 400:
-            self.track_error()
-        if res.get("metrics", {}).get("latencyMs"):
-            self.track_duration(res["metrics"]["latencyMs"])
-        if res.get("usage"):
-            self.track_tokens(_bedrock_to_token_usage(res["usage"]))
-        return res
-
     def track_tokens(self, tokens: TokenUsage) -> None:
         """
         Track token usage metrics.
 
         Records at most once per Tracker; further calls are ignored.
 
-        :param tokens: Token usage data from either custom, OpenAI, or Bedrock sources.
+        :param tokens: Token usage data.
         """
         if self._summary.tokens is not None:
             log.warning(
@@ -587,20 +565,6 @@ def get_summary(self) -> LDAIMetricSummary:
         return self._summary
 
 
-def _bedrock_to_token_usage(data: dict) -> TokenUsage:
-    """
-    Convert a Bedrock usage dictionary to a TokenUsage object.
-
-    :param data: Dictionary containing Bedrock usage data.
-    :return: TokenUsage object containing usage data.
-    """
-    return TokenUsage(
-        total=data.get("totalTokens", 0),
-        input=data.get("inputTokens", 0),
-        output=data.get("outputTokens", 0),
-    )
-
-
 class AIGraphTracker:
     """
     Tracks graph-level metrics for AI agent graph operations.

packages/sdk/server-ai/tests/test_tracker.py

@@ -174,82 +174,6 @@ def test_tracks_token_usage(client: LDClient):
     assert tracker.get_summary().tokens == tokens
 
 
-def test_tracks_bedrock_metrics(client: LDClient):
-    context = Context.create("user-key")
-    tracker = LDAIConfigTracker(
-        ld_client=client, run_id="test-run-id", config_key="config-key",
-        variation_key="variation-key", version=3, model_name="fakeModel",
-        provider_name="fakeProvider", context=context,
-    )
-
-    bedrock_result = {
-        "ResponseMetadata": {"HTTPStatusCode": 200},
-        "usage": {
-            "inputTokens": 220,
-            "outputTokens": 110,
-            "totalTokens": 330,
-        },
-        "metrics": {
-            "latencyMs": 50,
-        },
-    }
-    tracker.track_bedrock_converse_metrics(bedrock_result)
-
-    _btd = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key",
-            "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}
-    calls = [
-        call("$ld:ai:generation:success", context, _btd, 1),
-        call("$ld:ai:duration:total", context, _btd, 50),
-        call("$ld:ai:tokens:total", context, _btd, 330),
-        call("$ld:ai:tokens:input", context, _btd, 220),
-        call("$ld:ai:tokens:output", context, _btd, 110),
-    ]
-
-    client.track.assert_has_calls(calls)  # type: ignore
-
-    assert tracker.get_summary().success is True
-    assert tracker.get_summary().duration_ms == 50
-    assert tracker.get_summary().tokens == TokenUsage(330, 220, 110)
-
-
-def test_tracks_bedrock_metrics_with_error(client: LDClient):
-    context = Context.create("user-key")
-    tracker = LDAIConfigTracker(
-        ld_client=client, run_id="test-run-id", config_key="config-key",
-        variation_key="variation-key", version=3, model_name="fakeModel",
-        provider_name="fakeProvider", context=context,
-    )
-
-    bedrock_result = {
-        "ResponseMetadata": {"HTTPStatusCode": 500},
-        "usage": {
-            "totalTokens": 330,
-            "inputTokens": 220,
-            "outputTokens": 110,
-        },
-        "metrics": {
-            "latencyMs": 50,
-        },
-    }
-    tracker.track_bedrock_converse_metrics(bedrock_result)
-
-    _etd = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key",
-            "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}
-    calls = [
-        call("$ld:ai:generation:error", context, _etd, 1),
-        call("$ld:ai:duration:total", context, _etd, 50),
-        call("$ld:ai:tokens:total", context, _etd, 330),
-        call("$ld:ai:tokens:input", context, _etd, 220),
-        call("$ld:ai:tokens:output", context, _etd, 110),
-    ]
-
-    client.track.assert_has_calls(calls)  # type: ignore
-
-    assert tracker.get_summary().success is False
-    assert tracker.get_summary().duration_ms == 50
-    assert tracker.get_summary().tokens == TokenUsage(330, 220, 110)
-
-
 @pytest.mark.parametrize(
     "kind,label",
     [