Spaces:
Configuration error
Configuration error
File size: 5,266 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
# What is this?
## Unit test that rejected requests are also logged as failures
# 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
from typing import Literal
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_enterprise.enterprise_callbacks.secret_detection import (
_ENTERPRISE_SecretDetection,
)
from litellm.proxy.proxy_server import (
Depends,
HTTPException,
chat_completion,
completion,
embeddings,
)
from litellm.proxy.utils import ProxyLogging, hash_token
from litellm.router import Router
class testLogger(CustomLogger):
def __init__(self):
self.reaches_sync_failure_event = False
self.reaches_async_failure_event = False
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
],
):
raise HTTPException(
status_code=429, detail={"error": "Max parallel request limit reached"}
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.reaches_async_failure_event = True
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
self.reaches_sync_failure_event = True
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.parametrize(
"route, body",
[
(
"/v1/chat/completions",
{
"model": "fake-model",
"messages": [
{
"role": "user",
"content": "Hello here is my OPENAI_API_KEY = sk-12345",
}
],
},
),
("/v1/completions", {"model": "fake-model", "prompt": "ping"}),
(
"/v1/embeddings",
{
"input": "The food was delicious and the waiter...",
"model": "text-embedding-ada-002",
"encoding_format": "float",
},
),
],
)
@pytest.mark.asyncio
async def test_chat_completion_request_with_redaction(route, body):
"""
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 = [_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=route)
async def return_body():
import json
return json.dumps(body).encode()
request.body = return_body
try:
if route == "/v1/chat/completions":
response = await chat_completion(
request=request,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-12345", token="hashed_sk-12345", rpm_limit=0
),
fastapi_response=Response(),
)
elif route == "/v1/completions":
response = await completion(
request=request,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-12345", token="hashed_sk-12345", rpm_limit=0
),
fastapi_response=Response(),
)
elif route == "/v1/embeddings":
response = await embeddings(
request=request,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-12345", token="hashed_sk-12345", rpm_limit=0
),
fastapi_response=Response(),
)
except Exception:
pass
await asyncio.sleep(3)
assert _test_logger.reaches_async_failure_event is True
assert _test_logger.reaches_sync_failure_event is True
|