litellm/litellm/proxy/spend_tracking/spend_tracking_utils.py

import os
import re
import secrets
from datetime import datetime, timezone
from datetime import datetime as dt
from typing import Any, Final, Literal, cast

from pydantic import BaseModel

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
    LITELLM_PROXY_MASTER_KEY_ALIAS,
    LITELLM_TRUNCATED_PAYLOAD_FIELD,
    LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
    REDACTED_BY_LITELM_STRING,
)
from litellm.constants import (
    MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
)
from litellm.litellm_core_utils.core_helpers import (
    get_litellm_metadata_from_kwargs,
    reconstruct_model_name,
)
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.types.utils import (
    CallTypes,
    CostBreakdown,
    StandardLoggingGuardrailInformation,
    StandardLoggingMCPToolCall,
    StandardLoggingModelInformation,
    StandardLoggingPayload,
    StandardLoggingPayloadErrorInformation,
    StandardLoggingVectorStoreRequest,
    VectorStoreSearchResponse,
)
from litellm.utils import get_end_user_id_for_cost_tracking


def _get_max_string_length_prompt_in_db() -> int:
    """
    Resolve prompt truncation threshold at runtime so values loaded later via
    proxy config environment_variables are honored.
    """
    max_length_str: Final = os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB")
    if max_length_str is None:
        return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB
    try:
        return int(max_length_str)
    except (TypeError, ValueError):
        return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB


def _is_master_key(api_key: str | None, _master_key: str | None) -> bool:
    """
    Raw-only constant-time master-key comparison. The hashed form is never
    considered equivalent — only the raw master-key string matches.
    """
    if _master_key is None or api_key is None:
        return False
    return secrets.compare_digest(api_key, _master_key)


_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}")


def _is_non_secret_key_value(value: str) -> bool:
    return (
        value == LITELLM_PROXY_MASTER_KEY_ALIAS
        or is_valid_sha256_hash(value)
        or _HASHED_JWT_RE.fullmatch(value) is not None
    )


def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) -> str | None:
    if not isinstance(value, str) or not value:
        return None
    stripped: Final = re.sub(r"(?i)^bearer ", "", value)
    if not stripped:
        return None
    if already_redacted and _is_non_secret_key_value(stripped):
        return stripped
    return hash_token(stripped)


def _get_spend_logs_metadata(
    metadata: dict | None,
    applied_guardrails: list[str] | None = None,
    batch_models: list[str] | None = None,
    mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None,
    vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None,
    guardrail_information: list[StandardLoggingGuardrailInformation] | None = None,
    usage_object: dict | None = None,
    model_map_information: StandardLoggingModelInformation | None = None,
    cold_storage_object_key: str | None = None,
    litellm_overhead_time_ms: float | None = None,
    cost_breakdown: CostBreakdown | None = None,
    litellm_call_id: str | None = None,
    autorouter_savings: float | None = None,
) -> SpendLogsMetadata:
    if metadata is None:
        return SpendLogsMetadata(
            user_api_key=None,
            user_api_key_alias=None,
            user_api_key_team_id=None,
            user_api_key_project_id=None,
            user_api_key_project_alias=None,
            user_api_key_org_id=None,
            user_api_key_user_id=None,
            user_api_key_team_alias=None,
            spend_logs_metadata=None,
            requester_ip_address=None,
            additional_usage_values=None,
            applied_guardrails=None,
            status=None or "success",
            error_information=None,
            proxy_server_request=None,
            batch_models=None,
            mcp_tool_call_metadata=None,
            vector_store_request_metadata=None,
            model_map_information=None,
            usage_object=None,
            guardrail_information=None,
            internal_call_origin=None,
            eval_information=None,
            cold_storage_object_key=cold_storage_object_key,
            litellm_overhead_time_ms=None,
            attempted_retries=None,
            max_retries=None,
            cost_breakdown=None,
            compression_savings=None,
            autorouter_savings=autorouter_savings,
            litellm_call_id=litellm_call_id,
        )
    verbose_proxy_logger.debug(
        "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys()))
    )

    # Filter the metadata dictionary to include only the specified keys
    clean_metadata: Final = SpendLogsMetadata(**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__})
    _raw_key: Final = clean_metadata.get("user_api_key")
    _trusted_hash: Final = metadata.get("user_api_key_hash")
    _already_redacted: Final = (
        isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key
    )
    clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted)
    clean_metadata["applied_guardrails"] = applied_guardrails
    clean_metadata["batch_models"] = batch_models
    clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
    clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload(
        vector_store_request_metadata
    )
    clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(guardrail_information)
    clean_metadata["usage_object"] = usage_object
    clean_metadata["model_map_information"] = model_map_information
    clean_metadata["cold_storage_object_key"] = cold_storage_object_key
    clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
    clean_metadata["cost_breakdown"] = cost_breakdown
    clean_metadata["autorouter_savings"] = autorouter_savings
    clean_metadata["litellm_call_id"] = litellm_call_id

    return clean_metadata


BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost"


