feat(utils.py): support cost tracking for openai/azure image gen models

tests/test_keys.py

@@ -14,7 +14,11 @@ import litellm
 
 
 async def generate_key(
-    session, i, budget=None, budget_duration=None, models=["azure-models", "gpt-4"]
+    session,
+    i,
+    budget=None,
+    budget_duration=None,
+    models=["azure-models", "gpt-4", "dall-e-3"],
 ):
     url = "http://0.0.0.0:4000/key/generate"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
@@ -129,6 +133,39 @@ async def chat_completion(session, key, model="gpt-4"):
                 pass
 
 
+async def image_generation(session, key, model="dall-e-3"):
+    url = "http://0.0.0.0:4000/v1/images/generations"
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "Content-Type": "application/json",
+    }
+    data = {
+        "model": model,
+        "prompt": "A cute baby sea otter",
+    }
+
+    for i in range(3):
+        try:
+            async with session.post(url, headers=headers, json=data) as response:
+                status = response.status
+                response_text = await response.text()
+
+                print(response_text)
+                print()
+
+                if status != 200:
+                    raise Exception(
+                        f"Request did not return a 200 status code: {status}. Response: {response_text}"
+                    )
+
+                return await response.json()
+        except Exception as e:
+            if "Request did not return a 200 status code" in str(e):
+                raise e
+            else:
+                pass
+
+
 async def chat_completion_streaming(session, key, model="gpt-4"):
     client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000")
     messages = [
@@ -357,6 +394,40 @@ async def test_key_info_spend_values_streaming():
         assert rounded_response_cost == rounded_key_info_spend
 
 
+@pytest.mark.asyncio
+async def test_key_info_spend_values_image_generation():
+    """
+    Test to ensure spend is correctly calculated
+    - create key
+    - make image gen call
+    - assert cost is expected value
+    """
+
+    async def retry_request(func, *args, _max_attempts=5, **kwargs):
+        for attempt in range(_max_attempts):
+            try:
+                return await func(*args, **kwargs)
+            except aiohttp.client_exceptions.ClientOSError as e:
+                if attempt + 1 == _max_attempts:
+                    raise  # re-raise the last ClientOSError if all attempts failed
+                print(f"Attempt {attempt+1} failed, retrying...")
+
+    async with aiohttp.ClientSession(
+        timeout=aiohttp.ClientTimeout(total=600)
+    ) as session:
+        ## Test Spend Update ##
+        # completion
+        key_gen = await generate_key(session=session, i=0)
+        key = key_gen["key"]
+        response = await image_generation(session=session, key=key)
+        await asyncio.sleep(5)
+        key_info = await retry_request(
+            get_key_info, session=session, get_key=key, call_key=key
+        )
+        spend = key_info["info"]["spend"]
+        assert spend > 0
+
+
 @pytest.mark.asyncio
 async def test_key_with_budgets():
     """