diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py
index 0fb2ea1f7..0b5e9a4aa 100644
--- a/litellm/integrations/langfuse.py
+++ b/litellm/integrations/langfuse.py
@@ -300,8 +300,8 @@ class LangFuseLogger:
                 prompt=input,
                 completion=output,
                 usage={
-                    "prompt_tokens": response_obj["usage"]["prompt_tokens"],
-                    "completion_tokens": response_obj["usage"]["completion_tokens"],
+                    "prompt_tokens": response_obj.usage.prompt_tokens,
+                    "completion_tokens": response_obj.usage.completion_tokens,
                 },
                 metadata=metadata,
             )
@@ -526,8 +526,8 @@ class LangFuseLogger:
             if response_obj is not None and response_obj.get("id", None) is not None:
                 generation_id = litellm.utils.get_logging_id(start_time, response_obj)
                 usage = {
-                    "prompt_tokens": response_obj["usage"]["prompt_tokens"],
-                    "completion_tokens": response_obj["usage"]["completion_tokens"],
+                    "prompt_tokens": response_obj.usage.prompt_tokens,
+                    "completion_tokens": response_obj.usage.completion_tokens,
                     "total_cost": cost if supports_costs else None,
                 }
             generation_name = clean_metadata.pop("generation_name", None)
diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py
index d745a8d77..b4d12ad27 100644
--- a/tests/test_team_logging.py
+++ b/tests/test_team_logging.py
@@ -62,6 +62,7 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=
 
 
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_team_logging():
     """
     -> Team 1 logs to project 1