def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None:
    standard_logging_payload = kwargs.get("standard_logging_object")
    candidate_ids: Final = (
        response_obj.get("id"),
        standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None,
        kwargs.get("litellm_call_id"),
    )
    resolved_id: Final = next(
        (candidate for candidate in candidate_ids if isinstance(candidate, str) and candidate), None
    )
    if resolved_id is not None and call_type == CallTypes.aretrieve_batch.value:
        return f"{resolved_id}{BATCH_COST_REQUEST_ID_SUFFIX}"
    return resolved_id


def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> dict:
    """
    Extract usage information for OCR/AOCR calls.

    OCR responses use usage_info (with pages_processed) instead of token-based usage.

    Args:
        response_obj: The raw response object (can be dict, BaseModel, or other)
        response_obj_dict: Dictionary representation of the response object

    Returns:
        A dict with prompt_tokens=0, completion_tokens=0, total_tokens=0,
        and pages_processed from usage_info.
    """
    usage_info = None

    # Try to extract usage_info from dict
    if isinstance(response_obj_dict, dict) and "usage_info" in response_obj_dict:
        usage_info = response_obj_dict.get("usage_info")

    # Try to extract usage_info from object attributes if not found in dict
    if not usage_info and hasattr(response_obj, "usage_info"):
        usage_info = response_obj.usage_info
        if hasattr(usage_info, "model_dump"):
            usage_info = usage_info.model_dump()
        elif hasattr(usage_info, "__dict__"):
            usage_info = vars(usage_info)

    # For OCR, we track pages instead of tokens
    if usage_info is not None:
        # Handle dict or object with attributes
        if isinstance(usage_info, dict):
            result: Final = {
                "prompt_tokens": 0,  # OCR doesn't use traditional tokens
                "completion_tokens": 0,
                "total_tokens": 0,
            }
            # Add all fields from usage_info, including pages_processed
            for key, value in usage_info.items():
                result[key] = value
            # Ensure pages_processed exists
            if "pages_processed" not in result:
                result["pages_processed"] = 0
            return result
        else:
            return {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "pages_processed": 0,
            }
    else:
        return {}


def _sl_attribution_fallback(
    standard_logging_payload: StandardLoggingPayload | None,
    field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"],
) -> str:
    if standard_logging_payload is None:
        return ""
    return standard_logging_payload.get(field) or ""


