Added LiteLLM to the stack

Development/litellm/tests/local_testing/test_prometheus_service.py

@@ -0,0 +1,297 @@
+# What is this?
+## Unit Tests for prometheus service monitoring
+
+import json
+import sys
+import os
+import io, asyncio
+
+sys.path.insert(0, os.path.abspath("../.."))
+import pytest
+from litellm import acompletion, Cache
+from litellm._service_logger import ServiceLogging
+from litellm.integrations.prometheus_services import PrometheusServicesLogger
+from litellm.proxy.utils import ServiceTypes
+from unittest.mock import patch, AsyncMock
+import litellm
+
+"""
+- Check if it receives a call when redis is used 
+- Check if it fires messages accordingly
+"""
+
+
+@pytest.mark.asyncio
+async def test_init_prometheus():
+    """
+    - Run completion with caching
+    - Assert success callback gets called
+    """
+
+    pl = PrometheusServicesLogger(mock_testing=True)
+
+
+@pytest.mark.asyncio
+async def test_completion_with_caching():
+    """
+    - Run completion with caching
+    - Assert success callback gets called
+    """
+
+    litellm.set_verbose = True
+    litellm.cache = Cache(type="redis")
+    litellm.service_callback = ["prometheus_system"]
+
+    sl = ServiceLogging(mock_testing=True)
+    sl.prometheusServicesLogger.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
+    )
+
+    assert sl.mock_testing_async_success_hook > 0
+    assert sl.prometheusServicesLogger.mock_testing_success_calls > 0
+    assert sl.mock_testing_sync_failure_hook == 0
+    assert sl.mock_testing_async_failure_hook == 0
+
+
+@pytest.mark.asyncio
+async def test_completion_with_caching_bad_call():
+    """
+    - Run completion with caching (incorrect credentials)
+    - Assert failure callback gets called
+    """
+    litellm.set_verbose = True
+
+    try:
+        from litellm.caching.caching import RedisCache
+
+        litellm.service_callback = ["prometheus_system"]
+        sl = ServiceLogging(mock_testing=True)
+
+        RedisCache(host="hello-world", service_logger_obj=sl)
+    except Exception as e:
+        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
+
+
+@pytest.mark.asyncio
+async def test_router_with_caching():
+    """
+    - Run router with usage-based-routing-v2
+    - Assert success 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("gpt-4o-new-test"),
+                "tpm": 100,
+            },
+            {
+                "model_name": "azure/gpt-4",
+                "litellm_params": get_azure_params("gpt-4o-new-test"),
+                "tpm": 1000,
+            },
+        ]
+
+        router = litellm.Router(
+            model_list=model_list,
+            set_verbose=True,
+            debug_level="DEBUG",
+            routing_strategy="usage-based-routing-v2",
+            redis_host=os.environ["REDIS_HOST"],
+            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?"}]
+        response1 = await router.acompletion(model="azure/gpt-4", messages=messages)
+        response1 = await router.acompletion(model="azure/gpt-4", messages=messages)
+
+        assert sl.mock_testing_async_success_hook > 0
+        assert sl.mock_testing_sync_failure_hook == 0
+        assert sl.mock_testing_async_failure_hook == 0
+        assert sl.prometheusServicesLogger.mock_testing_success_calls > 0
+
+    except Exception as e:
+        pytest.fail(f"An exception occured - {str(e)}")
+
+
+@pytest.mark.asyncio
+async def test_service_logger_db_monitoring():
+    """
+    Test prometheus monitoring for database operations
+    """
+    litellm.service_callback = ["prometheus_system"]
+    sl = ServiceLogging()
+
+    # Create spy on prometheus logger's async_service_success_hook
+    with patch.object(
+        sl.prometheusServicesLogger,
+        "async_service_success_hook",
+        new_callable=AsyncMock,
+    ) as mock_prometheus_success:
+        # Test DB success monitoring
+        await sl.async_service_success_hook(
+            service=ServiceTypes.DB,
+            duration=0.3,
+            call_type="query",
+            event_metadata={"query_type": "SELECT", "table": "api_keys"},
+        )
+
+        # Assert prometheus logger's success hook was called
+        mock_prometheus_success.assert_called_once()
+        # Optionally verify the payload
+        actual_payload = mock_prometheus_success.call_args[1]["payload"]
+        print("actual_payload sent to prometheus: ", actual_payload)
+        assert actual_payload.service == ServiceTypes.DB
+        assert actual_payload.duration == 0.3
+        assert actual_payload.call_type == "query"
+        assert actual_payload.is_error is False
+
+
+@pytest.mark.asyncio
+async def test_service_logger_db_monitoring_failure():
+    """
+    Test prometheus monitoring for failed database operations
+    """
+    litellm.service_callback = ["prometheus_system"]
+    sl = ServiceLogging()
+
+    # Create spy on prometheus logger's async_service_failure_hook
+    with patch.object(
+        sl.prometheusServicesLogger,
+        "async_service_failure_hook",
+        new_callable=AsyncMock,
+    ) as mock_prometheus_failure:
+        # Test DB failure monitoring
+        test_error = Exception("Database connection failed")
+        await sl.async_service_failure_hook(
+            service=ServiceTypes.DB,
+            duration=0.3,
+            error=test_error,
+            call_type="query",
+            event_metadata={"query_type": "SELECT", "table": "api_keys"},
+        )
+
+        # Assert prometheus logger's failure hook was called
+        mock_prometheus_failure.assert_called_once()
+        # Verify the payload
+        actual_payload = mock_prometheus_failure.call_args[1]["payload"]
+        print("actual_payload sent to prometheus: ", actual_payload)
+        assert actual_payload.service == ServiceTypes.DB
+        assert actual_payload.duration == 0.3
+        assert actual_payload.call_type == "query"
+        assert actual_payload.is_error is True
+        assert actual_payload.error == "Database connection failed"
+
+
+def test_get_metric_existing():
+    """Test _get_metric when metric exists. _get_metric should return the metric object"""
+    pl = PrometheusServicesLogger()
+    # Create a metric first
+    hist = pl.create_histogram(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    # Test retrieving existing metric
+    retrieved_metric = pl._get_metric("litellm_test_service_test_type_of_request")
+    assert retrieved_metric is hist
+    assert retrieved_metric is not None
+
+
+def test_get_metric_non_existing():
+    """Test _get_metric when metric doesn't exist, returns None"""
+    pl = PrometheusServicesLogger()
+
+    # Test retrieving non-existent metric
+    non_existent = pl._get_metric("non_existent_metric")
+    assert non_existent is None
+
+
+def test_create_histogram_new():
+    """Test creating a new histogram"""
+    pl = PrometheusServicesLogger()
+
+    # Create new histogram
+    hist = pl.create_histogram(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    assert hist is not None
+    assert pl._get_metric("litellm_test_service_test_type_of_request") is hist
+
+
+def test_create_histogram_existing():
+    """Test creating a histogram that already exists"""
+    pl = PrometheusServicesLogger()
+
+    # Create initial histogram
+    hist1 = pl.create_histogram(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    # Create same histogram again
+    hist2 = pl.create_histogram(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    assert hist2 is hist1  # same object
+    assert pl._get_metric("litellm_test_service_test_type_of_request") is hist1
+
+
+def test_create_counter_new():
+    """Test creating a new counter"""
+    pl = PrometheusServicesLogger()
+
+    # Create new counter
+    counter = pl.create_counter(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    assert counter is not None
+    assert pl._get_metric("litellm_test_service_test_type_of_request") is counter
+
+
+def test_create_counter_existing():
+    """Test creating a counter that already exists"""
+    pl = PrometheusServicesLogger()
+
+    # Create initial counter
+    counter1 = pl.create_counter(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    # Create same counter again
+    counter2 = pl.create_counter(
+        service="test_service", type_of_request="test_type_of_request"
+    )
+
+    assert counter2 is counter1
+    assert pl._get_metric("litellm_test_service_test_type_of_request") is counter1