initial import

deepset-ai · Aug 12, 2024 · 6f412c1 · 6f412c1
1 parent 0451e6f
commit 6f412c1
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/...ock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/...ock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py
@@ -205,9 +205,9 @@ def run(
 
         # rename the meta key to be inline with OpenAI meta output keys
         for response in replies:
-            if response.meta is not None and "usage" in response.meta:
-                response.meta["usage"]["prompt_tokens"] = response.meta["usage"].pop("input_tokens")
-                response.meta["usage"]["completion_tokens"] = response.meta["usage"].pop("output_tokens")
+            if response.meta is not None:
+                response.meta["prompt_tokens"] = response.meta.pop("prompt_token_count")
+                response.meta["completion_tokens"] = response.meta.pop("generation_token_count")
 
         return {"replies": replies}