From 436e8eadb2bf600d632d83d3952fefc4263b7bde Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 29 Aug 2023 13:32:20 -0700 Subject: [PATCH] adding coverage for ai21 --- litellm/__pycache__/utils.cpython-311.pyc | Bin 73465 -> 74623 bytes litellm/llms/ai21.py | 2 +- litellm/tests/test_exceptions.py | 9 +++++-- litellm/utils.py | 31 ++++++++++++++++++++++ 4 files changed, 39 insertions(+), 3 deletions(-) diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index bf7b112ca47224d20d5d4bf242caba7ef2b62ddd..6bebd64510adead19d083727b21193adcfddcf48 100644 GIT binary patch delta 2314 zcma)7eN0nV6o2Qv*A|L^(3aAI^r_TR6eu9_5t*{c6d8^$qh_12rVkyGsWuWZuc#6&rTpX#*_`0=vtqz^98o8{<>3l<40Y;nlY@@lP{>v2IG zWBrJq?PPbNp-V~>?$NVdjCCs-?Y+U!))935qBohg?sL+7);5~(3Gca2Uy=gXMO#^H zS(0NGigE%!?+|AAF2X|Wf}Wr6gsllNpPiJD$dqL479vq`f(;m4g8{Mq*1Ap>SyG{ z7dTyN4e)u93Xt>XO)~KL6aDqv)Z|`Da{u2oGaY0~w-h*~q+^q(cd-(^>n%?-eN8)M z#|oo;hq%LPH0WKz6gf=D^qCUn`%0h~@xBIF?EcvogvNf>Y$+=2WLhgYp%fRnQ(G_C zcCI%X0S^x8mu$D0c{}>!%CZC>UJF;q zju$QJd#Nx$)&vS96lL#-qo_U9O)e@m8w4?^`+#CDfy`GTKevwLMvwE9qI5Y%KeJ+5G6~|2Gqr3X^p?MK~$3w zC@vM#O~I_PKz7rF9Z#~G0`m+3V?jv9>a;@?hG7CBiq&Wz%PD0BVcq3bRQeT_AxfE| z4JtLFcD0zh)StVGOth=L%38lfSt}}QQQcLIc`7+(mN0}F!!Usm#b~rxe|oN8U2;O< zQ!f?qSpJ&0N$_v#6ql0`lUg{VRCMBMf$4d&Oq40nmr~+I>_B-^EGJko!$Ti}D zwAZxg^#NnbRC7RIgrHxE4h|S+3b;3rGlMxNB&X7H{Ho&PcAu(3R8_=|uqcqdL9Bem zU%5`ilk5#>uwP9-jzdeXh1HE2jXKJwlQ)|fpYcVyc_HIpC}}piR|I81OQr6*>p#+T zveXTOFEel(t+~a?2T0ox(7W4j<0Z(?0~{cxhw>!Xr@lZza<7LLXipWhS_Yl|gZ#^|*D-p>-tn$)N^S zB!s*o5pSW5uq7eg!_~5Pr1*7cq0nGB8-4mPpiIRKMu>PErDFoShjr+iM{D4H^x)AM zP@(SO9q_SR_gFz={^Ic5OcE897Z;T`nd^7d7S=!#`Y>D%``iQJTxqFrktkV4M5L2+ zN}-C(YY7zKt$x-*9-AbqmJCI#NM{u;t&Q2PqYB81Mz_G-}MChK+v;N4($wH z{*s6z$31rv!NSr~nWrfg!VLI5wQ0}++UTc>$o`E$qy-~=H|)8thDxcTi2;3xdJ3&l JSAG`s{0r|Tg~I>< delta 1582 zcmbVKZA_C_6u#&Eq83~EX!+{W_?5w+Cv24;n3WrfzcpUG$60jr@;!&hD>be-TK%d;wf1>H$2nW% z1h;US&JCF8QlDM8#5rr}>aUYLIEU%|{=seFbhtS=^8|-_%17$00SPjklPqVA=EhxF z$VFCb0!_G-*4-t^*>34HQEYOKo5krdp0J!n=DK)S?s|%Bu5c6Cdc)}D3AxUGk@X=& z^>XFh&DBgVUDGqiY6)?7ZyL$uwtAWW>)IKxY+cSK+x{gNvA5y)zuB{E$KB4PI8JB;BbI08%pb3FSzR*=ni$62b zbAeLWNxumkfg+U~nE+Zc`XaqBdVoL2x6r$zH_bjEBq+9>YL+{cWHwBxXF%y#Rs0}h zlo3q%>BQLnl2P0v>=p^>_);R|JekU}`Gt%ScY+h`6Y*w$%8sv#2IDIN@f9KR7JYsE z1vspJJZ=@XE{Aj^$$r%^lvwJgJ0~lpU02TptmU+0@|7)-aLm4k&IL{Mx5@k%nKV_C z6D2cq86S!ho=ce{44xo_h+s~6s;50uOP2c5`byGYsc+Ep1AGnFAo6}uWSe6=9Soc3 zgXs`AEvkVB-*McGs9t>d8V?WY$Rn#hh^Imz>e{0sPCtu#gmks+j~RePn*F$zUE%S^ zJy1q%vv;6Uotb?DV1!nOx?rz*D|8fKg+4!@nf5np3kt^_PFcCj1oI^&53i5V#`#a6 zR-K+d3Hw%1e-+^`gkCK86G6|m`GSHP${GTG8D)gQH9o4tUa!K@PoKOEhw1cuhAKRL z2}f5i9x-ce>m5ovDjKMF@gTgT&Mk6aN@UIXpb^)&n6HX(7HwKK)3&e;|Dlbt82Qg={j zcI+ups*p`ZFD7Ll#)lAgYZ}g$vo2+kg$j@3w5J!(qM)v+3Av^wr5IW5x{O$ofslmo zITB`!(h&Tp`UQiaq0x1!z0T#)^NA4^ln1 z4N#Y=Jr=Q}z0KiLrcftxr5>9RjxwVcj8F$Z`1*}-mS@}jNvt#Ib0&d9;HD+tcnU1> rFzg$&Kqpv6v40F5{fwZ!Y3(5|_~Z=OD assert if the correct error (always the same one) is triggered # models = ["gpt-3.5-turbo", "chatgpt-test", "claude-instant-1", "command-nightly"] -test_model = "togethercomputer/CodeLlama-34b-Python" -models = ["togethercomputer/CodeLlama-34b-Python"] +test_model = "j2-light" +models = ["j2-light"] def logging_fn(model_call_dict): @@ -98,6 +98,9 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th elif model == "command-nightly": temporary_key = os.environ["COHERE_API_KEY"] os.environ["COHERE_API_KEY"] = "bad-key" + elif "j2" in model: + temporary_key = os.environ["AI21_API_KEY"] + os.environ["AI21_API_KEY"] = "bad-key" elif "togethercomputer" in model: temporary_key = os.environ["TOGETHERAI_API_KEY"] os.environ["TOGETHERAI_API_KEY"] = "84060c79880fc49df126d3e87b53f8a463ff6e1c6d27fe64207cde25cdfcd1f24a" @@ -138,6 +141,8 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th == "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1" ): os.environ["REPLICATE_API_KEY"] = temporary_key + elif "j2" in model: + os.environ["AI21_API_KEY"] = temporary_key elif ("togethercomputer" in model): os.environ["TOGETHERAI_API_KEY"] = temporary_key return diff --git a/litellm/utils.py b/litellm/utils.py index 994b4916c9..18c32d58b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1446,6 +1446,37 @@ def exception_type(model, original_exception, custom_llm_provider): message=f"HuggingfaceException - {original_exception.message}", llm_provider="huggingface", ) + elif custom_llm_provider == "ai21": + print(f"e: {original_exception}") + if hasattr(original_exception, "message"): + if "Prompt has too many tokens" in original_exception.message: + exception_mapping_worked = True + raise ContextWindowExceededError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21" + ) + if hasattr(original_exception, "status_code"): + print(f"status code: {original_exception.status_code}") + if original_exception.status_code == 401: + exception_mapping_worked = True + raise AuthenticationError( + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + ) + if original_exception.status_code == 422 or "Prompt has too many tokens" in original_exception.message: + exception_mapping_worked = True + raise InvalidRequestError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + ) + elif original_exception.status_code == 429: + exception_mapping_worked = True + raise RateLimitError( + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + ) elif custom_llm_provider == "together_ai": error_response = json.loads(error_str) if "error" in error_response and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"]: