From b1ec504421ec4bf6b56dc6f133e96eaa74599aab Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 9 Sep 2023 19:57:07 -0700 Subject: [PATCH] update budget manager --- litellm/__init__.py | 7 +- litellm/__pycache__/__init__.cpython-311.pyc | Bin 8070 -> 8083 bytes litellm/__pycache__/utils.cpython-311.pyc | Bin 90021 -> 90489 bytes litellm/budget_manager.py | 62 +++++++++++++----- litellm/tests/model_cost.json | 3 + ..._api_manager.py => test_budget_manager.py} | 23 ++++++- litellm/tests/user_cost.json | 9 +++ pyproject.toml | 2 +- 8 files changed, 79 insertions(+), 27 deletions(-) create mode 100644 litellm/tests/model_cost.json rename litellm/tests/{test_api_manager.py => test_budget_manager.py} (72%) create mode 100644 litellm/tests/user_cost.json diff --git a/litellm/__init__.py b/litellm/__init__.py index 920833c113..a166dcfcd5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -36,10 +36,6 @@ caching_with_models = False # if you want the caching key to be model + prompt cache: Optional[Cache] = None # cache object model_alias_map: Dict[str, str] = {} -####### BudgetManager ################### -from .budget_manager import BudgetManager -budget_manager = BudgetManager() - ############################################# def get_model_cost_map(): @@ -240,7 +236,7 @@ from .utils import ( get_max_tokens, register_prompt_template, validate_environment, - check_valid_key + check_valid_key, ) from .main import * # type: ignore from .integrations import * @@ -253,3 +249,4 @@ from .exceptions import ( ContextWindowExceededError, ) +from .budget_manager import BudgetManager \ No newline at end of file diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index f23105fe1b09a23d396ed6d1cdf7611985c50ccc..8878d126f808a7a9a6171d8cae23c10446075797 100644 GIT binary patch delta 1657 zcma)+O>Y`U7{_OLG4HnH7YN1%?AWm}*qGQYV1u!76XQhbC8?@TnRTsp3;qiWDs6wBNuVpo$bJQlA-bW7m7>YWe?%=l9GryUz@K zQ{k1+Z=sMM!4=QGrH{T3?e{bf@nA#|b*f&?M|>Ey5b-k?3q(bpyjYOBN-zQv0^(r- z<~($q@v{+Rx~sE3SDOh%bouFyMO0l`GC#F$*N&>8Vb;wNJ?K*p(`OhAE2 zTUOATAGVQ@MItsb+_TVS8yc}l)J8@vGG!x|EE2PkF^f!JAiW?Rx6sT56r>XtiQCAe zMG`g=wn)-OA{LoFPoV$M`pXuYJC7uZO5!0`pzZUpcMI^j3hiHnOy7U@eM)31h%B)o zI3I!Y=L1%V$vV?d5J@=@OFf6&Kg0&F3qfQVCZ^@Y%!P?FVxj}>BN~5#Xo8_HR`{?p zGP4S0R^OF5t&*ilSYVA}zU>V6T8Pg6ib)=P4!i(<1-uA8555V$0KNr&6}&_j=~Bx< z$}~ln!M7ovhIj=$1HJ>k0=^5r3SI@zg72|BD?Ek=U`45$Qnwz9lg&!bBG! zsk>75q*jHJ51{B86uo{{lx?ytR$;qrkCj-NS@UV`CqKDCueC7Qr|Wb>t!Xt+xRyyte8`m}5uMQO*?hMdW|0{0|<23)vm%s&n;;Y8QhW?K94KO_|B>S)A z`3-+&NP;Z?(*H5e@qs_GmII!`W{5pxZKF<^(NNTG7eh2D3z%Mo0#Obh>Eb9Je=*<94InXfZu?X!?u? zI%`m+-O`i%wap=?hmvFu?BcgF3qI;%a>rkfFn;odUA_6X| zG!4_OcA5?Kv~imvxzt~XBpi62Uysff&Efm}zLy4B_PgO4)DuNdEu^#8%S) delta 1686 zcmai!OKciP9L8sOd07m$V~m5rfE_zF4t_te!4FJglPHPQR;n~#e6=dT=jQyVss_tlCr5L;o|ZorFDfwDWzl|KIF4yR!>? z=gS@cGrwO&&>p?|mp*^uKWJe@5p`-kxYBTM>gi8t_+NT^xBBqB7Ph3 z2ts>EfMU|e++B%^q?q)xRyXoWq>=&ZBNwPj1|cA53+S~1hHS)dk&ul9EHeBiX(0qp z*k_@OZz4(hEfTho0gH^-$OVf;Y-G?P(Q^XpLxWogT4?kfosl7n#B3yFk+_WvTO?s4 z7cDZ@A)WO*w-A=7hm6CvUxIfx0eup-|1wNf-ufa&#LOxpQ_RCurrvYBMI$mNN=LP6 zJxP)nIS@;pz!g8pM-oD&Y3wl~GmaKwGhOHk1gB*1Y*%nx1Ut|HqKU_djxqFL?kw@d z3GvW7jm*Q6^If?MbR3pkga=xJz6^bZ;ucwj!5Z}S3Zj?3!(;<|0z3^q3BCz_8GH+T z3OoZo4W6Y*I#Y3wZF+^Kz;~cN3-w*_Iq*I3dGH+g0{A}oB6yx{u=M9}pKKG9|LhF3 zvhs9S>bBG!se4j$Qu9*xg_5ffwFFVuUPrMlmSMXr$MP)8wwZN3UH(=j2Xv)^$svu> zRjr^GNCAHGQM#sGe^2T|1}tC4+KwefSKOr zeS?)~wN^34_-{k;+h(6uZ8A;M=gOt>eNpcb7bzg7v(YrnzGAUdtJRyOX0={x6pMye zaPfa@`WmcR{FE6-^#(J1qTI_1p>bCLhWtTjkG}}5+>DC)sF)r&lKWQFI~!-j1rZx( zs(PI=tyr$Zg5v8pU9?(mnr^LLE@_=T+Y*zr?dWjafhTz;JT_sZMd^HNUnv<#d8n~% z<7)uoT{ZsU--IW}7f;du9Q@DlR4|q}Mg7MU$Bn1SQv~D0uPD);9{B~wJP6O6BB;JT eejkV(k3NY%jXw&s|Big&YLU!%{RSm zI`_-+&hX7_o9Sz4Yw*o#o8_C`Hk)~5doS_LX`3TVW|LcNpNq+uOp$~%k)z%g6Z19U zx{Uso(ckh8lWT>0wtJ?#Ma?(ojasSY%pqE8S>-Xct)`X+cMG*Onn&fV)N;{A^=pRL zhgWKOgMC`5-&(V&*`(zkFlnRQ6CSddOh2YSL4)@fs@Zi`k!PcCh}=AfsoNv5HH+6HYr)wHWFOM%*Fv4*vcT0K>09VTrY z?Hm40&;&``txbGL(6IDp{Ls8fYScasgA-2LmZeRmzT31Z?v`z4ZR!DQTgT8qNha-b z)N9k|Z#wo<(PkeAo7%Q(ZQ3QoW`{P1=Iqtx68f}8LccbT zFrdvR>|$mMXdtV$kOsO+Yoc1IUBUH>+(B(I4X}jT6SZby9wM$w4+vMcb}7;A)Gj03 zrESud(MY?s<@B^iTXDd2SzC|RLJwE7uve0>S6wXZwrjN2)V^0+LwK#nL~@#-O`u7% z5~HxTmTLEDqSfTstop3qOGs_%4u%3gwZ(Ry$cTPzQ$kdzI}#?Qj3Xh+@qh{H*9n_O zIjGo^xxyU`x;otTPN&E32|1l>)x5-&_DYadsXG(L3x|3!F|TMMRRpmS^W`M8RSZF^ z`di`>`(%Dpmn8j9_E;)9?6MX$4FHE~O0Hg9hvFpw=(vO}$$9{Dwo|DbKl=>!#@w1Z z1DUi>{^Dq_-Fa@2b0;U0YLg*%^17;8mpD)h@aD|$l(rV1NA$n`dIw9t& z_fuz-)KHrv$(Z&$aA02CtwPO68#|pBuox{%0L_4s~J^krZ=Sb;%xp&3Gg{l_Gh_!`&z9|=82DAKDx$sf}?dqXSry-Yls_{F8sw(UDLTt63)+IgukaMS7ZVd$8 zavjz6sNZDGr4^-Sw@(?d50_K3Yy&`|l2^wXvk{Mg{pv{e7-OW3>eiIP=-KRtMJ=e0 z3Ej5Skv07J-3p^|2nskH%>ZH<<%EpeBHV|;PLCp>oL+|DFaeAtB>+lRPhe;zHS4%4uA*H3D^$Up*H7Nl<;~W z*q+SgfsPK3zayriH>!K{?-P@wvqq%~yTMB>E2z0LfZll4t=%L<*@b660P>U=P|x-S zG`H8;?(%w_nmgq3c!Lg`jL+&SP%jcVqYD>Hfi z%uS+)UmYDYDStH;^>mNFt2@N&+3ITFAzRg7#>^<=Re+_Q?WOHqPhED;tEZ}on3ybf?Z-~a%UlHUQ`NTBC*1%igioZdi()4dZy;SFF->;U0S z1daqlF?cSq^%#>sWSU7-HxpGbfnZQNhH#Q%{Cl#7bT?M@w`gd>(Ij!LDwAYg*-mDp(4CSV}|M{GgcL1lXY9WtG&~0Xe=haWUbt;%4X)( z!nsLUwGjr!p@fLz5R%IajPhc&Y34TjK6>bBA<4FEbQBvBhdt_vnk(eS40MY>u#4rW zX9m05DPRQSItQT|MvNKgNpYbbL<=k9cI9sv?fM>`v0=KU!#yZ-_Krpj;NC7xBi2~d zh+X6@N#rfj?If*kZ^-Z(cY`aof|mYKLqX0VDvr=uHELE}LM6r4;024C&{Q1VHtUvz z@U`f3CjhE@s3*lx-kDQX>cUL*)PkJ&7Bo^-&SdV!tnLMu?*pL5dgeCrS>1!y$0o|+ zZ&s@oR?cRh`UB7olRO^div3T7s*W$Ls)WuxsE24D;jEN6&KaZg+QuJOMvE0D_B*QvKoRiKXWgvRN(67||?C zF~kn_?5Y*3kxuD0FS!e^5mlPeW(Ba^yHR2rG@Mw%F!eySVZ}R8;_c{E8&^-9fK8RX z0A3utlx}x*bx{zJ_o4d7sIhq6ep#xbs~f{CaSrsH0c-&cCj>0c5!(w1IZ_RGVs*)m zpl;u>3;#==2d{u9!z;in^^`%0upff(^o(}t3U-qB?h43|>RMAvxz+d9Ocl>XUt6=x zYDD12ROyD|3g*Y6e-#5;4V6s_xxH>5+4~;(Q+36L630u_D1Q#PxWn*pi`G->{tZ>) zzoVx%BqyY>HDOmMlhCi+Sd7W8J!aPJAzBYHXQn*%T(!bE!EV@%3On;#e~Mmi089II zD(N}!F|_f{U7k>J3&Qk5u>XhcEe2O{wHs86`_hU~x6Gw1l^Q2m-)lQn9T$Cn>-KtVeZ z7}KWN-s2f-`+$=&9|2Zc3v=fgNstkd*E99R;f( zL){v3?+)q7{%+q;<5%c(arbXV@-*h;WUl1+Efheyqcgpy%ntU&$%Dxxr4tf5rdxNo z_XLOH!a|JDq^biC7s1`>>Fr26y4=#3g3Kar{}TOUpjeD$pZGozwc7}Weo(g1->2(| zGGa=KR$tX5rdOJT>2i~n+-llwx!kl%I8s`SIKPxg<57!EiyRoMkWjbC$}ks=S%DOL zKt2_VlGvZY_(Epx^|VvmA2JF^o+nF3w@*-JiQ2nlz}FQrjv@bz>Nw|QuuLHcCQ*a2 zjpgOx!Q3&UojQ6ctsw6Uk;qX*`JkDov4{Ii=h91UIqwz}9!wi38QqsyIFM3$f7NeO z#`dL*eYa%diG===hNsK=OD;XQVxWBLiRS+DIel|i^_Q>iOC3FsTk}Xlf9|A%iw`dT zpu};)_VkjzlBUxoP5mWJ2g57Ql-HeD{908@cWvm;tLsa}??CFP!_)gx%jx%noT+`&ntxezI%jQv&f31z zwP%)HK9D`?{=I!mFF$87FBi$5n1G)#o=#laH#@PquVyI?m|531anlWGk^-32yeUcXW|buT32 zeXYAVL3F4Ouc<4oqG~Tv~Su9~*DKD!zuwostUxkFZ9q#PC6 z>ld@sdwbi*!(vDx-Trn%nds2M*`}VRkrD^?_$aY=dNfk<)#|EitHp;Za&7(kxnR=> zs0BCxGXWh~-7EmzbkZFPUXR zh;k9pd=y<5e$Qs4InCrbznsC*C=uPSqfsd{2UEUL88nr2@QTYvE+N_SbnVGqE51xyJlx?Zr2!u8FrwyFNq)j9Nv_63c7gpN6CLSBoWIchg(3J81-0M{ zlw6Rnx2_Sk%z=XHbEcIdX_5GZ|BHV1{pHq~r*p@i$tpZtdpfJKKdbVbNu(_jXUfJu zUE5dIbh@mmzpUxD)C=dFX6`uH|NQ4$RO^um;X2d()@t()#mYL<(FLXjwYHPhW-6bo z5h&L#ETo5*&Gtn}=9d#e^m5Vy^Wt>dsT5G1N)srjFUY2cSH{?v)S6$Z0o5zD3(U=R zwqJ|`)h{LplEZRdytkMBRG0TG9^%Z zy4=*tcO?&C?mQOyFdd!9yZz*AN6xDU1XlyDAvpGp%dfjkIn)+k49Asq%Kk+Ca(EnB zhF)ePM4`e_!Bx7YE6`=EoRu4gBqIe&a+51`vr4j_--@}ORr{XE4}TZMn?RX_66{A$ zXm@pm5Cr4)z&`Cb$nFLiAE)m^=?LoXK&=3^-E0sUKUF5VXm3_%?$EQz6lzqbeqN`tFXXAWn`>14%l1*1X{Nnd*JkY! zy9SflYZ7EJu|9pgSZ#Z`yzKwOFa1<`5lb#^LoB&L8(ey7vg2Ot9lK4|V8bM7g7W)l zxDW8JfUtVtRHgdcDVx}*es!wqQuufIFoABtvD`2Ti7k=G0C7`bzITCTe4~E^l69!t zuNq#dE#f^NLBq}0sh(FV>#qm-4*^F3aRXv~TZTHAlKe4BkE&n2QXj5h)i?AL1|yH7 zjSb^IJo1uHpl*cgH;e-$kE4|hgr{;lJ@QR}bMUs$(EAZrOA4Lt6k$R(5ZpEx7O@|t z2%0Z8hWKgOkv;6#%X%J{SX`dNOamN$0*+4sdI3)Zo)8r1jNE~D8;ttxs2X{bq3+wop%T#gi&J(8c`5@YgVUf~%eG0VociB$-@oqxBY0-+iVtDcSvW zezNwZOnk&nfM4rGIay_Gr<42ZyF}SIqH8g-L*tDQ-%7-otM~e5i~Xwf^rY})qBM@0 z|BKRjzy$(2@MY4oHlb$h*m@(ls$)g@9$NaIv2U#HgF)vf(M7CY7K-+dsiQsr+Jg)u7S z*zr-3hpt_M0|;0$^DLBDN^X>3k#zmSUHD|ijvx1KeEkwzGu^^L{1>L)p4iwnlBAVmsl7S{Vc!dX+Te1hJLcR^q2}qq* zt26HuRr2r`w~ZLSK&^Vaz%I{FucuVQ@23+(|LVR!?1+kOD=PY=CpsM;8jjSIYK9V+!J|A*p87 z@lT@<=3ZP{recPTK(nk zRbjRZ4k6eRJ<;hiK5USA)JIQSNS6+ahc4EQ@Hhp`eh=UQSwK8Bwt{>#Q2s|#1Hfa+ z1(*&*N>}Z4z6;7xsAYSZ7e(60hTgL#n5lwK%XanJR};su zWhD|p`QVT7E+P|TfV$}DudNy3e+RQZ@Ua;|Gd6^tg;%O@9`S2!9&pfzS#&%gw)^6p zmcIsDz+(y2b{Q4+nck%clUH!K@b)Y3ymFvy%DDtJE|@ zV+F6Ij(S3uYY+BMLh(lkAqSiR1HGF^iKj&33m7p$u8tHHis>owZ9NicD-;z*W?PR= zfX{mm6^ha3Fn1b!1;^rk88W#8z$*F(*20bmD#f2UNIfeaN@>;oI*j~+xH_O;Shz99i;*vl2F8tyi_XQwsRPIKXRl@6lJsXN(bj7G$q=?y2xu~)N4}Y zf6GL(7#*2cF1A`HQ(NR84v`;uu3Y><%#ZwFj2Kl0WgSpQ%sI}|csI1>VqVk6 zin;b#B=}%^#9t-ydv6&l#)%0$WBN95FcV@~uE*mBzz81i)zEwvL?+gXk0vgrSsK|B zFA#m`#t%6FJP-9}%p&QBRbN}GJRBHSaWQF^cs<+#(CVuTOt#Qt!~vk z?IHO*24Dlk(l`+6N%*=cH{_(Zt|KM%Ouc9l>w3Sc7e$HTD)hk4bNXg1G8fvwF)9b8 z8Gva3`g+kMZ2(>qd^La3C7tw4Uoo&Pv;EH}a^32uH$aA!FGVM|@>Nu#Fy=&?(-}5` z4{uK^XxxB}ph+-95|m9Cx()TeLTNKfTL7{6U>N05Ob20sLzkY`N>|vIQhM&fP4$q^ zIm~kge|kl_J5=nX8%mdPuD4MW>GV*-+ZmbEAjXKRde=0FK1&*#zeRRn;O&u#b3{Sv zMNXr2j;Iq4ME1`SE#W7KlOt>7r+n~bau)i}0W<i3x=e7&cWJU1 zJUISaIeuIR?yUeQmBX&5vdQ9qW7Ya?TbW_jlO9$J==TZ zBGGC~`w2v39^?`E@KSMiIGvQDLT^_x|4)}&g zf3#vG!^GIgafzne={(Zy4jN0)K)_~y1oY1lG#hsjA5u}b&?k=ZgAk<{yjDK%V01kr zKrgZ+9mkx5Z<_h0B+QGCEM6uSQ!qKSOpLMJKvn&{Pb?E(T19i@_$tvdk7u4Fe}-}S z0uM95XF@(OW&L4a!0zlio(+9yL`fo#@mXZ{YB9N@6TJ^&yk?ZXB~^^vvRahczDp!u z_dc;&WSPT%B}x*cm4sn9B5I?gvBjuiXU4#E=}TH<0Z}e#SwcVfn6!m_L7a$7U|_Hs zjPkK|hXoz>%bSRWZr1`Kw}19dBaHEZiFNBBX2LlXFZkDZWEpb4$NAg67~#>#)oVqy z_($Z(T2UZsBG0ZBug@x?Epm4lUMg5`ybj<`WN2ejsh6I25rLOp0&T;mj%kr2>qJ>o zI!ICj^zn+zQ_t7R65k)lBwCZ6J6Lc0htEsg6fZK-B{}vzPB4W1L8numL=#`Nm80}2dj17afyV~GJOBfp@dxPi{nbZ1A7yb6myPRT@H6v$VYVVFVk!jAY7CGlk7sHxf pwTY6f#q40#Ylo{igV9+H3sN+o$_4Yro7tqkRUC%kW+9 zpV>asKdXI~e|Gz9|D5(YOe51b*FUd)o-mnB9)ZR=jvpg%Pt=T*-yGhGN8`ZBGSs!lF#tw~Xp>bE5 zO^qfkXP-&S^-O)hVlq8Uzs65{E4PiKHpkc6R&(2UYRmgt+Zt`cK2!Tzl2E>HoxiQU z4J)EqCQ|Q7J)X8v(_F=klc}-b0W&c`zs8Syy_c4ziB5}FxKFgVYehtMgH}vNKO@61u@Y|gs#_CQ|(SotTs8X zR_Cm@s{^(}VoY?BT?tXB9#5E-;w0h8selUgUc%Mm%Bk3&vC|6z2i106@WI6@V$*sZ@ZU-3EDswr2l>DuQn2RQuj=`4! zW&oHW43n1wW)h6#mpK8fS8G#N8~o0W{xoH+5c8EaZB~96wKY*H_!9TGKa)0jVsGPN|lT1->H^ra%Go=C5n z!=eLSr^T3kIfxEbU{YOZhLB+QF$pbK6EjN0b!u_O!g`j7wV8gtE*D+~ll(@x@KOP8 zbt=P^#LVC^GB{r3s->B?j>mTEX}!|h6LM|!$W4KuN3Nl|?dpTf1+=Qltd7e@?Zs8p zEY|}dRmntIYufM#2&>~+g~l$lsh_5dkBaOEMLBvO6}oM!+t=qgYL+YLS+{ya9)A!n zr=Y-4j+JU~tW)e)UyLoL`6uPn^5)5Fsf*l5@C}_0sC#qJ`BYA|1M}i$Zp1ck- z)k_(WvO%Tht}NkQ=DqcRTqmF#&;{@UHmjYvMftoQ2)93@DbUsB?dgiC>37vrxp(p6 zI&6{MC1R|S<36$*MAa6@lxurH2G6>+kEAIBc=iJzSc#eSEPp`r_*@-spUsq7s7HpcyynwDVG2T^PPZh@MW9*)!}Znd#!{CZw<2;}+z z%>Pyl5gt7e@?m-qS}>7-r9i;Cq9^uB4^JB3rs0}|6^y&5A>ky&==Wz1D{-vq{pg{}`;){@HE&kA*rhsV)m;WH&{Kzn;|grr%w}OI zQpoG~7%IYRG&odm%xW2rP1bE4?v8Gcp|PD*(Xd8()xz01n-bPGtqJ;3{6q3t} zjPg=-{p?OVMA_d=l5M`qIl+)P3{y|k+#wG*&@DZIZ7fGUBiPqL(IXhwISAFTW^kY< z#fAD^w6HRI)eq)Obl-+&Y?yB8@(jrwZqbNnh&6hWThoX%7By-Yc}o&`OLRL)tIroQ z+{ceV0b4;!Po3+?zKM#XbXLutTa{2KsQz<{YFaWOdgI*hCxmxmi~>OQjtX`W5<(*{ zS%>~4C4Ot(XD2u;A^)4=+2Wypg0Y*q3k8#ER z2SQaZEiRc7XPu>aFGhZoFda*(#KvgDl2zu0>F5(r7s~M{0XG9~1>6UC6i^1J0ZapI z7>0f7_x1l`gslzHeGLP`{s@}3s=8&xzFSa)m!KfbQrV3PUPdcQdjT+Wa>DFa*bOiN zb`89M5q_L{>Y1MH9iHBhH_!tQfG8s$r}}@24lesFA&b=lTHrLxiS&m;lisrOO5|3$ z%}4ISXG9ftnO1aQxrb0<8#J6)!U*+1wPD4ZQR40BQroYbR*6lOKL_yQ;H7lCySJBu zh`bZkKSqtk=jo9X)oWMQhFRhq@i`T6p$}&UEY4Be3kf+>4R>O7%dVhq-?9yVCG){6 zU}ks);8IT+k_h`D2v1M%knUhNdGFqU45@>y<&8YnwOJ z?c%H;DA_;dqvH_(c0^Ai;zJJU9#ryP-KP>aR)2~*MK96h2O!b6rchDUQ zNh9>6X(Y)_0Xnp}`lOFtGlj6YJlF#7#WcKJ7SFFhHjQ1Ph7bzB2mhiMB>8u^Lfz8i z*4RwtPciTW;1uAbI@UMVwtM$(N?tzLSSl>hvF^i`S!L8lTXPa!{v8k-5wz0d#)3tQ zj396CF!F_SRQq#vLF=)?b~%NO`aOx}vHwAjdTOs*y8S_y-yQ1cmQSbyU6V*Z9_lK# zy+_Tf)L_@klp)^TePh=9^vqd4dC?xtqscNg+w0Gbd7Yr0Iv9Bdar_H!neAh$=u{WH zD=6ifzWJ~yW{H0eI$r?3B+yek0)3Q3hXPwXJwej2|JmYbf#7s&$g@49C-?OEha3Nj zL6;8ydiYLbN-m~K-rg)IYDC+7C(TZ1kDfe~K~fGOp;)?gi)Tl0I2hDnhQ+Er@L(Q% zoIa)lxkj%?8Z5{RqV}Dry?260fiUHHf}wQQwv^f;(Z=hRh*D>AGshw0C}a>tcdU!N z68=b&dPWw9xfts}J`(&&h<+$uklFgY9Tdlhl^jB4KImL_5oH#Zy;lbOy&>a#@gb^X zL}e04(v6HMm_!Z64iQ32!f+2X&_`(kc~h)b=2Fjl$WnuUrrx_v7qd(`xw}`Mb&Mah z=A2E)yG=WtP&AlObk0%nP~#cL<;PrS9QFIt&gM@$mT)G2&fwgpGx;kA6UU#;C_d77 zCZlq{ZNKe;qx7NrC+vfcg|9jmo^dSPpZ0dX^O)_4<>A5n##i$j&*V4Kh>LbpV%}|Q zi8lTJ#BR#0e#rA`#^q-+E(eW^iKe76KN@!?(>a)k|7R1k5B3cv=F|TR>C*;lmOR`3 zYI@6=^p?TImbV+%o*k2YWWnIlwHGbsWy1cU3HUMNtJbB1*22N!r8Hkgs&lZS>D82` zL2DC-Hg$Vw-DSwA^|a1_?C`jdJ_Y(h95VET;iDe>qn_yQ^}05>DIxZ%!air`?NrG@ zl6=B=$Q*QJ9+h;PyPtx-5qBn_BDI@RdibwScZWx2qbem7AYVm^lgsUum_X)(2z!!= zC`|$c$SXJ|v2td?wc&u^^Z8wd8OtK-^St`czQts6&aD#?M5n6XSv9Gcs!8?T+lNd^ zPo^+I7aN!SH+85KDm4gUkRaI-+5;iT!Tyi`8`ykBpc0qQ^ICKjsQgC@{{${_YU>qjaN>H?^F)e=}GWcIDZ9Y zUUc?%Ul5$=Q$jiwQ#b)-n8G4gsPPB#?bu1VQeAPNTs*3@16BF(DU~f2uA4ryrosk>1*z`2CC}!e zwilz%-qI+hVyPX-35I{*pDgtI(X$D(J9H^}y=kfGpyiqq(doC=*~LB4V5C7LYNpku zVpFJHS#B>CE49y^W?i! zlm`IM00L<41#AOsRg3Q2E>hK#cP7r-u_}`_d%yi9~cdkyK}{A7eY2f?g-n1j=LTvgqMdp}nEpd}<1Mohq+0 zH&)r6pNd}3R|=G?8Yj}jE2+t=CR$#}HLr45UUAsWt0qSO``&M-+V8`v9#y_aTkU;V z_jz^V(b*{`6uSTcWj}Ue*!#We zgJUHXOrsvnjTkg^%kxeB!)n1}Q^O(ZMahkSlSiYJ-RGA_G1g1qOr~^XN3Vxk`MTk5 zu*>8zhWpW>wy%eL;AmEMphG|4+XP3yc7GISCk1Jt;qTX|4j4~Uc0~O99 zuF@^NfnH;=?7V)AzMRk`dB_EMSjAXBZw9aDs9DeG=yp>!LqUG%ep`MYZEOgMya5Dy zLWjFIglH9xYc_kApJ0kZ=*MSNh0+n!9YpE7DD?u^X7-|V17J5`58!dYvjk3ul1-7o0=o5|^o%R|`cp+hn)-wl$?T6PV zY5D5h@f;OCezj;*ZyoOw`KtA~C)1vxmXa;(w&ash$FHO)J7~H2CGtqKXttrB61uG3 znHicp{A@CX<|)U?s8r5x2!j96W1rm386am z&nK^#{Qsh>+@}iiSZr}CVzCWb;o(!$ows4**!!_U8%9YC%G=R!2jEV?F7?T&V*3#k z!)ns=CExlN6z?O@EjZ~KCL*yR5-+7jmo}|UgXppb^XyjNd%ir6w|rC|Zr-C_ zdcL@NFS_3YxEByNDAu;+sDo+Ahfq4ICjGiPT*#(iXeSIv{upg+D10jCH9vy7Q7$}Y zEb=F4W&2=OZlOm$T)T!Y&A zBrAB_DB@nm*f`MPxDv%&j|2V{@H4DP&*oDWUF{mU+zd!y@r(`>bV{@EvEqaVC%5ewI(ntqkNDxo#sd^*oe z52y2!S5LK_E}Of$)b`qBJJq~aYNMLhrWR27OltDlJj}M=zG1JY5s60d+IrFMv-0pAparumfd2W9~ueb3p8zD*uK@I9T}=;O_*N zybRzlCLsTBI;JL4cgTG#zz{vVQ^0l6>tFkWsE0Q9r!1pvK2XpWkUmX9waN9zecyTN za@1m638Pj|<Bz9WV-LGF$wVs+}Cg227h)DqhQs^y`S2NC%^0(A_gL1y{ zu;^2-zi~nIs;7RJ7sjlVp2lZI3b&>KxjjfS@eGvq05kxMi0&kK2znbCB`Y7Pl_|Dn zI&5}(I<~kBCg|WTU#7`^Xd@<}(z{J($4`!~Ih!bu&%L=}=$J!ErMGhjy@?HZx+tvC zK~^SHy*-oYC;g$>DkiF{-l@%FPrwMZ{6Qn}Zilwg8mZq+dt+_FfKHFUOKw}JDSb`{s?$Sz5S;R;cw&@%^cJ(qQq|M z&nW#D;J*Ry0eHb!P~2gW0GarQsAcy5LZuM`gy4^t8ov2~aKKGeK{=HCIJCh|=x4SOXkc{$_oXi?pNpH}z} zP8n(F2ACXu;_udU_X{BQ3h3C3;24`i&m6vNh~EhFghQ6kqT}N*)Z$W&kS~HP;7DX* zx>zl02ChpNcUV%f6y4hEri;*y$hd4VC*?GH@_FjjNPD(eR=_qfmQVH4)#d5o+}73C zN4a(p@%_-ibJ=3GHH*~E6VPs~h-Z_|;;CwZdC zNK7j*2=J$YvI%0Mxq=4`y^Uh={R$Ge6~HQYKR9P!vnbLtQB0m%iKZDK!`dAiGx`F) z+!XokT#=V(e8citLgdYfV!x<~{C<)sabQ-CnDpWd1cUT!qR%SHj~qXeoU zpPZh(KSdg&a$}?9OqxG7CTFs!N=c+KayHniiL9I~Hkb0k%rdqd+687E4kZ?qB5zF= zPLUF^7YKYfK_NYIO@VmR&T-~^WPYK@%VLL=0?J2-B{K5LNOvKPnilzCp=cE85xr1s zqJ=Fe5?{_hj?i*S2fG_R=NT(tAgf2SY5t0mP~M9sRj(@SY9qB(^%=4o*xP%P?* z;$y{Pg~*Nktyt_FV&_nac%vCqWi4PKpdZr1It;Q_CM)aE6l0bmwLu!SO9_G&ny>)z zKk`H=%~FV2@*@Rh!aX$A;WF{;L`QUUo(@WUpC1#yVelkiQ?*pQ1}F^-nY zni{c0G!EQWLnak2j;#gxlCiXGlsKwnp)?aP1Mn5VOcNWTgm>j{ndH?}N3ZPJjM>5$ zQcvC5L%M4ics&NOfv*JLi7pq~T&}PYE_hQ~&}S21E&3!FA_>aHn6(Y{zecGYrE35& zk8YUeeP9QUoP&*?wwf-1zeSlgr&<&Rc$*_=Ukv*6K5% zijR#~+cg8{7m8{jvLjjKCgyY8vQtp%*rrJbhH~h(a;Vq?daVFxm(vbw?&gBMBxsSHy9DJs*!)i-t@WZTTn3(4{r*I?&W!(C4tNBk306qE@?_#a?LnJ` zzPh2bk7xQY@_MV7Uc?s&-@&X6pz{sMU!;7sm~6Y5x?CJ+S}iio;rFR0$Wx46aKb~uAss`KV_%vCWd>D5o@o=s zVftucM=!oQae&5nKptQM;AKEQ0d9-PmU(GgClg>?@hJ0I!66~``n diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index bdda4c1729..729ceb04c3 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -1,8 +1,36 @@ +import os, json import litellm from litellm.utils import ModelResponse + class BudgetManager: - def __init__(self): - self.user_dict = {} + def __init__(self, type: str): + self.type = type + ## load the data or init the initial dictionaries + self.load_data() + + def print_verbose(self, print_statement): + if litellm.print_verbose: + print(print_statement) + + def load_data(self): + if self.type == "local": + # Check if model dict file exists + if os.path.isfile("model_cost.json"): + # Load the model dict + with open("model_cost.json", 'r') as json_file: + self.model_dict = json.load(json_file) + else: + self.print_verbose("Model Dictionary not found!") + self.model_dict = {} + + # Check if user dict file exists + if os.path.isfile("user_cost.json"): + # Load the user dict + with open("user_cost.json", 'r') as json_file: + self.user_dict = json.load(json_file) + else: + self.print_verbose("User Dictionary not found!") + self.user_dict = {} def create_budget(self, total_budget: float, user: str): self.user_dict[user] = {"total_budget": total_budget} @@ -21,24 +49,22 @@ class BudgetManager: def update_cost(self, completion_obj: ModelResponse, user: str): cost = litellm.completion_cost(completion_response=completion_obj) - model = completion_obj.get("model", "") # if this throws an error try, model = completion_obj['model'] - self.store_llm_costs(cost, model) + model = completion_obj['model'] # if this throws an error try, model = completion_obj['model'] + self.model_dict[model] = cost + self.model_dict.get(model, 0) self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0) - return self.user_dict[user]["current_cost"] + return {"current_user_cost": self.user_dict[user]["current_cost"], "current_model_cost": self.model_dict[model]} def get_current_cost(self, user): return self.user_dict[user].get("current_cost", 0) - def store_llm_costs(self, cost, model): - try: - import json - with open("cost.json", 'r') as json_file: - cost_dict = json.load(json_file) - if model in cost_dict: - cost_dict[model] += cost - else: - cost_dict[model] = cost - with open("cost.json", 'w') as json_file: - json.dump(cost_dict, json_file, indent=4) # Indent for pretty formatting - except Exception as e: - return # [Non blocking] \ No newline at end of file + def save_data(self): + if self.type == "local": + import json + + # save the model dict + with open("model_cost.json", 'w') as json_file: + json.dump(self.model_dict, json_file, indent=4) # Indent for pretty formatting + + # save the user dict + with open("user_cost.json", 'w') as json_file: + json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting diff --git a/litellm/tests/model_cost.json b/litellm/tests/model_cost.json new file mode 100644 index 0000000000..e581433952 --- /dev/null +++ b/litellm/tests/model_cost.json @@ -0,0 +1,3 @@ +{ + "gpt-3.5-turbo-0613": 7.7e-05 +} \ No newline at end of file diff --git a/litellm/tests/test_api_manager.py b/litellm/tests/test_budget_manager.py similarity index 72% rename from litellm/tests/test_api_manager.py rename to litellm/tests/test_budget_manager.py index 18d507a7fa..619b4785f5 100644 --- a/litellm/tests/test_api_manager.py +++ b/litellm/tests/test_budget_manager.py @@ -8,7 +8,9 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm -from litellm import budget_manager, completion +from litellm import BudgetManager, completion + +budget_manager = BudgetManager(type="local") ## Scenario 1: User budget enough to make call def test_user_budget_enough(): @@ -25,11 +27,14 @@ def test_user_budget_enough(): messages = data["messages"] if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user): response = completion(**data) + print(budget_manager.update_cost(completion_obj=response, user=user)) else: response = "Sorry - no budget!" print(f"response: {response}") +test_user_budget_enough() + ## Scenario 2: User budget not enough to make call def test_user_budget_not_enough(): user = "12345" @@ -45,10 +50,22 @@ def test_user_budget_not_enough(): messages = data["messages"] if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user): response = completion(**data) - budget_manager.update_cost(completion_obj=response, user=user) + print(budget_manager.update_cost(completion_obj=response, user=user)) else: response = "Sorry - no budget!" print(f"response: {response}") -test_user_budget_not_enough() \ No newline at end of file +test_user_budget_not_enough() + +## Scenario 3: Saving budget to disk +def test_budget_save_to_disk(): + budget_manager.save_data() + +test_budget_save_to_disk() + +## Scenario 4: Loading budget from disk +def test_budget_load_from_disk(): + budget_manager_2 = BudgetManager(type="local") + +test_budget_load_from_disk() \ No newline at end of file diff --git a/litellm/tests/user_cost.json b/litellm/tests/user_cost.json new file mode 100644 index 0000000000..fb2db88e3c --- /dev/null +++ b/litellm/tests/user_cost.json @@ -0,0 +1,9 @@ +{ + "1234": { + "total_budget": 10, + "current_cost": 7.7e-05 + }, + "12345": { + "total_budget": 0 + } +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 907a4b7777..0907ff1acb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.583" +version = "0.1.582" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"