[REFACTOR] Replace CostCalculationHelper with litellm.cost_per_token

platform-service/src/unstract/platform_service/controller/platform.py

@@ -15,7 +15,6 @@
 from unstract.platform_service.helper.adapter_instance import (
     AdapterInstanceRequestHelper,
 )
-from unstract.platform_service.helper.cost_calculation import CostCalculationHelper
 from unstract.platform_service.helper.prompt_studio import PromptStudioRequestHelper
 
 platform_bp = Blueprint("platform", __name__)
@@ -213,23 +212,11 @@ def usage() -> Any:
     usage_type = payload.get("usage_type", "")
     llm_usage_reason = payload.get("llm_usage_reason", "")
     model_name = payload.get("model_name", "")
-    provider = payload.get("provider", "")
     embedding_tokens = payload.get("embedding_tokens", 0)
     prompt_tokens = payload.get("prompt_tokens", 0)
     completion_tokens = payload.get("completion_tokens", 0)
     total_tokens = payload.get("total_tokens", 0)
-    input_tokens = prompt_tokens
-    if usage_type == "embedding":
-        input_tokens = embedding_tokens
-    cost_in_dollars = 0.0
-    if provider:
-        cost_calculation_helper = CostCalculationHelper()
-        cost_in_dollars = cost_calculation_helper.calculate_cost(
-            model_name=model_name,
-            provider=provider,
-            input_tokens=input_tokens,
-            output_tokens=completion_tokens,
-        )
+    cost_in_dollars = payload.get("cost_in_dollars", 0.0)
     usage_id = uuid.uuid4()
     current_time = datetime.now()
     query = f"""

platform-service/src/unstract/platform_service/env.py

@@ -18,11 +18,6 @@ class Env:
     PG_BE_PASSWORD = os.environ.get("PG_BE_PASSWORD")
     PG_BE_DATABASE = os.environ.get("PG_BE_DATABASE")
     ENCRYPTION_KEY = EnvManager.get_required_setting("ENCRYPTION_KEY")
-    MODEL_PRICES_URL = EnvManager.get_required_setting("MODEL_PRICES_URL")
-    MODEL_PRICES_TTL_IN_DAYS = int(
-        EnvManager.get_required_setting("MODEL_PRICES_TTL_IN_DAYS")
-    )
-    MODEL_PRICES_FILE_PATH = EnvManager.get_required_setting("MODEL_PRICES_FILE_PATH")
     APPLICATION_NAME = EnvManager.get_required_setting(
         "APPLICATION_NAME", "unstract-platform-service"
     )

platform-service/src/unstract/platform_service/utils.py

@@ -38,23 +38,3 @@ def raise_for_missing_envs(cls) -> None:
                 cls.missing_settings
             )
             raise ValueError(ERROR_MESSAGE)
-
-
-def format_float_positional(value: float, precision: int = 10) -> str:
-    """Format floats to a string.
-
-    Formats a floating-point number to a string with the specified precision,
-    removing trailing zeros and the decimal point if not needed.
-
-    Args:
-        value (float): The floating-point number to format.
-        precision (int, optional): The number of decimal places to
-            include in the formatted output. Defaults to 10.
-
-    Returns:
-        str: The formatted string representation of the float,
-            with unnecessary trailing zeros and the decimal point
-            removed if the float is an integer.
-    """
-    formatted: str = f"{value:.{precision}f}"
-    return formatted.rstrip("0").rstrip(".") if "." in formatted else formatted

unstract/sdk1/src/unstract/sdk1/audit.py

@@ -1,12 +1,16 @@
+import logging
 from typing import Any
 
 import requests
+from litellm import cost_per_token
 from llama_index.core.callbacks import CBEventType, TokenCountingHandler
 from unstract.sdk1.constants import LogLevel, ToolEnv
 from unstract.sdk1.platform import PlatformHelper
 from unstract.sdk1.tool.stream import StreamMixin
 from unstract.sdk1.utils.common import TokenCounterCompat
 
+logger = logging.getLogger(__name__)
+
 
 class Audit(StreamMixin):
     """The 'Audit' class is responsible for pushing usage data to the platform service.
@@ -79,8 +83,31 @@ def push_usage_data(
         if event_type == "llm":
             llm_usage_reason = kwargs.get("llm_usage_reason", "")
 
-        if model_name is not None and model_name != "":
-            model_name = model_name.split("/", 1)[-1]
+        prompt_tokens = token_counter.prompt_llm_token_count
+        completion_tokens = token_counter.completion_llm_token_count
+        input_tokens = prompt_tokens
+        if event_type == "embedding":
+            input_tokens = token_counter.total_embedding_token_count
+            completion_tokens = 0
+
+        # Compute cost using the full model name (e.g. "azure/gpt-4o")
+        # before stripping the provider prefix for DB storage.
+        cost_in_dollars = 0.0
+        if model_name:
+            try:
+                prompt_cost, completion_cost = cost_per_token(
+                    model=model_name,
+                    prompt_tokens=input_tokens,
+                    completion_tokens=completion_tokens,
+                )
+                cost_in_dollars = prompt_cost + completion_cost
+            except Exception:
+                logger.debug(
+                    "Cost lookup failed for model %s, defaulting to 0", model_name
+                )
+
+        # Strip provider prefix for DB storage (e.g. "azure/gpt-4o" -> "gpt-4o")
+        display_model_name = model_name.split("/", 1)[-1] if model_name else ""
 
         data = {
             "workflow_id": workflow_id,
@@ -89,12 +116,13 @@ def push_usage_data(
             "run_id": run_id,
             "usage_type": event_type,
             "llm_usage_reason": llm_usage_reason,
-            "model_name": model_name,
+            "model_name": display_model_name,
             "provider": provider,
             "embedding_tokens": token_counter.total_embedding_token_count,
-            "prompt_tokens": token_counter.prompt_llm_token_count,
-            "completion_tokens": token_counter.completion_llm_token_count,
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
             "total_tokens": token_counter.total_llm_token_count,
+            "cost_in_dollars": cost_in_dollars,
         }
 
         url = f"{base_url}/usage"