From 80d77fed7123af22201158528feadb83ab9a95bc Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 11 Aug 2023 18:20:21 -0700
Subject: [PATCH] coverage for custom api scenario in load testing

---
 litellm/__pycache__/main.cpython-311.pyc    | Bin 22249 -> 22254 bytes
 litellm/__pycache__/timeout.cpython-311.pyc | Bin 5167 -> 5193 bytes
 litellm/__pycache__/utils.cpython-311.pyc   | Bin 37981 -> 38107 bytes
 litellm/main.py                             |   4 ++--
 litellm/tests/test_load_test_model.py       |   3 ++-
 litellm/timeout.py                          |   2 +-
 litellm/utils.py                            |   4 ++--
 pyproject.toml                              |   2 +-
 8 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index 3e2e932e1612bedebaf5ad4183cb304e5d7c0a4c..1884f5ce9e4485f93a3161c20d641cef69ed81f8 100644
GIT binary patch
delta 1576
zcmai!UuauZ7{GJ#=QdSSmzCCPlae)m)|<3RZHu;3+~(F~asRTUsj|VcyWG<(Gs(T4
zdsEt3H>RDogAL>qbXt^wPRr<nyD)qZHt@l>G1&CspbsJ_1({Dih(36}b6v$E^g@1o
z=lgfgch2`^Z52IOMZV8`zDFGVbN}$2y7-0fD-`Uyh8&LT{AYg^huW@osc6crJ~@x*
zM=01HQJ<Q39F3~ooG}#@9CE<wbbjk}RbO#h58QwC_>(-A3;<Fb6C_Um43CC)a+!Im
z4ZUfy#}!qrnpLf#)Jh7e7$gqf2|5`GOW;S+^ju`0-xgt=fnXG(J@HYxC3^OpNG9n-
ztLvsP%1$D@i{R}CZ?}nOOs*b-g<m7pN)@&v*wD3xfn|>!um=Lf_Ip7}0Gxm)0EanP
zOPOHB#47VN4W)uvd!G7Y#{-AJAqP0ZVL(tQM7DvVK;M<a9V;MC0Kx!qv>3e=>y^Z&
z5&BE)yj^b5YHV*qn`YcL${*(6j4SSqB_!zgeOES1)7P<nCMJpZF!(iP=pA`+v(T_)
ziWAV~cr>(r-|R7hUu$B+ESa-9=DDV!+zkCOeqr;3AKD48C$hUl!vA>>){1A!SciUD
zv7ec3U4PGY6Uot${us*AWBtROef$KhRXNinxb!NXB`x-eKIy-Nl5}X`MC2qW$^aFB
zIhaprn5T%QSIZ)1jNTl;s7?dP-uNWgz5kSHFt@H5rYMv7rZM52R;p_%t{G&Go=V1~
zHz4O(x|B2@O@k;Z8>VfkptKd%WAy1%5N)T2Q_s+gsX=KGltp?wl^hZEzs}dpT{Ur?
zl)!WX@ZZzg{VB?*(B8q}i8}av%E6pvjeHJ5%kEsZRLmbB1Z-l4eB!}Nx-whS6qU??
zn-@1!U8^=s=Hk;q=D^vmz*-rMbhJi60d?^4>!iPqKbL^6ml2b=d$R^Zmgfp4O*Kt^
z;Tt9)pFWDN(AD&Q^fCP-oj~8vXl5_EXPwTRMxM3jd7yQF-z`@`%-}A*B)b7G0(1_Z
zUEIkym(Of@IKHK6rPjndLTf|dmj%28*w67WVkiCiz|ipxT_pzmO>lh+T>W<JJD~6Z
z#8X@Z33^C04r?8+pTztht5G4BKzWsexp)VV1v-7OXU7Fl-T+ttFNf^iIGZw<#O^mi
z3Ie{i?i}>E&;q@mdqIM%<f_$mWDa@71$%h<<$O1~LBGr|NlReP(tKf$f0>KSUDotj
z;sNafy;eAcKA^uBM&!3a(F=GF&<(f;U;welPax&2{KWgn(|$fMAE3V#pKUd6NADJ1
SzN<{p)THcc<Jx1Bm;VN)JcjlF

delta 1593
zcmai!Urbw79LIZTdnv|>g}~Gf0u;Jk=qO=zLpEpU#-A~!P)bQ=vQm2w&<VYF=ichj
zZDbKQ(YVB%ne2}mUCaz2JeX@VCYtTR2cOImr*S8~7#|GL?8z7ZJb1pRib_mq(ocTp
z_wStFIlte%yGqtpN$~w(@E(u+t$q8AasGqgr=+3!I`Mc`<iGrz_87HBwLzwQ#siCl
zzfBryTa5=7JxAM&ha{s08T062x4z=*3h&xpk9*hmM`I{1V_7HAC6Q(w{P##tBq?R#
zFBGkkBdT+TQE&=oF*lpdv1yxiz&pk>tq~1=Y=AGc4u;Ai)`Jp)H1Wpv9)5q@nFUyc
zW-)Wla+Fc8V9>D)5Q5d`(36f-_hJz;Sz)?>?HIK!vuIP@Uk*5k0JVJ|q!>^EYyn0j
zL}i}QoI?%aFWI?iDr!?a7(KT2NjQuGDG3`vO--yBiZp*+kJK$g90wwRI@&gVBif>=
zP22d7(X-`pGj4I)ePq*&%{}sA{>ixP+gL)3f4BYeW@)z9t)EF~>OE|EO@sVPeQC4M
zh~}sh@HgArnv#DPCGKVPM2Xr?)|t1c%=H{BgZ%sUmo`s$wVZIJ<4{s1{Gay_RrF+@
zTIiSM@-vH?_4iyiu~D9gMad9<Dz>v;mnRTax#vraW}l(+tg?Kf_hPS+IPcyu*_wyK
z04OjeA%eN0Gs8@)kXJFi{G%O|6#3T9mi7tQ1Amoi3!i1$jw(|GpQTK8T6WGf=&a37
z^O?@5b_O{Y_{C0VU;ttgNB|E5yMWz55AX<31E5#g9$+u<D6kKBjHlzfNH;$hPiU_|
zk>NMvoe6bWr{r4r3J#rPlQ89g|CUy6W>xbHZ|T~ZF2d)!gs8}y42>!(_m3RPj;Ds1
z4VxOcpz4mSm7AY6a|WA-TYwe~%PbTf;gzGrPQ$reuDjgTT2~c?0!7PNt>-`Y?diaX
z=tOAJy)=u79Fq#6O_dz^Ivd6&)qj*+;j8_J$vgbl{tog5ZyPvBZo9_^juZFJ;0<rH
z8rA*S<Z<8ta6&@nRkzX4KN${h8G`ySF#IoDlK(W^eJlyJ@-2ULuzV3%l+Z&=#l8$h
z5K!0sDkL<O>LgLsH}^D^pRnRIy8z`i3E`D3!It<7M;d$2LFob?0s#p<uyI)1VN&0l
zk0CVxpSxcj@q5V<zcaE=LsoXpZ5Ta80_qa|GX1-$Fmd^3sf*ehFc0yGu>+yEq$qrO
z)0$^~XqWiP*c0R`|7|RxpNFCaxCDfO+kgPn9zQ@Dbtlr7h`;7+crna>8Q)#SZr0sO
SWo}JUp2+Clpwv!gF8m3y_J)7}

