Merge pull request #3173 from BerriAI/litellm_ssl_caching_fix

fix(caching.py): fix redis url parsing logic to work with ssl urls

litellm/_redis.py

@@ -32,6 +32,25 @@ def _get_redis_kwargs():
     return available_args
 
 
+def _get_redis_url_kwargs(client=None):
+    if client is None:
+        client = redis.Redis.from_url
+    arg_spec = inspect.getfullargspec(redis.Redis.from_url)
+
+    # Only allow primitive arguments
+    exclude_args = {
+        "self",
+        "connection_pool",
+        "retry",
+    }
+
+    include_args = ["url"]
+
+    available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
+
+    return available_args
+
+
 def _get_redis_env_kwarg_mapping():
     PREFIX = "REDIS_"
 
@@ -98,20 +117,31 @@ def _get_redis_client_logic(**env_overrides):
 def get_redis_client(**env_overrides):
     redis_kwargs = _get_redis_client_logic(**env_overrides)
     if "url" in redis_kwargs and redis_kwargs["url"] is not None:
-        redis_kwargs.pop(
-            "connection_pool", None
-        )  # redis.from_url doesn't support setting your own connection pool
-        return redis.Redis.from_url(**redis_kwargs)
+        args = _get_redis_url_kwargs()
+        url_kwargs = {}
+        for arg in redis_kwargs:
+            if arg in args:
+                url_kwargs[arg] = redis_kwargs[arg]
+
+        return redis.Redis.from_url(**url_kwargs)
     return redis.Redis(**redis_kwargs)
 
 
 def get_redis_async_client(**env_overrides):
     redis_kwargs = _get_redis_client_logic(**env_overrides)
     if "url" in redis_kwargs and redis_kwargs["url"] is not None:
-        redis_kwargs.pop(
-            "connection_pool", None
-        )  # redis.from_url doesn't support setting your own connection pool
-        return async_redis.Redis.from_url(**redis_kwargs)
+        args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
+        url_kwargs = {}
+        for arg in redis_kwargs:
+            if arg in args:
+                url_kwargs[arg] = redis_kwargs[arg]
+            else:
+                litellm.print_verbose(
+                    "REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(
+                        arg
+                    )
+                )
+        return async_redis.Redis.from_url(**url_kwargs)
     return async_redis.Redis(
         socket_timeout=5,
         **redis_kwargs,

litellm/caching.py

@@ -149,18 +149,19 @@ class RedisCache(BaseCache):
         if password is not None:
             redis_kwargs["password"] = password
 
+        ### HEALTH MONITORING OBJECT ###
+        if kwargs.get("service_logger_obj", None) is not None and isinstance(
+            kwargs["service_logger_obj"], ServiceLogging
+        ):
+            self.service_logger_obj = kwargs.pop("service_logger_obj")
+        else:
+            self.service_logger_obj = ServiceLogging()
+
         redis_kwargs.update(kwargs)
         self.redis_client = get_redis_client(**redis_kwargs)
         self.redis_kwargs = redis_kwargs
         self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs)
 
-        if "url" in redis_kwargs and redis_kwargs["url"] is not None:
-            parsed_kwargs = redis.connection.parse_url(redis_kwargs["url"])
-            redis_kwargs.update(parsed_kwargs)
-            self.redis_kwargs.update(parsed_kwargs)
-            # pop url
-            self.redis_kwargs.pop("url")
-
         # redis namespaces
         self.namespace = namespace
         # for high traffic, we store the redis results in memory and then batch write to redis
@@ -172,8 +173,15 @@ class RedisCache(BaseCache):
         except Exception as e:
             pass
 
-        ### HEALTH MONITORING OBJECT ###
-        self.service_logger_obj = ServiceLogging()
+        ### ASYNC HEALTH PING ###
+        try:
+            # asyncio.get_running_loop().create_task(self.ping())
+            result = asyncio.get_running_loop().create_task(self.ping())
+        except Exception:
+            pass
+
+        ### SYNC HEALTH PING ###
+        self.redis_client.ping()
 
     def init_async_client(self):
         from ._redis import get_redis_async_client
@@ -601,15 +609,72 @@ class RedisCache(BaseCache):
             print_verbose(f"Error occurred in pipeline read - {str(e)}")
             return key_value_dict
 