def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload:
    if kwargs is None:
        kwargs = {}

    if response_obj is None:
        response_obj = {}
    elif not isinstance(response_obj, BaseModel) and not isinstance(response_obj, dict):
        response_obj = {"result": str(response_obj)}
    # standardize this function to be used across, s3, dynamoDB, langfuse logging
    litellm_params: Final = kwargs.get("litellm_params", {})
    metadata: Final = get_litellm_metadata_from_kwargs(kwargs)
    completion_start_time: Final = kwargs.get("completion_start_time", end_time)
    call_type: Final = kwargs.get("call_type")
    cache_hit: Final = kwargs.get("cache_hit", False)

    # Convert response_obj to dict first
    if isinstance(response_obj, dict):
        response_obj_dict = response_obj
    elif isinstance(response_obj, BaseModel):
        response_obj_dict = response_obj.model_dump()
    else:
        response_obj_dict = {}

    # Handle OCR responses which use usage_info instead of usage
    usage: dict = {}
    if call_type in ["ocr", "aocr"]:
        usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict)
    else:
        # Use response_obj_dict instead of response_obj to avoid calling .get() on Pydantic models
        _usage: Final = response_obj_dict.get("usage", None) or {}
        if isinstance(_usage, litellm.Usage):
            usage = dict(_usage)
        elif isinstance(_usage, dict):
            usage = _usage

    # A request that failed mid-stream has no usable response_obj usage, but the
    # streaming handler may have recovered the usage from the chunks already
    # delivered. Honor that override so the partial usage lands in spend tracking.
    _combined_usage: Final = kwargs.get("combined_usage_object")
    if not usage and isinstance(_combined_usage, litellm.Usage):
        usage = _combined_usage.model_dump()

    id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs)
    standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None))

    end_user_id = get_end_user_id_for_cost_tracking(litellm_params)

    api_key = metadata.get("user_api_key", "")

    standard_logging_prompt_tokens: int = 0
    standard_logging_completion_tokens: int = 0
    standard_logging_total_tokens: int = 0
    if standard_logging_payload is not None:
        standard_logging_prompt_tokens = standard_logging_payload.get("prompt_tokens", 0)
        standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0)
        standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0)
    _trusted_hash = metadata.get("user_api_key_hash")
    _key_already_redacted = (
        isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == api_key
    )
    api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted) or ""

    if (
        standard_logging_payload is not None
    ):  # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data
        api_key = (
            api_key
            or _redact_logged_api_key(
                standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True
            )
            or ""
        )
        end_user_id = end_user_id or standard_logging_payload["metadata"].get("user_api_key_end_user_id")
    request_tags = safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) else "[]"
    if (
        standard_logging_payload is not None and standard_logging_payload.get("request_tags") is not None
    ):  # use 'tags' from standard logging payload instead
        request_tags = safe_dumps(standard_logging_payload["request_tags"])

    _model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback(
        standard_logging_payload, "model_id"
    )
    _model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback(
        standard_logging_payload, "model_group"
    )
    _api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback(
        standard_logging_payload, "api_base"
    )

    # Extract overhead from hidden_params if available
    litellm_overhead_time_ms = None
    if standard_logging_payload is not None:
        hidden_params: Final = standard_logging_payload.get("hidden_params", {})
        litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms")

    # clean up litellm metadata
    clean_metadata = _get_spend_logs_metadata(
        metadata,
        applied_guardrails=(
            standard_logging_payload["metadata"].get("applied_guardrails", None)
            if standard_logging_payload is not None
            else None
        ),
        batch_models=(
            standard_logging_payload.get("hidden_params", {}).get("batch_models", None)
            if standard_logging_payload is not None
            else None
        ),
        mcp_tool_call_metadata=(
            standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None)
            if standard_logging_payload is not None
            else None
        ),
        vector_store_request_metadata=(
            standard_logging_payload["metadata"].get("vector_store_request_metadata", None)
            if standard_logging_payload is not None
            else None
        ),
        usage_object=(
            standard_logging_payload["metadata"].get("usage_object", None)
            if standard_logging_payload is not None
            else None
        ),
        model_map_information=(
            standard_logging_payload["model_map_information"] if standard_logging_payload is not None else None
        ),
        guardrail_information=(
            standard_logging_payload.get("guardrail_information", None)
            if standard_logging_payload is not None
            else (metadata.get("standard_logging_guardrail_information", None) if metadata is not None else None)
        ),
        cold_storage_object_key=(
            standard_logging_payload["metadata"].get("cold_storage_object_key", None)
            if standard_logging_payload is not None
            else None
        ),
        litellm_overhead_time_ms=litellm_overhead_time_ms,
        cost_breakdown=(
            standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None
        ),
        autorouter_savings=(
            standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None
        ),
        litellm_call_id=cast(
            str | None,
            kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
        ),
    )

    special_usage_fields: Final = ["completion_tokens", "prompt_tokens", "total_tokens"]
    additional_usage_values = {}
    for k, v in usage.items():
        if k not in special_usage_fields:
            if isinstance(v, BaseModel):
                v = v.model_dump()
            additional_usage_values.update({k: v})
    prompt_tokens_details = additional_usage_values.get("prompt_tokens_details")
    if not isinstance(prompt_tokens_details, dict):
        usage_object: Final = clean_metadata.get("usage_object")
        if isinstance(usage_object, dict):
            prompt_tokens_details = usage_object.get("prompt_tokens_details")
    if isinstance(prompt_tokens_details, dict):
        if "cache_read_input_tokens" not in additional_usage_values:
            cached_tokens: Final = prompt_tokens_details.get("cached_tokens")
            if isinstance(cached_tokens, int) and cached_tokens > 0:
                additional_usage_values["cache_read_input_tokens"] = cached_tokens
        if "cache_creation_input_tokens" not in additional_usage_values:
            cache_write_tokens: Final = prompt_tokens_details.get("cache_write_tokens") or prompt_tokens_details.get(
                "cache_creation_tokens"
            )
            if isinstance(cache_write_tokens, int) and cache_write_tokens > 0:
                additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens
    clean_metadata["additional_usage_values"] = additional_usage_values

    if litellm.cache is not None:
        cache_key = litellm.cache.get_cache_key(**kwargs)
    else:
        cache_key = "Cache OFF"
    if cache_hit is True:
        import time

        id = f"{id}_cache_hit{time.time()}"  # SpendLogs does not allow duplicate request_id

    mcp_namespaced_tool_name = None
    mcp_tool_call_metadata: Final[StandardLoggingMCPToolCall | None] = clean_metadata.get("mcp_tool_call_metadata")
    if mcp_tool_call_metadata is not None:
        mcp_namespaced_tool_name = mcp_tool_call_metadata.get("namespaced_tool_name", None)

    # Extract agent_id for A2A requests (set directly on model_call_details)
    agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id")
    custom_llm_provider: Final = (
        kwargs.get("custom_llm_provider")
        or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider")
        or None
    )
    raw_model: Final = cast(str, kwargs.get("model") or "")
    model_name: Final = (
        standard_logging_payload.get("model") if standard_logging_payload is not None else None
    ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})

    try:
        payload: Final[SpendLogsPayload] = SpendLogsPayload(
            request_id=str(id),
            call_type=call_type or "",
            api_key=str(api_key),
            cache_hit=str(cache_hit),
            startTime=_ensure_datetime_utc(start_time),
            endTime=_ensure_datetime_utc(end_time),
            completionStartTime=_ensure_datetime_utc(completion_start_time),
            model=model_name,
            user=metadata.get("user_api_key_user_id", "") or "",
            team_id=metadata.get("user_api_key_team_id", "") or "",
            organization_id=metadata.get("user_api_key_org_id") or "",
            metadata=safe_dumps(clean_metadata),
            cache_key=cache_key,
            spend=kwargs.get("response_cost", 0),
            total_tokens=usage.get("total_tokens", standard_logging_total_tokens),
            prompt_tokens=usage.get("prompt_tokens", standard_logging_prompt_tokens),
            completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens),
            request_tags=request_tags,
            end_user=end_user_id or "",
            api_base=_api_base,
            model_group=_model_group,
            model_id=_model_id,
            mcp_namespaced_tool_name=mcp_namespaced_tool_name,
            agent_id=agent_id,
            requester_ip_address=clean_metadata.get("requester_ip_address", None),
            custom_llm_provider=custom_llm_provider or "",
            messages=_get_messages_for_spend_logs_payload(
                standard_logging_payload=standard_logging_payload, metadata=metadata
            ),
            response=_get_response_for_spend_logs_payload(payload=standard_logging_payload, kwargs=kwargs),
            proxy_server_request=_get_proxy_server_request_for_spend_logs_payload(
                metadata=metadata, litellm_params=litellm_params, kwargs=kwargs
            ),
            session_id=_get_session_id_for_spend_log(
                kwargs=kwargs,
                standard_logging_payload=standard_logging_payload,
            ),
            request_duration_ms=_get_request_duration_ms(start_time, end_time),
            status=_get_status_for_spend_log(
                metadata=metadata,
            ),
        )

        verbose_proxy_logger.debug(
            "SpendTable: created payload - request_id: %s, model: %s, spend: %s",
            payload.get("request_id"),
            payload.get("model"),
            payload.get("spend"),
        )

        # Explicitly clear large intermediate objects to reduce memory pressure
        del response_obj_dict, usage, clean_metadata, additional_usage_values

        return payload
    except Exception as e:
        spend_log_error("Error creating spendlogs object - %s", str(e), exc=e)
        raise e


