From 0b436cf76741adb173a2c5c15cd0f7df9ce8337d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 14 Sep 2023 17:16:43 -0700 Subject: [PATCH] fix streaming objects --- litellm/__pycache__/utils.cpython-311.pyc | Bin 107576 -> 107651 bytes litellm/tests/test_streaming.py | 27 +++++++++++++++------- litellm/utils.py | 7 ++++-- pyproject.toml | 2 +- 4 files changed, 25 insertions(+), 11 deletions(-) diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index ac9c4d897e842f568e4dbe023085aaaaf057ca3d..c4098f31b1535a8d4ef4736b71e1619e5a69dc6a 100644 GIT binary patch delta 2838 zcmb7GeNa@_6@TZwefz<(?DGA&tbDsD5iGD&X^<}y5l9WTRw^Q}E3$n0c#)5S*k~g< zu~lyH1){Jdq&8|A&CVQiIgCB!!`xxhV<+JMnRRL8wsGejb#=%sAovCE)U}qbc zgPfgna1K=@ovpNkY(H}l?5Y!nO1_`zgi3VG&i)=7anSDsB^TJ3oOp} z@ODR&SVIo^50QPw;0U1@jaIer2zJ6Y3i%X@NvTaAX9kxN3G&2B3PAf)}opI)`nS4GB7R4Fh)5JiZ#4t{yPsXUr^X*z!Ao(aFRHHjV zw-iL)iOYuoQgbDgTqx0zxOf{faqP`dp2JaAf;|ut)T{QT(d3?m37dPdv#R z_@n$;cy`dv^rA#Au9l8iw(`>R_cCeBk>!OV4mE`!eh7Lvl!mc&ARO4jSCEw=FD*ys?Yk(MO!Z9oz=R# zHg4i-^J7Yl%sbgp1%b>>6QF!E(63O--Hc8sv27dx$kc4g##X~y+(8( z5e^U6(%PYf^qAV8s>4gCy}GWZsj0T6u2C4Ip1UZVryx;yp284?Cn$_iIE>)YHMQ1P zJB~P7MByU2T#y3S$lZc8)0)MP5@TV6X^eWkNZ}<4%Jl! z%&7XN0R~CGv&FZ}e>;=Z3W?8?l*R;PWM|`^RgSMV#^XL_8cPEyB_3UKSG9Nu{{-xU z%bnJAPYse=*Ph2>-F7_H!@}$2R~@gzfV->n&l1BP-L|L0^f0L_m+l~PFv8wX*vSN$@iw6Nd3a}$E#C*_+jp< z;7RxYiATjeZ%$M~DuI3pm_aGyZJO%?3QN`AKOYrE3G#BNJW#)&>Li76GBBTlNpH;G z*CD6Fe>ty0PBT6aMcZSa-}6&@gSglK>j>|Y#}_AKVL)n*1p!<#($tUum*tWINMo_@ zEP_vrP%KyN2CE8)eDYhciv>~spb~;K^^f2MYQw}!&2G9H_(NR!75sTyugV_OqjJXh zbh}@qPh_fTk(Ab>s$=&u-Yd0YsqcUerjy&tbbvKj3aEv`VES?)e2TE!3f^b`QK`BX z#=?~Wg{^Hk>U;lqO^U07M6CSwI@qeekAL&Wm<86Rib1lf zv+aGeZPsk-FQ3cK?JuA8k9TFv`X|l$e`PK>ZpMG@_e*RG8VrGbdlII8x-hsbkGU0A zvL3tN=t&f`0e$I#qb*I4a)P`oE%9>WW-N$VOR8}7(^8sQ7T z>h5TUVx-~+&5#Pem2Wh|BDYk3>8-<<@?$wbfGRZ*si_mpAW5e?;R-yEv%8>!NB0K-&EinHZzfkOde|WT0j(G`U;6F3j?Nk5& delta 2758 zcmZ`*4Nw%<9pC>hyLWKh2M65A$8ny=5e@|;5t<>W;Q%!VB#1SQ)llHB;GNvzvI1I= z4>PUJ#7WHqSBMOEkdQAYjak|?ZIfoS$pD%Rc5G@+O>LWL$INuZHm$~))V{ap@zJL5 z=6>)0{_p)i_WyhD?H-=t&P;QLW7gz^N9XY0;L2tq385RjU2y@3uD99z3E>gwjY*@Kw2pbjP6wliwM zit=rilqI%Z;9W>YFr$P!NjmJp7ylAk*kvAEAZQF=T zvwl?y#(j3PJ~!%xqTE^`55a?02JZrNp(TUE)zkN0y@q3UY&aQ)a z#M%oYMOF$I_2__yi4XLiLxV20nEtPibV5`#wmUH;DiUH?l%Ht_3#!YsfdOr(w16pE zZ5v}2N=N`bq5OX1^}>4e(bZIRv_Pf84EzM18$x$09k4*sH=7xwW)G>Qn~r)$ZRMZF zE_oj5T`SN|=T`K(^I8JmhIL`J$M3Ee_OyD!fKKE*4L#^~jw9`IishvBNaZ!N^^!UJ z_vSUfHLtm7&bwsJJF-PeUmXiNMthJBx?l3mQ$&(V*W@$3E~v_)&Zechim0=Nid+;| zmdl)@>}A=^IXjDI&tVl0Sa#pFadIQt6P+FY38vhym|HqpcMM}Dv-xve0p>6A4up+GH2MG)jI8I7{al81I_iG4v$~p$55y10eBmYxpEEv zm3IhTb*;!7Csr>Jc$vT?0Ukrx==XHiw+8S{SI>Le+XDV}4?iKTD%+=qH__nMO3hEn zQo=Pqj&5wd!bP3|%~qvB5dEdvTD{hwB%*}>C0RWo^;f@2 zbFUEjl$8I(6kYl{k^3>=i&xoqi)`#7Qwx+h*+`{3i1tGQ2MJ6_nRN-Va1s^m>e4L9 z0ll-U944fH@7kdS5q;b1k3Pzb*TE@;A3L1cK;WQ>${6%8?s39N)CRT?n*R1)7^Zyt1{OYTim#wiB2{ho`b&0`*rX zOTV6aEO8m@iPPoF#*NdxFpmB-?Odic&5S_^{d}g}y0{^fd=~iPBSfNdnvyx@;{{Go*7(-wD!IG&&QGEXfzltK~_B`?RDS=3M zwEB8#B&c$TddSSup!L@u;K+$>H=_5CCZoOqR+_y2ek>-T`HSN*W`(1e2uI2h4-f|v zQqs*Uaa<&V^v|z4*r>ek-gE{Gip3egN2)5H&4BH8I3=51kW1sfq8vWg!zMXrJ2+!t zNZ$S(sHI`8Jh>B$DtkQs3fO={+;0Pb-=XXwcJUqRG>X#$j!p5NLZ$@bC@7o`664A#+;HZGudg5VtkKBbvAH?bk=m(}l4sZ;4Y+!4W#o=b? zWceW)U#!*^sDZPxzXi^NzKF%H7obht;)So(Bwu`%MX}lk7C0^Xd{6{4^0*If0WRWT zD^%cQ!>y15b296Pd(8bZi*Y>meO-Qvhus_uiE!Ex-2+ybRDRyMKu7p;H|A5sX5&yz&Ut(a#0qIXDDA0~~IL2y5V&d{~4e8bacUVR(d7_F_={bQpfDQI7n>WJ?aqts^k2 ZUfV%~W1jkEyrt;BTsq6k%@dFg{{c2?<6!^* diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 43308aa168..7a20c6cb2e 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -40,7 +40,8 @@ def test_completion_cohere_stream(): # Add any assertions here to check the response for chunk in response: print(f"chunk: {chunk}") - complete_response += chunk["choices"][0]["delta"]["content"] + if "content" in chunk["choices"][0]["delta"]: + complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") print(f"completion_response: {complete_response}") @@ -79,7 +80,8 @@ def test_openai_text_completion_call(): for chunk in response: chunk_time = time.time() print(f"chunk: {chunk}") - complete_response += chunk["choices"][0]["delta"]["content"] + if "content" in chunk["choices"][0]["delta"]: + complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") except: @@ -98,15 +100,15 @@ def ai21_completion_call(): for chunk in response: chunk_time = time.time() print(f"time since initial request: {chunk_time - start_time:.5f}") - print(chunk["choices"][0]["delta"]) - complete_response += chunk["choices"][0]["delta"]["content"] + print(chunk) + if "content" in chunk["choices"][0]["delta"]: + complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") except: print(f"error occurred: {traceback.format_exc()}") pass -ai21_completion_call() # test on openai completion call def test_openai_chat_completion_call(): try: @@ -118,14 +120,16 @@ def test_openai_chat_completion_call(): for chunk in response: chunk_time = time.time() print(f"time since initial request: {chunk_time - start_time:.5f}") - print(chunk["choices"][0]["delta"]) - complete_response += chunk["choices"][0]["delta"]["content"] + print(chunk) + if "content" in chunk["choices"][0]["delta"]: + complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") except: print(f"error occurred: {traceback.format_exc()}") pass +# test_openai_chat_completion_call() async def completion_call(): try: response = completion( @@ -139,7 +143,8 @@ async def completion_call(): chunk_time = time.time() print(f"time since initial request: {chunk_time - start_time:.5f}") print(chunk["choices"][0]["delta"]) - complete_response += chunk["choices"][0]["delta"]["content"] + if "content" in chunk["choices"][0]["delta"]: + complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") except: @@ -205,6 +210,8 @@ def test_together_ai_completion_call_replit(): ) if complete_response == "": raise Exception("Empty response received") + except KeyError as e: + pass except: print(f"error occurred: {traceback.format_exc()}") pass @@ -232,6 +239,8 @@ def test_together_ai_completion_call_starcoder(): print(complete_response) if complete_response == "": raise Exception("Empty response received") + except KeyError as e: + pass except: print(f"error occurred: {traceback.format_exc()}") pass @@ -281,6 +290,8 @@ async def ai21_async_completion_call(): complete_response += chunk["choices"][0]["delta"]["content"] if complete_response == "": raise Exception("Empty response received") + except KeyError as e: + pass except: print(f"error occurred: {traceback.format_exc()}") pass \ No newline at end of file diff --git a/litellm/utils.py b/litellm/utils.py index e3547dc1c6..20206461a1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -103,7 +103,7 @@ class Choices(OpenAIObject): self.message = message class StreamingChoices(OpenAIObject): - def __init__(self, finish_reason="stop", index=0, delta=Delta(), **params): + def __init__(self, finish_reason=None, index=0, delta: Optional[Delta]={}, **params): super(StreamingChoices, self).__init__(**params) self.finish_reason = finish_reason self.index = index @@ -2493,7 +2493,10 @@ class CustomStreamWrapper: model_response.choices[0].delta = completion_obj return model_response except Exception as e: - raise StopIteration + model_response = ModelResponse(stream=True) + model_response.choices[0].finish_reason = "stop" + return model_response + # raise StopIteration async def __anext__(self): try: diff --git a/pyproject.toml b/pyproject.toml index dd424cb157..5e62451466 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.632" +version = "0.1.633" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"