fix(aiml): attribute aimlapi.com traffic, read its documented key env var, and repair its cost map

README.md

@@ -270,8 +270,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
 
 | Provider                                                                            | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` |
 |-------------------------------------------------------------------------------------|---------------------|-------------|--------------|---------------|----------------------|-------------------------|-----------------|----------------|-----------|-----------|
+| [aimlapi.com (`aiml`)](https://docs.litellm.ai/docs/providers/aiml) | ✅ | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  |
 | [Abliteration (`abliteration`)](https://docs.litellm.ai/docs/providers/abliteration) | ✅ |  |  |  |  |  |  |  |  |  |
-| [AI/ML API (`aiml`)](https://docs.litellm.ai/docs/providers/aiml) | ✅ | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  |
 | [AI21 (`ai21`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [AI21 Chat (`ai21_chat`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Aleph Alpha](https://docs.litellm.ai/docs/providers/aleph_alpha) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |

litellm/llms/aiml/chat/transformation.py

@@ -1,7 +1,7 @@
 from typing import Final
 
+from litellm.llms.aiml.common_utils import get_aiml_api_base, get_aiml_api_key
 from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
-from litellm.secret_managers.main import get_secret_str
 
 
 class AIMLChatConfig(OpenAIGPTConfig):
@@ -12,9 +12,6 @@ def custom_llm_provider(self) -> str | None:
     def _get_openai_compatible_provider_info(
         self, api_base: str | None, api_key: str | None
     ) -> tuple[str | None, str | None]:
-        # AIML is openai compatible, we just need to set the api_base
-        api_base = (
-            api_base or get_secret_str("AIML_API_BASE") or "https://api.aimlapi.com/v1"  # Default AIML API base URL
-        )
-        dynamic_api_key: Final = api_key or get_secret_str("AIML_API_KEY")
-        return api_base, dynamic_api_key
+        resolved_api_base: Final = get_aiml_api_base(api_base)
+        dynamic_api_key: Final = get_aiml_api_key(api_key)
+        return resolved_api_base, dynamic_api_key

litellm/llms/aiml/common_utils.py

@@ -0,0 +1,47 @@
+import types
+from collections.abc import Mapping
+from typing import Final
+from urllib.parse import urlparse
+
+from litellm.secret_managers.main import get_secret_str
+
+AIML_DEFAULT_API_BASE: Final = "https://api.aimlapi.com/v1"
+
+AIML_ATTRIBUTION_HOSTS: Final = frozenset({"api.aimlapi.com"})
+
+AIML_ATTRIBUTION_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType(
+    {
+        "HTTP-Referer": "https://github.com/BerriAI/litellm",
+        "X-Title": "LiteLLM",
+        "X-AIMLAPI-Partner-ID": "part_O0eykPA6gQNIFEYaUBojBbU4",
+        "X-AIMLAPI-Source": "agent/litellm",
+    }
+)
+
+_NO_ATTRIBUTION: Final[Mapping[str, str]] = types.MappingProxyType({})
+
+
+def get_aiml_api_key(api_key: str | None = None) -> str | None:
+    return (
+        api_key or get_secret_str("AIML_API_KEY") or get_secret_str("AIMLAPI_API_KEY") or get_secret_str("AIMLAPI_KEY")
+    )
+
+
+def get_aiml_api_base(api_base: str | None = None) -> str:
+    return api_base or get_secret_str("AIML_API_BASE") or AIML_DEFAULT_API_BASE
+
+
+def aiml_attribution_headers(api_base: str | None) -> Mapping[str, str]:
+    host: Final = urlparse(get_aiml_api_base(api_base)).hostname
+    if host in AIML_ATTRIBUTION_HOSTS:
+        return AIML_ATTRIBUTION_HEADERS
+    return _NO_ATTRIBUTION
+
+
+def with_aiml_attribution(
+    headers: Mapping[str, str] | None, api_base: str | None
+) -> dict[str, str]:  # mutable-ok: the OpenAI SDK extra_headers slot this feeds is typed as a plain dict
+    return {  # mutable-ok: a fresh dict per request leaves the caller's headers and AIML_ATTRIBUTION_HEADERS unmutated
+        **aiml_attribution_headers(api_base),
+        **(headers or _NO_ATTRIBUTION),
+    }

litellm/llms/aiml/image_generation/transformation.py

@@ -2,6 +2,7 @@
 
 import httpx
 
+from litellm.llms.aiml.common_utils import aiml_attribution_headers, get_aiml_api_key
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
@@ -129,15 +130,16 @@ def validate_environment(
         api_key: str | None = None,
         api_base: str | None = None,
     ) -> dict:
-        final_api_key: Final[str | None] = (
-            api_key or get_secret_str("AIML_API_KEY") or get_secret_str("AIMLAPI_KEY")  # Alternative name
-        )
+        final_api_key: Final[str | None] = get_aiml_api_key(api_key)
         if not final_api_key:
-            raise ValueError("AIML_API_KEY or AIMLAPI_KEY is not set")
-
-        headers["Authorization"] = f"Bearer {final_api_key}"
-        headers["Content-Type"] = "application/json"
-        return headers
+            raise ValueError("AIML_API_KEY, AIMLAPI_API_KEY or AIMLAPI_KEY is not set")
+
+        return {  # mutable-ok: a fresh dict leaves the caller's header mapping unmutated; the base signature returns dict
+            **aiml_attribution_headers(api_base or self.DEFAULT_BASE_URL),
+            **headers,
+            "Authorization": f"Bearer {final_api_key}",
+            "Content-Type": "application/json",
+        }
 
     def transform_image_generation_request(
         self,

litellm/main.py

@@ -2534,6 +2534,11 @@ def _complete_custom_openai(
 
     headers = headers or litellm.headers
 
+    if custom_llm_provider == "aiml":
+        from litellm.llms.aiml.common_utils import with_aiml_attribution
+
+        headers = with_aiml_attribution(headers, api_base)
+
     # Add GitHub Copilot headers (same as /responses endpoint does)
     if custom_llm_provider == "github_copilot":
         from litellm.llms.github_copilot.authenticator import Authenticator

litellm/model_prices_and_context_window_backup.json

@@ -143,26 +143,20 @@
         "output_cost_per_token": 7e-07,
         "supports_system_messages": true
     },
-    "aiml/dall-e-2": {
+    "aiml/blackforestlabs/flux-pro-1.1": {
         "litellm_provider": "aiml",
-        "metadata": {
-            "notes": "DALL-E 2 via AI/ML API - Reliable text-to-image generation"
-        },
         "mode": "image_generation",
-        "output_cost_per_image": 0.026,
-        "source": "https://docs.aimlapi.com/",
+        "output_cost_per_image": 0.052,
+        "source": "https://api.aimlapi.com/v1/models?include=pricing",
         "supported_endpoints": [
             "/v1/images/generations"
         ]
     },
-    "aiml/dall-e-3": {
+    "aiml/blackforestlabs/flux-pro-1.1-ultra": {
         "litellm_provider": "aiml",
-        "metadata": {
-            "notes": "DALL-E 3 via AI/ML API - High-quality text-to-image generation"
-        },
         "mode": "image_generation",
-        "output_cost_per_image": 0.052,
-        "source": "https://docs.aimlapi.com/",
+        "output_cost_per_image": 0.078,
+        "source": "https://api.aimlapi.com/v1/models?include=pricing",
         "supported_endpoints": [
             "/v1/images/generations"
         ]
@@ -190,7 +184,7 @@
     "aiml/flux-pro/v1.1-ultra": {
         "litellm_provider": "aiml",
         "mode": "image_generation",
-        "output_cost_per_image": 0.063,
+        "output_cost_per_image": 0.078,
         "supported_endpoints": [
             "/v1/images/generations"
         ]

litellm/provider_endpoints_support_backup.json

@@ -32,6 +32,24 @@
     }
   },
   "providers": {
+    "aiml": {
+      "display_name": "aimlapi.com (`aiml`)",
+      "url": "https://docs.litellm.ai/docs/providers/aiml",
+      "endpoints": {
+        "chat_completions": true,
+        "messages": true,
+        "responses": true,
+        "embeddings": true,
+        "image_generations": true,
+        "audio_transcriptions": false,
+        "audio_speech": false,
+        "moderations": false,
+        "batches": false,
+        "rerank": false,
+        "a2a": true,
+        "interactions": true
+      }
+    },
     "a2a": {
       "display_name": "A2A (Agent-to-Agent) (`a2a`)",
       "url": "https://docs.litellm.ai/docs/providers/a2a",
@@ -66,24 +84,6 @@
         "a2a": false
       }
     },
-    "aiml": {
-      "display_name": "AI/ML API (`aiml`)",
-      "url": "https://docs.litellm.ai/docs/providers/aiml",
-      "endpoints": {
-        "chat_completions": true,
-        "messages": true,
-        "responses": true,
-        "embeddings": true,
-        "image_generations": true,
-        "audio_transcriptions": false,
-        "audio_speech": false,
-        "moderations": false,
-        "batches": false,
-        "rerank": false,
-        "a2a": true,
-        "interactions": true
-      }
-    },
     "ai21": {
       "display_name": "AI21 (`ai21`)",
       "url": "https://docs.litellm.ai/docs/providers/ai21",

litellm/proxy/public_endpoints/provider_create_fields.json

@@ -1,7 +1,7 @@
 [
   {
     "provider": "AIML",
-    "provider_display_name": "AI/ML API",
+    "provider_display_name": "aimlapi.com",
     "litellm_provider": "aiml",
     "credential_fields": [
       {

litellm/proxy/spend_tracking/budget_reservation.py

@@ -1202,7 +1202,7 @@ def _estimate_image_generation_cost(
 
     The "output" vs "input" cost-per-image naming is inconsistent across
     providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while
-    aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed.
+    aiml/blackforestlabs/flux-pro-1.1 uses ``output_cost_per_image`` — so both are summed.
     """
     # Gate strictly on `mode`. Several chat and embedding models carry
     # ``input_cost_per_image`` / ``output_cost_per_image`` to price multimodal

model_prices_and_context_window.json

@@ -143,26 +143,20 @@
         "output_cost_per_token": 7e-07,
         "supports_system_messages": true
     },
-    "aiml/dall-e-2": {
+    "aiml/blackforestlabs/flux-pro-1.1": {
         "litellm_provider": "aiml",
-        "metadata": {
-            "notes": "DALL-E 2 via AI/ML API - Reliable text-to-image generation"
-        },
         "mode": "image_generation",
-        "output_cost_per_image": 0.026,
-        "source": "https://docs.aimlapi.com/",
+        "output_cost_per_image": 0.052,
+        "source": "https://api.aimlapi.com/v1/models?include=pricing",
         "supported_endpoints": [
             "/v1/images/generations"
         ]
     },
-    "aiml/dall-e-3": {
+    "aiml/blackforestlabs/flux-pro-1.1-ultra": {
         "litellm_provider": "aiml",
-        "metadata": {
-            "notes": "DALL-E 3 via AI/ML API - High-quality text-to-image generation"
-        },
         "mode": "image_generation",
-        "output_cost_per_image": 0.052,
-        "source": "https://docs.aimlapi.com/",
+        "output_cost_per_image": 0.078,
+        "source": "https://api.aimlapi.com/v1/models?include=pricing",
         "supported_endpoints": [
             "/v1/images/generations"
         ]
@@ -190,7 +184,7 @@
     "aiml/flux-pro/v1.1-ultra": {
         "litellm_provider": "aiml",
         "mode": "image_generation",
-        "output_cost_per_image": 0.063,
+        "output_cost_per_image": 0.078,
         "supported_endpoints": [
             "/v1/images/generations"
         ]

provider_endpoints_support.json

@@ -32,6 +32,24 @@
     }
   },
   "providers": {
+    "aiml": {
+      "display_name": "aimlapi.com (`aiml`)",
+      "url": "https://docs.litellm.ai/docs/providers/aiml",
+      "endpoints": {
+        "chat_completions": true,
+        "messages": true,
+        "responses": true,
+        "embeddings": true,
+        "image_generations": true,
+        "audio_transcriptions": false,
+        "audio_speech": false,
+        "moderations": false,
+        "batches": false,
+        "rerank": false,
+        "a2a": true,
+        "interactions": true
+      }
+    },
     "a2a": {
       "display_name": "A2A (Agent-to-Agent) (`a2a`)",
       "url": "https://docs.litellm.ai/docs/providers/a2a",
@@ -66,24 +84,6 @@
         "a2a": false
       }
     },
-    "aiml": {
-      "display_name": "AI/ML API (`aiml`)",
-      "url": "https://docs.litellm.ai/docs/providers/aiml",
-      "endpoints": {
-        "chat_completions": true,
-        "messages": true,
-        "responses": true,
-        "embeddings": true,
-        "image_generations": true,
-        "audio_transcriptions": false,
-        "audio_speech": false,
-        "moderations": false,
-        "batches": false,
-        "rerank": false,
-        "a2a": true,
-        "interactions": true
-      }
-    },
     "ai21": {
       "display_name": "AI21 (`ai21`)",
       "url": "https://docs.litellm.ai/docs/providers/ai21",