fix(utils.py): fix replicate completion cost calculation

BerriAI · May 18, 2024 · 1cecdc4 · 1cecdc4
1 parent c8a1cf6
commit 1cecdc4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
@@ -2301,6 +2301,8 @@ def test_completion_azure_deployment_id():
 
 # test_completion_azure_deployment_id()
 
+import asyncio
+
 
 @pytest.mark.parametrize("sync_mode", [False, True])
 @pytest.mark.asyncio

diff --git a/litellm/utils.py b/litellm/utils.py
@@ -3853,7 +3853,7 @@ def get_replicate_completion_pricing(completion_response=None, total_time=0.0):
     )
     if total_time == 0.0:  # total time is in ms
         start_time = completion_response["created"]
-        end_time = completion_response["ended"]
+        end_time = getattr(completion_response, "ended", time.time())
         total_time = end_time - start_time
 
     return a100_80gb_price_per_second_public * total_time / 1000