From c6be8326dbd0da26fcb2164f1a140929cd14fd92 Mon Sep 17 00:00:00 2001 From: Peter Muller Date: Mon, 1 Jul 2024 16:00:42 -0700 Subject: [PATCH 1/8] Allow calling SageMaker endpoints from different regions --- litellm/llms/sagemaker.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 8e75428bb7c..079951b935b 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -185,7 +185,8 @@ def completion( # I assume majority of users use .env for auth region_name = ( get_secret("AWS_REGION_NAME") - or "us-west-2" # default to us-west-2 if user not specified + or aws_region_name # get region from config file if specified + or "us-west-2" # default to us-west-2 if region not specified ) client = boto3.client( service_name="sagemaker-runtime", @@ -439,7 +440,8 @@ async def async_streaming( # I assume majority of users use .env for auth region_name = ( get_secret("AWS_REGION_NAME") - or "us-west-2" # default to us-west-2 if user not specified + or aws_region_name # get region from config file if specified + or "us-west-2" # default to us-west-2 if region not specified ) _client = session.client( service_name="sagemaker-runtime", @@ -506,7 +508,8 @@ async def async_completion( # I assume majority of users use .env for auth region_name = ( get_secret("AWS_REGION_NAME") - or "us-west-2" # default to us-west-2 if user not specified + or aws_region_name # get region from config file if specified + or "us-west-2" # default to us-west-2 if region not specified ) _client = session.client( service_name="sagemaker-runtime", @@ -661,7 +664,8 @@ def embedding( # I assume majority of users use .env for auth region_name = ( get_secret("AWS_REGION_NAME") - or "us-west-2" # default to us-west-2 if user not specified + or aws_region_name # get region from config file if specified + or "us-west-2" # default to us-west-2 if region not specified ) client = boto3.client( service_name="sagemaker-runtime", From d9e9a8645bf05c6dbd7519034ddf1ce3f5f21d85 Mon Sep 17 00:00:00 2001 From: Peter Muller Date: Tue, 2 Jul 2024 15:30:39 -0700 Subject: [PATCH 2/8] Add tests for SageMaker region selection --- litellm/llms/sagemaker.py | 19 +--- .../tests/test_provider_specific_config.py | 100 ++++++++++++++++++ 2 files changed, 103 insertions(+), 16 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 079951b935b..0e0fa8006e1 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -9,6 +9,9 @@ from litellm.utils import ModelResponse, EmbeddingResponse, get_secret, Usage import sys from copy import deepcopy import httpx # type: ignore +import boto3 +import aioboto3 +import io from .prompt_templates.factory import prompt_factory, custom_prompt @@ -25,10 +28,6 @@ class SagemakerError(Exception): ) # Call the base class constructor with the parameters it needs -import io -import json - - class TokenIterator: def __init__(self, stream, acompletion: bool = False): if acompletion == False: @@ -160,8 +159,6 @@ def completion( logger_fn=None, acompletion: bool = False, ): - import boto3 - # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) aws_access_key_id = optional_params.pop("aws_access_key_id", None) @@ -416,10 +413,6 @@ async def async_streaming( aws_access_key_id: Optional[str], aws_region_name: Optional[str], ): - """ - Use aioboto3 - """ - import aioboto3 session = aioboto3.Session() @@ -484,10 +477,6 @@ async def async_completion( aws_access_key_id: Optional[str], aws_region_name: Optional[str], ): - """ - Use aioboto3 - """ - import aioboto3 session = aioboto3.Session() @@ -639,8 +628,6 @@ def embedding( """ Supports Huggingface Jumpstart embeddings like GPT-6B """ - ### BOTO3 INIT - import boto3 # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) diff --git a/litellm/tests/test_provider_specific_config.py b/litellm/tests/test_provider_specific_config.py index 08a84b5604e..e79b5769f92 100644 --- a/litellm/tests/test_provider_specific_config.py +++ b/litellm/tests/test_provider_specific_config.py @@ -512,6 +512,106 @@ def sagemaker_test_completion(): # sagemaker_test_completion() + +def test_sagemaker_default_region(mocker): + """ + If no regions are specified in config or in environment, the default region is us-west-2 + """ + mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + try: + response = litellm.completion( + model="sagemaker/mock-endpoint", + messages=[ + { + "content": "Hello, world!", + "role": "user" + } + ] + ) + except Exception: + pass # expected serialization exception because AWS client was replaced with a Mock + assert mock_client.call_args.kwargs["region_name"] == "us-west-2" + +# test_sagemaker_provided_region() + + +def test_sagemaker_environment_region(mocker): + """ + If a region is specified in the environment, use that region instead of us-west-2 + """ + expected_region = "us-east-1" + os.environ["AWS_REGION_NAME"] = expected_region + mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + try: + response = litellm.completion( + model="sagemaker/mock-endpoint", + messages=[ + { + "content": "Hello, world!", + "role": "user" + } + ] + ) + except Exception: + pass # expected serialization exception because AWS client was replaced with a Mock + del os.environ["AWS_REGION_NAME"] # cleanup + assert mock_client.call_args.kwargs["region_name"] == expected_region + +# test_sagemaker_environment_region() + + +def test_sagemaker_config_region(mocker): + """ + If a region is specified as part of the optional parameters of the completion, including as + part of the config file, then use that region instead of us-west-2 + """ + expected_region = "us-east-1" + mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + try: + response = litellm.completion( + model="sagemaker/mock-endpoint", + messages=[ + { + "content": "Hello, world!", + "role": "user" + } + ], + aws_region_name=expected_region, + ) + except Exception: + pass # expected serialization exception because AWS client was replaced with a Mock + assert mock_client.call_args.kwargs["region_name"] == expected_region + +# test_sagemaker_config_region() + + +def test_sagemaker_config_and_environment_region(mocker): + """ + If both the environment and config file specify a region, the environment region is expected + """ + expected_region = "us-east-1" + unexpected_region = "us-east-2" + os.environ["AWS_REGION_NAME"] = expected_region + mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + try: + response = litellm.completion( + model="sagemaker/mock-endpoint", + messages=[ + { + "content": "Hello, world!", + "role": "user" + } + ], + aws_region_name=unexpected_region, + ) + except Exception: + pass # expected serialization exception because AWS client was replaced with a Mock + del os.environ["AWS_REGION_NAME"] # cleanup + assert mock_client.call_args.kwargs["region_name"] == expected_region + +# test_sagemaker_config_and_environment_region() + + # Bedrock From 47c97e1fa28345859da11de920c533681f0e1aa1 Mon Sep 17 00:00:00 2001 From: Peter Muller Date: Tue, 2 Jul 2024 15:38:15 -0700 Subject: [PATCH 3/8] Fix test name typo in comment --- litellm/tests/test_provider_specific_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_provider_specific_config.py b/litellm/tests/test_provider_specific_config.py index e79b5769f92..c1d5362ec12 100644 --- a/litellm/tests/test_provider_specific_config.py +++ b/litellm/tests/test_provider_specific_config.py @@ -532,7 +532,7 @@ def test_sagemaker_default_region(mocker): pass # expected serialization exception because AWS client was replaced with a Mock assert mock_client.call_args.kwargs["region_name"] == "us-west-2" -# test_sagemaker_provided_region() +# test_sagemaker_default_region() def test_sagemaker_environment_region(mocker): From d8fc8252fa6d4371a46643cf255a948837563bb9 Mon Sep 17 00:00:00 2001 From: Peter Muller Date: Tue, 2 Jul 2024 19:09:22 -0700 Subject: [PATCH 4/8] Revert imports changes, update tests to match --- litellm/llms/sagemaker.py | 14 ++++++++++++-- litellm/tests/test_provider_specific_config.py | 8 ++++---- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 0e0fa8006e1..6892445f080 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -9,8 +9,6 @@ from litellm.utils import ModelResponse, EmbeddingResponse, get_secret, Usage import sys from copy import deepcopy import httpx # type: ignore -import boto3 -import aioboto3 import io from .prompt_templates.factory import prompt_factory, custom_prompt @@ -159,6 +157,8 @@ def completion( logger_fn=None, acompletion: bool = False, ): + import boto3 + # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) aws_access_key_id = optional_params.pop("aws_access_key_id", None) @@ -413,6 +413,10 @@ async def async_streaming( aws_access_key_id: Optional[str], aws_region_name: Optional[str], ): + """ + Use aioboto3 + """ + import aioboto3 session = aioboto3.Session() @@ -477,6 +481,10 @@ async def async_completion( aws_access_key_id: Optional[str], aws_region_name: Optional[str], ): + """ + Use aioboto3 + """ + import aioboto3 session = aioboto3.Session() @@ -628,6 +636,8 @@ def embedding( """ Supports Huggingface Jumpstart embeddings like GPT-6B """ + ### BOTO3 INIT + import boto3 # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) diff --git a/litellm/tests/test_provider_specific_config.py b/litellm/tests/test_provider_specific_config.py index c1d5362ec12..c20c44fb134 100644 --- a/litellm/tests/test_provider_specific_config.py +++ b/litellm/tests/test_provider_specific_config.py @@ -517,7 +517,7 @@ def test_sagemaker_default_region(mocker): """ If no regions are specified in config or in environment, the default region is us-west-2 """ - mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + mock_client = mocker.patch("boto3.client") try: response = litellm.completion( model="sagemaker/mock-endpoint", @@ -541,7 +541,7 @@ def test_sagemaker_environment_region(mocker): """ expected_region = "us-east-1" os.environ["AWS_REGION_NAME"] = expected_region - mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + mock_client = mocker.patch("boto3.client") try: response = litellm.completion( model="sagemaker/mock-endpoint", @@ -566,7 +566,7 @@ def test_sagemaker_config_region(mocker): part of the config file, then use that region instead of us-west-2 """ expected_region = "us-east-1" - mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + mock_client = mocker.patch("boto3.client") try: response = litellm.completion( model="sagemaker/mock-endpoint", @@ -592,7 +592,7 @@ def test_sagemaker_config_and_environment_region(mocker): expected_region = "us-east-1" unexpected_region = "us-east-2" os.environ["AWS_REGION_NAME"] = expected_region - mock_client = mocker.patch("litellm.llms.sagemaker.boto3.client") + mock_client = mocker.patch("boto3.client") try: response = litellm.completion( model="sagemaker/mock-endpoint", From 606d04b05b7f44ac044f1bd8f409d4f28d6fc60b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 3 Jul 2024 09:48:19 -0700 Subject: [PATCH 5/8] fix(_service_logging.py): only trigger otel if in service_callback Fixes https://github.com/BerriAI/litellm/issues/4511 --- litellm/_service_logger.py | 18 ++++----- litellm/proxy/_new_secret_config.yaml | 4 ++ litellm/proxy/proxy_server.py | 2 + litellm/tests/test_caching.py | 56 +++++++++++++++++++++++---- 4 files changed, 64 insertions(+), 16 deletions(-) diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 0aa0312b804..be8d7cf895a 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -75,16 +75,16 @@ class ServiceLogging(CustomLogger): await self.prometheusServicesLogger.async_service_success_hook( payload=payload ) + elif callback == "otel": + from litellm.proxy.proxy_server import open_telemetry_logger - from litellm.proxy.proxy_server import open_telemetry_logger - - if parent_otel_span is not None and open_telemetry_logger is not None: - await open_telemetry_logger.async_service_success_hook( - payload=payload, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) + if parent_otel_span is not None and open_telemetry_logger is not None: + await open_telemetry_logger.async_service_success_hook( + payload=payload, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + ) async def async_service_failure_hook( self, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index f135922ea11..dba1a46826f 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -4,6 +4,10 @@ model_list: model: "openai/*" mock_response: "Hello world!" +litellm_settings: + callbacks: ["otel"] + cache: True + general_settings: alerting: ["slack"] alerting_threshold: 10 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1ca18072239..b451617ebd8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1465,6 +1465,8 @@ class ProxyConfig: open_telemetry_logger = OpenTelemetry() imported_list.append(open_telemetry_logger) + + litellm.service_callback.append("otel") elif isinstance(callback, str) and callback == "presidio": from litellm.proxy.hooks.presidio_pii_masking import ( _OPTIONAL_PresidioPIIMasking, diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py index 2f0f1dbfe6d..c5e9c7f1fc2 100644 --- a/litellm/tests/test_caching.py +++ b/litellm/tests/test_caching.py @@ -1,6 +1,9 @@ -import sys, os, uuid +import os +import sys import time import traceback +import uuid + from dotenv import load_dotenv load_dotenv() @@ -9,12 +12,15 @@ import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import pytest -import litellm -from litellm import embedding, completion, aembedding -from litellm.caching import Cache +import asyncio +import hashlib import random -import hashlib, asyncio + +import pytest + +import litellm +from litellm import aembedding, completion, embedding +from litellm.caching import Cache # litellm.set_verbose=True @@ -656,6 +662,7 @@ def test_redis_cache_completion(): assert response1.created == response2.created assert response1.choices[0].message.content == response2.choices[0].message.content + # test_redis_cache_completion() @@ -877,6 +884,7 @@ async def test_redis_cache_acompletion_stream_bedrock(): print(e) raise e + def test_disk_cache_completion(): litellm.set_verbose = False @@ -925,7 +933,7 @@ def test_disk_cache_completion(): litellm.success_callback = [] litellm._async_success_callback = [] - # 1 & 2 should be exactly the same + # 1 & 2 should be exactly the same # 1 & 3 should be different, since input params are diff if ( response1["choices"][0]["message"]["content"] @@ -1569,3 +1577,37 @@ async def test_redis_semantic_cache_acompletion(): ) print(f"response2: {response2}") assert response1.id == response2.id + + +def test_caching_redis_simple(caplog): + """ + Relevant issue - https://github.com/BerriAI/litellm/issues/4511 + """ + litellm.cache = Cache( + type="redis", url=os.getenv("REDIS_SSL_URL") + ) # passing `supported_call_types = ["completion"]` has no effect + + s = time.time() + x = completion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello, how are you? Wink"}], + stream=True, + ) + for m in x: + print(m) + print(time.time() - s) + + s2 = time.time() + x = completion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello, how are you? Wink"}], + stream=True, + ) + for m in x: + print(m) + print(time.time() - s2) + + captured_logs = [rec.message for rec in caplog.records] + + assert "LiteLLM Redis Caching: async set" not in captured_logs + assert "ServiceLogging.async_service_success_hook" not in captured_logs From c6962bda803fc5cf2cd6b762aa1b8aa986f55197 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 3 Jul 2024 10:02:50 -0700 Subject: [PATCH 6/8] fix(user_api_key_auth.py): expose `/health` as an info route this still requires key access. Allows users to check model health --- litellm/proxy/_super_secret_config.yaml | 8 ++++---- litellm/proxy/_types.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_super_secret_config.yaml b/litellm/proxy/_super_secret_config.yaml index ede853094ed..eb725656a75 100644 --- a/litellm/proxy/_super_secret_config.yaml +++ b/litellm/proxy/_super_secret_config.yaml @@ -2,10 +2,10 @@ model_list: - model_name: claude-3-5-sonnet litellm_params: model: anthropic/claude-3-5-sonnet -- model_name: gemini-1.5-flash-gemini - litellm_params: - model: vertex_ai_beta/gemini-1.5-flash - api_base: https://gateway.ai.cloudflare.com/v1/fa4cdcab1f32b95ca3b53fd36043d691/test/google-vertex-ai/v1/projects/adroit-crow-413218/locations/us-central1/publishers/google/models/gemini-1.5-flash +# - model_name: gemini-1.5-flash-gemini +# litellm_params: +# model: vertex_ai_beta/gemini-1.5-flash +# api_base: https://gateway.ai.cloudflare.com/v1/fa4cdcab1f32b95ca3b53fd36043d691/test/google-vertex-ai/v1/projects/adroit-crow-413218/locations/us-central1/publishers/google/models/gemini-1.5-flash - litellm_params: api_base: http://0.0.0.0:8080 api_key: '' diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1f1aaf0eeac..b3ecd19640b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -218,6 +218,7 @@ class LiteLLMRoutes(enum.Enum): "/v2/model/info", "/v2/key/info", "/model_group/info", + "/health", ] # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend From d18dc2648c5825c27c9fe4f5300f3516e3fcf1c3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 3 Jul 2024 10:09:46 -0700 Subject: [PATCH 7/8] fix(replicate.py): correctly handly invalid replicate response --- litellm/llms/replicate.py | 10 +++++++++- litellm/tests/test_streaming.py | 1 - 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm/llms/replicate.py b/litellm/llms/replicate.py index 56549cfd4aa..77dc52aae86 100644 --- a/litellm/llms/replicate.py +++ b/litellm/llms/replicate.py @@ -295,7 +295,15 @@ def handle_prediction_response_streaming(prediction_url, api_token, print_verbos response_data = response.json() status = response_data["status"] if "output" in response_data: - output_string = "".join(response_data["output"]) + try: + output_string = "".join(response_data["output"]) + except Exception as e: + raise ReplicateError( + status_code=422, + message="Unable to parse response. Got={}".format( + response_data["output"] + ), + ) new_output = output_string[len(previous_output) :] print_verbose(f"New chunk: {new_output}") yield {"output": new_output, "status": status} diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 1f1b253a064..0dd81e3b34b 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -1415,7 +1415,6 @@ def test_bedrock_claude_3_streaming(): "gpt-3.5-turbo", "databricks/databricks-dbrx-instruct", # databricks "predibase/llama-3-8b-instruct", # predibase - "replicate/meta/meta-llama-3-8b-instruct", # replicate ], ) @pytest.mark.asyncio From 3d61a316cb29488acf58c5db9e72fc3fead14e27 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 3 Jul 2024 12:00:00 -0700 Subject: [PATCH 8/8] fix(router.py): bump azure default api version Allows 'tool_choice' to be passed to azure --- litellm/__init__.py | 2 ++ litellm/router.py | 2 +- litellm/utils.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 29b5bc360ab..e00e4f804ed 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -113,6 +113,8 @@ ssl_verify: bool = True ssl_certificate: Optional[str] = None disable_streaming_logging: bool = False in_memory_llm_clients_cache: dict = {} +### DEFAULT AZURE API VERSION ### +AZURE_DEFAULT_API_VERSION = "2024-02-01" # this is updated to the latest ### GUARDRAILS ### llamaguard_model_name: Optional[str] = None openai_moderations_model_name: Optional[str] = None diff --git a/litellm/router.py b/litellm/router.py index ac61ec729fb..8d9884d1f01 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3437,7 +3437,7 @@ class Router: if azure_ad_token.startswith("oidc/"): azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) if api_version is None: - api_version = "2023-07-01-preview" + api_version = litellm.AZURE_DEFAULT_API_VERSION if "gateway.ai.cloudflare.com" in api_base: if not api_base.endswith("/"): diff --git a/litellm/utils.py b/litellm/utils.py index 82e3ca17127..6c1814629c7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3634,7 +3634,7 @@ def get_model_region( model=_model, api_key=litellm_params.api_key, api_base=litellm_params.api_base, - api_version=litellm_params.api_version or "2023-07-01-preview", + api_version=litellm_params.api_version or litellm.AZURE_DEFAULT_API_VERSION, timeout=10, mode=mode or "chat", )