From bb3c2fefdb5818d7de8547478868d988cf39fdcb Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 26 Aug 2023 19:00:53 -0700
Subject: [PATCH] anthropic and together ai streaming no longer waits till
 completion to return response

---
 litellm/__pycache__/main.cpython-311.pyc  | Bin 28957 -> 28867 bytes
 litellm/__pycache__/utils.cpython-311.pyc | Bin 67694 -> 66078 bytes
 litellm/main.py                           |  80 ++++++++++++----------
 litellm/tests/test_streaming.py           |  71 +++++++++----------
 litellm/utils.py                          |  39 -----------
 pyproject.toml                            |   2 +-
 6 files changed, 78 insertions(+), 114 deletions(-)

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index d7813a16bcb4b74aec01a510ebbce7242f055b96..7048c99f86c3b8e4c7c90fd8496fe708db2efd7e 100644
GIT binary patch
delta 2663
zcmai0U2IfU5WZ)-KkaURD8*7>Eu}3#m(qXQZEY##r%L5Vgkq&omfPLC?WMbWFX!H^
zlm~uHV~BrI2MkI;qYpr&54a>I8Zkh8&<GOYLGK%=QPGDs!h<9lXU>-H`e3@*d^=~(
z%$#$+Gc#vT-(}bCu)MQ*d0q)W$Hwo4%{TMrv6hQVtC*Wf4oUUEWO+$8lfngKQlAV*
z*iyC}&M_sgDqMn^72$b^)llqMmsb<^qg)$)9I-ALaSSoW;M**Ja-p+_xhimhn&bi3
zu4SwmZdYI0P=g9yk09spWhBc4%L%FoRuBlIT9WDz?1GpcR--{BswrkLrWog)ybEsD
zyvk(guU*)@feg1G$PV5^GAU#y-$pWpw+nTy4Qg>E8sw^J=y6l!+u`HdzGdXaw~~1l
z?;<%fR)zvr*1gEe;gh;rHXDAaTf!<KzkWgYw7Z?H7+TO8HPop<w%g!(%&;EqOT>!M
zUOb^8T??tE?X=Vb7FTH1e&F?%e>aJH2!aTv1EH*pmy==u4h1S^ZDc5hXk1AUfu@E)
z{%RCuCJx)f7sJ5@bG|sUFRBmXnFoXVz&_bCbvp1gR`|Nf`DubGXlmT*PVoe<G(P8>
zri|}^KN|PWxFz3r=`F$AD;wGq6jog63$(B73|7T`9-#SAf{a~|CF3>VY`Qh$imU|Y
zHWU9H%}tBatmblce?m2_)Ezi%=RQ-9^DS_^`NoVPecO@hDPG}1VEfDtMutOo%hs74
zUdnLzv88urhxaoa##%p}(IId>?SORt7A8&|c2<H%xlVGH?bhN(!ZO8BTnlpk9$b2&
zd!e9_YR;ZLT`1~-kZM>OlDG%_ZGQG7w6@i(>_sE45>V0bT!Lp1>>TA_f~&z{bu?9R
z*#?e6N;HahaH{Pj3&6_uR+fNxdy9`Umiq}t;bMC;I|{$Hm$M=8ty0+=uzytr130m2
ziTgO2oPaB<#{BQl8Xk-p)<Tx;Q?!QqplAvvKLtuhV3s(7DBc$GYK1c$Rg3JLcp?@Y
zP?ea;*OGc2!Fqya1Tx&~Xg*v@qR7kXl(aol#bakrHw<xLQNcwatn8>t<fQvafTCO`
z5H<7`O14w8)EM7SN>1=4DQ78hYl!QH77B)<ni_}mt5*b$()3p}?Fi}oItjCf5~6-1
zN(h-GON-wivzG}(=lK9~J4<{(<0-CcJB#NCu5w3Ij{PM?5^}V&v$XgQh0PyOOx053
zk3yb<*E)9;lv50GB?K)n*;&H2CB19DE|{Ay%`2NIl8VYEC2vvweS%}_CQ6dOZ+(!>
zI^fREfJ{EvZv~MHDbimEwuh@0A62>O#8DZ3csZ=;d))Om4j(|Mua8wGFZbPNPP#T=
zt;O)o?zU#?e=@V($a^AcVMpJK3fmn?SeRe`?TI*XuyR1LLPJ3)Ybj-3*xfJpP;11#
za|M6+=LC--*v_bCTDIFv48$}`oYZ!TTyiUhp~l0qH@$u0h=Rsc(^Ljklb?Yf`m69&
zbM2}26I1rAltDTVcrNzAhsiB_mby1p)1v7Ijb1SR6haRdGaG5ViC{AVJ{Re7qDtb=
zlGc&Rq8_*sF7}@#wWz%3N!poV&;@^n>vkO>&C~<-5t&de%SGHra=K(iwTd(r3yCg$
z5l!rzuGoMY4&!6NljJf=@CCs*98yaUUm)==!7YMI1Rew=P1K#p`1G;GiN$07Gg<yb
zAadgYN~RZqvLW>m&zqg2PH&(Pi7d!MyUh+sIwDQvuXZIDMD}JCjiaCK!`uyGY7gl6
zgx`eMwZ*B_6+0&AgR5Eu6HDG#$kvNoa`BXBNRLOf!JwgF^%bCHD&6T+$TnrI?VB=D
z_@FtRX>e$Gu4^7{N`ZHWd-Ftg)$`9NKrj3=+`3>6&hhy)zD7ZuA>A0|E*g)+w&>Nv
z;zgWG;}V+3XITjc!#bAW2(KWGD}t(2PI9UZkxQI-NL>Uip88uriwnXYX18PaOGMQU
zF5@hICG#MXHj#yYzP^`=Z>GP=)3MvE)WIAR(*J=MKGb7u1PYCswIg2UoscqcJ(Ju;
j6J0DrC8j4m4u|8ugfQun9NA;V6JB9rz)|C*lXm(q@T7?r

delta 2766
zcmb7FX>3$g6n=MRI$egA((06!ii4%gSj)7fJEc?>0R<VbiXEh{Z)R?19=w@1-1mm1
z{D2^C2{GzHB{3wJXpAeGaf`;d#keJiTVkGxdo-G8N`iY_&bd?Q_^WX;^Ub+ux%b?2
z&+_`0?D$u#{GsylG6g?B9sD6~-&tPEhM#I~WtFLxNJ}ah2{Of{@Zuep!aR>3GOU=j
z!d}+QE`t+HD{qU>KdeNS#A^{lP~lot-X5<<x+8uOVrOp5m1LXUPF;fAs%aQnN>c^5
zGG^N1s$~nVrJTzAziOMr(bEYhV4J9)<U-YSHD;zV2Df!HJ#By*xv$+Ltf&!vZqE%C
zZEIm6IM(*GuN@g;DMGLmUL2}xkn0X&b|N@)Q)ZkSs%Ge#rKYsZ!yd5~jxITZ7g`j$
zc;S_#7$$3%xQ=LI1U+Je=vEjD)v#+}cPP?AI${e6-C`}#g_Z>j_$_pZEr$K=A$9>g
z+TL`aiEKPhwVB~*yqPHyT3Vmf=x*|0xv5j4Wpvfvm*LafJKl`I7})20lP4r`E2lV?
zH1!y_9FNWS+9F22sCfZK*5+}iC~Ij6?o@AC6W*{M!X1(NF;W^Q&=4#aLLeX(6C)0v
z_cl~+WJq#ale;1bzjyRjT#jUr$xw`lW*F+U>ROTI_zW|Vzz0#y=#F6VbTr`I&IaFA
zWW0f(4Sw$od#9tgY3X&oJV~Oh@Z{3*SzYpdG~*IPx;h8<kS(1r_8|mI3SL3J_9Th+
z6BKwsnoP9ATU{T{dLpZXPrI8~6S#W1z4>^<@}4>thU<Gy&1&r%!KhA$5Hh{j%x<u&
zz@WS@JiEay1qKiGZJpiV{sMz_{V&XF(0gje03-K(oWkdF+*x6WT|~Pbub$3iZA&iY
z3t)KfhrI*qFP0Ib;&ZA@p<{k%hRXNGct&RWID9cs&(^_T153I_QAs2SsGh_eg3Sm{
zsWzDvT;0j{$%l4IGPsIqh9SbBE<48h;kRY|Y!7q}_W39|#X^GpkR9w{H^O6sm#`i1
z-XLdpLHF_ob~lVJZ}Q$lD)+)o%MY{rplqm)9f!uD-u(Ms3C2+CBC(QGRv|d0>1;}k
zVt=y4YGSS-xB{LZ>Z$7@vYQ}4AhRu>Xz=$?jaS|z2@6)#HujVJ6#|*_N0D+oy3JF<
zBBo97CXB8a?!A%7cS-Jwnc_Vns*+h*TN7FguWsl(ZHv>Sc7#ASp$Fjfiuv7l(>i}S
zvMY-vHCvUv&hZc}?j!b`-CSUs$P2`AxM+Bx|7#LdL^X@sJbmtrM_}{twJowE$e<Do
z;LrM`Mo|tbKZltzr|P!8BYg)~Pq{YdP7S{^w`Pt~8<?7>R0qx|rPUQ@32t3A<<FfA
z|6Ia`;la_~D%u9SX#oGkB7#5Rr%`_co#WILH2FQcox~Y|lYm>|AC&$A&5;Oe${mWF
zWgcp^u&culcigb7XB-zn=6G?UWDOhJb`%+ocOq+J?)<m4p<r@xRI_7A6*{^WG9LuC
z2S;Q>E)lPzRJ=@ZA#LC>bjx<URyLZ_ZJAh(C$7hkrDd>l$Ae`vx1i|JrYUYwr?tch
zIKI7=x#8{YZT0eQm3f7k8;S~SXMg6Z#+tncLUib51Z@N~ro?)L5g{cTX?-=pCIp-Z
zGo?aRBCa8>E1z66vL-9f57*tAk>vvc@i6JhGK&zkt-xZf0?)@o2M!TydXha(I#jvA
zYT+Y#ra)zF%G{Pmw~{o6;FPXOMR`1qQ$RdJOxbu}CF&%J=hL*KLC(*q<x|iSPms<x
zI|)A~c$T0T0Rt>sh%B0!5XtDsIDJQ&-x34}enqg#5CZ2va-un<=g$<?81%~A`#kNg
zmTvhl=T$6;sA+NW0c_D5=gO3q-V<ztoAn7yyC3y=Y<wIU!J_<yd@(aUp(oUgW+%m5
z6y$r)H>Bc{sqLIo(L};3%tL41v~58)EI`MiFt&4Rxy;6;;w9R06kglezrc?&QBUhP
z5U`rf%sx>>>yzL&UOONMAy%a#<`dKs)FIGZ*5Yd1#IoKk8i`XhfvkKs%c_!DD5H{p
zFmgkT(*|<vQ!M@{8N@D<HTaNdMJaw2>vp%IOt~3jpD6y*xs9pMSf_`%rj-8!6|mk+
zu`E1iE?Jo^V`WoH0j_3>w|Z&~D_}|SnPQjAbyh(*Q>3^`4*RFdq)Hl^GsisS^S|l>
ByaE6K

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index 29b9675c41a5aa40c319cd5315ece028bf37b35c..ef3e1e402e981cc3246fe1ba05c5e51af6d700fc 100644
GIT binary patch
delta 231
zcmaDifn{C`3*T~HUM>b8&|3Q{WwPKzJ_*JT8`Y06@q45wwlGHdrFf<&18M(B%p#Ko
zm|0jj8B)9^i!u8$dQVPa_7?R?v1nmTj7hO<VTg)Nu>z8DDZZ1BFxyM|rMLq1rKY$6
z#nMvTfn<6LV=#lJ|7InYk8d}7{5Z+XC_P!|k2R|j69a?I=Ab{7OpF<u&-{DBqyZ8s
yk_S4fNCQNu1BqW8Ho5sJr8%i~MW#SmMxan}<@URbj45ogeA<j69~dwRux<cy2tT<1

delta 1433
zcmbW0ZAe>J7{|}K_vYo^m}@lVWhEx@CDG^_*BNbnTf4>7t*z~#FxQt=g6FC!YUCzG
ztqBn(gAG5Z_d^47Yv{%*Y%oW@^vl2!Ka{~(A{Pw8V1tgW;>vX|C42d>bJWz9m3`RB
z&F}u7bDrlp4^N(aILkboXZYKCJqPf5IJ1Rh%K~qOkG`l|qe@>xRkQHoS#%23;P<&j
zx<WG048sC+S~AnU@Qj4$)23c@Y!(JDqqbRa@h#Mjc?F%7*68Dob0~o01auPb4Wb^*
zA@nNdFoNfS=k>IW+4<Q_k(>9I8oz10@#~26^@DbnLNHFVAVx{fElV9812Om+^$88Y
zr}(?KeBf<}z%=~^xIw?e#F-c)UEb=_sbVVe8je(@RMJAoBpq)pRmasabxOT!zBgQI
zL#pSBNkCi^Eq!c&t6Hh#rgbo`jcJkkzzJ!d(AonXXAz4yq+6omT#SuzsH79;LTaRs
zaUnPaL$D9v@*o6r6Ey&W3t(<~lD+^Y;9Q~=NS}XWUfgbmlr--%ODQ8Sy}50A%HpNe
zc1faoI2s*~v^O;c$3xy|cu0s23!-;&>UvXXY)}w|u>rwDi621##3h3J7}wP(7!9^N
zMFtlnb`X$mi{BI22zCe@#J`%TKd|y}B8R<19l@>QQ9>*s_K(EY6SaHO17h;n^yoI(
ztC8?nKe=$P00Z%QpE#5eRZd{U*wpPhMxtV9Y-p1u0sSa65ZyEmgh$6mglGt-5rxQj
zcq}4_f8zuZ9RLa7JXm5KE_Tyd>_UGq)c-^*&MX8+adOdtuj|4y6{*1DC;XWK-}0sn
zzGj`T$?-0kcg^(VIq!^Lah9p<dE=oQor|3*df9WQZ@x2U49LcS0(iC(BeBb8<6cGH
zUcEH3VQ*fyH|Ojvvb{wCdOZ&NFU+Iwj1~px$|@C#GganYp4*phU0Si^Ty3(eEy?BG
zFKoD5*WIlvlR5W^gzP@CS|+>wN!<^%^&g*Fp3c>F$h93wHt#)phn2mrtR9iQfz?6T
z+ncpE=3nf(SDp3r;wQ=HeLYD&HI=g-$vf(odsePw8@sa&fvo*x*80>_P$-BivD%VM
zp-@mPq&5|Bv}Wxs35DKk#C{8QNMC;2Z_cyZQf%G-J|#62Z!nM_p_s`DO3hMSaY_?P
zL*IEk3pefE7VzaEe*=BjqV}7ayOlK2RWCp9=b4NiVk5&-{&G#mpds2?j{P$YoZqL(
z_-LY!LZXimy=Jj>H`D2oU%sMYm9+7DhpHZr&p_JpXo`ZN^g`i2b(l<+NVY`>`QOwN
vkfHoj+c_FW*G7bw#14FlhZjWNO@HtgQXe+dK)V7S;|y$s6$<z_c9Z=BgC&NG

diff --git a/litellm/main.py b/litellm/main.py
index 3ff22ecc31..f0238a6e71 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -35,8 +35,6 @@ from litellm.utils import (
 )
 from litellm.utils import (
     get_ollama_response_stream,
-    stream_to_string,
-    together_ai_completion_streaming,
 )
 
 ####### ENVIRONMENT VARIABLES ###################
@@ -542,46 +540,56 @@ def completion(
             logging.pre_call(input=prompt, api_key=TOGETHER_AI_TOKEN)
 
             print(f"TOGETHER_AI_TOKEN: {TOGETHER_AI_TOKEN}")
-
-            res = requests.post(
-                endpoint,
-                json={
-                    "model": model,
-                    "prompt": prompt,
-                    "request_type": "language-model-inference",
-                    **optional_params,
-                },
-                headers=headers,
-            )
-
             if "stream_tokens" in optional_params and optional_params["stream_tokens"] == True:
+                res = requests.post(
+                    endpoint,
+                    json={
+                        "model": model,
+                        "prompt": prompt,
+                        "request_type": "language-model-inference",
+                        **optional_params,
+                    },
+                    stream=optional_params["stream_tokens"],
+                    headers=headers,
+                )
                 response = CustomStreamWrapper(
                     res.iter_lines(), model, custom_llm_provider="together_ai"
                 )
                 return response
-            ## LOGGING
-            logging.post_call(
-                input=prompt, api_key=TOGETHER_AI_TOKEN, original_response=res.text
-            )
-            # make this safe for reading, if output does not exist raise an error
-            json_response = res.json()
-            if "output" not in json_response:
-                raise Exception(
-                    f"liteLLM: Error Making TogetherAI request, JSON Response {json_response}"
+            else:
+                res = requests.post(
+                    endpoint,
+                    json={
+                        "model": model,
+                        "prompt": prompt,
+                        "request_type": "language-model-inference",
+                        **optional_params,
+                    },
+                    headers=headers,
                 )
-            completion_response = json_response["output"]["choices"][0]["text"]
-            prompt_tokens = len(encoding.encode(prompt))
-            completion_tokens = len(encoding.encode(completion_response))
-            ## RESPONSE OBJECT
-            model_response["choices"][0]["message"]["content"] = completion_response
-            model_response["created"] = time.time()
-            model_response["model"] = model
-            model_response["usage"] = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": prompt_tokens + completion_tokens,
-            }
-            response = model_response
+                ## LOGGING
+                logging.post_call(
+                    input=prompt, api_key=TOGETHER_AI_TOKEN, original_response=res.text
+                )
+                # make this safe for reading, if output does not exist raise an error
+                json_response = res.json()
+                if "output" not in json_response:
+                    raise Exception(
+                        f"liteLLM: Error Making TogetherAI request, JSON Response {json_response}"
+                    )
+                completion_response = json_response["output"]["choices"][0]["text"]
+                prompt_tokens = len(encoding.encode(prompt))
+                completion_tokens = len(encoding.encode(completion_response))
+                ## RESPONSE OBJECT
+                model_response["choices"][0]["message"]["content"] = completion_response
+                model_response["created"] = time.time()
+                model_response["model"] = model
+                model_response["usage"] = {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens,
+                }
+                response = model_response
         elif model in litellm.vertex_chat_models:
             # import vertexai/if it fails then pip install vertexai# import cohere/if it fails then pip install cohere
             install_and_import("vertexai")
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 214e62d683..b6deb20984 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -62,27 +62,22 @@ messages = [{"content": user_message, "role": "user"}]
 
 
 # test on anthropic completion call
-try:
-    response = completion(
-        model="claude-instant-1", messages=messages, stream=True, logger_fn=logger_fn
-    )
-    complete_response = ""
-    start_time = time.time()
-    time_since_initial_request = []
-    for chunk in response:
-        chunk_time = time.time()
-        time_since_initial_request.append(chunk_time - start_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 complete_response == "": 
-        raise Exception("Empty response received")
-    print(f"set(time_since_initial_request): {set(time_since_initial_request)}")
-    if len(set(time_since_initial_request)) == 1: 
-        raise Exception("All time since initial request is the same")
-except:
-    print(f"error occurred: {traceback.format_exc()}")
-    pass
+# try:
+#     response = completion(
+#         model="claude-instant-1", messages=messages, stream=True, logger_fn=logger_fn
+#     )
+#     complete_response = ""
+#     start_time = time.time()
+#     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"]
+#     if complete_response == "": 
+#         raise Exception("Empty response received")
+# except:
+#     print(f"error occurred: {traceback.format_exc()}")
+#     pass
 
 
 # # test on huggingface completion call
@@ -104,23 +99,23 @@ except:
 #     pass
 
 # test on together ai completion call
-# try:
-#     start_time = time.time()
-#     response = completion(
-#         model="Replit-Code-3B", messages=messages, logger_fn=logger_fn, stream= True
-#     )
-#     complete_response = ""
-#     print(f"returned response object: {response}")
-#     for chunk in response:
-#         chunk_time = time.time() 
-#         print(f"time since initial request: {chunk_time - start_time:.2f}")
-#         print(chunk["choices"][0]["delta"])
-#         complete_response += chunk["choices"][0]["delta"]["content"] if len(chunk["choices"][0]["delta"].keys()) > 0 else "" 
-#     if complete_response == "": 
-#         raise Exception("Empty response received")
-# except:
-#     print(f"error occurred: {traceback.format_exc()}")
-#     pass
+try:
+    start_time = time.time()
+    response = completion(
+        model="Replit-Code-3B", messages=messages, logger_fn=logger_fn, stream= True
+    )
+    complete_response = ""
+    print(f"returned response object: {response}")
+    for chunk in response:
+        chunk_time = time.time() 
+        print(f"time since initial request: {chunk_time - start_time:.2f}")
+        print(chunk["choices"][0]["delta"])
+        complete_response += chunk["choices"][0]["delta"]["content"] if len(chunk["choices"][0]["delta"].keys()) > 0 else "" 
+    if complete_response == "": 
+        raise Exception("Empty response received")
+except:
+    print(f"error occurred: {traceback.format_exc()}")
+    pass
 
 
 # # test on azure completion call
diff --git a/litellm/utils.py b/litellm/utils.py
index 348304dca3..77ce504cfb 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1593,45 +1593,6 @@ async def stream_to_string(generator):
     return response
 
 
-########## Together AI streaming ############################# [TODO] move together ai to it's own llm class
-async def together_ai_completion_streaming(json_data, headers):
-    session = aiohttp.ClientSession()
-    url = "https://api.together.xyz/inference"
-    # headers = {
-    #     'Authorization': f'Bearer {together_ai_token}',
-    #     'Content-Type': 'application/json'
-    # }
-
-    # data = {
-    #     "model": "togethercomputer/llama-2-70b-chat",
-    #     "prompt": "write 1 page on the topic of the history of the united state",
-    #     "max_tokens": 1000,
-    #     "temperature": 0.7,
-    #     "top_p": 0.7,
-    #     "top_k": 50,
-    #     "repetition_penalty": 1,
-    #     "stream_tokens": True
-    # }
-    try:
-        async with session.post(url, json=json_data, headers=headers) as resp:
-            async for line in resp.content.iter_any():
-                # print(line)
-                if line:
-                    try:
-                        json_chunk = line.decode("utf-8")
-                        json_string = json_chunk.split("data: ")[1]
-                        # Convert the JSON string to a dictionary
-                        data_dict = json.loads(json_string)
-                        completion_response = data_dict["choices"][0]["text"]
-                        completion_obj = {"role": "assistant", "content": ""}
-                        completion_obj["content"] = completion_response
-                        yield {"choices": [{"delta": completion_obj}]}
-                    except:
-                        pass
-    finally:
-        await session.close()
-
-
 def completion_with_fallbacks(**kwargs):
     response = None
     rate_limited_models = set()
diff --git a/pyproject.toml b/pyproject.toml
index badea01008..20675dce82 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.489"
+version = "0.1.490"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"