def _get_session_id_for_spend_log(
    kwargs: dict,
    standard_logging_payload: StandardLoggingPayload | None,
) -> str:
    """
    Get the session id for the spend log.

    This ensures each spend log is associated with a unique session id.

    """
    from litellm._uuid import uuid

    if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
        return str(standard_logging_payload.get("trace_id"))

    # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs
    if kwargs.get("litellm_trace_id") is not None:
        return str(kwargs.get("litellm_trace_id"))

    # Ensure we always have a session id, if none is provided
    return str(uuid.uuid4())


def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None:
    """Compute request duration in milliseconds from start and end times."""
    try:
        return int((end_time - start_time).total_seconds() * 1000)
    except Exception:
        return None


def _ensure_datetime_utc(timestamp: datetime) -> datetime:
    """Helper to ensure datetime is in UTC"""
    timestamp = timestamp.astimezone(timezone.utc)
    return timestamp


async def get_spend_by_team(
    start_date: dt,
    end_date: dt,
    team_id: str | None,
    prisma_client: PrismaClient,
):
    sql_query: Final = """
    WITH SpendByModelApiKey AS (
        SELECT
            date_trunc('day', sl."startTime") AS group_by_day,
            COALESCE(tt.team_alias, 'Unassigned Team') AS team_name,
            sl.model,
            sl.api_key,
            SUM(sl.spend) AS model_api_spend,
            SUM(sl.total_tokens) AS model_api_tokens
        FROM
            "LiteLLM_SpendLogs" sl
        LEFT JOIN
            "LiteLLM_TeamTable" tt
        ON
            sl.team_id = tt.team_id
        WHERE
            sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
            AND sl."startTime" <  (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
            AND ($3::text IS NULL OR sl.team_id = $3)
        GROUP BY
            date_trunc('day', sl."startTime"),
            tt.team_alias,
            sl.model,
            sl.api_key
    )
        SELECT
            group_by_day,
            jsonb_agg(jsonb_build_object(
                'team_name', team_name,
                'total_spend', total_spend,
                'metadata', metadata
            )) AS teams
        FROM (
            SELECT
                group_by_day,
                team_name,
                SUM(model_api_spend) AS total_spend,
                jsonb_agg(jsonb_build_object(
                    'model', model,
                    'api_key', api_key,
                    'spend', model_api_spend,
                    'total_tokens', model_api_tokens
                )) AS metadata
            FROM
                SpendByModelApiKey
            GROUP BY
                group_by_day,
                team_name
        ) AS aggregated
        GROUP BY
            group_by_day
        ORDER BY
            group_by_day;
    """

    db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id)
    if db_response is None:
        return []

    return db_response


