wip: fixing tests

integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py

@@ -206,16 +206,16 @@ def run(
         # rename the meta key to be inline with OpenAI meta output keys
         for response in replies:
             if response.meta is not None:
-                if 'usage' not in response.meta:
+                if "usage" not in response.meta:
                     if "prompt_token_count" in response.meta:
                         response.meta["prompt_tokens"] = response.meta.pop("prompt_token_count")
-                    if 'generation_token_count' in response.meta:
+                    if "generation_token_count" in response.meta:
                         response.meta["completion_tokens"] = response.meta.pop("generation_token_count")
                 elif "usage" in response.meta:
-                    if "input_tokens" in response.meta['usage']:
-                        response.meta["usage"]['prompt_tokens'] = response.meta['usage'].pop("input_tokens")
-                    if "output_token" in response.meta['usage']:
-                        response.meta["usage"]['completion_tokens'] = response.meta['usage'].pop("output_token")
+                    if "input_tokens" in response.meta["usage"]:
+                        response.meta["usage"]["prompt_tokens"] = response.meta["usage"].pop("input_tokens")
+                    if "output_token" in response.meta["usage"]:
+                        response.meta["usage"]["completion_tokens"] = response.meta["usage"].pop("output_token")
                 else:
                     print("DEBUG", response.meta)
 

integrations/amazon_bedrock/tests/test_chat_generator.py

@@ -358,13 +358,13 @@ def test_default_inference_params(self, model_name, chat_messages):
         assert "paris" in first_reply.content.lower(), "First reply does not contain 'paris'"
         assert first_reply.meta, "First reply has no metadata"
 
-        if first_reply.meta and 'usage' in first_reply.meta:
-            assert 'prompt_tokens' in first_reply.meta['usage']
-            assert 'completion_tokens' in first_reply.meta['usage']
+        if first_reply.meta and "usage" in first_reply.meta:
+            assert "prompt_tokens" in first_reply.meta["usage"]
+            assert "completion_tokens" in first_reply.meta["usage"]
 
-        if first_reply.meta and 'usage' not in first_reply.meta:
-            assert 'prompt_tokens' in first_reply.meta
-            assert 'completion_tokens' in first_reply.meta
+        if first_reply.meta and "usage" not in first_reply.meta:
+            assert "prompt_tokens" in first_reply.meta
+            assert "completion_tokens" in first_reply.meta
 
     @pytest.mark.parametrize("model_name", MODELS_TO_TEST)
     @pytest.mark.integration