-    async def ping(self):
+    def sync_ping(self) -> bool:
+        """
+        Tests if the sync redis client is correctly setup.
+        """
+        print_verbose(f"Pinging Sync Redis Cache")
+        start_time = time.time()
+        try:
+            response = self.redis_client.ping()
+            print_verbose(f"Redis Cache PING: {response}")
+            ## LOGGING ##
+            end_time = time.time()
+            _duration = end_time - start_time
+            self.service_logger_obj.service_success_hook(
+                service=ServiceTypes.REDIS,
+                duration=_duration,
+                call_type="sync_ping",
+            )
+            return response
+        except Exception as e:
+            # NON blocking - notify users Redis is throwing an exception
+            ## LOGGING ##
+            end_time = time.time()
+            _duration = end_time - start_time
+            self.service_logger_obj.service_failure_hook(
+                service=ServiceTypes.REDIS,
+                duration=_duration,
+                error=e,
+                call_type="sync_ping",
+            )
+            print_verbose(
+                f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}"
+            )
+            traceback.print_exc()
+            raise e
+
+    async def ping(self) -> bool:
         _redis_client = self.init_async_client()
+        start_time = time.time()
         async with _redis_client as redis_client:
             print_verbose(f"Pinging Async Redis Cache")
             try:
                 response = await redis_client.ping()
-                print_verbose(f"Redis Cache PING: {response}")
+                ## LOGGING ##
+                end_time = time.time()
+                _duration = end_time - start_time
+                asyncio.create_task(
+                    self.service_logger_obj.async_service_success_hook(
+                        service=ServiceTypes.REDIS,
+                        duration=_duration,
+                        call_type="async_ping",
+                    )
+                )
+                return response
             except Exception as e:
                 # NON blocking - notify users Redis is throwing an exception
+                ## LOGGING ##
+                end_time = time.time()
+                _duration = end_time - start_time
+                asyncio.create_task(
+                    self.service_logger_obj.async_service_failure_hook(
+                        service=ServiceTypes.REDIS,
+                        duration=_duration,
+                        error=e,
+                        call_type="async_ping",
+                    )
+                )
                 print_verbose(
                     f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}"
                 )

litellm/integrations/prometheus.py

@@ -67,7 +67,7 @@ class PrometheusLogger:
 
             # unpack kwargs
             model = kwargs.get("model", "")
-            response_cost = kwargs.get("response_cost", 0.0)
+            response_cost = kwargs.get("response_cost", 0.0) or 0
             litellm_params = kwargs.get("litellm_params", {}) or {}
             proxy_server_request = litellm_params.get("proxy_server_request") or {}
             end_user_id = proxy_server_request.get("body", {}).get("user", None)

litellm/integrations/prometheus_services.py

@@ -30,7 +30,6 @@ class PrometheusServicesLogger:
                 raise Exception(
                     "Missing prometheus_client. Run `pip install prometheus-client`"
                 )
-            print("INITIALIZES PROMETHEUS SERVICE LOGGER!")
 
             self.Histogram = Histogram
             self.Counter = Counter

litellm/tests/test_caching.py

@@ -435,6 +435,7 @@ async def test_embedding_caching_azure_individual_items_reordered():
 @pytest.mark.asyncio
 async def test_embedding_caching_base_64():
     """ """
+    litellm.set_verbose = True
     litellm.cache = Cache(
         type="redis",
         host=os.environ["REDIS_HOST"],
@@ -453,6 +454,8 @@ async def test_embedding_caching_base_64():
         caching=True,
         encoding_format="base64",
     )
+    await asyncio.sleep(5)
+    print("\n\nCALL2\n\n")
     embedding_val_2 = await aembedding(
         model="azure/azure-embedding-model",
         input=inputs,
@@ -1139,10 +1142,6 @@ def test_custom_redis_cache_params():
             port=os.environ["REDIS_PORT"],
             password=os.environ["REDIS_PASSWORD"],
             db=0,
-            ssl=True,
-            ssl_certfile="./redis_user.crt",
-            ssl_keyfile="./redis_user_private.key",
-            ssl_ca_certs="./redis_ca.pem",
         )
 
         print(litellm.cache.cache.redis_client)
@@ -1150,7 +1149,7 @@ def test_custom_redis_cache_params():
         litellm.success_callback = []
         litellm._async_success_callback = []
     except Exception as e:
-        pytest.fail(f"Error occurred:", e)
+        pytest.fail(f"Error occurred: {str(e)}")
 
 
 def test_get_cache_key():

litellm/tests/test_prometheus_service.py