async def get_spend_by_team_and_customer(
    start_date: dt,
    end_date: dt,
    team_id: str,
    customer_id: str,
    prisma_client: PrismaClient,
):
    sql_query: Final = """
    WITH SpendByModelApiKey AS (
        SELECT
            date_trunc('day', sl."startTime") AS group_by_day,
            COALESCE(tt.team_alias, 'Unassigned Team') AS team_name,
            sl.end_user AS customer,
            sl.model,
            sl.api_key,
            SUM(sl.spend) AS model_api_spend,
            SUM(sl.total_tokens) AS model_api_tokens
        FROM 
            "LiteLLM_SpendLogs" sl
        LEFT JOIN 
            "LiteLLM_TeamTable" tt 
        ON 
            sl.team_id = tt.team_id
        WHERE
            sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
            AND sl."startTime" <  (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
            AND sl.team_id = $3
            AND sl.end_user = $4
        GROUP BY
            date_trunc('day', sl."startTime"),
            tt.team_alias,
            sl.end_user,
            sl.model,
            sl.api_key
    )
        SELECT
            group_by_day,
            jsonb_agg(jsonb_build_object(
                'team_name', team_name,
                'customer', customer,
                'total_spend', total_spend,
                'metadata', metadata
            )) AS teams_customers
        FROM (
            SELECT
                group_by_day,
                team_name,
                customer,
                SUM(model_api_spend) AS total_spend,
                jsonb_agg(jsonb_build_object(
                    'model', model,
                    'api_key', api_key,
                    'spend', model_api_spend,
                    'total_tokens', model_api_tokens
                )) AS metadata
            FROM 
                SpendByModelApiKey
            GROUP BY
                group_by_day,
                team_name,
                customer
        ) AS aggregated
        GROUP BY
            group_by_day
        ORDER BY
            group_by_day;
    """

    db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id, customer_id)
    if db_response is None:
        return []

    return db_response


def _get_messages_for_spend_logs_payload(
    standard_logging_payload: StandardLoggingPayload | None,
    metadata: dict | None = None,
) -> str:
    if _should_store_prompts_and_responses_in_spend_logs():
        if standard_logging_payload is not None:
            call_type: Final = standard_logging_payload.get("call_type", "")
            if call_type == "_arealtime":
                messages: Final = standard_logging_payload.get("messages")
                if messages is not None:
                    try:
                        return safe_dumps(messages)
                    except Exception:
                        return "{}"
    return "{}"


_SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"})