diff --git a/litellm/__pycache__/timeout.cpython-311.pyc b/litellm/__pycache__/timeout.cpython-311.pyc
index 3817a2602c5a15c756f6cff1856d6a4a2aa976a7..14a356abc300ab281fe54fc47cfb5e804879a83a 100644
GIT binary patch
delta 101
zcmZ3laZ-bKIWI340}!;{y_QnCk++A1anj^jENcV=IvG=#BtbG@(74%+HJX`QhLM5c
z$B!Qk3?En~Z{(1fe1gM+Yem=w|0}vK7jy$A3vjOGQWJBuWjJVS>}0U{IA;bI0M88`
ATL1t6

delta 73
zcmX@9v0j6BIWI340}xC#yPD#=k++A1v1#%wmNlEhSfiP_g%}wae*E~+!0>@(@<9%n
d$u~GWCTnr7<&qO~G-WtwYV4%3`5|Wp7XVd47lHr)

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index acbc91552ba6f7d989ef951f7e45c36b470fecc6..a01a26e20832148de53b3d844c981eba8e8c467e 100644
GIT binary patch
delta 1855
zcmY+Ed2AF_9LML)?#^!a=*dklplrL^mZMz?6fG@cOAtz1=p|v;?wi8WJ<8izuw)lh
zs+I~Ac!~a_e~=iI07e{)*%*@27)Z1x974=MV?eG1qTm@p(BE&?11Ix7@ArPc-@Cr=
z&5V5@UOp_k-gmj|0z4PTF42MguE}KW@#42g+EX;UpLAByr}=5JWJf8TsLZ94piG7`
z1<F*Vn&v9i-A<a<Z|Ph_^ZSL)XJ`TNSy~9Jq0@ki>2zQ%odK+)WwZ#!Tta6;>pAKM
zE)|7(ffn}*wB(Q|2uGl#p3di~2jcaVv<RMZ?PalJ_C`I!AB(DyknHOX$X&ju!uCV5
zo)(UUWWO&MjOwoL2=gnl8VD(onCh4Q84YSH6cd!6T8SXM;e1=5WY9^)gcLV1F>V@k
z#?4Th6ViZ_Ppilw68BaV5G#LFG1Fp86I2TyHjT@U3vp|l#KpMfsQHiu7Hr`MD`vUU
z;%RZ)Hjpu4kK6eebaY@x&;R^Xq}I5DO7EEcVuwJ5gk!s-<?mWsUP(x<g_Xhzbz4Y@
zMtvKUC_4g8b`<d;U+3*Da~TOwxKRLh%t!#}c2$Y0DJ@u%-}g>;d<yQU3s6|dbLUNW
zd=BOuP#~I4__BE$-S-5ck2Hu-wTlW?mTvLB!JiGLJpAOmfvHusUg4D2URPo|Rg+&g
z!!(jdz#f@lsVmXeSAAn8XZV7e-po;$fqjMe8gYRqYX<YrfT?Fim9WbCWR>~gkeE+Z
zQh66N-{SrZa-QF8td&k<r%_(EYA3nEkFT0#OM=5gU^};0c}bGzH%%j#`Qj#z`!~$|
z9dJ0M8+(Y3O^6_3B4UgWH7(7=F?5>}rg98!4SUG1H&G(;WzA*v*D#A=Jk!6?r^>yN
zsJbGhKl=`2R}lpW+?j6m`+C(FQ*?7I7NEKXS;>a^$>zMQo0yOmjBJpVEx6Si{7&-}
za*InX^^Q9@!d=8|-rmwclG?$RsToqr%}2WCc)rJ=q24cOokt8Ien7xM2<%6M&U5{9
z=i`>x1H@&-ABabYD~MkaV~EFyKM?{PJ2Uau{j+z7XqpjO7->QnGuzN7A<__b#J>tL
zD>e;L3~d)+x9b4ifqRa|{C*IQ-9Wb!Q&A6=gK%N<7FzECxR)+0e?kx=4KuP4;}GK!
z>4;K{Cbf6zPAUBa*els)phsk%@baE>BzYiMDLF9Fp)Cw%6EenEgj7<bjfA*FN_eYU
z0^jluwVHVNXR0Ts1XIfpZbUJFXY3hoH7r)F<;3dDMABOOs!WohftU06fa5<oB<@a>
zk_EgaQ9`^to@kURaiB`=T;h^RYDBwHJNC*bA#Hry&M#am&|@r_uy(DfzmJf1?aS9D
zn52(!XD9fr0dI91CLO@G0+w!*<*@RSD$A@EeF@yKnYCjm-5&|{1{F0B3CoeLO>8}H
z+m%OpwAij$qyfKd-7d?a2#p05S=ODhygB9zhJ7Juu?`5c7ZErM>p|ek*d)XR#A@`d
z=646X$ZD<e&2mDj`1ak4q($hj=EJ+kftnY0cag2y{5{XWFKZ<}9uuO0xnT%hq7hXI
zZ-MN=y~PEFa7(baAIBfY44WG0Q6oJ{nAP&LdrxP6i;=5{R>Ua(XlPQAae0g@^8z~C
z5XQYV?p`4_C4Oya&V2l(u{DSQVgtgs5JfmE9xz*nKI52-^D#ELl`r3yXMYvK@Cwn#
o8ryeMbf0Zrb5F=30dlrGkTT%I+fHiX*N9$dbckz3?&lBw0{%$#Z2$lO

delta 1752
zcmY+EYitx%6vyYzPNzG&6xtFfrP9(uw=5J&yOi>jK6oWU?Y2@vV_3Iy2Poa{mb<eE
zWETu<w6r3Hi}-~Peo&MYUmZ=qC<#P~shYG#N!f--K}dX{Bw!R2;{DH>sFR)FocllL
z-nr-8von59zV^19^;TAvOA^<~u`4wDL{>g|=7X}+M4Lx*lccSNdiiD2aA-cwGpgte
zap#MBrm=E|lNKa1+E&pgl2Y4hI!o|LS}3@N775nU*@AVnSg@YXrzK*o20BNK*3!9x
z>ttzzL`#zr^_-O@X+YfRt8G0k6aEd9v`F3sKR|5s%2vPZU);n_h*3%l$GQS~FcgaF
zN_vh>RnQ^sd(A|WUUI!AQPS(8azcsAOpeRFuDC-y9SNn!wS!Pc(rFUDXZ=jx>B}J*
zJmD*qRgLc{%$ppS;u&#rUO6jP8*uQid=*)l@yxh-KvXH=j61ouq=;yk;r*XZWtth+
zsCq`x<!us`652tn<)5eZ{1QTP9jshzFy-trqS4?^BgzK&sp{@IZd+yp!e?j5&LhfY
z8c{QyCmZ6GOG~tmgdBAX@a;=W+?Rw=DN>Ezmw4aOU6V!>V^lG|Q4C@5{F<JE>N=lv
z(dVl7I4-U(tS@$4Dj>qXRGd-ou|BRD%Os!iSJ!lB4~uo!7r@uRD4$l_D_pPF=Bc(i
zYJF0hoI*x;eN&zKDME(%(WWCL#k2es&LJUqiP+#;e>EB62mHn4DnH}*dVWLX?*ivC
z(&;^bup0;gdB8aT$-gceYowe;m<D1|gRuv^wwV%%zuCO7{5ql-{Mq`hpc&|nM9t0V
z>ahv<z5|MYVg9&zCb_`#TV|Dij~QAhvNK@pHNqyliGZmPrv%Jh%m~te9toSqK9k+z
z?JXPJ=ri^saGQVFvXKl~k6H?*sA-uO>XqJ~;k0XY57r299JmAg0^9}0_&&XA8B$_@
z09S#Bz@I=0xDSj2e*uqx#{l6s^u>oHajHxPav=Hv`^_pe1;_-Pz;$?0A(nyR1T4F@
z%VFIRNV!pnXiV2d;nJmWArwWji^AL(-oo$=0e*t6SoEKQK+i!$HjoQU1ttM#Xf|X$
zphwingTiiPdj!2Qvv_aksOX4iLrYW*k(zZnluO7sKi6fF600!GRpQ};<~*YCadRap
z<9RV}wg;gLfVn^^Z;IusrM78BV)Y79tik;?4%IGzkN+C)asOKjg&$6olN$bBVjl7F
zvBcAA6=qgh3lCm#sGDJLwzR`v60(I~IWnB}6cl^Egtc1lB=-~2YL&e>-J!mZv<7*@
z(dunm5Ofk_U05k+AP^RP69}+6XbEJiuvWMx>yfVRkYRR2!huNpZnhnwhAs=xpc#xR
zB6ha8*mTV7uttwo5K_TOU!ixa2xZ#<yo7ZC*ch7u*zx$Rv1Y!pubnho1AU7KsphGd
zR~4*)zFZ)sgd#!88ZZOB&J<pGtewQHQ^!`5X^zImM(;-E5ss8fBc>7F%a-#yuap(p
z^{c}wI%Zu!gwyQkG$Wlxn62Ty<CiyGgU9}_wZgKy@j9#t09}@DEW4-dPTB-*D`2;?
z-IB90RQS67%4PU^u`NIcuoJM`qy!tW4}3c``vC2ev-6Mfcl&3#UJx<jCq-{s-}c{@
U-JegrE9H<5GV+|xC-}p^0SHabsQ>@~

diff --git a/litellm/main.py b/litellm/main.py
index 6e54d658e..4c3d75bf5 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -36,7 +36,7 @@ async def acompletion(*args, **kwargs):
 
 @client
 # @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(2), reraise=True, retry_error_callback=lambda retry_state: setattr(retry_state.outcome, 'retry_variable', litellm.retry)) # retry call, turn this off by setting `litellm.retry = False`
-@timeout(60) ## set timeouts, in case calls hang (e.g. Azure) - default is 60s, override with `force_timeout`
+@timeout(600) ## set timeouts, in case calls hang (e.g. Azure) - default is 60s, override with `force_timeout`
 def completion(
     messages, model="gpt-3.5-turbo",# required params
     # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create
@@ -44,7 +44,7 @@ def completion(
     temperature=1, top_p=1, n=1, stream=False, stop=None, max_tokens=float('inf'),
     presence_penalty=0, frequency_penalty=0, logit_bias={}, user="", deployment_id=None,
     # Optional liteLLM function params
-    *, return_async=False, api_key=None, force_timeout=60, azure=False, logger_fn=None, verbose=False,
+    *, return_async=False, api_key=None, force_timeout=600, azure=False, logger_fn=None, verbose=False,
     hugging_face = False, replicate=False,together_ai = False, custom_llm_provider=None, custom_api_base=None
   ):
   try:
diff --git a/litellm/tests/test_load_test_model.py b/litellm/tests/test_load_test_model.py
index 1ff74d580..373da1a6d 100644
--- a/litellm/tests/test_load_test_model.py
+++ b/litellm/tests/test_load_test_model.py
@@ -4,5 +4,6 @@ sys.path.insert(0, os.path.abspath('../..'))  # Adds the parent directory to the
 import litellm
 from litellm import load_test_model
 
-result = load_test_model(model="gpt-3.5-turbo", num_calls=5)
+model="gpt-3.5-turbo"
+result = load_test_model(model=model, num_calls=5)
 print(result)
\ No newline at end of file
diff --git a/litellm/timeout.py b/litellm/timeout.py
index 37bbbffc1..81d99e7de 100644
--- a/litellm/timeout.py
+++ b/litellm/timeout.py
@@ -38,7 +38,7 @@ def timeout(
             thread.start()
             future = asyncio.run_coroutine_threadsafe(async_func(), thread.loop)
             local_timeout_duration = timeout_duration
-            if "force_timeout" in kwargs:
+            if "force_timeout" in kwargs and kwargs["force_timeout"] is not None:
                 local_timeout_duration = kwargs["force_timeout"]
             try:
                 result = future.result(timeout=local_timeout_duration)
diff --git a/litellm/utils.py b/litellm/utils.py
index ca0c6d83e..65cd96a8e 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -302,7 +302,7 @@ def get_optional_params(
     return optional_params
   return optional_params
 
-def load_test_model(model: str, custom_llm_provider: str = None, prompt: str = None, num_calls: int = None):
+def load_test_model(model: str, custom_llm_provider: str = None, custom_api_base: str = None, prompt: str = None, num_calls: int = None, force_timeout: int = None):
   test_prompt = "Hey, how's it going"
   test_calls = 100
   if prompt:
@@ -312,7 +312,7 @@ def load_test_model(model: str, custom_llm_provider: str = None, prompt: str = N
   messages = [[{"role": "user", "content": test_prompt}] for _ in range(test_calls)]
   start_time = time.time()
   try:
-    litellm.batch_completion(model=model, messages=messages, custom_llm_provider=custom_llm_provider)
+    litellm.batch_completion(model=model, messages=messages, custom_llm_provider=custom_llm_provider, custom_api_base = custom_api_base, force_timeout=force_timeout)
     end_time = time.time() 
     response_time = end_time - start_time
     return {"total_response_time": response_time, "calls_made": 100, "status": "success", "exception": None}
diff --git a/pyproject.toml b/pyproject.toml
index 78f9d4712..d0094c107 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.382"
+version = "0.1.383"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"