From e1c91903da19532f595083c3b21343ecb355b63e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 12 Sep 2023 13:01:55 -0700
Subject: [PATCH] add reset budget on duration end to budget manager

---
 litellm/__pycache__/main.cpython-311.pyc  | Bin 32862 -> 32824 bytes
 litellm/__pycache__/utils.cpython-311.pyc | Bin 100211 -> 99701 bytes
 litellm/budget_manager.py                 |  38 +++++++++++++++++++---
 litellm/tests/test_budget_manager.py      |  37 +++++++++++++++++++--
 pyproject.toml                            |   2 +-
 5 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index a7c59fbcb796b7fd9e9705e7315135d0311e46b9..d4386ac9ceb4fc5dc2a7fbe5c9601b5ab284f08c 100644
GIT binary patch
delta 3967
zcmcIndu$xV8Q;B^&py9m$8r3;bNtFWu^syy+j0DM9(53D0;Vb3B**%8d|rF^b~C%@
z*yeG8q^*DyK?f<S+P|WT5E0SR11SlhzUdzbO4UC~w=JkrRkc+U1@%P{s``Di;cOy7
zkg8sCzx&O6-^_fEnQzALobp{h<tw>RQc@(qCwfI7uRc{$BmVWN_8(_OK`l_jdZQ7Q
zqnsB4kI6C4i-0%DO`J!7Z;`ifUJSfhZsxoMc#GV^c`5K#xs~%W;B9gn=jFiL<#x_D
z0N*Na<-9`AHX4#FK|@eQy+iJR-%5R(v0dKI`!?!9qf_oQy5ugSTkbY?$U8*ACy<VX
z9&`miYSlBMAUp?O?<4oZ9GmpbMvYu!)XKF+T#g%cavdL6xX|~J-ZMa+6IhuR4#QH_
zX9RgC`%sj`epc$=%tn1Z%77XrU1|)zO=LjbLi*Ka__mNiwUz8r+km&jcPklEC6ITJ
zZZc5w(d>e{?UwfTTiTr@&ffAh*X>rjVeR3RkP_4y$OZnZRYh<TE$PnwU#*yHRC~@8
z$Rp~I+6X?Nzvx?>49tbqzB562l)d8LUpb}@fQ{qoAmD_$3viNc3Y-#y>|EgUxqk;L
z{o+)vBwR06cR&E#h>|o7ODB$Irs;0ziRZ=(3w>fMI}quMOh5;nM3~Aw8u?Fa>F0od
zAa>IY*ym$o6Ajr`pl-l+OpD6c`z35|n4{!GhNP1#35%o^-C3bWQFsi6B`UFWbHyNO
zC!wkIZWK%-90ZVx=s0q|^d4mU5cVT*r72_w5bgtTSx@h#5TD{0yR*05qtk4)_wYu~
za8SDFcY|8m$~bg_z1th@nZeNp0$+vCh{`jP=tguimD19wggLt)1*pljSa-ZywAlXm
z;x;}EM@S*h^0mqMJZm8{*ze*C4eg+M(MPw!k6U06+g4J<_CYHaqa94`>lXL2r~BH(
zVfJ!g-O>E)ZdkFjgtKB1DahAxLl%WhIk(vD79{6QElF%Q;E-j9J_Iw<2l!%1N)(5v
zZZLyu(#@=QXO|bA88)}Gd0Prq=MgZ%ZWyv<s^Ga76bgNkkKzxVW-soXY~exX-adw%
zk0X414tkc=_HPOC5KuPRf2i6EA5P_8A1!vV{|k}4jM}Vd;EDSCaY76ieGniYBkagB
zbEvr*s);~PVS5+=jN=xaX4ePw_?mG%3!iKu4!>#O*9L1Nd`odBD|6o)42ZY2A~L^b
z#n*@KzODYzudeCi+8Y7f;)&_rp2WnV>BJp-@1<4fP2sxue7g1A+B%%8o?RarzO4m2
z{`fHqFwv@IYH4R59YJG)D$yNfwZPPM#ZbCwE&Imq=Wjb}=Wah(a0<=f*$C3p$ax7k
zg6xY3{6s^dSZ?@ul;;Qo?pByogF~k139_6d7B4>SY{y7_=sq0P%<dYQ+*OXW7$Jl3
z2m(*;Br?pQTd1sNC`l}m6*p`FNh-Qd4fe*!vtpbrjczO7jH+D-e7l}yFO2qz73>eA
zQV8omm09K3QSr;n8tV^LW5;=RX>7At$$mPvCv*Wjo@Hg@XT(eF>G76Lm^+>=H=sF$
z(wj+mbG%cmWgm|Bm%oGx1pv}CPbV9hh)P%>=`2FNKGGyM^A(Y**c@l?s}1b@#6(*k
zG~I%mg^2Rquy)XfP$&|5nw3u0MBc!WKS%f_>z(Y%{tD^asNzeS^qbiF7Q#E&Rpyu}
z;>;78P&5o0{XL32H&I@coL8KLVZw%R1IXzm>@Hp+6pD){px1EBcd@@PMbd;Sj!E4z
zN^Hwa+w9%R^3V;`_#6Aj<b!1$NO}Ab4&)x5I$T?av2Y704ApEXZ*JIuD<Y}E9q}5w
zc5F)pszT|-DA1h<MXc%gCb5GJ9&ZSiih?l5<l~P;$5F(OGT#t<Yw!!>APS~){=_+7
zt%BUUt&m~Hq%wLB{IKsRReL@|Zt}a<+suJ==kz>InPMfgKWX?4PV9I6y6w<=fxDt+
z`{6gn%l!4(PVoY(Pu^3DXXVeBZoFy?WjS%uq>F^wbOz3H4f}R-cSXA=4M9jVx&viy
zD|;_l8RDVLvLaOyUu1o%)WnZnF}(un^ksyr2(KU%BGj>m)aq7lV*|9^ke!(|G>50d
z4Jj5}%&HX2@6-LCAr|8}?<a>GrF|?$n#ypyVv5lNn{H%>$e>um9wkyey5W}3Do%1f
zvF_;sk5a+Py-A+*i(!_~s;3^n$<6>k0A9c`C39L@v*#0VtJ&~YftRbWcH7afpi1$b
zW^(Z^ZzmEcI#;yzY%|WY0hKXD^mT;S0HpA`L~_^R{VRBftrD19x@H}|1?c-Q$PNQQ
zO8JL!SIP`nbG;{BF!(a0-u^r2{wbGOxDw1_orN9EcPBofy!Q)#*vvt5hHBW!B@4(F
z5p)2!HXlLGH#gtrY3%iR?~)TJ;>GbKvgLIm+_kG62hG@$qHFA&)pGb8N^jo6mr)%n
zu~b1{MgFN}mVO^p+`_%aEK&Ib;fFZnDl0zGnB}_<J*7WFz{t{{Art^e#a@j0P0zjK
z5jcSII}msT4kN>=M*oSxkNv-Z*=W?wSg#9sA<5ej7a`8S$Xz;7;XiU3lwpgn*bXr$
z{FNfm8R(IGuHS-}DK9DhbtMIbYX=m^mf`zpeIf^7k}6h}nT5C6a>fy_vTK<hucGnG
z1R=^gmqx_l+~U#(aSuPZL4FL~NYYHtX{m&zIP<g=G`%W_k%H`DRm-a=G-c;sJM0H1
ztAk^(ax?qw$s;ZNjQjTS-7Zbhw{Rj%xa$W`=}Xvrj~!VauG^0Eb%bsd2a=|>LW9`8
z%AQ}Ij<^9`H*8t~J<PY#%XY0qvse{qBSI7*hR_7y7Q-!($Ug_qVt+IC!`oO<6RHXI
zdWrwB;0EWQ%i$S!{qUr41GA<Hmxrn)y+T8=UQ;~X6b}ZCBVc6d;}`^ru6mC?bd?{e
z55mK61!Z(HN)Lk+{tNbWNDyxLMNxd4HLbR>*lHEqzxwNaA<=h3_+OzkXWjpoKmgbO
EFG$s(761SM

delta 4119
zcmcInYit}>6`t9rpYbz}AF<=CH?}8R$L~iP+wm)Dolw%enr4!CHQu>ik2ABgxijl{
zUD7qwKt&U&?fp>%>7OE%M5vY86=^~WLR66c2o+RSA0t#K5`qFk0Tl=pRGf3i_PUYU
z@P|p}+jH-^_sn_SbFQy0dEQ_06o0q4xIlpC{L2D)>#5=z@zzr<uVst|xj_`vpqi&|
zF&gDYPUi!?Ro=?!0-!^3h|`5YH_1($E&@6%hdEsg^fq}Lr#ArIEH`ty1n3sIh0~=#
zx5}-YE(2PUC7?4wwHyZ9<?S#KR8fz}5%{gp?=jlsHa=IW`;8s)4x?ReH#+1Fqf_n_
z1&=@?b6sc-KI+EjMM3x;yzV1+gN!P@+NhChj9R(Y=#hJjI=PPT%b)9AUH81e9%(IN
zx5S{>$8L*~*v~?qNo7FYLfX|vc(;-PHAMQ=CV0bSr@D>Yt2P7O0&gqXrAolxPCCdy
z&8kR2jodN5=Z<k3>0viLO?886`?4Sp#f7+_)<6L8U#%*FW&Nzc+n)KqYBAHGb}k3y
zVf9|M0c=2f(Y9_fFq5bDEc@l%?5g*^iV?LRbR1O&0LRpwfa9#(w<HD_^BpKJn<rG8
zNk^4fj96CmG)ZSa@Kt)niR?duo5f8L@QG8X#7rZp6I(MAbP#5GvQmD&M{H)@g*}C1
zFhR!=CbG2flWmzvpg$JdX*te$5FSD}hVU@LaReEm8Nl&bwwa_SaC{QSWiynVNs&Y>
z9ZiyiqT6YD3Yk;LELKTUH`4}5*ilWTN06}(;Q)YCKu3l9V8C_SgVQh)==}(M)dW)e
z5flKYJZD%mrsz6|yzFy@PO-kOmVhgt#kvkwx+cP|%U;K?C6Xzdj<MIfwoD#HaR~(e
zS^f!Jd_)qRLfwqVwM0B>PR~g`np|Dwl$lfmhN4F)v65!OBJ>RF=sq3cf>4GOa23gk
z=aO^oPOzVK&(*iW(#swy!H*L(AcjhuSm^b}MjBxc^t6k6*}0zP%n*w7ITa@QREwH4
zt|YWYg~xR5n@*mR)S`AeNhCl2x)VrJh`oJ>eNHeoYif`&j?X3wHqC%=dXB3QqeQWZ
z>iAQr3B^O@lpWf`LcJ^AJjxFAHbrJ}hlYR=a`GVDrV5U7eQ3d)<e~UO_pxVt$HF`i
z+!asY<avZo%b*vSw=Wdn{;^o3@4-!L-j=-UUCEyAd#GxSM{N+b*fVa<5=z*ufflx{
z|K#Q`p?q{LeH38LpE$9^+IKZ=f-^4AFXNay*5iI)nZ47mKb+H*N9?l|#_g`YKtVF5
z;3ODP*W*!Di`n1<p0?Z(mDz6(_{6)aR;cAvV;}B3eAf*pAJ1(-jW(h(MWd5lozc+;
zC!_oK9H&(<CFQSk;^*7M)!mG;>Y0+NWdG@Isi${yX~*wBr3R*inl!b9y_XK7Mt+s(
zwz3#BbzL!(c3Q{2vg^gW8rL@H1vPSk^3!G1i&vXmG2+DoV#qEIPOY5zr-!P=R#r2#
zIe?XdhS{#6u`Ojt79q?dJchty%!|VqyEs%MCfJpsOJWa;4o6Bc`)NA@PuL6WTf<#q
z1^fB16sX176YSr^C&jO^<GcH|RO7^V0Gx8$j1wF35zZ-|M<~j)lD)WlB5(=UTxPd-
zFN-g*b0gs@j1`ZL<I`+H>59`|8EF&i*zZRAim%{;Ab>Op<~`b3FRc|`8bjV%-K8oH
z*D5U>hBv0F*%PCC4;)5DFMtzVDQA?&KX?9iAyXvWw@$^nqPkYx-ar|@KzNJQjqR9y
z7s>Z<n<r+{i#U1~;nz4-vg&{+cmS>cEwWtWkzWv-RqUu?LPj_~r0EMdT?D-h%8|>U
zf5AO3v6seL0)IvFBliB-xsnKyJR%7DvvcDIYsDCQ^;C^`gS~w!B;hnvYBZPjAry)r
zI!K94(}Y?+m}-OvYU|didVdKp<7|BD8|=-g%0k|84&bJJ**{M`<H<}To$G@b9jWB=
z+u>cJ;x;4{i{Al7bHxU8hu6dy93=iQ&-B&$UjZq39j|WLyy-fkW_e-GHR7G=Ht`bk
z$Bx!w9s3Y-0Na+KEJV+m^fXvZkMcA8cx<E(3nFj>5Y?29#ff*Fz)Z?c!4!TK9Rojf
z31Bfit=O^IC_6aVAby|KtCGxXWfA=WE_oH<I>HYT@)6KMPMNzcmqivo7PJAG0V_3a
zXf_X_6Ht=S%~Yw;P3`dilG>QnrRxHPpY3J&WNQhgGX@%Mv*<>)lMDoEA-+~4+c4f#
z&kb`*?^KS$m038GeT6*fJ&}*8lCM~nX4y)@zU*;))YJ*@jP$D@9$fY^ex-OuOK8?C
z^nS&HI~Cltg<KD%U&kdy516sjhgR<fN7!>(OC}5qT8@h`sPrbn4FD-`T^it>io$(e
za1UuE>a-!}OuUKEf593n4*-&YKdx~l27EJ`TT9zv^EF5o>mykEhwS}x*ZdhQ#*la1
zg_A(L*D37cG!Aft&LPP;r;*YT3;^iDk0CvQFpgm2tjE0-QJmmabpfgLIu@q*;`T3^
zjq{4Gv87}<^G)QgH1%t^9J{1cPJe{-XV+)s{usc?-(yS@mESjhf|HnrPBmVQu;P$a
zrdvrYH(&pb3tvOH!LK{7BN+scid^sWD+RZk`({7#zktAfa|kJ{lk}emyuRE9%0hKl
zqOdL$`D;qh9tgsl7=)8>!j^<v8+K-&JLC5rTZWBr9@3Ug3<`gH3G^t;NFK+Vg!?qF
zW8QUj44Juuk9$k-?zTSM10bY|ePU0GLvZG8@j81u)#;YZ^RR(8u|`%kKO_!ikIa{e
z5nd4duKNmOW@1K*N0W*@OE=(hRDq(~AjW;{$+Ok<SR0(Om4)l_Yy>_WW<Nc9EX*r}
zXD>ez(ggh_iomRJykIE3g2TVDk%hs!HY9H%bRgRoGm~lR$MJRc%)(?L#a`|BbltGn
zpBB12a22UdhuCPkB11#SDuNyxUAuN(z@2!$@XD^JQPs3bVxIrC<M?M_%BC1_#|sw{
z$2V=7(0TBN<{F7@fENOOVtB+TI@4|YXq|hRMC)h>z3_3~AoQ=4PKN0rV8egGIvfy$
rTV7EV-(_WswXAfp4&DHZEWXD+SPZe>EH=SgAeLtLJ@QAN4`uxq4_xJW

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index a12a9697700203e04ee6663b0541d45c9e185120..ea63bbb384b0f62f53bcf6905ef1dfbcb4f96158 100644
GIT binary patch
delta 2390
zcmZ`)dr(wW7{9-LT$Wv4CJ2&>Y$mQRK*1t~5*08@Q|!SuxvZSKxJj&?yDVyvmN>r3
z!MW3%rmR7lOf9DeHP@PSa{8mGsietfigHHnVWm!J51E)A=R0fknilTweBXDz_xaAf
zKkU?h+OB6jo1X5~<geg0%}BhvD<cp2FBMI92EuG&CxVS^C7UGIFl%Cy<+?Ie`}2=6
zHbqLBtc}imSE|MgYqxVABGyUgOf*~Bt|Cd#am_AX#by|b*#aZZYUEmLERJ|sU8g(P
z#v*KCC&CGKW2Y8uXNx4RWl?LdYtz`f*<FUgRv23&E*6ry+gPLA-Of<21y=Vib9+9+
z<JO%gi$^in^X|PaO?ys0$saVPTcvd)(|sMz290SQ&Ig<<#X53wbm2cF^_TbB(XOqP
zn(YX9tO=(IOTMLc1YEpUE<8ezAeaVCTV5mOONk2-j39W(YCBa?{iWPJ;np;Bo0M(T
z26QcdeBadSzN~rlpE-YH-GD=P=bZ3l9{1#T=e(Hr<~4n{H}rXeA9{kvJwZuWXHS(_
zgQv!fA8^8*GdMVCW=mIF^ZZ(mcdpm@a=L`0ZueZT)pWWa-&s{(e2U)Qyf2@26ne+Y
zUqE4FV`GwY^0P7}evaV0b>i!ZzQ6MsXQpDHclMe2Zd|f9o~yu+_4>I=r1lP;OLO@y
zlF<(YKU$d=CTIOb<j(|`2nMZ%7uMzL4oTKi2Yx4^Ml1*%*5?=Uvy^5E@k6woX6YCG
zNbe0?d@=`__S-2q?Z8EGz>7`D6w`h117fKU)nii0<}>Lx;7J}qE(zYHK%>M9KGf(1
zBq*?le8?SvA$#Olyx~Hb-B*ZJE)0o^iKxn^Yy%Gd3&p0a0$yGwc1%R+jl<dYXYp{<
zG(wRW51H|Byv>w{tCF1|4f<2ogx9si!i-Z9CXE+$nU_!^e#*da>qVI4&?||lw95im
z<G=#@nJGvBwPMUPRAQkWoQ4dCt{OAeZY{>=z!I^j6yx%iQ507s+RFH0Dp%*z$oF#l
z!BPxi?h<l)3&B!?T~vvS;44ygaWd!i1WD(T+V8yVf~2?EhiBqNiq=|=QL-<a%W++f
zvKltv21N|p$IH=$0u_;wt$Z$zxIt`$AS|T6(f00ZFdNw<4nUdjcl??6*=tb?tQK-D
z5Ie3zjwq<W1Y1&d&Mf6z`Aya&zNkdu6xD8Zmz2k}WL4CTf0f4bL~#`Y2-^)+Fg+V`
zC?)0YLsDKKK^}K2NoNzCpPU_aSLeyD^}AY)wQ;_H6mBO_0~L}~zb6!;y-<kHAMWu=
zQi&5ZlS)dgExx8LZt@z^h&nWn%#X+xQC};MdAZnCi>vf@sk?rRI9iNs`-57Xb`~Va
z#X53vMSVQBX0ExewP|78;Nf^Q*2-(eu|?>)t&Wry5!_6m(xc3{%7PN0kTMuSj6h{Z
z(Rc?z5w)O{Kuu>g$*S$QZ$b)CZ6CWC(}AnSZ;LTUpG<oYH!MLrHrt1mV3rd<itm@9
zq+s}#%SQUCW%@laX0(`mkSg_ws(O@c_*K$@EPjxJ{r?+)g4{|fYG8Ya`h?(9g5e9!
zHxb(@LC*dbvZ1i*z9sMBR+4R#VCqtduAewxk8#*;r!B_{H@1k?4fxu>oz!Ro_z{9M
zm3t!G(q@>+5T-DLvIkp4PY@IIoyx0tFNkJ5YM0%KhuugKM;kE<LJT&_@qWrK3*%!4
z9<jX)8NgmK+Q9c{5UCMqbHKhgA}3MrCO>;bcQbOaUmR*igT9YM`$X<4sogD7qL`!~
zB+)@p9>td2=Ot~*N0yR2DdnQu{wRt=m;H5#U8~Wh?;)!_c2NsTbo7WjIHr_$kxZu!
zsn+f#_5cBS3HXwqD)rIoA?aa)BjON8vHmpC`^9;V3cO_(nFzztq**Yt?WzP`ak!63
zg&FpxwRjx4Z%DdkabO4UR*ht_xY3%BGqZ0!{65vZ<J9Kg#mDPWhgtTx`*1xlTCBJq
zL%N!Y-$bYbQFMuSIxuwueH!_91ax8f08wfN(uq|s^Y_VN8L@r>^+r`wyqQ>azK;_1
z6v1DGt9GjfC(lCe_PbWbVlDhFQlz4Ik)M1IUyffqbL3RFwHQ_WjNF#FXy&W-6AwV>
j<GsjvZ03L_mwhAE@-MKLZN@hkm%x>{VnF*Zs3?B`Pz+b;

delta 2746
zcmZ`*eNa@_6~AZq?PFnK#mK@!39!bZD-fWFsUIXk=%NrYW}8f5rY`H=$2TnAeS8Lk
z4a7iVHRL7Uq@xk-7<FQ%v2^+}Hg;^%KT>Tw*pgJ&Iz!@=Db<NfO)%9a?LBXUR%d#5
z_jk`Zdp_?wXYYRTf_#0qY(J8a5J$w?^>d<(pB=S3A^XJA2hHvP%^rf<I!fs>vHPh)
zmy3OCESZ1NdYU670d0plO^PM7LEC38gGO!2Tm(&8lC)INYo#^$&Gb7;HT|BVQjgdx
zl<ImbT|Z>0?WB!#!w~qpXyp*8?WR=%duWUHrZnhor9V^@x<whTm#9zZj?g+`XQKWg
zsrCiBRcO9Q!Asn2;~tCkBK%g{G?kxrTzb<?$XntQT|lVCOCBJqq@A8h&5^x?lI`Ya
z%-x{M^E!kqA<VVitj6Avs5V~M3~Bt0<zP>G7Uqq*NsqaPbb~ua`)tZp;5M@gVZ-VX
zx&(aR<CfWG<Z2NXB0P(rApBW-<dan;mw}K=7D7TV3fT^FJ~eY9cilu{8NL^k?^qu(
z$rjfo>w-yZ+L+@V=jn$gtX1c%Rg=~#L1-sGDcInO5U@nylHIv~(QBJuPu`b&&Ym@C
z&zi7ieJNJNEHuO(N)<tdQ(WX=Txq8HM1p`{r<bOgf0Jerq_(grSu6SU1DMg^@&)+C
z_{GaVH9@ZS%yhEtA0kb51>vf8U^@FQequTou8n^+U2cInZQa#Xa8v8QS_0PblUL&<
z+bkOW6XCjcF0wrFb0oh&m_xXs*=BpwEG9vE@l?J>;eQYa%0PQ#HZ9T6l##!Q-`KPd
zXBR{K_?6lH4zTMZG4QDgW_fEI41%3U;$SflzQ_ja(yeH7QTTORqt7N81!pjjgKx8e
zN6tV&hJMNhDGTAI{!u!-Cqa?^vI`m|xXEvQ8y-%=Y~3dIulcOpajb}M%!b0s+u4Q_
znZGHd`09hq7gGIdXUHu@lPweh!Z98G-qxU>GAtq_iTGJ!u>wrQg&DZ?53^yJNnVT0
zT77OAbeN!AZ^(fz01qFnf<=5b7nZ|%Jv|TXCfPWhbiE)S{t8gd6AB^2xfw%A^-b-R
zox+-BmWKhndRZadg3@aAy$RuQgi$O<LiighOAO7~69`e~f;wjY#sx~B)CY>-9gJ2`
z49VhfDvDvX!?3z-5N+)`eQz<m1er#}0NRSrH&Z~7RS@d9@RzD@c?i}(5{lhmm>)CE
z=l%9~p#q?pi#?NXTn!H0y&780iH2XpWwacAxCC4|MsFCCW;lzF7KA5pPgN|1&z682
z>hy0t3?VBd@jV`JKpQ{ofvm^NF}J4>j6T^SsIb-NYgc+ypO2N_J~cn8gb?acB_`OZ
zc66#C=0S}{lzWkC;f@L@0WbGfzyng3kn@3DaOy8tz?3<&6ZLw~{ywi7>?jTOwg)z-
z3iGQ?!FFam_(M-spi(8mV+clWm=-fq!06F<El9N^7-?ZF){n3h&!7;&=(tjpmFgpp
zK@32tKDrSe1X#jPSHoJ_jo%S&+63J&r1x!tm1elkPdotynYSlb%z~IO5|-P7RA~*d
z8(3%DUN003{9BOjMD{ku`u{_XF`h&fqx}vcbspgY!tI%5gUAjE5S_deZ7`}aLebHB
z9%aJ<Lip>A6Q!j$Te<EP-8rKF(F<EFu#@N2z%-2Tq*}-V#*1r3u{$}fg(dQj4ST-3
z78>9ceXbUESwQB4b+8il@{@I<E%)nlb?}i1cIm$lfF0m4zewR4{D7ZTgv~L%s9v;D
z#sTy*CLp{Z;BWW&DfP7}p^(2pamTQuysrUj<RhqagkNY7<_Gxi8ey4y97V_ZjYfDr
z<t;(uLlad*pONA7fc|O|ybH3y?BU)PcuB@MY)qeNfdUy0@-zlH#d}aD<KY;mdKlSb
z2x#cGMQ<eIZaRq4_Yh9;J_h;nA*7G;pD|bk@9Q%R{9q5q`2zt_fUiaA4S;sAIVr1O
zq{c}<QN=#Ld?H|U#S(3_Ct0g3Naj;2NCVVNrG2IbpqAtE-8b-dpfPv@1F0d@>Ycwp
zW2;#VroT=23rngJQkh@yGx$H`u4^?<NX4}(u}W0OOEzpTYxJvS!M2W81&1aq2L_np
zR~6BvABf_@3Fywplu$>oJ*52aA%_=sDXgi!*LSB}c1D=25Me|3Ln7{DN9J>$q3Vb!
z&f>V_%$RiM{l&TBPtFzR?{iJ1W*+rSBo>Kx#B55oM`V(g9<h<s;`~`6r^bDbFq{yv
zlf<N9^~;C}ERORFl7>5nJMk*E3je#aSCU_v*615lIBT+eEJBy*son4sNcokZJ=N}A
zY)sS{mT6E`+PlPDp6P+bHluGRg+2S4-{^t$P^>@N3y%Oe`HnugB^!hH75~xG&;)z<
zm8T(h0RLvN8H8=ZGrNYAF=+9~8kgW7&|wj>E(GHyGJYwBkR3sI1z|tJofb11+i2xG
zV<2q}23y(tsEFsuKIWtSaDUuqf);^EXhMItAGn-l1=nju5hAvui;V3a(1U~UceuX~
S;=wT{G{kl<Wbm5jApO6e1<73i

diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py
index ddfae14d192..ff7de9b7be3 100644
--- a/litellm/budget_manager.py
+++ b/litellm/budget_manager.py
@@ -1,8 +1,8 @@
-import os, json
+import os, json, time
 import litellm 
 from litellm.utils import ModelResponse
 import requests, threading
-from typing import Optional
+from typing import Optional, Union, Literal
 
 class BudgetManager:
     def __init__(self, project_name: str, client_type: str = "local", api_base: Optional[str] = None):
@@ -41,8 +41,18 @@ class BudgetManager:
             else:
                 self.user_dict = response["data"]
 
-    def create_budget(self, total_budget: float, user: str):
-        self.user_dict[user] = {"total_budget": total_budget}
+    def create_budget(self, total_budget: float, user: str, duration: Literal["daily", "weekly", "monthly", "yearly"], created_at: float = time.time()): 
+        if duration == 'daily':
+            duration_in_days = 1
+        elif duration == 'weekly':
+            duration_in_days = 7
+        elif duration == 'monthly':
+            duration_in_days = 28
+        elif duration == 'yearly':
+            duration_in_days = 365
+        else:
+            raise ValueError('Invalid duration')
+        self.user_dict[user] = {"total_budget": total_budget, "duration": duration_in_days, "created_at": created_at, "last_updated_at": created_at}
         return self.user_dict[user]
     
     def projected_cost(self, model: str, messages: list, user: str):
@@ -84,6 +94,26 @@ class BudgetManager:
         self.user_dict[user]["model_cost"] = {}
         return {"user": self.user_dict[user]}
     
+    def reset_on_duration(self, user: str):
+        # Get current and creation time
+        last_updated_at = self.user_dict[user]["last_updated_at"]
+        current_time = time.time()
+
+        # Convert duration from days to seconds
+        duration_in_seconds = self.user_dict[user]["duration"] * 24 * 60 * 60
+        
+        # Check if duration has elapsed
+        if current_time - last_updated_at >= duration_in_seconds:
+            # Reset cost if duration has elapsed and update the creation time
+            self.reset_cost(user)
+            self.user_dict[user]["last_updated_at"] = current_time
+            self._save_data_thread()  # Save the data
+    
+    def update_budget_all_users(self):
+        for user in self.get_users():
+            if "duration" in self.user_dict[user]:
+                self.reset_on_duration(user)
+
     def _save_data_thread(self):
         thread = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution
         thread.start()
diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py
index 1ca3d6dbbd7..c94e729e625 100644
--- a/litellm/tests/test_budget_manager.py
+++ b/litellm/tests/test_budget_manager.py
@@ -36,8 +36,6 @@ def test_user_budget_enough():
     except Exception as e:
         pytest.fail(f"An error occurred - {str(e)}")
 
-test_user_budget_enough()
-
 ## Scenario 2: User budget not enough to make call
 def test_user_budget_not_enough():
     try: 
@@ -78,4 +76,37 @@ def test_get_users():
         response = budget_manager.get_users()
         print(response)
     except:
-        pytest.fail(f"An error occurred") 
\ No newline at end of file
+        pytest.fail(f"An error occurred") 
+
+
+## Scenario 5: Reset budget at the end of duration 
+def test_reset_on_duration():
+    try:
+        # First, set a short duration budget for a user
+        user = "123456"
+        budget_manager.create_budget(total_budget=10, user=user, duration="daily")
+
+        # Use some of the budget
+        data = {
+            "model": "gpt-3.5-turbo",
+            "messages": [{"role": "user", "content": "Hello!"}]
+        }
+        if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user):
+            response = litellm.completion(**data)
+            print(budget_manager.update_cost(completion_obj=response, user=user))
+
+        assert budget_manager.get_current_cost(user) > 0, f"Test setup failed: Budget did not decrease after completion"
+
+        # Now, we need to simulate the passing of time. Since we don't want our tests to actually take days, we're going
+        # to cheat a little -- we'll manually adjust the "created_at" time so it seems like a day has passed.
+        # In a real-world testing scenario, we might instead use something like the `freezegun` library to mock the system time. 
+        one_day_in_seconds = 24 * 60 * 60
+        budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds
+
+        # Now the duration should have expired, so our budget should reset
+        budget_manager.update_budget_all_users()
+
+        # Make sure the budget was actually reset
+        assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired" 
+    except Exception as e:
+        pytest.fail(f"An error occurred - {str(e)}")
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 306a0a866d9..3407318b35b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.605"
+version = "0.1.606"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"