@@ -65,23 +65,18 @@ async def test_completion_with_caching_bad_call():
     - Assert failure callback gets called
     """
     litellm.set_verbose = True
-    sl = ServiceLogging(mock_testing=True)
+
     try:
-        litellm.cache = Cache(type="redis", host="hello-world")
+        from litellm.caching import RedisCache
+
         litellm.service_callback = ["prometheus_system"]
+        sl = ServiceLogging(mock_testing=True)
 
-        litellm.cache.cache.service_logger_obj = sl
-
-        messages = [{"role": "user", "content": "Hey, how's it going?"}]
-        response1 = await acompletion(
-            model="gpt-3.5-turbo", messages=messages, caching=True
-        )
-        response1 = await acompletion(
-            model="gpt-3.5-turbo", messages=messages, caching=True
-        )
+        RedisCache(host="hello-world", service_logger_obj=sl)
     except Exception as e:
-        pass
+        print(f"Receives exception = {str(e)}")
 
+    await asyncio.sleep(5)
     assert sl.mock_testing_async_failure_hook > 0
     assert sl.mock_testing_async_success_hook == 0
     assert sl.mock_testing_sync_success_hook == 0
@@ -144,64 +139,3 @@ async def test_router_with_caching():
 
     except Exception as e:
         pytest.fail(f"An exception occured - {str(e)}")
-
-
-@pytest.mark.asyncio
-async def test_router_with_caching_bad_call():
-    """
-    - Run completion with caching (incorrect credentials)
-    - Assert failure callback gets called
-    """
-    try:
-
-        def get_azure_params(deployment_name: str):
-            params = {
-                "model": f"azure/{deployment_name}",
-                "api_key": os.environ["AZURE_API_KEY"],
-                "api_version": os.environ["AZURE_API_VERSION"],
-                "api_base": os.environ["AZURE_API_BASE"],
-            }
-            return params
-
-        model_list = [
-            {
-                "model_name": "azure/gpt-4",
-                "litellm_params": get_azure_params("chatgpt-v-2"),
-                "tpm": 100,
-            },
-            {
-                "model_name": "azure/gpt-4",
-                "litellm_params": get_azure_params("chatgpt-v-2"),
-                "tpm": 1000,
-            },
-        ]
-
-        router = litellm.Router(
-            model_list=model_list,
-            set_verbose=True,
-            debug_level="DEBUG",
-            routing_strategy="usage-based-routing-v2",
-            redis_host="hello world",
-            redis_port=os.environ["REDIS_PORT"],
-            redis_password=os.environ["REDIS_PASSWORD"],
-        )
-
-        litellm.service_callback = ["prometheus_system"]
-
-        sl = ServiceLogging(mock_testing=True)
-        sl.prometheusServicesLogger.mock_testing = True
-        router.cache.redis_cache.service_logger_obj = sl
-
-        messages = [{"role": "user", "content": "Hey, how's it going?"}]
-        try:
-            response1 = await router.acompletion(model="azure/gpt-4", messages=messages)
-            response1 = await router.acompletion(model="azure/gpt-4", messages=messages)
-        except Exception as e:
-            pass
-
-        assert sl.mock_testing_async_failure_hook > 0
-        assert sl.mock_testing_async_success_hook == 0
-        assert sl.mock_testing_sync_success_hook == 0
-
-    except Exception as e:
-        pytest.fail(f"An exception occured - {str(e)}")

litellm/tests/test_router_caching.py

@@ -15,6 +15,61 @@ from litellm import Router
 ## 2. 2 models - openai, azure - 2 diff model groups, 1 caching group
 
 
+@pytest.mark.asyncio
+async def test_router_async_caching_with_ssl_url():
+    """
+    Tests when a redis url is passed to the router, if caching is correctly setup
+    """
+    try:
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "gpt-3.5-turbo",
+                    "litellm_params": {
+                        "model": "gpt-3.5-turbo-0613",
+                        "api_key": os.getenv("OPENAI_API_KEY"),
+                    },
+                    "tpm": 100000,
+                    "rpm": 10000,
+                },
+            ],
+            redis_url=os.getenv("REDIS_SSL_URL"),
+        )
+
+        response = await router.cache.redis_cache.ping()
+        print(f"response: {response}")
+        assert response == True
+    except Exception as e:
+        pytest.fail(f"An exception occurred - {str(e)}")
+
+
+def test_router_sync_caching_with_ssl_url():
+    """
+    Tests when a redis url is passed to the router, if caching is correctly setup
+    """
+    try:
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "gpt-3.5-turbo",
+                    "litellm_params": {
+                        "model": "gpt-3.5-turbo-0613",
+                        "api_key": os.getenv("OPENAI_API_KEY"),
+                    },
+                    "tpm": 100000,
+                    "rpm": 10000,
+                },
+            ],
+            redis_url=os.getenv("REDIS_SSL_URL"),
+        )
+
+        response = router.cache.redis_cache.sync_ping()
+        print(f"response: {response}")
+        assert response == True
+    except Exception as e:
+        pytest.fail(f"An exception occurred - {str(e)}")
+
+
 @pytest.mark.asyncio
 async def test_acompletion_caching_on_router():
     # tests acompletion + caching on router