tPj"dZddlZddlmZddlmZddlmZmZmZddl m Z ddl Z ddl m Z m Z ddlmZdd lmZdd lmZdd lmZdd lmZmZdd lmZmZmZmZmZmZddl m!Z!m"Z"ddl#m$Z$ddl%m&Z&ddl'm(Z(erddl'm)Z)ddZ*GddeZ+y)zI Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting N)Callable)datetime) TYPE_CHECKINGFinalLiteral) HTTPException) ModelResponseRouter)verbose_proxy_logger) DualCache) CustomLogger)UserAPIKeyAuth)ProxyRateLimitErrormap_v3_rate_limit_type)RateLimitDescriptor"RateLimitDescriptorRateLimitObjectRateLimitResponse$_PROXY_MaxParallelRequestsHandler_v3claim_request_stash_for_dataget_or_create_request_stash)convert_priority_to_percent#resolve_llm_provider_for_rate_limit)InternalUsageCache)ModelGroupInfo)CallTypesLiteralPriorityReservationSettingscDtj}| ddlm}|S|S)z Get the priority reservation settings, guaranteed to be non-None. The settings are lazy-loaded in litellm.__init__ and always return an instance. This helper provides proper type narrowing for mypy. rr)litellmpriority_reservation_settingslitellm.types.utilsr)settingsrs p/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/proxy/hooks/dynamic_rate_limiter_v3.py_get_priority_settingsr$*s&;;HC*,, Oc eZdZdZ d"dedegefdzfdZdefdZ de fd Z d e de dzfd Z d"d e dzd edzdefdZdede dzfdZd edee effdZ d"de d e dzdee efd edzdeee ff dZde dedefdZde ded e dzdeefdZ d#de dede defdZde deded e dzdeddf dZdededededee zezdzf dZdedefd Z d!Z!y)$ _PROXY_DynamicRateLimitHandlerV3aM Saturation-aware priority-based rate limiter using v3 infrastructure. Key features: 1. Model capacity ALWAYS enforced at 100% (prevents over-allocation) 2. Priority usage tracked from first request (accurate accounting) 3. Priority limits only enforced when saturated >= threshold 4. Three-phase checking prevents partial counter increments 5. Reuses v3 limiter's Redis-based tracking (multi-instance safe) How it works: - Phase 1: Read-only check of ALL limits (no increments) - Phase 2: Decide enforcement based on saturation - Phase 3: Increment counters only if request allowed - When under-saturated: priorities can borrow unused capacity (generous) - When saturated: strict priority-based limits enforced (fair) - Uses v3 limiter's atomic Lua scripts for race-free increments Ninternal_usage_cache time_providerc^t||_t|j||_y)N) dual_cache)r))rr(r v3_limiter)selfr(r)s r#__init__z)_PROXY_DynamicRateLimitHandlerV3.__init__Ns' %7BV$W!>t?X?Xhuvr% llm_routerc||_yN)r/)r-r/s r#update_variablesz1_PROXY_DynamicRateLimitHandlerV3.update_variablesVs $r%returnc*tjS)zHGet the configurable TTL for local cache when reading saturation values.)r$saturation_check_cache_ttl)r-s r#_get_saturation_check_cache_ttlz@_PROXY_DynamicRateLimitHandlerV3._get_saturation_check_cache_ttlYs%'BBBr% counter_keyc|K|j}|jj|dd|d{S7w)a Get saturation value with configurable local cache TTL. Uses DualCache with configurable TTL for local cache storage. TTL is configurable via litellm.priority_reservation_settings.saturation_check_cache_ttl Args: counter_key: The cache key for the saturation counter Returns: Counter value as string, or None if not found NF)keylitellm_parent_otel_span local_onlyttl)r6r(async_get_cache)r-r7local_cache_ttls r# _get_saturation_value_from_cachezA_PROXY_DynamicRateLimitHandlerV3._get_saturation_value_from_cache]sJ "&!E!E!G..>>%) ?    s 3<:<priority model_infoc\tj}tj|tjvrt j d|S|\tjLt jddt jd|Stj|}t||}|S)zEGet the weight for a given priority from litellm.priority_reservationNz4Priority Reservation not set for the given priority.LITELLM_LICENSEzPREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this. Get a license: https://docs.litellm.ai/docs/proxy/enterprise.) r$default_priorityrpriority_reservationr debugosgetenverrorr)r-r@rAweightvalues r#_get_priority_weightz5_PROXY_DynamicRateLimitHandlerV3._get_priority_weightvs.0AA  ' ' /87C_C_3_ & &'] ^  !g&B&B&Nyy*D19$**V   ';;HE4UJG r%user_api_key_dictcd}|j|jjdd}||jjdd}|S)a Get priority from user_api_key_dict. Checks team metadata first (takes precedence), then falls back to key metadata. Args: user_api_key_dict: User authentication info Returns: Priority string if found, None otherwise Nr@) team_metadatagetmetadata)r-rMr@s r#$_get_priority_from_user_api_key_dictzE_PROXY_DynamicRateLimitHandlerV3._get_priority_from_user_api_key_dictsT $  * * 6(66:::tLH  (1155j$GHr%c^tjiSi}tjjD]\}}t||||<t |j }|dkDr?|jDcic] \}}|||z  }}}t jd|||S|Scc}}w)z Normalize priority weights if they sum to > 1.0 Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57} Converts absolute rpm/tpm values to percentages based on model capacity. g?z.Normalized over-allocated priorities: %s -> %s)rrEitemsrsumvaluesr rF)r-rAweightskv total_weight normalizeds r#_normalize_priority_weightsz<_PROXY_DynamicRateLimitHandlerV3._normalize_priority_weightss  ' ' /I,.00668 DDAq4Q CGAJ D"'.."23 # AH QAA $4!4 QJ Q & &'WY`bl m  !Rs;B)modelnormalized_weightsc|duxr&tjduxr|tjv}|r/|-|j||j||}|d|}||fSt j }|d}||fS)a! Get priority weight and pool key for a given priority. For explicit priorities: returns specific allocation and unique pool key For default priority: returns default allocation and shared pool key Args: model: Model name priority: Priority level (None for default) normalized_weights: Pre-computed normalized weights model_info: Model configuration (optional, for fallback conversion) Returns: tuple: (priority_weight, priority_key) N: :default_pool)rrErPrLr$rD)r-r]r@r^rAhas_explicit_prioritypriority_weight priority_keys r#_get_priority_allocationz9_PROXY_DynamicRateLimitHandlerV3._get_priority_allocations0 D  9,,D8 9G888  !X%9044Xt?X?XYacm?noO#WAhZ0L ,, 56GGO#WM2L ,,r%model_group_infoc PK d}|j|jdkDr|jjd|d}|j|d{}|St |}||jz }t ||}t jd|d |d |jd |d d |j|jdkDr|jjd|d}|j|d{}|St|}||jz } t || }t jd|d|d |jd | d d t jd|d|d |S7 7z#t$r!} t jd|| Yd} ~ yd} ~ wwxYww)aC Check current saturation by directly querying v3 limiter's cache keys. Reuses v3 limiter's Redis-based tracking (works across multiple instances). Reads counters WITHOUT incrementing them. Returns: float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over) gNrmodel_saturation_checkrequestsr9rKrate_limit_type)r7zModel z RPM: /z (.1%)tokensz TPM: z overall saturation: z$Error checking saturation for %s: %s) rpmr,create_rate_limit_keysr?intmaxr rFtpmfloat ExceptionrI) r-r]rfmax_saturationr7 counter_valuecurrent_requestsrpm_saturationcurrent_tokenstpm_saturationes r#_check_model_saturationz8_PROXY_DynamicRateLimitHandlerV3._check_model_saturations2  N ##/4D4H4H14L"ooDD0$.E '+&K&KXc&K&d d  ,.1-.@$,-?qAQAUAU@VVXYghkXllmn  ##/4D4H4H14L"ooDD0$,E '+&K&KXc&K&d d  ,,1-,@N,:=M=Q=Q,QN%(%HN(.. vn-=Q?O?S?S>TTVWefiVjjkl ! & &w6KN[^K_'` a! !A!e&!e  & &'MuVW X sUF&AE9E4B'E9<E7=A6E93F&4E97E99 F#FF&F##F&cg}tj|S|jj|}||S|j |}|j ||||\}}i} |j t|j |z} | | d<|jt|j|z} | | d<| r6|jj| d<|jtd|| |S)aU Create rate limit descriptors with normalized priority weights. Uses normalized weights to handle over-allocation scenarios. For explicit priorities: each priority gets its own pool (e.g., prod gets 75%) For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%)  model_group)r]r@r^rAtokens_per_unitrequests_per_unit window_sizepriority_modelr9rK rate_limit) rrEr/get_model_group_infor\rertrrrpr,rappendr) r-r]rMr@ descriptorsrfr^rcrdrate_limit_config reserved_tpm reserved_rpms r#"_create_priority_based_descriptorszC_PROXY_DynamicRateLimitHandlerV3._create_priority_based_descriptors!s$9;  ' ' / :>9]9]jo9]9p  # %)$D$DEU$V(,(E(E1' )F) %HJ    +"%&6&:&:_&L"ML3? / 0    +"%&6&:&:_&L"ML5A 1 2 /3/J/J m ,   #(&0 r%high_limit_multiplierctd||jr|j|znd|jr|j|znd|jjdS)aW Create a descriptor for tracking model-wide usage. Args: model: Model name model_group_info: Model configuration with RPM/TPM limits high_limit_multiplier: Multiplier for limits (use >1 for tracking-only) Returns: Rate limit descriptor for model-wide tracking rhN)rrrr)rrprtr,r)r-r]rfrs r#!_create_model_tracking_descriptorzB_PROXY_DynamicRateLimitHandlerV3._create_model_tracking_descriptorZs_"#(VfVjVj&6&:&:=R&RptTdThTh$4$8$8;P$Pnr#::  r% saturationc Kddl}tj}||k\}g} |j||d} | j | |j |||} | r| j | | g} | r|r| j | ddd} |jj| | Dcgc]}| c}|jd{}tjd|j|d t |d d k(rt|\}}|d D]n}|dd k7r |d}|dk(rtdd|d|d|dd|j |j ndd|j" |j"ndd|d it%|jj&t%|d|xsddt)|d|||dk(stjd|d |d!d"|tdd#|d$|d|dd|j |j ndd|j" |j"ndd|dd%|d!it%|jj&t%|d|xsd|d&d't)|d||t+d(|d Dd}tj,d)|td*|r|dnd+|rt%|dnd+d,t)|r|dndt%|jj&|xsdd-||.| rY|sW|jj/| |jd/0d{}t1|d |d |d z1t3_y|t3_ycc}w77Cw)2ad Check rate limits using THREE-PHASE approach to prevent partial increments. Phase 1: Read-only check of ALL limits (no increments) Phase 2: Decide which limits to enforce based on saturation Phase 3: Increment ALL counters atomically (model + priority) This prevents the bug where: - Model counter increments in stage 1 - Priority check fails in stage 2 - Request blocked but model counter already incremented Key behaviors: - All checks performed first (read-only) - Only increment counters if request will be allowed - Model capacity: Always enforced at 100% - Priority limits: Only enforced when saturated >= threshold - Both counters tracked from first request (accurate accounting) Args: model: Model name model_group_info: Model configuration user_api_key_dict: User authentication info priority: User's priority level saturation: Current saturation level Raises: HTTPException: If any limit is exceeded rN)r]rfr)r]rMr@)riro)r incrementsparent_otel_spanz#Atomic check+increment response: %s)indentdefault overall_code OVER_LIMITstatusescodedescriptor_keyrhrIzModel capacity reached for z . Priority: z, Rate limit type: rkz , Model TPM: znot configuredz , Model RPM: z , Remaining: limit_remainingr) retry-afterrkx-litellm-priority)detailheadersrkr] llm_providerrzEnforcing priority limits for z, saturation: rmz , priority: z+Priority-based rate limit exceeded. Model: z , Priority: z, Model saturation: z.2%)rrkrzx-litellm-saturationc32K|]}|ddk(s |yw)rrN).0ss r# zF_PROXY_DynamicRateLimitHandlerV3._check_rate_limits..sUq1V9 ;TUs ujDynamic rate limiter: OVER_LIMIT response with unknown descriptor_key(s) — refusing request. response=%szRate limit exceededunknown)rIrrk)rr)rrkrr]rF)rr read_only)rr)jsonr$saturation_thresholdrrrextendr,atomic_check_and_increment_by_nrr rFdumpslistrrrtrpstrrrnextrIshould_rate_limitrrrate_limit_response)r-r]rfrMr@rrrshould_enforce_prioritydescriptors_to_checkmodel_wide_descriptorpriority_descriptorsenforced_descriptorsper_request_increment_atomic_responseresolved_modelrstatusr offendingpriority_tracking_responses r#_check_rate_limitsz3_PROXY_DynamicRateLimitHandlerV3._check_rate_limitsusJ &<&>&S&S)37K)KBD(,'M'M-"#(N(  ##$9:'+&M&M/'N'  ' '(< =CXAX $; ' '(< =R (,'V'V,7KL!-L.??(W( "  "" 14::oVWae:3f  > *l :+Nu+U (NL)*5/ &>\1!'(8!9!%==-#'B5'J))1 30067H0I/JK*BRBVBVBb*:*>*>hx)yz*BRBVBVBb*:*>*>hx)yz**01B*C)D &F ,/t/J/J+K/26:K3L/M2:2Gi! )?vFW?X(Y,%1!$"%55(..8~jY\M]]ijrist.#(&&+W-))1 30067H0I/JK*BRBVBVBb*:*>*>hx)yz*BRBVBVBb*:*>*>hx)yz**01B*C)DE11;C0@&B ,/t/J/J+K/26:K3L/M2:2Gi7A#6F ! )?vFW?X(Y,%1'7/ n $UOJ7U I ! & &| &2FOy1A'BU^MVI6G,H(I\e !7W`yAR7Sfj k#&t'B'B#C*2*?i%)  2 (?6:oo6W6W0!2!C!C7X71 & AR,^<(47QR\7]]A ' ) = AP ' ) =[M" J1s9B#N% M9 .NM>C?NE8N=N>ANNcachedata call_typec Kd|vryt||d}|j|}|jj|}|t j d|y |j ||d{}tj} t j d|d|dd | dd ||j||||| d{y7Z7#t$rt$r } t jd | Yd} ~ yd} ~ wwxYww) u/ Saturation-aware pre-call hook for priority-based rate limiting. Flow: 1. Check current saturation level 2. THREE-PHASE rate limit check: - PHASE 1: Read-only check of ALL limits (no increments) - PHASE 2: Decide which limits to enforce based on saturation - PHASE 3: Increment ALL counters atomically if request allowed This three-phase approach ensures: - Model capacity is NEVER exceeded (always enforced at 100%) - Priority usage tracked from first request (accurate metrics) - Counters only increment when request will be allowed (prevents phantom usage) - When under-saturated: priorities can borrow unused capacity (generous) - When saturated: fair allocation based on normalized priority weights (strict) Example with 100 RPM model, 60% priority allocation, 80% threshold: - Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only) - Saturation >= 80%: Priority limited to 60 RPM (both limits enforced) Prevents bugs where: - Model counter increments but priority check fails → model over-capacity - Priority counter increments but not enforced → inaccurate metrics Args: user_api_key_dict: User authentication and metadata cache: Dual cache instance data: Request data containing model name call_type: Type of API call being made Returns: None if request is allowed, otherwise raises HTTPException r]NrMrz,No model group info for %s, allowing requestz[Dynamic Rate Limiter] Model=z , Saturation=rmz , Threshold=z , Priority=)r]rfrMr@rz3Error in dynamic rate limiter: %s, allowing request) rrRr/rr rFr~r$rrrrvrI) r-rMrrrr]r@rfrrr}s r#async_pre_call_hookz4_PROXY_DynamicRateLimitHandlerV3.async_pre_call_hook4sAR $ $T*G}CCVgCh:>9]9]jo9]9p  # & &'UW\ ] &*&B&B5JZ&[ [J*@*B*W*W & &/wmJsCST1#6k(M ))!1"3!% *  ;!\     & &'\^_ ` sUAD  C5C6AC C CD CCD&D<D DD cK |jj|||d{t|trc|j |}t |dij dixsi}|xsd|d<d|d <t|dsi|_||jd<|S7y#t$r"}tjd ||cYd}~Sd}~wwxYww) z Post-call hook to add rate limit headers to response. Leverages v3 limiter's post-call hook functionality. )rrMresponseNr_hidden_paramsadditional_headersrrv3zx-litellm-rate-limiter-versionz3Error in dynamic rate limiter v3 post-call hook: %s) r,async_post_call_success_hook isinstancer rRgetattrrPhasattrrrvr exception)r-rrMrr@rr}s r#rz=_PROXY_DynamicRateLimitHandlerV3.async_post_call_success_hooks  //>>->?   (M2"&"K"K^o"K"p-4H>NPR,S,W,WXlnp,q,wuw"lowhighmediumrREDACTEDzI[Dynamic Rate Limiter] Incremented tokens by %s for model=%s, priority=%sz/Error in dynamic rate limiter success event: %s)'litellm.litellm_core_utils.core_helpersr)litellm.proxy.common_utils.callback_utilsrlitellm.types.cachingrr!rr rFrPr,get_rate_limit_typerr rcompletion_tokens prompt_tokens total_tokensrqrrrrE,async_increment_tokens_with_ttl_preservationrvr)r-kwargs response_obj start_timeend_timerrrrr:rstandard_logging_metadatarr key_priorityrrk_usagermodel_token_keyrbrdpriority_token_keySAFE_PRIORITIESlogged_priorityr}s r#async_log_success_eventz8_PROXY_DynamicRateLimitHandlerV3.async_log_success_events   J-a a & &'Z [.OPV.W $.4ZZ8Q-R-XVX #/F/J/J:/V/\Z\ %"EV!LK1J0M0MNj0k0qoq &.H.L.LZ.XLL%)__%H%H%JO, 6 ' gt Dj7&(2'-'?'? (G3'-';'; (G3'-':': q QS &*__%K%K,! (&L&O & &/'$033 D(A00<A G$@$@@ " %)A"-a ~> "-m< (,(N(N$" ()O)  & &/*$033 #ooRR(;%=S *N9E9X^h$**_ # # a * *+\^_ ` ` as[I(AH<-I(.B'H<I(C9H<H:(H<8I(:H<< I%I I( I%%I(r1)r)"__name__ __module__ __qualname____doc__r rrr.r r2rrr6rr?rrurLrrRdictr\tuplerer~rrrrrrrvrrrrr%r#r'r':s_,8<w'w H -4w%6%CC   t 2 S4Z ^VZEZ fk nY\_cYc0ncSXjIY:-1 (-(-*(-!e, (- #T) (- ucz  (-T@@)@  @D77*7* 7 ! " 7z&'   )  #    6}P}P)}P* }P * }P  }P }P~U)UU U $ U S4 $ & UntP^Brar%r')r3r),rrGcollections.abcrrtypingrrrfastapirrr r litellm._loggingr litellm.caching.cachingr "litellm.integrations.custom_loggerr litellm.proxy._typesr1litellm.proxy.common_utils.proxy_rate_limit_errorrr/litellm.proxy.hooks.parallel_request_limiter_v3rrrrrr&litellm.proxy.hooks.rate_limiter_utilsrrlitellm.proxy.utilsrlitellm.types.routerrr!rrr$r'rr%r#r sj $00!)1-;/3/0?  d a|d ar%