diff --git a/litellm/tests/test_secret_detect_hook.py b/litellm/tests/test_secret_detect_hook.py index a1bf10eba..cb1e01810 100644 --- a/litellm/tests/test_secret_detect_hook.py +++ b/litellm/tests/test_secret_detect_hook.py @@ -21,15 +21,20 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest +from fastapi import Request, Response +from starlette.datastructures import URL import litellm from litellm import Router, mock_completion from litellm.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.enterprise.enterprise_hooks.secret_detection import ( _ENTERPRISE_SecretDetection, ) +from litellm.proxy.proxy_server import chat_completion from litellm.proxy.utils import ProxyLogging, hash_token +from litellm.router import Router ### UNIT TESTS FOR OpenAI Moderation ### @@ -214,3 +219,82 @@ async def test_basic_secret_detection_embeddings_list(): ], "model": "gpt-3.5-turbo", } + + +class testLogger(CustomLogger): + + def __init__(self): + self.logged_message = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + print(f"On Async Success") + + self.logged_message = kwargs.get("messages") + + +router = Router( + model_list=[ + { + "model_name": "fake-model", + "litellm_params": { + "model": "openai/fake", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "api_key": "sk-12345", + }, + } + ] +) + + +@pytest.mark.asyncio +async def test_chat_completion_request_with_redaction(): + """ + IMPORTANT Enterprise Test - Do not delete it: + Makes a /chat/completions request on LiteLLM Proxy + + Ensures that the secret is redacted EVEN on the callback + """ + from litellm.proxy import proxy_server + + setattr(proxy_server, "llm_router", router) + _test_logger = testLogger() + litellm.callbacks = [_ENTERPRISE_SecretDetection(), _test_logger] + litellm.set_verbose = True + + # Prepare the query string + query_params = "param1=value1¶m2=value2" + + # Create the Request object with query parameters + request = Request( + scope={ + "type": "http", + "method": "POST", + "headers": [(b"content-type", b"application/json")], + "query_string": query_params.encode(), + } + ) + + request._url = URL(url="/chat/completions") + + async def return_body(): + return b'{"model": "fake-model", "messages": [{"role": "user", "content": "Hello here is my OPENAI_API_KEY = sk-12345"}]}' + + request.body = return_body + + response = await chat_completion( + request=request, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-12345", + token="hashed_sk-12345", + ), + fastapi_response=Response(), + ) + + await asyncio.sleep(3) + + print("Info in callback after running request=", _test_logger.logged_message) + + assert _test_logger.logged_message == [ + {"role": "user", "content": "Hello here is my OPENAI_API_KEY = [REDACTED]"} + ] + pass