From 400653992c2f1b37b3a035191b38d1d43712cfa1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 7 Aug 2024 16:43:12 -0700 Subject: [PATCH 1/5] feat(router.py): allow using .acompletion() for request prioritization allows /chat/completion endpoint to work for request prioritization calls --- docs/my-website/sidebars.js | 2 +- litellm/router.py | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 414838280e..f84a43c56f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -193,12 +193,12 @@ const sidebars = { "vertex_ai" ], }, + "scheduler", { type: "category", label: "🚅 LiteLLM Python SDK", items: [ "routing", - "scheduler", "set_keys", "completion/token_usage", "sdk_custom_pricing", diff --git a/litellm/router.py b/litellm/router.py index 74562566db..fb9af96188 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -654,7 +654,12 @@ class Router: timeout = kwargs.get("request_timeout", self.timeout) kwargs.setdefault("metadata", {}).update({"model_group": model}) - response = await self.async_function_with_fallbacks(**kwargs) + if kwargs.get("priority", None) is not None and isinstance( + kwargs.get("priority"), int + ): + response = await self.schedule_acompletion(**kwargs) + else: + response = await self.async_function_with_fallbacks(**kwargs) return response except Exception as e: @@ -1097,6 +1102,10 @@ class Router: _response = await self.acompletion( model=model, messages=messages, stream=stream, **kwargs ) + _response._hidden_params.setdefault("additional_headers", {}) + _response._hidden_params["additional_headers"].update( + {"x-litellm-request-prioritization-used": True} + ) return _response except Exception as e: setattr(e, "priority", priority) From 4c375d13d6fe19e8672dcc7f06ef44dbb38a2ad0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 7 Aug 2024 16:26:56 -0700 Subject: [PATCH 2/5] fix test_drop_params_parallel_tool_calls --- litellm/tests/test_optional_params.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_optional_params.py b/litellm/tests/test_optional_params.py index b2b0a0a2a4..2cd5c11492 100644 --- a/litellm/tests/test_optional_params.py +++ b/litellm/tests/test_optional_params.py @@ -345,7 +345,7 @@ def test_drop_params_parallel_tool_calls(model, provider, should_drop): response = litellm.utils.get_optional_params( model=model, custom_llm_provider=provider, - response_format="json", + response_format={"type": "json"}, parallel_tool_calls=True, drop_params=True, ) From 472a1df76f32f7ce987e0e6cc24a2a5ad2949c84 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 7 Aug 2024 16:41:00 -0700 Subject: [PATCH 3/5] fix test for wildcard routing --- .circleci/config.yml | 1 + litellm/tests/test_completion.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 385c913981..4fbd58c003 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -209,6 +209,7 @@ jobs: -e MISTRAL_API_KEY=$MISTRAL_API_KEY \ -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \ -e GROQ_API_KEY=$GROQ_API_KEY \ + -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \ -e COHERE_API_KEY=$COHERE_API_KEY \ -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \ -e AWS_REGION_NAME=$AWS_REGION_NAME \ diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 45c9c64437..9367b98db5 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -23,7 +23,7 @@ from litellm import RateLimitError, Timeout, completion, completion_cost, embedd from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.prompt_templates.factory import anthropic_messages_pt -# litellm.num_retries = 3 +# litellm.num_retries=3 litellm.cache = None litellm.success_callback = [] user_message = "Write a short poem about the sky" From 75bfe817485c7d139f242c9490c7fc6b19ddf6ce Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 7 Aug 2024 17:52:40 -0700 Subject: [PATCH 4/5] fix all optional param tests --- litellm/tests/test_optional_params.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_optional_params.py b/litellm/tests/test_optional_params.py index 2cd5c11492..d961190c29 100644 --- a/litellm/tests/test_optional_params.py +++ b/litellm/tests/test_optional_params.py @@ -301,7 +301,7 @@ def test_dynamic_drop_params(drop_params): optional_params = litellm.utils.get_optional_params( model="command-r", custom_llm_provider="cohere", - response_format="json", + response_format={"type": "json"}, drop_params=drop_params, ) else: @@ -309,7 +309,7 @@ def test_dynamic_drop_params(drop_params): optional_params = litellm.utils.get_optional_params( model="command-r", custom_llm_provider="cohere", - response_format="json", + response_format={"type": "json"}, drop_params=drop_params, ) pytest.fail("Expected to fail") @@ -389,7 +389,7 @@ def test_dynamic_drop_additional_params(drop_params): optional_params = litellm.utils.get_optional_params( model="command-r", custom_llm_provider="cohere", - response_format="json", + response_format={"type": "json"}, additional_drop_params=["response_format"], ) else: @@ -397,7 +397,7 @@ def test_dynamic_drop_additional_params(drop_params): optional_params = litellm.utils.get_optional_params( model="command-r", custom_llm_provider="cohere", - response_format="json", + response_format={"type": "json"}, ) pytest.fail("Expected to fail") except Exception as e: From 4e02425d8b5374be2c084d66dfc62eba752e3fa6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 7 Aug 2024 18:07:14 -0700 Subject: [PATCH 5/5] fix(vertex_ai_partner.py): pass model for llama3 param mapping --- litellm/llms/vertex_ai_partner.py | 6 ++++-- litellm/utils.py | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai_partner.py b/litellm/llms/vertex_ai_partner.py index 378ee7290d..24586a3fe4 100644 --- a/litellm/llms/vertex_ai_partner.py +++ b/litellm/llms/vertex_ai_partner.py @@ -96,11 +96,13 @@ class VertexAILlama3Config: def get_supported_openai_params(self): return litellm.OpenAIConfig().get_supported_openai_params(model="gpt-3.5-turbo") - def map_openai_params(self, non_default_params: dict, optional_params: dict): + def map_openai_params( + self, non_default_params: dict, optional_params: dict, model: str + ): return litellm.OpenAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, - model="gpt-3.5-turbo", + model=model, ) diff --git a/litellm/utils.py b/litellm/utils.py index 98c8b01841..a20e961727 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3190,6 +3190,7 @@ def get_optional_params( optional_params = litellm.VertexAILlama3Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, + model=model, ) elif custom_llm_provider == "vertex_ai" and model in litellm.vertex_mistral_models: supported_params = get_supported_openai_params(