[feat] add support for custom model pricing and update cost calculation logic

Author: Mounir-charef

edenai_apis/llmengine/clients/litellm_client/litellm_client.py

@@ -187,7 +187,8 @@ def completion(
             call_params["model_list"] = model_list
         if user is not None:
             call_params["user"] = user
-        # See if there's a custom pricing here
+        # See if there's custom pricing (model_pricing for extended pricing, or legacy per-token pricing)
+        model_pricing = kwargs.pop("model_pricing", None)
         custom_pricing = {}
         if kwargs.get("input_cost_per_token", None) and kwargs.get(
             "output_cost_per_token", None
@@ -200,6 +201,13 @@ def completion(
             if drop_invalid_params == True:
                 litellm.drop_params = True
             kwargs.pop("moderate_content", None)
+            # Register custom model pricing in litellm's registry for extended pricing support
+            if model_pricing:
+                # Merge with existing pricing to preserve other fields (max_tokens, mode, etc.)
+                if model_name in litellm.model_cost:
+                    litellm.model_cost[model_name].update(model_pricing)
+                else:
+                    litellm.model_cost[model_name] = model_pricing
             provider_start_time = time.time_ns()
             c_response = completion(**call_params, **kwargs)
             provider_end_time = time.time_ns()
@@ -216,7 +224,8 @@ def generate_chunks():
                     "completion_response": c_response,
                     "call_type": "completion",
                 }
-                if len(custom_pricing.keys()) > 0:
+                # Use model_pricing via registry lookup, or fall back to legacy custom_cost_per_token
+                if not model_pricing and len(custom_pricing.keys()) > 0:
                     cost_calc_params["custom_cost_per_token"] = custom_pricing
                 response = {
                     **c_response.model_dump(),
@@ -807,7 +816,8 @@ async def acompletion(
             call_params["model_list"] = model_list
         if user is not None:
             call_params["user"] = user
-        # See if there's a custom pricing here
+        # See if there's custom pricing (model_pricing for extended pricing, or legacy per-token pricing)
+        model_pricing = kwargs.pop("model_pricing", None)
         custom_pricing = {}
         if kwargs.get("input_cost_per_token", None) and kwargs.get(
             "output_cost_per_token", None
@@ -820,6 +830,13 @@ async def acompletion(
             if drop_invalid_params == True:
                 litellm.drop_params = True
             kwargs.pop("moderate_content", None)
+            # Register custom model pricing in litellm's registry for extended pricing support
+            if model_pricing:
+                # Merge with existing pricing to preserve other fields (max_tokens, mode, etc.)
+                if model_name in litellm.model_cost:
+                    litellm.model_cost[model_name].update(model_pricing)
+                else:
+                    litellm.model_cost[model_name] = model_pricing
             provider_start_time = time.time_ns()
             c_response = await acompletion(**call_params, **kwargs)
             provider_end_time = time.time_ns()
@@ -834,9 +851,10 @@ async def generate_chunks():
             else:
                 cost_calc_params = {
                     "completion_response": c_response,
-                    "call_type": "completion",
+                    "call_type": "acompletion",
                 }
-                if len(custom_pricing.keys()) > 0:
+                # Use model_pricing via registry lookup, or fall back to legacy custom_cost_per_token
+                if not model_pricing and len(custom_pricing.keys()) > 0:
                     cost_calc_params["custom_cost_per_token"] = custom_pricing
                 response = {
                     **c_response.model_dump(),

edenai_apis/llmengine/utils/calculate_cost.py

@@ -5,18 +5,6 @@
 
 def calculate_cost(
     completion_response: Union[ModelResponse, dict],
-    model: str,
-    call_type: Literal[
-        "completion",
-        "embedding",
-        "image_generation",
-        "moderation",
-        "acompletion",
-        "aembedding",
-        "aimage_generation",
-        "amoderation",
-        "arerank",
-    ] = "completion",
     input_cost_per_token: Optional[float] = None,
     output_cost_per_token: Optional[float] = None,
 ) -> float:
@@ -35,7 +23,7 @@ def calculate_cost(
     """
     cost_calc_params = {
         "completion_response": completion_response,
-        "call_type": call_type,
+        "call_type": "acompletion",  # For now, we only support completion cost calculation
     }
 
     if input_cost_per_token is not None and output_cost_per_token is not None:
@@ -44,4 +32,4 @@ def calculate_cost(
             "output_cost_per_token": output_cost_per_token,
         }
 
-    return completion_cost(**cost_calc_params, model=model)
+    return completion_cost(**cost_calc_params)