From d542066d4bae6d599ee8b4d356f0a56234a26fbd Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 28 Aug 2023 09:38:40 -0700
Subject: [PATCH] fixes to streaming for ai21, baseten, and openai text
 completions

---
 litellm/__init__.py                          |   1 +
 litellm/__pycache__/__init__.cpython-311.pyc | Bin 7422 -> 7483 bytes
 litellm/__pycache__/main.cpython-311.pyc     | Bin 29079 -> 29196 bytes
 litellm/__pycache__/utils.cpython-311.pyc    | Bin 66013 -> 68280 bytes
 litellm/llms/ai21.py                         | 127 ++++++++++++++
 litellm/llms/baseten.py                      |   2 +-
 litellm/main.py                              |  57 ++++---
 litellm/tests/test_streaming.py              | 164 +++++++++----------
 litellm/utils.py                             |  39 +++++
 9 files changed, 273 insertions(+), 117 deletions(-)
 create mode 100644 litellm/llms/ai21.py

diff --git a/litellm/__init__.py b/litellm/__init__.py
index baa3aa4af8b..227a7d66250 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -20,6 +20,7 @@ azure_key: Optional[str] = None
 anthropic_key: Optional[str] = None
 replicate_key: Optional[str] = None
 cohere_key: Optional[str] = None
+ai21_key: Optional[str] = None
 openrouter_key: Optional[str] = None
 huggingface_key: Optional[str] = None
 vertex_project: Optional[str] = None
diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc
index df608813db25262329c7815d96ed31611665a2b1..b15b88021181c0a14e70994b95401895e58eeb7e 100644
GIT binary patch
delta 1147
zcmY+BNla5g7{_NyOP8b?S%k8brBGT*DP8CSr64LG`zGjkK*BfF>MNMXBN1YPBp6NP
zpuTXDlLw69qUjOuM!mQs9(wj@Ja|$M#%~@4>`cGkfBy4*Gt-&(H+IlrePOj)2(}+@
zzr)TGtBrns-Sm@gP*%<L+-MUit#K2S#c>3M%noK&gIbGGmm;j58@Q1h7j=2n<)S)8
zzy=jiiM?vKK&4QP&W`=i)j%!O9lQDyw}GXO6N|1P&ozP*n!uIY#+~Py^PC4-u+p2W
zY|V2%bY5u7<@^wUAcQoxL7^RO7$Ui36k>SX(SaS%`QI(#=(@0LcVZ45=_iobBGApr
z&4$C*BV@0DkbR89#u{~mH}0X?Kbj)U2|ZhcH34O=GYB_EGlmDF1=O7a&qk6!-*d`>
z&<_Ji9SgxF7({Ky`5~MSqYk4+P)AUssH3Pc-pYOJco1*X=+_v~7}OZn7}D6TF`_Z5
zF($}5lEVbXaJ$Zv?YN%@c$i0dj0bs$=a$0wu7$;6B1u^nOv03uNGDi=5*$xU-P1xZ
zOX>r?JDMJ1L%I23ZGKuBK_xIFji$%gSbCg|rzhA%dXi11r`QzM(kbGeRln)q7)<l2
zfIsuh;47f(>IYMAWgGs@mqhlX$UXtj5!O_fxusA`cvNPdwRth))@4u6DK}(34|5C4
zOA5cE$X@lmxw1qo&YDudmz31vnjAyBss1+m!(s=wNfuvKZWM4L)=wpGE8dVd7`Wwa
z3;Tm|t;p01)CoA$>y~nBr5M&A+y>dL-nR@Cy2MOV=8NS~VQWH+y9Ig#&I$AioEPX9
zxFB#*pf9t1CQj+J`mSi1R;jV#KAmh+?-oZK6`0Oi(#yalywZx2HAu?~^U~o%M8w)y
z=4<g&S~4gKjS62;QcKCKVO2><E9z!xKqt?ruS%WqSy6jM;Hm%<xF*mpaQwo5OmRGH
sRpJSG3lkJ$<yWTYY%49Uu-`L1bUhAkdp{C8{yq_Vrn~IGpcjk(0Xi5K3IG5A

delta 1011
zcmY+A%TE(g6vk)zqOB5v0=3jaq2<xi7xW3Vyi^|YrU-fy1B7#_HJy?Sg#d}7k)0d$
zZds8C5;vM|jJou%XmqC=6E$&T;sQN)D8$aq{J!(u^O(8!efLRM`B8bfQNZVG;u~x}
zE4PSm-?aV`cXh-D22u|?(txi$Es#d4rDe2iRb}Dj%0u;Fj$@P-;4)aC4wrRb%|y6@
z&W7vM)k6c=&s}4QYbtS$Gq<aZaY8e+Kx=VKTZwa(I5)IorKecwEpa|{9`F})0qB5E
z2o_xkT^J%o7ljypjepP2PtgV4|DB=-T^yHkCeox5h6PCM3D8T0dk(gwHrt#kfjSTI
zK3<((50W%P-=08PKx9wVwc3;b{Riww-5!G>7*4535Jq4WHH4WP2O<fhj-f_S$5ErG
z6R0ujrM?{;$S!IuyI~&Dpu&*Ch{CAC7$@n%&PnX-KI^1@8lW9COrtbLJ86&>J7H?y
zNP1v8C6YMIz-%TVB}hUPFivKA=eW*g`lNo+FAb0ZX^;%ELybLDEErP?l9_Q-0d8g{
zq)9R<O_3>SnoLVGWJa1Lvsf!81=l<~(R|2%)VkH;Hv6d`wEFQ!j|Dk^BKtVpXV_xn
z29HijxY+~4g2jUwm#T1SF?Uy{%dohzmd(+voa|vg4AyGiT+pY1W^?J)Z8?VaA+s4f
zs`vz#UUs0!PL5`_WULN1@UW4?&e4Q$sdp;)L@z9Frn4()Zy*qq>v*=FXSH%0J2sB!
zTKP^({<raoPLsfxJ*up$7~-2(Ij(VB&%dsWi(-;}uUZqW?2c)Oy)=bH3;S#`Ur>%1
z;bpP>U(<8ZGKz7bEJxRK>1?W?-N>ag>+%?GvKKX->{E>+KF_NcIBs#=<{%s)j`Q35
gqmN5#u#us?@?%Vh_^yA+_tg5t{C)E~t%eu=14v-~>Hq)$

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index 58ce173b3b7815a0a0f2ee000a679a9f9e9e6f9f..ed8d32b9652700b7db2bf3cf763de23dc0972d7c 100644
GIT binary patch
delta 8414
zcmcgx4R9ORalXSp34$aD0w4(Dm%kt_QY0l(6eUuWL{YL#N|t2Nf=xppA4q`$K==U4
zl6nZKSatoQx>B}NJGR`Rl2}eFy0*$Du05HFHEn(-ZDYz-MhY^rlTn=1si$RXyVE#z
z)7=9=0Q#})nKVb<-rl=?@9pl}{r27AscZa+pY!!E=yZh~9(Vuq`|P(*aNISz$sGDL
z^TEzj_51C7sK{HyIHm~c!o}WVo<p7wm3T|=stlbvW1I;ay+)~76*7g*UURs@TM@Q+
zE#XRUW!UPq;xjeVgsQ?euTA0$LiVu3>yY@ukTdM^x+Gp3a)+zE)yUHnHQpM$7KLiV
zb>2FjlW~EX$@+|rOm~*&xN~@NtEDz?0}2;kr+~-nAtj^%OYm8vw^15MJ=p|$g=q#K
zlpg22cLYjt^fIPj9#a}}f}7x28R|}d2F9bz3Ra#kgOj3lL6YOgOGt|8@ia&rk6alu
zG8)HpTF1-*&$x_L99N&k@y~G_JUn9T@-{OI>d?X}k*-W@V$Gig`x}@QN8`tptm-%~
z^@n34=5lW<vthqBW=Go2_}!exvA|?iyi(n_t*sS~%4hhu;bZx};&_2%ltd<#Xui2c
z5m#p0M8%{kb0^Br%1}c*Q%6E6%I4H?#%h#vq7tSQPF@2i6-{m0v8hNra5#R*FT|2X
zV^h&1zI01cE=)&~B}jb1h%azBFg6{Z5<DvS7e%EbDIbd;PU^>}BV)dYg7E{HPBGHJ
z>4f(c+Y3F&w%{3?g+67+;72M>Yr0hEzF1gusjy~t)71*cocwF*ua{q}Xued@yr?iP
zskk!JWqs|N`r0|Rc*pvS`kqVro<;2q|B4<sqpaX{@S5_7S``(h!qGUisa6?vIGq+x
zIi4S^#aMfai}Nr$QgK(pR=|l>6Wl?1#0f=l=}J7uu}6ZOX!{!fBsyQ@QATo`P$XIy
zpO#lK8Q#fH%O>~)5C5Su!9~?-wT&s(VSl_gLyfvYYZr}zoGBBghyd@XEk+en=Pv`(
zSU7m6+6d36>kKz}3U8@vJ8x`PTO>V+Bsf+enlc(PmB=091&wH8DrnJICeRXAD9V6m
zT2YR)NK_!z;ga*LSX3b`5!FabMGaEDSb(%lEJSJ$wMffF9a5QCjMONWAT@~`QnP46
zS|L^fdMV+6mo)7Ima`a;=az#q;fx!W)#OC<r05Wx^Oa{6D1~S0NVrf}WlH~(yzn@~
zU@S60Z^3%Fv#<nqSak5Mq5}A3wGRHKU?pD#eHJTxT(Fb3!S+HM(qN&7cfm1cfmw@D
zl3Oe_!OshwO8aT~yaTlQdRVD-^N+yKELNTKv|M!1$|l?q3$$uohDS&<<~l9Qw2Lk{
zRoTM3;X>gi3ytB<*O=!{mPQUjcag(bozgp|WHsn_4)x#2YEl0RS+ojVC8oM)Y22Ez
z6yvHc9|WR2t4~ym)$mf$t9*lGPea<C^SV}bHS?U7-C#lEEi7od&VmY+s2APr4m875
z5#(=TMM<O-eo|~RtRUUdhGxkInV^$wpfzZbqlTzmGsp||(5CCi(Xs(bO8fa1(ASvx
zmEb6^mRrzs=XGYh>hYN=y%rtH*3$~fiecVD%FN8SoyBB3x14aa6metLO_R$mn=H??
zGFm1JYtI^qf$KQL33k!JO0wy|R;5#+EZ&vr%W<Mx<dzeUFah(%ExIwgU2suvR-(qD
zJ6Go!()osUM|ZBynF@0^bFhvqpQsUQMEflzJF}8(4WR{&+Ldj?o?Yv>8Eoc|HKhgC
z<_WBoTI~t<$)d>nnEteUYD}Yo<uv+sygQ>BhFw8^*kO&-FfL?I)W&<VQj^XMKj-4$
zE7o$p8}=DI{5p8SVMG3m!KHI&L=tuJ0eHh;wcT_8)=dhsRBWuL*r>O#9?=ry6ZMEO
zv}}c9fmqLaB}0hYAaPRsxGZ9wST7c^K8pVQP~kbTcCH4_lv`#tE-$F0RHET?1w;cn
zCBGm_@xTxl#D>k>h&tiHpz$2VU`e=7Xhb)zs%tJMNa3npTSA)FEbHoq-Ns&w@t+xO
z@Un4N{b!5;Xs!3ON~@l?=#65Ni04uA<R;8{JIgM+HsXedO<m@jdVkJEWu4y)@0t4K
zuK9jgVO}v~kKdnJJDde^STxLUIjhDr#4~jy?4Qen_MDWG21sWVg=Jyn3nCX^#(1oV
zKQ`-i1Iyw0U>?G3h2Dy5`fd3^8ZCx`$|}Tv^D?@d)Z0^Q270|Ki|N)}MBJ2zw|9ZL
z(&ie>A34z?HqXk~P@-9E!RwAh3)KOuR0l+CV>?gFSk9QKdSG_X-+g^7p43GC8!lHm
z9J$H1D}OdhHZWRS9J}+|qs4ie?tw9DcN$Ia$wgDdrq=oG@EfaR`1c2ax$7k%pcQ8k
zZ(<>0lLpRI_5VSlS{*qImu(L1-ZX^2#P$MjcYIoudxo#`s{K}wCgsGv@Pxeyb7CL7
zYHzt|PI%91FemU#9m_f4yhWKQ&nq*S6KwDNefe>clNYMUHxQ=pPuW+FC3hsfZTMJ1
z<h6;_@QT9*Pdkcvp9H)633jvLmFQr8HY#>J%RMJUFzsYxi8X1MW!D8w2L_GLX_E)U
z4)BkbSFmvjcZboA@oX_+tKgiow#C9G*a6yCJdb;{+~;EFd@z6aljw?f!S9^y*86W^
z=jyrja9(SyT_Y9j?6{rtlh^sBD^-2<7O_L@n(WB<eEx2D$<<{%D6YkwY+3wRrX6uG
z1fRGp^#^nG|CD+RFIIK^LE=^sc^M|EZJt~d2<OLTy5w$5bmKxt;Kk~@4O98`2PLF>
z2|aZ+nVpo<!Z=*UgEbG9SjKsLCrRyGu@cd)r4OHqH8-NtXWxBL<uTq+c}~hBUxj0}
zpWU7Z;d`}%x3Y&9@&i+{=vdvh|38buH&`T6O_D|aI=?1O?#x6M?nnNN#S$h_A8pL7
zT9bEYB%+^+D~5RfY9}_vpNHq{X7ff7vF5s{-z08=UbFS4=yR)z{Mpebc8jb`c|eEg
zGo3;onDTU{>&>Cu71Cl<v4_Mp_{QJ>b)(B-xlaGh;%3}8o7iFKYjkLjq_O%*yg%6Z
zsaVY)g;yF===*Nt=i8BdcKZ+C+WyEVsP!ME^o>=*sU{0A(o0$yF8@EVL}DO=!D&SO
zLatkMlxzcRhRsT^WUKFJvRT*+!u>IPLHuHVDkI=a3ED(`+*x2-?thaP3%KB_S;HRA
zBHQPAi=aHl-6p;_<n17F%_C*;7c()8_<3Bk!S7r~e78Neg7#IXDy6r?tq8fVtf=Py
z7=E(C)%qtxLKF5$DyCxr!6Qp5gsD&<sg}M1MdHbl@DvM#eEv|-AM=I%QK235%}>Ki
z%_i+C>?W*1;wgZ^-b%A{PapFH0(W5plD|?2JAiLFfuZ+gi=%@QOc+MukqINjQ-en#
z+($fpD@y98gy2Ll;t%<RKrA{HL7nc0Yb|%^lO;QL_7C-K^Y!iA=G)k}yB|tcHaE~%
z!af>AF6<yacZQq>2&~+%6naqtez?+tueJZWvIUDpMQa`32)(V=f#rdc8h<qCiyw&w
zGDCLKxM`)OA*o>xdg#DZa4Zl@D(L%}FiBFq@YU9M)i|*eL=GT{(bpuWQ)r-ugYbH1
zb?ILG>M4ZizL7{LjJI`a?nKtZ3-`gfwpgX)3vFm>0(~*zn;M<;sMGern)Yh#F6y?M
zNGpuD?^LDy2j6Pnt6g@b@Br9W?Z2%u+N)V*;Crju)+K0InqR=TQqH5v@vF49v-F6n
zw2-3(N%F!fSiSnK+s?>W!`2S+`&7r~wq=JVi-f>~(}7q#?fax+GBy<vw$emt=YLt^
z;oXjJ-*yh|1bRK~V)*yY-M8)FKA!D>(<BxCU|Valh(5te5x2d6A9Q#1-F6)9H?!lw
zce`G@ZHMk_*$z1IdkhI?p?q2h2@er3PpX2E=yW_L-Ab$B;gbsHkNe>_YqvNg4^rwS
z%a(H|aCj^bjR(<x6YxNHm1+Y%7Aj!2+p}svu@NGaHbN1Rdyymy{70vSfbU@7NGcOG
zQM?=TheE<IT<r#JF-g+L6k#JATIVSeoYYGm3Pgkiyu5CWmI7F?5jh4Qtn1((gU0o4
zJ_>{D1N@g^VSP3KEL>h+r#ej{&w;$>c+(jgLX8u#;8;AVC7=5wyoq5N(hX@tZ{dOk
z&i8nBNw!KUzKI$)6X_?ig~$LADPXn{C+V}DxVw-f3nJ5D-)O)ejtPU*xSq%m+}Z1%
zxemA}za?A~x`<m##7$(D$a2&t)#?09YL>Bze3~?j;RoqZfc&3Dv5nOCABjjA`V0vw
zg7H9DI8sQh^x`l40511#=zNUWACiD<Y)W{On6d-Y5)_R4$8b?X!9XM~yhUPPB0~8g
zyZ}QRDzwiMoA&U78%$dy4;LK@2nbRs*9z~VNZQw_k&*=yfk;5`<L7<h7xb`r)E^7P
z1Ch*3FTxue?rD$!Oj=(e<i1#61;@>DAMN6~qO!-<oLDovc}X@3cl9}4SG0zc!s9)U
z_B=g4H~P%v)00<R^)GK(tmvkH^SVd3E`<8NBX7|=7fWN$cV4P)$LlwaE^gd?v2)L*
z&OLaycyDY;&gtk^0Z09xZaJ~zimP_P*S}xJkHPB$?e<I7_C@RH-`9R$g%|qY<xepu
zH9sve!`nN|{8iXF<X9-*+Nm49;%-<ha{RQ^yu``$ZC4GImkiZQa%A5%82^0jQ$1hm
znQOjiXuV`;T~gzu9~E*|_f?DQ>uX*&qnNe$*Bp`$i9EV}NyFJamu>BD+S*}o*ix|S
z*PO<V4>j<k9UlG={BB2u61A?VTNv7VC(j?j2HV0nhyRYBiK5d^>iY&=+*gJ<x?HZ_
z;-Kr>@_q|lui3Vh)Ajwf{&u?lXWgbQy8f4QXa!v_^cVNFNX(X=jZTT#=Ge?j%%H2V
zZ>Um!VX&#OZ^$UWFl1EH{TR-fKH$o}Ier@)@pgJ>39;}>VO@>9@V`isY9SC8js%35
z0`Hu72o3OKuT3Sx8ljcL|9JQDjSGA4`;b>q62cp+O~6LF%c#~NCaEwDrbBV8?Ljy{
zVo;?k-Va7Ryc0eg@%S1^P9uC5FTy_%F(FARLcv%(sftaHhJ$elKoH#s4^;kWG!S8)
z!mOGUn(?s%zHlHG^G{&$#ow<oxfX-?0}Z?io_k<L6-}KiP03~@gfg_Fbm12d)T?IN
zY5tW&T8T(I7J5V2BS=s7P~$yB_9DTpKo&9yEd?*N%2K&Jg8b3a3<mfP`9srzwN7E4
z<n|L8B|__v>9K=@6M>fd9;K%A#^pT85D<9`0{w)Q>M6}CVYGzS($rFIeg#F61^waC
z0AsjC5MHG&$B2A~$O4$g?K6K%>@SI2A@Uj#H4^lwgrpJ@WgQ{;K=Q_~Nb)@*Qh<Dn
zT#QhPLLzlY+nHoRU|9h(y0%u1a)UC$gr(wM<-)@Fez{=*`z5uACq6`>(J9=}{RrwN
zYfGf)kTfO|1@GiI2G@m2J>Mk7lTt|a#Es~9aKab$#}5c4D48lTWEFa|8a_E#VUt!{
zpBBtnZv26odkES>!#XMQRtXoV|9<$3P?r;bQ|1IK-G85m9Z6C_zkdq!lPWs!SD{ly
z1))$l)=ZU1Fv3G&_ly*<MH6UDATl;3c|Jknh=W+8Xbjx6`I(PRA$%SZJk+8b$EH+Z
z3zSGw6bj_8)Pj+gsqG?ZDLM9I%{nZD)1g3*uo9o*KNtH;bgpVi&hz{ga7CQZ6Y*?W
zx5)h-{^3mI-%9O^T;4xCK6SQ4%l9sE%cKf=qb79`S6q6fxD@t9TloQaBI@kCqSYbX
sX?0i1jrd@xpC8cjg-cux-NJJ!!_q+h19&HDQb|I);J4AI6=ePY0bv$ho&W#<

delta 8202
zcmcIp3vgWJb-s6>TJ1}#*Gk%5?Mtt<EXlHL>*a@JOUR0`WMhG0{aCMcujG~YA>UmY
z%iN9Z#tnwLAm=m%CZ;hoF-a}jFlMGnGY$0+8kiQ4!Bd>g&<2uBn=(n|0@EQaq38U&
zEA3i#3~9O2J^#7?<NW8G|D5mt@7*&$=O6nSU;m8JSjgek{rtZR&p*U*Kck22Rg_`g
z?>Jk3U`7=-M2vnT&k3BshfRJHmg=x6V)mOOMgF3Q#czoe`->wb{*p+kzm)aUgv%o3
z{_=>`Z;e#=D<YNt%1D*J3fr_oLD&|t`|XU^g&mP<e>LL^!_J7y?_%Tmu-os(dGuk=
zNW|;+@|=Pb-4iug4F$smp5rd!oxc3F*x7i)q;>u}nwSYRVOzbwp2^WpG(c-X`3yFg
zPjP;qXhz`{p(tPEI44+8Tly;&c<ix4D9P`EQ=&Z~>c@|nsFzTLcQNC5<Vpn&rIejB
z2<2kkm_o3g(q2G)FQ#<{t)jTCQK&!_Rtl9!n=*=2<xhhmiiIi^$B(H6+bJIYhj&`+
z*uO<^AYCQ!yEtF<C0SwT;m?)R{3GB{-KUZj=T*?_FRS5X<pl*!ig)_@5~}@x@*eIn
z<(E}Bk{|M^L8)#mHcL~{KqwjzkBOsGiJ0V5!2xxZI<CTD@Pv9>p%<BEyyN@e*XngE
z7BpOO`CFFSH!QX1g>Us;v9w;bw9f07{i|Bxh^B%!z{8qp&}-XF29%-4TZ#Akr{GSl
zOKsP2vJH-Ct%VPUINAOU{t<o`%xkUis&=i$F7WG71$uNA__PkWOi~I87%p&>C<Rr%
zae^9V3O3xd{`G>|_T~Lmm8eDfN^*imE<em1B^d;z%pK*W0=ZmJLZhy795oZPvI2Oa
zKvp8v$tt9UvKpyg)*v;=TBJs~0I5mVAvMc|NQ-1WQj2UvS}dE8mdIwLr80-KOfE-S
zE>{9toven7y0!u91syIK-WkoDiIObOaWW_Ka<%N7tGIw`bCK4ET}c<}sz~cU%S#WV
zj~qq)Fi~jYE8!=FHh9gbgC`3&@KyM97>*U1;USZmmCu-LybB)HA5q)SDd!yED6WT<
z1`q!r3>$35>T^oDn${%gj#fgG!Cmqo=^!}IDYEsl3!XApz<=1Ryc=FHmqVb$2tU=k
zDoLIzUr(N!Es8!3?MCNXcUs%HTJWH^In=sc@Fq&L7I9pF0)<_&^`t_!M)e651WQT{
zNssKo3Qv?5YT&Z*dweZCU~(XR!sKJ7RKQP6m0FKbcTTa)p!!=F)Nq4AHl<tep*fXh
z%S|vx?}A^M9VIJBaWtSYW5D}@7u`O<lUB{}NudSynsl(;)U>GKM@12SB?MPm`6igL
zc$Hps-lKXemKJbZ%3-i%i@N!oifO%@_mVry=UOgcx?H4`VHd81J5dTx6_>)E;>t4b
z+^YOMNV?gISk)+0c4x)Tkl1CW$Znx^DGLO34!?SdPu9pavg4LL+m?C?YcO3ol)84Q
z4o9wC&&{ylbqJ-I4(swdtYfu~q<hvF-HNG5e$`?~YZgQLg+yn@ek=q{K29Q>s|7CX
zNY*AuQ0AI&X1Vh&4mRWd)DB-S_3<6>E0+U#z1!XI&UQ)GCHk{Gs#=#n!Og`A?KGQ^
zUeHcj(0Z@XC3{1BvK~>O`Y0zA$~w7T=w{{-x1Moq(cNCTPOg`A!UhV{{59q|x%Rvl
z&g#o%IJtgq;|(j^fOtdFtL19Ju;fT#QwG6`6N5RnG@!fkwgHR~_byK8!Lp%`+ow(X
zWVh@)ZsfS6`%HoC#`SicFH1{+Z=20#(zIq#R~I}{z6pb~sXA8wDWd{nR$nWNgS;?U
zAvel+A15#NVq!Xkr9jX}zX-2bJFT}^P}<=5&CpWOt8~q6fp|sJj3co($7)1Lxn!>I
zf);ZJ@AQ>)d?wF1WbS63lPOROqbZq7EFv)OQqD??v41gOk7U5$4g`Y)d{i0Z2jGwF
z4t^`c;~iBx1<U<r#oQpgT{W1&+gy0s0@GD?IBYY!x8+A;vPEt_p%Cs&Hp?woZcnxl
zGTRt3%UCB2ol^)ol9mwK>Yf|EG16vr(GhUiozC0@+mSybB^d0nx9~g32H0vh!(ZE*
zGElZF7s?QITG66S4(D#t$1AK_0@lyud-pATPrHrPhO@{wu?3Mw4gc!s|J<Wf8@&n+
zr&B+Y$+*7~Mqt3{{9V!NyF1^flhrpzmGdqZ+xNf~XCub;ZurF6a?{wp_d)^2Hr_YH
zw(AzCP@0Dd7~8_g+<p1ckVXa;-`c1J0ob3m&x5<%p3E-y8-}%A$+Zbvc4|!Dz3@7M
z>OO*2c+FJ=RqpDNfN(%w_ciWG1p;Eb5KMMtpz>gD!0hzO>*faY(~Rd5(>bz7(^B%?
zhao925ZsOC3b`N%*WI-(UO^PbXs`b!9+-05<@UMp{6k5yGcgGLp0+ao&D;Th>~Xqt
zrM1r;y1{9k=?&L{mdKqG>#~~94a0|?PHRZ+!oy^7;$*fSy)_|s!J|T1eU}i*(fpQM
zF`V>vf9@fzjvC<E8iy|z1`aRtIqkEXlie5tVR*l0XG!G7mCLMgG<S{FUV%eUc&_${
z**nJLF+Y8z#BGL>xffgIn&rUs=?@*Wy_7B6US^fiPvA#&pWdPv80)v+$|8u?)@09X
z8GC+IKlu6CQ@P9@mX^w>2ENx|+xcGJdFaSZ*J)>Rb6$>0`Mu=Lf$Zyk_7D@#U9~2!
zym?}=$MPJ-)<pz3e-|BAwjtmzJC1IVw+M-xD9=4)PUY0+A2?ZLY|GQ&{~jZboBLD|
zH6qj<h5zzx$ZTo%=lVg9@hKdWJ3>63%fGt9UFp3c|NG>=T^vf7hSyg(^~W<?!Xx<n
z(Td->CGZbGe`6Zt<Bgx~0Q3I>^3jh`YdNiNyb7EvE0xLgOBM{TR2#pT9LQpQCaY(1
z{bHoN9S}|kC)s<}iz^!)Uf~38D}R75eqYMRh3rd76OC~%VdUcYo4lHw--Iz$3P0Fq
z>JUyX1<Gf68zEr8Gb+A+<efY5{1z=vJevc^e<<4tkns)H+-${{$fwLU<yHjn{$?-#
zWe7F9%T5oa%#oNNh6BNHC>Re!f|F7ce82hedsS0$QA!yFaWWj6j)>7jAS9&Jtgl2L
zO{7*NKAqG+ypPBVB8@~=5@DTMh+9R3Hjr<tq&wi7EoU}T|CBi>2%$tM77gO;(s<mb
zOlidEcql4LTWQ1)5}!gEA)YFHD(NobeGpq^#m^}Zu38B{T-Dq_{iGe#K`GrqeD3I+
z3aDy5pi5~-W6=bf12e6aya%3WZQ)DcmDW0Zzg=jx4K4~x6$B?kfyDHrm=T-OOiHoH
zWFo5<lRSp=p~yqA(5M(ssS@I`gcKwl_Q74N6LvbfOZ$lgki_XrjLRi?sNx{Jy3Sj4
zCy(qXJ=YK!gR-`EQxCE}o*6qL)xt<yyoxzYAC8Tqd&UE?gA+b&##VTz&8w#uBk2o7
z8o|6~hbH4Eam_vYG)ql)!{4tta9fr1Z!W0>lWSM6pQ5pBf$vAoSCHdlwz|iteVRy)
z7SxlMn&F?<zIodj@ng>Qb>#nw_P*6iUNuPKkts2r$gDw1H4%?RrHwRECNj=49v*3b
z;kI+=_o44Ij)uoNcHMS>=lIe9oF=6ThE}(x484P^TLZm21_Rsr?}O($I&Uk-J)18_
zKZ-FF7Afp8UW8yG2-`Y)Z!1Rs7fWK`pF02MwgbA}S{i`*|58BWQ>2trQdo)+kI5E_
zPEIA_?2#=huGrVH2@1RKaI$r1#-~(qObTJtVUmi+M#V|eT-pz@Zo7t3TC&2~Zr_^S
z#EL{H5heN}Dh(q^>4L|nBr$MUoMthUDwxEl(O@_%4Z*LwL4&X7oYVnl*840eC4R7#
zOf*Uv7NgP>TwdRy@1YKMBGaJSu#P_o-5WgoA=tk`<j=wnH+Xq~f7?)}`4j5$RVdwf
zYQ+=CrnESb5E@OS^t5sT23&FJYpf~#MOc(7;Khx;8D=WOj82l#MWmYu;j6TP2#dN+
z#PtyAC9)YwN*A4q1P+S9NL<=Nm3>6$lgq7fv58vAvr-xyVf&_%rfPa-nY5ThDQza*
zQw5S3pTu|()6N}(i<>HIJ7~lYh_DR*3UVn`C?Q6qBUC97c?o{Gsi*xUv9D0OVl*bb
zN=)&gDTX9t!BJd|a0u~GdX2h0Ok{}2({P}tA_T;yU5i|*aP&|xk@jUuMZCm(Z8$1Q
z2z}{>w2zI<>siH9x^WR<J&0f7rFTe>>0mG}A`E9I{}%j9&)p3knvJbB5wbAeT)}ZC
zxDR*moT2!kjx!x6`W6&>VQ;VN(v!U}DVuBB=Dizd&9mm02CnvvEGRi+*+K#5sGYa3
znk|?um|wT^L#!(~GrXYTJaw1O_8(C2JK#SD+REOtwY_0$`)=*ant5B>72EJt+wg3`
zrRdfz#=5r+)sG*&VyL@nsGB#`T{BouZ==6o*lN#5uh`aJwXHqfcjito-s^%Fb{S`$
zUq5fKUorGtHT29IdVXqmog043-u8yQ?OM6*n$@wO<{WFTb4rKl10tulqk`U5-!WeG
zZ2az#tDfG!tGntMd>P5kd4m(xG?u<)aJ*q~Tr0O<tFWIlT=T45&~nD=>zvAnB9JUF
zV&XRQck`!*U|^&f&fTrMbpJhlJP13gFa3D;^ZY^DaAqyNLmKYsy*?`5HuP6g`EK1{
zDV6Va_qS1by=6-$l^^<cvhwnPce9HsE`LGQ$1`R~Q`kGSNqKpwXhrW(7nWUWeEy1m
zj^6~&`rCc9TDUv2Aa)`zeS{>XmBfTJElTVM7Z+ZV2aJ0hnvb#Z69uf=b01%G>ES&e
z@c8v;X5UETev-ncG#E5f;RNpd2S|!$Y$|~Wa>L&F+kHO16dLyXcGJFzo6==0q`xIX
zpZHR$a44QgY2s4{BcTKfDFj<W15I#pQj7||!X<4f7!#w10ueDD4~~m*>CaFCd}n_H
z&%=-RH`ysFQ$^`MOXw{bc#czD@&xKNGc2W-jz8?YGk~K;B-XMGS?Nw9+mYaLVhJb-
zBc+{Gt4QZDopz2FW%DL*G#H)|yIj&U)Q^Gey~N#@GvW>oUKU$sPEb|m0CbW15XAUO
zB>KG}J(A7JkeOjLo7zXM&ml?a`XdKLJdj1lrAt&fP2>e4&r)+m4BZw8h5}-?TNYC1
zsL#8HiM~nXc_Q>}5`D~|mEqHpi<oyBsnair{G14joDY$U7b2lpNnh-!l+rEEDOu3z
zm6zTbiz-W)T6#p00HV*q7^34m9Gs{%GZ10g6A|D~C&n<&KAy1f2Uvuu(+1&2932ad
z2PT7wLz0=4V8HV(a<~>0#i?u;I-KUlU&@jhn-dNAt1NdIqT$^}7Lcvd_etb#cq80t
z!=I2jsfwOoB(bW|*yOaNrsrp&F>)6EF;YFlqSi2u)`-#37+VH*@Ul`j#6Z+09198o
zA%>`VRPs`ddJL=5(JfI@v5;ggO0OR2ls<iuPRy6TDOis?;Z#`MDEY7z|GD^aC5C$_
zI^5MV&;6eN@Z0En@ZYg-n6>=s1#Xd)!e367t4y5P0{<}C%J;y>ldc^IaeCvm(z0u&
yqJ=GdkDf1F;Bsgm&uK~)`j<6a*D4f>4>%;()tsW>RQW<78`TGsM;=#^f&U96R2h)~

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index 773d47432605c370372ec1bb4f67920ee3531192..8d8cbae60af951f7ba332fd4843415cba9ed8752 100644
GIT binary patch
delta 3660
zcmai1eQ*=U6~8^*Svp&?EbG&jjb&^ZNnk_lgfD}s02@dgV`?ziG<Aee#C9B+tYktG
zHKvf6x=kl_SsFBxF!i)eYC=Aoxg^j(G70&kLpy1D@*JSJD8_uGkitYvNb7-2(zhoY
z1fl79?{~Xz-+R07?R#(K<X-ORZ@An)8VuPCJodj`7N2yyl{>Z)$z*b6IeLJM*4VOf
zg@~ePCxc5#Xq9bcR1@JSclrUcA#)tn6wFL9ElkK6Wiw}i+Zf@<$f~095Gw)$eajfc
zw5h#anl|PMbck)j%R$K@4f$eKA&(u<C&KEk>=&UkwVO$?Ci3*^l7XbTP%&4f7{OF=
z#!)i7P<D8sCCYsxt)pvR>xh>xiwOy5`Rh5dvj$qStZL**rEG~@v}Edxz4&FPY@Y`$
zSvqf|RVlqUmX#=R#|$Z+aa5*sjH75Mk}%tc?wv}pNN}iK7)(ut?XZizb$*wRImSBt
zF6NjMQFO6>x26-p<jliErv%<Fhsf6Y1wt`+kp&0X$zSWM_cc_q_z0<OXsdk#WWrAY
zpK6*Qm=e~Kq4r^mY+W#71AsQzkTlv9qc6qqhRTG+8QUpa+|bCyh7v8$SSx7XvZl<O
z2Ws|K`Wfc!Y=C17>*uA$H3PcpeBO~_?gE}{-jrhQ0uEwoc96Vgi{xz%3apNkhT6&y
zdYru7egL^i<JN9I70rz>#AUXSkGGZ;k424Q{-Z2xfWg&KHfj>NWzdiE&_fdUqvti~
z35j&fVuf<wpfRS4>q@23z{?u+q0}1uAwq9UhqjmK(OKy)xXXbmrRwJn+tKIJr@wnn
ziyo7DUOk6Um^8iiA$nXgz5Wb~c1uHV{tTfZ$??{kJc>!j-(8%Aj!5V|e-1h!tvcSD
z1!AKgODyV_#P|~$^c*>@yhRQg1^znh>i~JP#-1%Qk80E%-$UN0agebS)rBE}#-qsm
z1hp|CO*$f@sD+s}wfUr-%ry(Bg^YID$)8`ali!{+Y4L72M6Zr~dU7Fgo$|N+0nR?W
z3pDR$lG<EFTRy^#Y3Iq>d8sT;C`?+56l?V<>+&({@<g8Vss`DLuQ33pi@hlgV=xak
zOr<y^6eg?%gAMc#;x1vd4jPi=@25QaunsiuD__JM;g*w!Z1%hcY7i4-+L6e#vtkaJ
zop9TA&ob?TC;(=Q*?>7BSmD)^#>9CvA>H%I*BW$=R3&FQK7+lFbKAQ^BBmMU&BEpI
z!Iva4c?0>+!7M&c;Q|FhrzR;}rMw)b()gbgeox^dz#+}F21@uB8J_qDS=ekvN6Ciq
zDul?d*VxF~_f1+>V4|$Zl2D6nK;&+f4UxaqPgjSC7S@`6#H=WKkvYh|0?H4=n|}77
z^htTOO_=U@Jg2kou^qj^NOxZ-jA>NjgLF5O(7Rr_DHPb&8+1kbTss0d9CYD_gW(;(
z7_4%w$K8)~hXTDWJWX2anlbXGedF1!f!_XLJ;r?){~OeXX%M+wE<6Xm<EE}aNbC)E
z1iGv5>FDh04?X&nn|w6k%)3RHr18;7zL~H+oU}Oh-yiFbTilApJt(ADjnIk`=Azg`
zih1scA1d6)PQ|=%6rfJl*Qw>G6(#eF<l>FTHvvG4=WkN-H_7@<v_9p)9u61ydFB{z
z@fT>wuU@t3+ktj~0e>SAk<oe^+0~*a>@4#D_r1>Kn{X{iru|1<hp&URUW9MZ4}0sB
z!_>nw_LoiD-?t+e3Uqfwg1aIew@zc)Ba2q`ha-I?K4Ba9KE6xTQFWl4mi@G69I*n$
zTs91zsTc;QRLT0P|L+u<F?-B5<^_vS)2S9IAkc>cuQ2rZzWfN02TxbcMCTlcOtoCU
zn;CpGUT+0w1hf))zpVnV5H+_?L7l(vPC**z<jEdAN$k(f!j_rXFUgF3s52R{uIeOc
zk&S$^-!K>JZeIg#PhQ<(*&Q-`!Hmp+Dsx*tMU1N5D*Ufp617k%bkmunf2>ao_IC93
zM|Si_(pjYD)0dz^jfZq4(4{NE%_PC^L5!>z=fv*L2!4t3en^2v&RufGG95G90^wjJ
z7@7_e4^11(ZL}|FKo}AEcEUPvN9rVU3Sw<?P6f22u^{FhZc++s-q{#8E>(<6Q;gQo
zj*jj;ULdywly%$W6||jr<Z&54onoNcKfqoC&8U%u(SvMDdbE-KieXyWb{2|V7urgR
z^aGL!yK7*0mSSHpYE$fskD3&Ft(;ezG&y7IhC@oxeeZ<hrdq{RORwy91l%gKoVOrh
zDjqJ7O+IL^KXugmL8G#C!-<yG`25Yv{LMh5tVN2wO3tggU~|c(TjUL`6ic_nZEcFJ
zP0njOW3~=eV~-R3*$M8v`Oq*{Lphft`_9y$Dl$Admu&moO3qB?=^^*k8+izY{zP1p
zR^)_fFX@;RI3N+yGkG7nA4o%V2%={SNYOfhY&`3Ce;s9dvSyaSi!(=}Zu-uzR2RE&
zwz4AhUM9mG?80sAd!06N<E*_R^vw5&WM_&X1NrIU^6VL7LMzmbzO+u@Am!%KJs>0H
zaLjOsIU<<ejUIs)qS`=apYzq_g4@P9T94HNH_k(ZM*`kF^>-&;1Eyk|$`XR&g(qmt
zUZ(I0ow-H^&J{Ic7P0_;O{u@2@C^Xm6!hvG&&?=zecO8QC72r*)Iul*B7ramOmJPK
z?+dJD$@)(U4RooFj&Epu&u7W?bL&u?tUm9}jqBZs+<YasBE?vRE#&F*+hOAD`6Zeh
z7|pt{a6kqIgGn*CQ;bg7f=0>}-^vfV$9(H$-+D-qm_;d=AGa(}EDHwf6P|geJpM6{
z{{wg2vsUq}9c+p%ffX`z@HcgAAM-WJzQ()is2lUt9qo^M>XkZA{cTEq*?w@5;`GY~
z+U_dbJmzVJsE&I!E1u2d!3%DrCzH9>hDqv2)z8;x?>`h&2UFjs{jH@=r$4H(hXTEC
z@o73s;RyBJE9CtP9UHRvR+M5gFppCy8K^h#wiI&{W|i=ZQ_OcEz0hfC*=H9~mI;<G
z&60c<J309A^!(*HEYFf7m)D}t_Fb`}i^O*2$Jy#(qaLhd$un1$@v}jsMtc9svl?`X
ztiD>tyMfRq^<MRIIdiBr^j{r3K+tR(dG{+PUj?*>rITOvusQj3+DySBJur2Gtx+wd
zd%#pHrkXJ2zr!Qfb<3b;_xA?tuwQDq@pI16hmbSHd<VOvYu`4Z9~jVFaKsE$Y)1$@
Uad*6`(ywLCZqyvugV5Ih0uh`Z-T(jq

delta 1890
zcmZuxYfMx}6rOYM-eq@LR+h&;crEe}7X=gq!B<6G(4+!twN+D*xS$9w_OfcL>!Y<v
z)wU_^NrSC5K}{@;STVN|e>AoBN1HaWQoG=WLIbTL0=7#Pi;viz*;QkrleyoW^Y~`w
zoHH}$K34VaQyD%F3)3-bgI^6fOO0m?mlwco>|c-)S>X}f-0-MvWAaMu8!%0fXFI_I
zB<6U)30ADgjmG^&dOW#IA4l34j}Xv5=CK_#+9kRfmKRxcPPUx!kHp<YaoAiGZ4yi6
zx#ZHV-dyjjs%x-|vV}zarYHt`i#F=ckubiO`Gs&STWH(UoY!kk@{5p?>}LemB4fOf
znSQ1*CiPit&3PRbJN1_8wwGJ$du{10d0ty~TfWyehbUQdiIO$9FD1QYWm`#W#V;vy
z4r#elf!CJXVafG}Fl#!o16~0b)3F45i^&^5t+Ch_JF#&|U%ENA(LL@LKx^?unj2>{
z&KMsbZ-}5h?kltwX0r<u4IHzD!kPHrLNok?X2&%1EXtTbxykkjlreG+HY_c&@U*GO
zWW{nvHtu#zH&6mf$S(81Wec8htTLVg#u`2$IgkjQ)`&|NMqyU5CQTid;fsnrw5=m5
zmAWr3sd?sSWv6zY%kyR}_%`30^<qaH@8~GDg@v%_Ndc#TrnWdaI=F95VLZEF)Nza#
z3X_zh#d|c_M#<u5&w)vylF`qe12I@x7O&KoxwZO-B4bUghV3$X4A_=qQlwQPKpkpV
zH$fsEU0ov$dkk~}`nN}6@|xtZi=GM2h+2UsMb5yqH3`bjHQPkkto*Y!LBOl)B9)BF
z4@J13oOdn*P?U&jn;!a=Y<Z&vrYeUIo;1S^rR3PVAyB9M`o*6BGPa-m4z?&eT6YVu
zN4fjuI{+UjZ?>J5ph=m2HeU^Al}%q4hJsgV{h>}xH0#fb0B+@C$9548qQiRzdsb)#
z6CiXZ!wWnJoZFe5PbZtzQgpOYI^}eXwR{|GRAY9SIfWzQ2`bk~`hx8(pitLciozqE
zW~{xWRS`{(dvl|NW(_Xg8;v8K#!S0dx<j+PYQ0?LPE6YopO)yVcPCcYZ*@5nYg~y9
zLBr-sSLKE(*{(%ZSDwCGx3aQsYn4Nm>t&qPRj58L(J6T!-|UK9F{Wk_R<AMhim~9b
zvB0M{kBSgsz0C-Etf_vHX(#lB8XI3aIB)+vU!-;al4}uhzHrB(*&X4IU?)Ok2@r#j
z?m}1zTtjcyWJ#V(QQ=yXDSQSEnphQEO*vUDI5n7cB~_xF(DrJtT!lWR-1k6)tGJ{m
z!E&9l(XFbkaXIC~M6;{q6#B_S%D$faP<WNB@-+?v9FUv(IgD~&mCV`lAcx}|ZV=EV
z3C_S3BY6Jz-|+N6wD6vQ-w)*At--07dOgG9y6a(^XpUfWN{`N14Ek@_L-{kMnnNa+
z7mr=fg~g<ug5qG6;M#(9gX>f05n271EYEyea%vt;H8nu8-N*=cYMnX{)W8P%C<&qH
z8%ocaP4sSwo;7!+T_1Snf_K|Q+wNAms$O@?eLRAj98Qq4Jcjd!<8aHc1%~jzaJ+CI
z(B+E><E|AI{k+1pfR~0BLkC(%?1m0~s?QMNHB9m|y><zfk5rSPeWXAXNV+wWhc|B~
zfEN1=(KG*evd|bG{mbD#2OD3PmsL*V(93lk+WBe?xZ$sg1e28SXU{^A)vK8Q9t5>S
z7x=!UNTcgjRP~y%G~q7bM`Mf3f#ORITk2}u@{fG-U=?w6Y^`C4<Zjp`Kj8X3So%sF
zzHui`GSOUN%BOeU7AXz3d&!cO6gf)CJ-Z5G@wJDur6f|!QqDZwB!sGID2H&+v`+%)
z6ruwm@cqbqBAJurq^58c#6xa-rV0+IEG~dpKYIqN+b?OL6y_UXil04!$@2jy0nf+L
Qfd2N=)<r_ot2$WoKMpH9n*aa+

diff --git a/litellm/llms/ai21.py b/litellm/llms/ai21.py
new file mode 100644
index 00000000000..9b856be4c7d
--- /dev/null
+++ b/litellm/llms/ai21.py
@@ -0,0 +1,127 @@
+import os, json
+from enum import Enum
+import requests
+import time
+from typing import Callable
+from litellm.utils import ModelResponse
+
+class AI21Error(Exception):
+    def __init__(self, status_code, message):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(
+            self.message
+        )  # Call the base class constructor with the parameters it needs
+
+
+class AI21LLM:
+    def __init__(
+        self, encoding, logging_obj, api_key=None
+    ):
+        self.encoding = encoding
+        self.completion_url_fragment_1 = "https://api.ai21.com/studio/v1/"
+        self.completion_url_fragment_2 = "/complete"
+        self.api_key = api_key
+        self.logging_obj = logging_obj
+        self.validate_environment(api_key=api_key)
+
+    def validate_environment(
+        self, api_key
+    ):  # set up the environment required to run the model
+        # set the api key
+        if self.api_key == None:
+            raise ValueError(
+                "Missing Baseten API Key - A call is being made to baseten but no key is set either in the environment variables or via params"
+            )
+        self.api_key = api_key
+        self.headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "Authorization": "Bearer " + self.api_key,
+        }
+
+    def completion(
+        self,
+        model: str,
+        messages: list,
+        model_response: ModelResponse,
+        print_verbose: Callable,
+        optional_params=None,
+        litellm_params=None,
+        logger_fn=None,
+    ):  # logic for parsing in - calling - parsing out model completion calls
+        model = model
+        prompt = ""
+        for message in messages:
+            if "role" in message:
+                if message["role"] == "user":
+                    prompt += (
+                        f"{message['content']}"
+                    )
+                else:
+                    prompt += (
+                        f"{message['content']}"
+                    )
+            else:
+                prompt += f"{message['content']}"
+        data = {
+            "prompt": prompt,
+            # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg
+            **optional_params,
+        }
+
+        ## LOGGING
+        self.logging_obj.pre_call(
+            input=prompt,
+            api_key=self.api_key,
+            additional_args={"complete_input_dict": data},
+        )
+        ## COMPLETION CALL
+        response = requests.post(
+            self.completion_url_fragment_1 + model + self.completion_url_fragment_2, headers=self.headers, data=json.dumps(data)
+        )
+        if "stream" in optional_params and optional_params["stream"] == True:
+            return response.iter_lines()
+        else:
+            ## LOGGING
+            self.logging_obj.post_call(
+                input=prompt,
+                api_key=self.api_key,
+                original_response=response.text,
+                additional_args={"complete_input_dict": data},
+            )
+            print_verbose(f"raw model_response: {response.text}")
+            ## RESPONSE OBJECT
+            completion_response = response.json()
+            if "error" in completion_response:
+                raise AI21Error(
+                    message=completion_response["error"],
+                    status_code=response.status_code,
+                )
+            else:
+                try:
+                    model_response["choices"][0]["message"]["content"] = completion_response["completions"][0]["data"]["text"]
+                except:
+                    raise ValueError(f"Unable to parse response. Original response: {response.text}")
+
+            ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. 
+            prompt_tokens = len(
+                self.encoding.encode(prompt)
+            ) 
+            completion_tokens = len(
+                self.encoding.encode(model_response["choices"][0]["message"]["content"])
+            )
+
+            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,
+            }
+            return model_response
+
+    def embedding(
+        self,
+    ):  # logic for parsing in - calling - parsing out model embedding calls
+        pass
diff --git a/litellm/llms/baseten.py b/litellm/llms/baseten.py
index cc0fcec8d14..49753d67b34 100644
--- a/litellm/llms/baseten.py
+++ b/litellm/llms/baseten.py
@@ -62,7 +62,7 @@ class BasetenLLM:
         data = {
             "prompt": prompt,
             # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg
-            # **optional_params,
+            **optional_params,
         }
 
         ## LOGGING
diff --git a/litellm/main.py b/litellm/main.py
index fa8cc847b53..f20f173ff73 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -22,6 +22,7 @@ from litellm.utils import (
 from .llms.anthropic import AnthropicLLM
 from .llms.huggingface_restapi import HuggingfaceRestAPILLM
 from .llms.baseten import BasetenLLM
+from .llms.ai21 import AI21LLM
 import tiktoken
 from concurrent.futures import ThreadPoolExecutor
 
@@ -302,7 +303,11 @@ def completion(
                     headers=litellm.headers,
                 )
             else:
-                response = openai.Completion.create(model=model, prompt=prompt)
+                response = openai.Completion.create(model=model, prompt=prompt, **optional_params)
+            
+            if "stream" in optional_params and optional_params["stream"] == True:
+                response = CustomStreamWrapper(response, model)
+                return response
             ## LOGGING
             logging.post_call(
                 input=prompt,
@@ -661,32 +666,34 @@ def completion(
             model_response["model"] = model
             response = model_response
         elif model in litellm.ai21_models:
-            install_and_import("ai21")
-            import ai21
-
-            ai21.api_key = get_secret("AI21_API_KEY")
-
-            prompt = " ".join([message["content"] for message in messages])
-            ## LOGGING
-            logging.pre_call(input=prompt, api_key=ai21.api_key)
-
-            ai21_response = ai21.Completion.execute(
+            custom_llm_provider = "ai21"
+            ai21_key = (
+                api_key
+                or litellm.ai21_key
+                or os.environ.get("AI21_API_KEY")
+            )
+            ai21_client = AI21LLM(
+                encoding=encoding, api_key=ai21_key, logging_obj=logging
+            )
+            
+            model_response = ai21_client.completion(
                 model=model,
-                prompt=prompt,
+                messages=messages,
+                model_response=model_response,
+                print_verbose=print_verbose,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                logger_fn=logger_fn,
             )
-            completion_response = ai21_response["completions"][0]["data"]["text"]
-
-            ## LOGGING
-            logging.post_call(
-                input=prompt,
-                api_key=ai21.api_key,
-                original_response=completion_response,
-            )
-
+            
+            if "stream" in optional_params and optional_params["stream"] == True:
+                # don't try to access stream object,
+                response = CustomStreamWrapper(
+                    model_response, model, custom_llm_provider="ai21"
+                )
+                return response
+            
             ## RESPONSE OBJECT
-            model_response["choices"][0]["message"]["content"] = completion_response
-            model_response["created"] = time.time()
-            model_response["model"] = model
             response = model_response
         elif custom_llm_provider == "ollama":
             endpoint = (
@@ -725,7 +732,7 @@ def completion(
             if "stream" in optional_params and optional_params["stream"] == True:
                 # don't try to access stream object,
                 response = CustomStreamWrapper(
-                    model_response, model, custom_llm_provider="huggingface"
+                    model_response, model, custom_llm_provider="baseten"
                 )
                 return response
             response = model_response
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index f6a676177a5..306a317eb90 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -18,30 +18,85 @@ score = 0
 
 
 def logger_fn(model_call_object: dict):
-    return
     print(f"model call details: {model_call_object}")
 
 
 user_message = "Hello, how are you?"
 messages = [{"content": user_message, "role": "user"}]
 
+# test on baseten completion call
+try:
+    response = completion(
+        model="wizard-lm", messages=messages, stream=True, logger_fn=logger_fn
+    )
+    print(f"response: {response}")
+    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 openai completion call
-# try:
-#     response = completion(
-#         model="gpt-3.5-turbo", 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
+try:
+    response = completion(
+        model="text-davinci-003", messages=messages, stream=True, logger_fn=logger_fn
+    )
+    complete_response = ""
+    start_time = time.time()
+    for chunk in response:
+        chunk_time = time.time()
+        print(f"chunk: {chunk}")
+        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 ai21 completion call
+try:
+    response = completion(
+        model="j2-ultra", messages=messages, stream=True, logger_fn=logger_fn
+    )
+    print(f"response: {response}")
+    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 openai completion call
+try:
+    response = completion(
+        model="gpt-3.5-turbo", 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 azure completion call
@@ -63,25 +118,6 @@ messages = [{"content": user_message, "role": "user"}]
 #     pass
 
 
-# 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()
-    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
 # try:
 #     start_time = time.time()
@@ -123,7 +159,7 @@ except:
     print(f"error occurred: {traceback.format_exc()}")
     pass
 
-# test on together ai completion call - starcoder
+# # test on together ai completion call - starcoder
 try:
     start_time = time.time()
     response = completion(
@@ -148,57 +184,3 @@ try:
 except:
     print(f"error occurred: {traceback.format_exc()}")
     pass
-
-
-# # test on azure completion call
-# try:
-#     response = completion(
-#         model="azure/chatgpt-test", messages=messages, stream=True, logger_fn=logger_fn
-#     )
-#     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"])
-#         response += chunk["choices"][0]["delta"]
-#     if response == "":
-#         raise Exception("Empty response received")
-# except:
-#     print(f"error occurred: {traceback.format_exc()}")
-#     pass
-
-
-# # test on anthropic completion call
-# try:
-#     response = completion(
-#         model="claude-instant-1", messages=messages, stream=True, logger_fn=logger_fn
-#     )
-#     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"])
-#         response += chunk["choices"][0]["delta"]
-#     if response == "":
-#         raise Exception("Empty response received")
-# except:
-#     print(f"error occurred: {traceback.format_exc()}")
-#     pass
-
-
-# # # test on huggingface completion call
-# # try:
-# #     response = completion(
-# #         model="meta-llama/Llama-2-7b-chat-hf",
-# #         messages=messages,
-# #         custom_llm_provider="huggingface",
-# #         custom_api_base="https://s7c7gytn18vnu4tw.us-east-1.aws.endpoints.huggingface.cloud",
-# #         stream=True,
-# #         logger_fn=logger_fn,
-# #     )
-# #     for chunk in response:
-# #         print(chunk["choices"][0]["delta"])
-# #     score += 1
-# # except:
-# #     print(f"error occurred: {traceback.format_exc()}")
-# #     pass
diff --git a/litellm/utils.py b/litellm/utils.py
index 7e85decf5f7..fb136a3301f 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -648,6 +648,7 @@ def get_optional_params(  # use the openai defaults
         optional_params["top_k"] = top_k
     elif custom_llm_provider == "baseten":
         optional_params["temperature"] = temperature
+        optional_params["stream"] = stream
         optional_params["top_p"] = top_p
         optional_params["top_k"] = top_k
         optional_params["num_beams"] = num_beams
@@ -1561,6 +1562,35 @@ class CustomStreamWrapper:
             else:
                 return ""
         return ""
+    
+    def handle_ai21_chunk(self, chunk):
+        chunk = chunk.decode("utf-8")
+        data_json = json.loads(chunk)
+        try:
+            return data_json["completions"][0]["data"]["text"]
+        except:
+            raise ValueError(f"Unable to parse response. Original response: {chunk}")
+    
+    def handle_openai_text_completion_chunk(self, chunk):
+        try:
+            return chunk["choices"][0]["text"]
+        except:
+            raise ValueError(f"Unable to parse response. Original response: {chunk}")
+    
+    def handle_baseten_chunk(self, chunk):
+        chunk = chunk.decode("utf-8")
+        data_json = json.loads(chunk)
+        if "model_output" in data_json:
+            if isinstance(data_json["model_output"], dict) and "data" in data_json["model_output"] and isinstance(data_json["model_output"]["data"], list):
+                return data_json["model_output"]["data"][0]
+            elif isinstance(data_json["model_output"], str):
+                return data_json["model_output"]
+            elif "completion" in data_json and isinstance(data_json["completion"], str):
+                return data_json["completion"]
+            else:
+                raise ValueError(f"Unable to parse response. Original response: {chunk}")
+        else:
+            return ""
 
     def __next__(self):
         completion_obj = {"role": "assistant", "content": ""}
@@ -1584,6 +1614,15 @@ class CustomStreamWrapper:
         elif self.custom_llm_provider and self.custom_llm_provider == "huggingface":
             chunk = next(self.completion_stream)
             completion_obj["content"] = self.handle_huggingface_chunk(chunk)
+        elif self.custom_llm_provider and self.custom_llm_provider == "baseten": # baseten doesn't provide streaming
+            chunk = next(self.completion_stream)
+            completion_obj["content"] = self.handle_baseten_chunk(chunk)
+        elif self.custom_llm_provider and self.custom_llm_provider == "ai21": #ai21 doesn't provide streaming
+            chunk = next(self.completion_stream)
+            completion_obj["content"] = self.handle_ai21_chunk(chunk)
+        elif self.model in litellm.open_ai_text_completion_models:
+            chunk = next(self.completion_stream)
+            completion_obj["content"] = self.handle_openai_text_completion_chunk(chunk)
         # return this for all models
         return {"choices": [{"delta": completion_obj}]}