From 7762ae7762bfc067e1086d2102bbecce6f677649 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 31 Oct 2023 22:32:29 -0700 Subject: [PATCH] feat(utils.py): accept context window fallback dictionary --- litellm/main.py | 14 +++----------- litellm/tests/test_completion_with_retries.py | 2 ++ litellm/tests/test_exceptions.py | 11 ++++++++++- litellm/utils.py | 17 +++++++++++++++++ 4 files changed, 32 insertions(+), 12 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index e3df4f80892..b1ed2eb0dcb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -255,9 +255,10 @@ def completion( fallbacks = kwargs.get('fallbacks', None) headers = kwargs.get("headers", None) num_retries = kwargs.get("num_retries", None) + context_window_fallback_dict = kwargs.get("context_window_fallback_dict", None) ######## end of unpacking kwargs ########### openai_params = ["functions", "function_call", "temperature", "temperature", "top_p", "n", "stream", "stop", "max_tokens", "presence_penalty", "frequency_penalty", "logit_bias", "user", "request_timeout", "api_base", "api_version", "api_key"] - litellm_params = ["metadata", "acompletion", "caching", "return_async", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries"] + litellm_params = ["metadata", "acompletion", "caching", "return_async", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries", "context_window_fallback_dict"] default_params = openai_params + litellm_params non_default_params = {k: v for k,v in kwargs.items() if k not in default_params} # model-specific params - pass them straight to the model/provider if mock_response: @@ -1326,18 +1327,9 @@ def completion( return response except Exception as e: ## Map to OpenAI Exception - try: - raise exception_type( + raise exception_type( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=args, ) - except Exception as e: - if num_retries: - if (isinstance(e, openai.error.APIError) - or isinstance(e, openai.error.Timeout) - or isinstance(e, openai.error.ServiceUnavailableError)): - return completion_with_retries(num_retries=num_retries, **args) - else: - raise e def completion_with_retries(*args, **kwargs): diff --git a/litellm/tests/test_completion_with_retries.py b/litellm/tests/test_completion_with_retries.py index f5096be77c2..30ec9ec8c7b 100644 --- a/litellm/tests/test_completion_with_retries.py +++ b/litellm/tests/test_completion_with_retries.py @@ -49,3 +49,5 @@ def test_completion_with_num_retries(): pass except Exception as e: pytest.fail(f"Unmapped exception occurred") + +# test_completion_with_num_retries() \ No newline at end of file diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py index 0bed38418b8..b7cb8375ac2 100644 --- a/litellm/tests/test_exceptions.py +++ b/litellm/tests/test_exceptions.py @@ -44,7 +44,16 @@ def test_context_window(model): with pytest.raises(ContextWindowExceededError): completion(model=model, messages=messages) -test_context_window(model="command-nightly") +@pytest.mark.parametrize("model", models) +def test_context_window_with_fallbacks(model): + ctx_window_fallback_dict = {"command-nightly": "claude-2"} + sample_text = "how does a court case get to the Supreme Court?" * 1000 + messages = [{"content": sample_text, "role": "user"}] + + completion(model=model, messages=messages, context_window_fallback_dict=ctx_window_fallback_dict) + +# test_context_window(model="command-nightly") +test_context_window_with_fallbacks(model="command-nightly") # Test 2: InvalidAuth Errors @pytest.mark.parametrize("model", models) def invalid_auth(model): # set the model key to an invalid key, depending on the model diff --git a/litellm/utils.py b/litellm/utils.py index 19c6705a3d7..76d79429faa 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -824,6 +824,23 @@ def client(original_function): result._response_ms = (end_time - start_time).total_seconds() * 1000 # return response latency in ms like openai return result except Exception as e: + call_type = original_function.__name__ + if call_type == CallTypes.completion.value: + num_retries = kwargs.get("num_retries", None) + context_window_fallback_dict = kwargs.get("context_window_fallback_dict", {}) + + if num_retries: + if (isinstance(e, openai.error.APIError) + or isinstance(e, openai.error.Timeout) + or isinstance(e, openai.error.ServiceUnavailableError)): + kwargs["num_retries"] = num_retries + return litellm.completion_with_retries(*args, **kwargs) + elif isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict and model in context_window_fallback_dict: + if len(args) > 0: + args[0] = context_window_fallback_dict[model] + else: + kwargs["model"] = context_window_fallback_dict[model] + return original_function(*args, **kwargs) traceback_exception = traceback.format_exc() crash_reporting(*args, **kwargs, exception=traceback_exception) end_time = datetime.datetime.now()