Use this to update the cache with new user spend. Put any alerting logic in here.
(
token: Optional[str],
user_id: Optional[str],
end_user_id: Optional[str],
team_id: Optional[str],
response_cost: Optional[float],
parent_otel_span: Optional[Span], # type: ignore
tags: Optional[List[str]] = None,
)
| 2618 | |
| 2619 | |
| 2620 | async def update_cache( |
| 2621 | token: Optional[str], |
| 2622 | user_id: Optional[str], |
| 2623 | end_user_id: Optional[str], |
| 2624 | team_id: Optional[str], |
| 2625 | response_cost: Optional[float], |
| 2626 | parent_otel_span: Optional[Span], # type: ignore |
| 2627 | tags: Optional[List[str]] = None, |
| 2628 | ): |
| 2629 | """ |
| 2630 | Use this to update the cache with new user spend. |
| 2631 | |
| 2632 | Put any alerting logic in here. |
| 2633 | """ |
| 2634 | |
| 2635 | values_to_update_in_cache: List[Tuple[Any, Any]] = [] |
| 2636 | |
| 2637 | ### UPDATE KEY SPEND ### |
| 2638 | async def _update_key_cache(token: str, response_cost: float): |
| 2639 | # Fetch the existing cost for the given token |
| 2640 | if isinstance(token, str) and token.startswith("sk-"): |
| 2641 | hashed_token = hash_token(token=token) |
| 2642 | else: |
| 2643 | hashed_token = token |
| 2644 | verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token) |
| 2645 | existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) |
| 2646 | verbose_proxy_logger.debug(f"_update_key_cache: existing_spend_obj={existing_spend_obj}") |
| 2647 | if existing_spend_obj is None: |
| 2648 | return |
| 2649 | |
| 2650 | existing_spend = existing_spend_obj.spend or 0.0 |
| 2651 | # Calculate the new cost by adding the existing cost and response_cost |
| 2652 | new_spend = existing_spend + response_cost |
| 2653 | |
| 2654 | ## CHECK IF USER PROJECTED SPEND > SOFT LIMIT |
| 2655 | if ( |
| 2656 | existing_spend_obj.soft_budget_cooldown is False |
| 2657 | and existing_spend_obj.soft_budget is not None |
| 2658 | and ( |
| 2659 | _is_projected_spend_over_limit( |
| 2660 | current_spend=new_spend, |
| 2661 | soft_budget_limit=existing_spend_obj.soft_budget, |
| 2662 | ) |
| 2663 | is True |
| 2664 | ) |
| 2665 | ): |
| 2666 | projected_spend, projected_exceeded_date = _get_projected_spend_over_limit( |
| 2667 | current_spend=new_spend, |
| 2668 | soft_budget_limit=existing_spend_obj.soft_budget, |
| 2669 | ) # type: ignore |
| 2670 | soft_limit = existing_spend_obj.soft_budget |
| 2671 | call_info = CallInfo( |
| 2672 | token=existing_spend_obj.token or "", |
| 2673 | spend=new_spend, |
| 2674 | key_alias=existing_spend_obj.key_alias, |
| 2675 | max_budget=soft_limit, |
| 2676 | user_id=existing_spend_obj.user_id, |
| 2677 | projected_spend=projected_spend, |
no test coverage detected