uPj |dZddlZddlmZddlmZddlmZmZm Z m Z m Z m Z ddl mZmZedGd d Zd ed ed efdZdededed efdZddedej(dzd efdZdedeed efdZe e fdededeededed ef dZe e dfdeeefdeeefdededej(dzd ef dZy) uu Thompson sampling and prior initialization for the adaptive router bandit. Each (router, request_type, model) cell is a Beta(alpha, beta) posterior. - alpha = pseudo-successes - beta = pseudo-failures - mean = alpha / (alpha + beta) - total samples = alpha + beta - COLD_START_MASS (informative prior, not data) Hot path: thompson_sample() — pure function, no I/O. N) dataclass)Final)BASE_TIER_WEIGHTCOLD_START_MASSDEFAULT_COST_WEIGHTDEFAULT_QUALITY_WEIGHT SAMPLE_CAPSTRENGTH_BONUS)AdaptiveRouterPreferences RequestTypeT)frozencReZdZUdZeed<eed<edefdZedefdZ y) BanditCellz@Posterior state for a single (router, request_type, model) cell.alphabetareturnc`|j|jz}|dkDr|j|z SdS)Nr?rr)selftotals s/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/router_strategy/adaptive_router/bandit.pymeanzBanditCell.mean#s.zzDII-%*QYtzzE!7C7chtdt|j|jztz S)Nr)maxintrrr)rs r total_sampleszBanditCell.total_samples(s&1c$**tyy0?BCDDrN) __name__ __module__ __qualname____doc__float__annotations__propertyrrrrrrrsFJ L K 8e88EsEErrprefs request_typerc.|jtvr*tt}td|jd|t|j}||jvrt nd}t d||z}|tz}d|z tz}t||S)a" Cold-start prior for a (model, request_type) cell. mean = base_tier_weight[tier] + (STRENGTH_BONUS if request_type in strengths else 0) capped at 0.95 to avoid an over-confident prior. Total mass = COLD_START_MASS so that ~10 real observations can move it noticeably. z quality_tier=z# is not supported; valid tiers are ggffffff??r) quality_tierrsorted ValueError strengthsr minrr)r'r(validbasebonusrrrs r initial_cellr3-s !11./=););(<<_`e_fghh"5#5#56D%1U__%D>#EdD5L)D/)E:0D E --rcell delta_alpha delta_betact|j|z}|j|z}||ztkDr|St||S)u Apply a learning update to a cell, enforcing the sample cap. SAMPLE_CAP is a HARD cap on (alpha + beta). When the cap would be exceeded, we drop the update. (D5: hard cap, no rescaling — keep v0 simple.) r)rrr r)r4r5r6 new_alphanew_betas r apply_deltar:@s?zzK/Iii*,H8j( IH 55rrngcb||nt}|j|j|jS)zKDraw a sample from Beta(alpha, beta). Returns a quality estimate in [0, 1].)random betavariaterr)r4r;rs rthompson_sampler@Ns&os6A ==TYY //r model_cost all_costsc\|syt|t|}}||k(ryd||z ||z z z S)z Map a raw $/1k-token cost into [0, 1] where 0 = most expensive, 1 = cheapest. Returns 0.5 when there's no spread. rr*)r/r)rArBlohis rnormalized_costrFTs=  ^S^B Rx :?rBw/ 00rquality_samplequality_weight cost_weightc0t||}||z||zzS)z Multi-objective score. V0 is a weighted linear sum of (quality, normalized_cost). Higher is better. Both inputs are in [0, 1]. )rF)rGrArBrHrI cost_scores rscorerLas%( I>J N *[:-E EErcells model_costsc|s tdt|j}d}td}|j D].\}} t | |} t | |||||} | |kDs+| }|}0|J|S)z Sample once per model, score each, return the model with highest score. cells: {model_name: BanditCell} model_costs: {model_name: $/1k tokens} zpick_best called with no modelsNz-inf)r;)r-listvaluesr#itemsr@rL) rMrNrHrIr;rB best_model best_scoremodelr4qss r pick_bestrXps :;;K..01I!JvJ{{} t Dc * !['NK P z>JJ   !! ! r)N)r"r= dataclassesrtypingr.litellm.router_strategy.adaptive_router.configrrrrr r litellm.types.routerr r rr3r#r:Randomr@rPrFrLdictstrrXr&rrr`sw !H $ E E E .1..Q[.& 6j 6u 6% 6J 60*06==4+?050 1 1$u+ 1% 1"3, F F FE{ F F  F  F$3, $  Z c5j!      r