forked from phoenix/litellm-mirror
4d1b4beb3d
* use folder for caching * fix importing caching * fix clickhouse pyright * fix linting * fix correctly pass kwargs and args * fix test case for embedding * fix linting * fix embedding caching logic * fix refactor handle utils.py * fix test_embedding_caching_azure_individual_items_reordered
308 lines
9 KiB
Python
308 lines
9 KiB
Python
# What is this?
|
|
## This tests the llm guard integration
|
|
|
|
import asyncio
|
|
import os
|
|
import random
|
|
|
|
# What is this?
|
|
## Unit test for presidio pii masking
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from datetime import datetime
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import os
|
|
|
|
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.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 ###
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_secret_detection_chat():
|
|
"""
|
|
Tests to see if secret detection hook will mask api keys
|
|
|
|
|
|
It should mask the following API_KEY = 'sk_1234567890abcdef' and OPENAI_API_KEY = 'sk_1234567890abcdef'
|
|
"""
|
|
secret_instance = _ENTERPRISE_SecretDetection()
|
|
_api_key = "sk-12345"
|
|
_api_key = hash_token("sk-12345")
|
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
|
local_cache = DualCache()
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
test_data = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Hey, how's it going, API_KEY = 'sk_1234567890abcdef'",
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": "Hello! I'm doing well. How can I assist you today?",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "this is my OPENAI_API_KEY = 'sk_1234567890abcdef'",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "My hi API Key is sk-Pc4nlxVoMz41290028TbMCxx, does it seem to be in the correct format?",
|
|
},
|
|
{"role": "user", "content": "i think it is +1 412-555-5555"},
|
|
],
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
|
|
await secret_instance.async_pre_call_hook(
|
|
cache=local_cache,
|
|
data=test_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type="completion",
|
|
)
|
|
print(
|
|
"test data after running pre_call_hook: Expect all API Keys to be masked",
|
|
test_data,
|
|
)
|
|
|
|
assert test_data == {
|
|
"messages": [
|
|
{"role": "user", "content": "Hey, how's it going, API_KEY = '[REDACTED]'"},
|
|
{
|
|
"role": "assistant",
|
|
"content": "Hello! I'm doing well. How can I assist you today?",
|
|
},
|
|
{"role": "user", "content": "this is my OPENAI_API_KEY = '[REDACTED]'"},
|
|
{
|
|
"role": "user",
|
|
"content": "My hi API Key is [REDACTED], does it seem to be in the correct format?",
|
|
},
|
|
{"role": "user", "content": "i think it is +1 412-555-5555"},
|
|
],
|
|
"model": "gpt-3.5-turbo",
|
|
}, "Expect all API Keys to be masked"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_secret_detection_text_completion():
|
|
"""
|
|
Tests to see if secret detection hook will mask api keys
|
|
|
|
|
|
It should mask the following API_KEY = 'sk_1234567890abcdef' and OPENAI_API_KEY = 'sk_1234567890abcdef'
|
|
"""
|
|
secret_instance = _ENTERPRISE_SecretDetection()
|
|
_api_key = "sk-12345"
|
|
_api_key = hash_token("sk-12345")
|
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
|
local_cache = DualCache()
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
test_data = {
|
|
"prompt": "Hey, how's it going, API_KEY = 'sk_1234567890abcdef', my OPENAI_API_KEY = 'sk_1234567890abcdef' and i want to know what is the weather",
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
|
|
await secret_instance.async_pre_call_hook(
|
|
cache=local_cache,
|
|
data=test_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type="completion",
|
|
)
|
|
|
|
test_data == {
|
|
"prompt": "Hey, how's it going, API_KEY = '[REDACTED]', my OPENAI_API_KEY = '[REDACTED]' and i want to know what is the weather",
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
print(
|
|
"test data after running pre_call_hook: Expect all API Keys to be masked",
|
|
test_data,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_secret_detection_embeddings():
|
|
"""
|
|
Tests to see if secret detection hook will mask api keys
|
|
|
|
|
|
It should mask the following API_KEY = 'sk_1234567890abcdef' and OPENAI_API_KEY = 'sk_1234567890abcdef'
|
|
"""
|
|
secret_instance = _ENTERPRISE_SecretDetection()
|
|
_api_key = "sk-12345"
|
|
_api_key = hash_token("sk-12345")
|
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
|
local_cache = DualCache()
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
test_data = {
|
|
"input": "Hey, how's it going, API_KEY = 'sk_1234567890abcdef', my OPENAI_API_KEY = 'sk_1234567890abcdef' and i want to know what is the weather",
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
|
|
await secret_instance.async_pre_call_hook(
|
|
cache=local_cache,
|
|
data=test_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type="embedding",
|
|
)
|
|
|
|
assert test_data == {
|
|
"input": "Hey, how's it going, API_KEY = '[REDACTED]', my OPENAI_API_KEY = '[REDACTED]' and i want to know what is the weather",
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
print(
|
|
"test data after running pre_call_hook: Expect all API Keys to be masked",
|
|
test_data,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_secret_detection_embeddings_list():
|
|
"""
|
|
Tests to see if secret detection hook will mask api keys
|
|
|
|
|
|
It should mask the following API_KEY = 'sk_1234567890abcdef' and OPENAI_API_KEY = 'sk_1234567890abcdef'
|
|
"""
|
|
secret_instance = _ENTERPRISE_SecretDetection()
|
|
_api_key = "sk-12345"
|
|
_api_key = hash_token("sk-12345")
|
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
|
local_cache = DualCache()
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
test_data = {
|
|
"input": [
|
|
"hey",
|
|
"how's it going, API_KEY = 'sk_1234567890abcdef'",
|
|
"my OPENAI_API_KEY = 'sk_1234567890abcdef' and i want to know what is the weather",
|
|
],
|
|
"model": "gpt-3.5-turbo",
|
|
}
|
|
|
|
await secret_instance.async_pre_call_hook(
|
|
cache=local_cache,
|
|
data=test_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type="embedding",
|
|
)
|
|
|
|
print(
|
|
"test data after running pre_call_hook: Expect all API Keys to be masked",
|
|
test_data,
|
|
)
|
|
assert test_data == {
|
|
"input": [
|
|
"hey",
|
|
"how's it going, API_KEY = '[REDACTED]'",
|
|
"my OPENAI_API_KEY = '[REDACTED]' and i want to know what is the weather",
|
|
],
|
|
"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
|