From 606d04b05b7f44ac044f1bd8f409d4f28d6fc60b Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 3 Jul 2024 09:48:19 -0700
Subject: [PATCH] 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 0aa0312b8..be8d7cf89 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 f135922ea..dba1a4682 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 1ca180722..b451617eb 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 2f0f1dbfe..c5e9c7f1f 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