def _sanitize_request_body_for_spend_logs_payload(
    request_body: dict,
    visited: set | None = None,
    max_string_length_prompt_in_db: int | None = None,
) -> dict:
    """
    Recursively sanitize request body to prevent logging large base64 strings or other large values.
    Truncates strings longer than MAX_STRING_LENGTH_PROMPT_IN_DB characters and handles nested dictionaries.

    Also strips keys listed in _SENSITIVE_REQUEST_BODY_KEYS (e.g. secret_fields
    which contains raw HTTP headers including Authorization tokens).
    """
    from litellm.constants import (
        LITELLM_TRUNCATED_PAYLOAD_FIELD,
        LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
    )

    if visited is None:
        visited = set()
    if max_string_length_prompt_in_db is None:
        max_string_length_prompt_in_db = _get_max_string_length_prompt_in_db()

    # Get the object's memory address to track visited objects
    obj_id: Final = id(request_body)
    if obj_id in visited:
        return {}
    visited.add(obj_id)

    def _sanitize_value(value: Any) -> Any:
        if isinstance(value, dict):
            return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db)
        elif isinstance(value, list):
            return [_sanitize_value(item) for item in value]
        elif isinstance(value, str):
            if len(value) > max_string_length_prompt_in_db:
                # Keep 35% from beginning and 65% from end (end is usually more important)
                # This split ensures we keep more context from the end of conversations
                start_ratio: Final = 0.35
                end_ratio: Final = 0.65

                # Calculate character distribution
                start_chars: Final = int(max_string_length_prompt_in_db * start_ratio)
                end_chars = int(max_string_length_prompt_in_db * end_ratio)

                # Ensure we don't exceed the total limit
                total_keep: Final = start_chars + end_chars
                if total_keep > max_string_length_prompt_in_db:
                    end_chars = max_string_length_prompt_in_db - start_chars

                # If the string length is less than what we want to keep, just truncate normally
                if len(value) <= max_string_length_prompt_in_db:
                    return value

                # Calculate how many characters are being skipped
                skipped_chars: Final = len(value) - total_keep

                # Build the truncated string: beginning + truncation marker + end
                truncated_value: Final = (
                    f"{value[:start_chars]}"
                    f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. "
                    f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..."
                    f"{value[-end_chars:]}"
                )
                return truncated_value
            return value
        return value

    return {k: _sanitize_value(v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS}


# Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by
# ``:``. Covers JSON bodies and Python dict-reprs in provider error strings.
# ``prompt`` is included for ``/v1/completions``-style payloads where the user
# input lives under a top-level ``prompt`` key rather than ``messages``.
_ERROR_MESSAGE_PROMPT_LEAK_KEYS: Final = ("input", "messages", "prompt")


# Assignment-style keys: Pydantic v2 validation errors render the offending
# value as ``input_value=<repr>`` inside ``[type=..., input_value=...,
# input_type=...]``. The same prompt body that would appear under an
# ``"input"`` JSON key is echoed here as a Python repr, so we redact it
# under the same store_prompts_in_spend_logs gate.
_ERROR_MESSAGE_ASSIGN_LEAK_KEYS: Final = ("input_value",)


_SENSITIVE_KEY_START_PATTERN: Final = re.compile(
    r"(?:"
    r"['\"](?:" + "|".join(_ERROR_MESSAGE_PROMPT_LEAK_KEYS) + r")['\"]\s*:\s*"
    r"|"
    r"\b(?:" + "|".join(_ERROR_MESSAGE_ASSIGN_LEAK_KEYS) + r")\s*=\s*"
    r")"
)


def _scan_quoted_string_end(text: str, start: int, quote: str) -> int:
    """
    Given ``text[start] == quote`` (``'`` or ``"``), return the index just
    past the matching close quote, honoring backslash escapes. Returns
    ``-1`` if unterminated.
    """
    n: Final = len(text)
    i = start + 1
    while i < n:
        c = text[i]
        if c == "\\":
            i += 2
            continue
        if c == quote:
            return i + 1
        i += 1
    return -1


def _scan_balanced_value_end(text: str, start: int) -> int:
    """
    Given ``text[start]`` is ``[``, ``{``, ``'`` or ``"``, return the index
    just past the matching close, accounting for nested brackets and
    quoted strings (with escape sequences). Returns ``-1`` if the
    structure is unterminated.

    Implemented iteratively (no self-recursion): the bracket scanner
    inlines a quote-skip helper rather than re-entering itself, since
    JSON-style values cannot contain another bracket *as a first char*
    inside a quoted string — only the quote-skip case can occur.
    """
    n: Final = len(text)
    if start >= n:
        return -1
    first: Final = text[start]
    if first in ("'", '"'):
        return _scan_quoted_string_end(text, start, first)
    if first == "[":
        close = "]"
    elif first == "{":
        close = "}"
    else:
        return -1
    depth = 0
    i = start
    while i < n:
        c = text[i]
        if c in ("'", '"'):
            end = _scan_quoted_string_end(text, i, c)
            if end == -1:
                return -1
            i = end
            continue
        if c == first:
            depth += 1
        elif c == close:
            depth -= 1
            if depth == 0:
                return i + 1
        i += 1
    return -1


def _redact_prompt_leaks_in_error_string(text: str) -> str:
    """
    Strip echoed request input from provider error strings.

    Provider validation errors (e.g. OpenAI ``RateLimitError`` carrying 178
    pydantic validation errors, each with its own ``'input': [...]`` field)
    embed the full request body in their message. When prompts must not be
    stored in spend logs, that echo is a back-door leak.

    Two leak shapes are handled:

    - Quoted-key form — ``"<key>": <value>`` where ``key`` is ``input``,
      ``messages`` or ``prompt`` (covers JSON bodies, Python dict-reprs,
      and ``/v1/completions`` payloads).
    - Assignment form — ``input_value=<value>`` from Pydantic v2 validation
      errors, which render the offending value as a Python repr inside
      ``[type=..., input_value=..., input_type=...]``.

    The value scan understands nested ``[]`` / ``{}`` and quoted strings,
    so multi-modal payloads (``'messages': [{'content': [{...}]}]``) and
    user text containing brackets (``"secret[123"``) are handled correctly.
    """
    if not text:
        return text
    redaction: Final = f'"{REDACTED_BY_LITELM_STRING}"'
    out: Final[list[str]] = []
    n: Final = len(text)
    pos = 0
    while pos < n:
        m = _SENSITIVE_KEY_START_PATTERN.search(text, pos)
        if not m:
            out.append(text[pos:])
            break
        out.append(text[pos : m.end()])
        v_start = m.end()
        if v_start >= n:
            break
        first = text[v_start]
        if first in ("[", "{", "'", '"'):
            v_end = _scan_balanced_value_end(text, v_start)
            if v_end == -1:
                # Unterminated value — redact through the rest of the string
                # so a malformed leak can't slip past.
                out.append(redaction)
                pos = n
                break
            out.append(redaction)
            pos = v_end
        else:
            # Unquoted scalar (number, null, bare identifier) — not a leak
            # carrier, leave intact and resume after the key match.
            pos = v_start
    return "".join(out)


def _sanitize_guardrail_information_for_spend_logs(
    guardrail_information: list[StandardLoggingGuardrailInformation] | None,
) -> list[StandardLoggingGuardrailInformation] | None:
    """
    When ``store_prompts_in_spend_logs`` is False, redact prompt-carrying fields
    (``guardrail_request``, ``guardrail_response``, ``match_details``,
    ``classification``) before they land in ``LiteLLM_SpendLogs.metadata``.

    Guardrail hooks may echo the LLM request payload back into
    ``guardrail_response``, and two first-party hooks
    (``block_code_execution``, ``litellm_content_filter``) inline user-prompt
    substrings into ``match_details`` / ``classification`` too, so the flag
    must cover all four fields. Every other typed field on the entry (name,
    provider, mode, status, timings, action, violation_categories, risk_score,
    masked_entity_count, ...) is preserved so guardrail dashboards keep
    working.

    ``guardrail_information`` is typed ``Optional[List[...]]`` but at least
    one writer (``xecguard``) assigns a bare dict, so normalize to a list
    here to match OTEL's defensive read pattern; otherwise iteration would
    yield the dict's keys and crash the whole spend-log write.
    """
    if guardrail_information is None or _should_store_prompts_and_responses_in_spend_logs():
        return guardrail_information
    entries: Final = [guardrail_information] if isinstance(guardrail_information, dict) else guardrail_information
    return [_redact_prompt_fields_in_guardrail_entry(entry) for entry in entries if isinstance(entry, dict)]


_PROMPT_CARRYING_GUARDRAIL_FIELDS: Final = (
    "guardrail_request",
    "guardrail_response",
    "match_details",
    "classification",
)

_NUMERIC_COMPRESSION_STAT_KEYS: Final = (
    "tokens_before",
    "tokens_after",
    "tokens_saved",
    "compression_ratio",
)


def _numeric_compression_stats_from_guardrail_response(
    guardrail_response: object,
) -> dict[str, int | float] | None:
    if not isinstance(guardrail_response, dict):
        return None
    stats: Final = {
        key: value
        for key, value in guardrail_response.items()
        if key in _NUMERIC_COMPRESSION_STAT_KEYS and isinstance(value, (int, float)) and not isinstance(value, bool)
    }
    return stats or None


def _redact_prompt_fields_in_guardrail_entry(
    entry: StandardLoggingGuardrailInformation,
) -> StandardLoggingGuardrailInformation:
    """
    Replace prompt-carrying fields with the redaction marker. Purely numeric
    compression stats inside ``guardrail_response`` (e.g. Headroom's
    ``tokens_saved``) cannot carry prompt content, so they are preserved as a
    stats-only dict; spend aggregation reads them via
    ``extract_compression_saved_tokens``.
    """
    preserved_stats: Final = _numeric_compression_stats_from_guardrail_response(entry.get("guardrail_response"))
    redacted: Final[StandardLoggingGuardrailInformation] = {
        **entry,
        **{key: REDACTED_BY_LITELM_STRING for key in _PROMPT_CARRYING_GUARDRAIL_FIELDS if key in entry},
    }
    if preserved_stats is None:
        return redacted
    return {**redacted, "guardrail_response": preserved_stats}


def _sanitize_error_information_for_spend_logs(
    error_information: StandardLoggingPayloadErrorInformation | None,
) -> StandardLoggingPayloadErrorInformation | None:
    """
    Sanitize ``error_information`` before it lands in ``LiteLLM_SpendLogs.metadata``.

    Provider errors are stored verbatim via ``str(original_exception)``; those
    strings can echo the full request body, producing multi-megabyte spend-log
    rows.

    - Always: cap ``error_message`` and ``traceback`` with the existing
      ``MAX_STRING_LENGTH_PROMPT_IN_DB`` DB-storage safeguard.
    - When ``store_prompts_in_spend_logs`` is False: additionally redact
      ``'input'`` / ``'messages'`` / ``'prompt'`` values *and* Pydantic v2
      ``input_value=...`` assignments inside both ``error_message`` and
      ``traceback`` so prompts cannot leak through either field.

    Scoped to the spend-log path — OTEL/Datadog/etc. callbacks still receive
    the untruncated error per ``LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE``.
    """
    if error_information is None:
        return None

    sanitized = cast(dict, {**error_information})

    if not _should_store_prompts_and_responses_in_spend_logs():
        for field in ("error_message", "traceback"):
            value = sanitized.get(field)
            if isinstance(value, str):
                sanitized[field] = _redact_prompt_leaks_in_error_string(value)

    sanitized = _sanitize_request_body_for_spend_logs_payload(sanitized)
    return cast(StandardLoggingPayloadErrorInformation, sanitized)


def _convert_to_json_serializable_dict(obj: Any, visited: set | None = None, max_depth: int = 20) -> Any:
    """
    Convert object to JSON-serializable dict, handling Pydantic models safely.

    This avoids pickle-based deepcopy which fails on Pydantic v2 models
    containing _thread.RLock objects.

    Args:
        obj: Object to convert (dict, list, Pydantic model, or primitive)
        visited: Set of object IDs to track circular references
        max_depth: Maximum recursion depth to prevent infinite recursion

    Returns:
        JSON-serializable version of the object
    """
    if max_depth <= 0:
        # Return a placeholder if max depth is exceeded
        return "<max_depth_exceeded>"

    if visited is None:
        visited = set()

    # Get the object's memory address to track visited objects
    obj_id: Final = id(obj)
    if obj_id in visited:
        # Circular reference detected, return placeholder
        return "<circular_reference>"

    # Only track mutable objects (dict, list, objects with __dict__)
    if isinstance(obj, (dict, list)) or hasattr(obj, "__dict__"):
        visited.add(obj_id)

    try:
        if isinstance(obj, BaseModel):
            # Use Pydantic's model_dump() instead of pickle
            result: Final = obj.model_dump()
            # Recursively process the dumped dict
            return _convert_to_json_serializable_dict(result, visited, max_depth - 1)
        elif isinstance(obj, dict):
            return {k: _convert_to_json_serializable_dict(v, visited, max_depth - 1) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [_convert_to_json_serializable_dict(item, visited, max_depth - 1) for item in obj]
        elif hasattr(obj, "__dict__"):
            # Handle objects with __dict__ attribute
            return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1)
        else:
            # Primitives (str, int, float, bool, None) pass through
            return obj
    finally:
        # Remove from visited set when done processing this object
        if obj_id in visited:
            visited.remove(obj_id)


def _get_proxy_server_request_for_spend_logs_payload(
    metadata: dict,
    litellm_params: dict,
    kwargs: dict | None = None,
) -> str:
    """
    Only store if _should_store_prompts_and_responses_in_spend_logs() is True

    If turn_off_message_logging is enabled, redact messages in the request body.
    """
    if _should_store_prompts_and_responses_in_spend_logs():
        _proxy_server_request: Final = cast(dict | None, litellm_params.get("proxy_server_request", {}))
        if _proxy_server_request is not None:
            _request_body = _proxy_server_request.get("body", {}) or {}

            if kwargs is not None:
                realtime_tools: Final = kwargs.get("realtime_tools")
                if realtime_tools:
                    _request_body = dict(_request_body)
                    _request_body["tools"] = realtime_tools

            # Apply message redaction if turn_off_message_logging is enabled
            if kwargs is not None:
                from litellm.litellm_core_utils.redact_messages import (
                    perform_redaction,
                    should_redact_message_logging,
                )

                # Build model_call_details dict to check redaction settings
                model_call_details: Final = {
                    "litellm_params": litellm_params,
                    "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"),
                }

                # If redaction is enabled, convert to serializable dict before redacting
                if should_redact_message_logging(model_call_details=model_call_details):
                    _request_body = _convert_to_json_serializable_dict(_request_body)
                    perform_redaction(model_call_details=_request_body, result=None)

            _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body)
            _request_body_json_str: Final = safe_dumps(_request_body)
            if LITELLM_TRUNCATED_PAYLOAD_FIELD in _request_body_json_str:
                verbose_proxy_logger.info(
                    "Spend Log: request body was truncated before storing in DB. %s",
                    LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
                )
            return _request_body_json_str
    return "{}"


