2024-01-04 12:54:51 +08:00
#### What this tests ####
# This tests using caching w/ litellm which requires SSL=True
import sys , os
import traceback
from dotenv import load_dotenv
load_dotenv ( )
import os , io
# this file is to test litellm/proxy
sys . path . insert (
0 , os . path . abspath ( " ../.. " )
) # Adds the parent directory to the system path
2024-01-04 20:58:18 +08:00
import pytest , logging , asyncio
2024-01-04 12:54:51 +08:00
import litellm
from litellm import embedding , completion , completion_cost , Timeout
from litellm import RateLimitError
# Configure logging
logging . basicConfig (
level = logging . DEBUG , # Set the desired logging level
format = " %(asctime)s - %(levelname)s - %(message)s " ,
)
# test /chat/completion request to the proxy
from fastapi . testclient import TestClient
from fastapi import FastAPI
from litellm . proxy . proxy_server import (
router ,
save_worker_config ,
initialize ,
) # Replace with the actual module where your FastAPI router is defined
# Your bearer token
2024-01-27 12:15:35 +08:00
token = " sk-1234 "
2024-01-04 12:54:51 +08:00
headers = { " Authorization " : f " Bearer { token } " }
@pytest.fixture ( scope = " function " )
def client_no_auth ( ) :
# Assuming litellm.proxy.proxy_server is an object
from litellm . proxy . proxy_server import cleanup_router_config_variables
cleanup_router_config_variables ( )
filepath = os . path . dirname ( os . path . abspath ( __file__ ) )
config_fp = f " { filepath } /test_configs/test_cloudflare_azure_with_cache_config.yaml "
# initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
2024-01-04 20:58:18 +08:00
asyncio . run ( initialize ( config = config_fp , debug = True ) )
2024-01-04 12:54:51 +08:00
app = FastAPI ( )
app . include_router ( router ) # Include your router in the test app
return TestClient ( app )
def generate_random_word ( length = 4 ) :
import string , random
letters = string . ascii_lowercase
return " " . join ( random . choice ( letters ) for _ in range ( length ) )
2024-02-29 06:55:57 +08:00
@pytest.mark.skip ( reason = " AWS Suspended Account " )
2024-01-04 12:54:51 +08:00
def test_chat_completion ( client_no_auth ) :
global headers
try :
user_message = f " Write a poem about { generate_random_word ( ) } "
messages = [ { " content " : user_message , " role " : " user " } ]
# Your test data
test_data = {
" model " : " azure-cloudflare " ,
" messages " : messages ,
" max_tokens " : 10 ,
}
print ( " testing proxy server with chat completions " )
response = client_no_auth . post ( " /v1/chat/completions " , json = test_data )
print ( f " response - { response . text } " )
assert response . status_code == 200
response = response . json ( )
print ( response )
content = response [ " choices " ] [ 0 ] [ " message " ] [ " content " ]
2024-01-04 13:41:08 +08:00
response1_id = response [ " id " ]
2024-01-04 12:54:51 +08:00
print ( " \n content " , content )
assert len ( content ) > 1
2024-01-04 13:41:08 +08:00
print ( " \n making 2nd request to proxy. Testing caching + non streaming " )
response = client_no_auth . post ( " /v1/chat/completions " , json = test_data )
print ( f " response - { response . text } " )
assert response . status_code == 200
response = response . json ( )
print ( response )
response2_id = response [ " id " ]
assert response1_id == response2_id
litellm . disable_cache ( )
2024-01-04 12:54:51 +08:00
except Exception as e :
pytest . fail ( f " LiteLLM Proxy test failed. Exception - { str ( e ) } " )