def _get_vector_store_request_for_spend_logs_payload(
    vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None,
) -> list[StandardLoggingVectorStoreRequest] | None:
    """
    If user does not want to store prompts and responses, then remove the content from the vector store request metadata
    """
    if _should_store_prompts_and_responses_in_spend_logs():
        return vector_store_request_metadata

    # if user does not want to store prompts and responses, then remove the content from the vector store request metadata
    if vector_store_request_metadata is None:
        return None
    for vector_store_request in vector_store_request_metadata:
        vector_store_search_response: VectorStoreSearchResponse = (
            vector_store_request.get("vector_store_search_response") or VectorStoreSearchResponse()
        )
        response_data = vector_store_search_response.get("data", []) or []
        for response_item in response_data:
            for content_item in response_item.get("content", []) or []:
                if "text" in content_item:
                    content_item["text"] = REDACTED_BY_LITELM_STRING
    return vector_store_request_metadata


def _get_response_for_spend_logs_payload(
    payload: StandardLoggingPayload | None,
    kwargs: dict | None = None,
) -> str:
    if payload is None:
        return "{}"
    if _should_store_prompts_and_responses_in_spend_logs():
        response_obj: Any = payload.get("response")
        if response_obj is None:
            return "{}"

        if kwargs is not None:
            realtime_tool_calls: Final = kwargs.get("realtime_tool_calls")
            if realtime_tool_calls and isinstance(response_obj, dict):
                response_obj = dict(response_obj)
                response_obj["tool_calls"] = realtime_tool_calls

        # Apply message redaction if turn_off_message_logging is enabled
        if kwargs is not None:
            from litellm.litellm_core_utils.redact_messages import (
                perform_redaction,
                should_redact_message_logging,
            )

            litellm_params: Final = kwargs.get("litellm_params", {})
            model_call_details: Final = {
                "litellm_params": litellm_params,
                "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"),
            }

            # If redaction is enabled, convert to serializable dict before redacting
            if should_redact_message_logging(model_call_details=model_call_details):
                response_obj = _convert_to_json_serializable_dict(response_obj)
                response_obj = perform_redaction(model_call_details={}, result=response_obj)

        sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload({"response": response_obj})

        sanitized_response: Final = sanitized_wrapper.get("response", response_obj)

        if sanitized_response is None:
            return "{}"
        if isinstance(sanitized_response, str):
            result_str = strip_null_bytes(sanitized_response)
        else:
            result_str = safe_dumps(sanitized_response)
        if LITELLM_TRUNCATED_PAYLOAD_FIELD in result_str:
            verbose_proxy_logger.info(
                "Spend Log: response was truncated before storing in DB. %s",
                LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
            )
        return result_str
    return "{}"


def _should_store_prompts_and_responses_in_spend_logs() -> bool:
    from litellm.proxy.proxy_server import general_settings
    from litellm.secret_managers.main import get_secret_bool

    # Check general_settings (from DB or proxy_config.yaml)
    store_prompts_value: Final = general_settings.get("store_prompts_in_spend_logs")

    # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null
    if store_prompts_value is True:
        return True
    elif isinstance(store_prompts_value, str):
        # Case-insensitive string comparison
        if store_prompts_value.lower() == "true":
            return True

    # Also check environment variable
    return get_secret_bool("STORE_PROMPTS_IN_SPEND_LOGS") is True


def _get_status_for_spend_log(
    metadata: dict,
) -> Literal["success", "failure"]:
    """
    Get the status for the spend log.

    It's only a failure if metadata.get("status") is "failure"
    """
    _status: Final[str | None] = metadata.get("status", None)
    if _status == "failure":
        return "failure"
    return "success"