File size: 6,174 Bytes
395201c |
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 194 |
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
import pytest
import litellm
from litellm import embedding, completion, completion_cost, Timeout
from litellm import RateLimitError
# test /chat/completion request to the proxy
from fastapi.testclient import TestClient
from fastapi import FastAPI
from litellm.proxy.proxy_server import router # Replace with the actual module where your FastAPI router is defined
app = FastAPI()
app.include_router(router) # Include your router in the test app
client = TestClient(app)
def test_chat_completion():
try:
# Your test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "hi"
},
],
"max_tokens": 10,
}
print("testing proxy server")
response = client.post("/v1/chat/completions", json=test_data)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
# Run the test
# test_chat_completion()
def test_chat_completion_azure():
try:
# Your test data
test_data = {
"model": "azure/chatgpt-v-2",
"messages": [
{
"role": "user",
"content": "write 1 sentence poem"
},
],
"max_tokens": 10,
}
print("testing proxy server with Azure Request")
response = client.post("/v1/chat/completions", json=test_data)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
assert len(result["choices"][0]["message"]["content"]) > 0
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
# Run the test
# test_chat_completion_azure()
def test_embedding():
try:
test_data = {
"model": "azure/azure-embedding-model",
"input": ["good morning from litellm"],
}
print("testing proxy server with OpenAI embedding")
response = client.post("/v1/embeddings", json=test_data)
assert response.status_code == 200
result = response.json()
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
# Run the test
# test_embedding()
def test_add_new_model():
try:
test_data = {
"model_name": "test_openai_models",
"litellm_params": {
"model": "gpt-3.5-turbo",
},
"model_info": {
"description": "this is a test openai model"
}
}
client.post("/model/new", json=test_data)
response = client.get("/model/info")
assert response.status_code == 200
result = response.json()
print(f"response: {result}")
model_info = None
for m in result["data"]:
if m["id"]["model_name"] == "test_openai_models":
model_info = m["id"]["model_info"]
assert model_info["description"] == "this is a test openai model"
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")
# test_add_new_model()
from litellm.integrations.custom_logger import CustomLogger
class MyCustomHandler(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-API Call")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
assert kwargs["user"] == "proxy-user"
assert kwargs["model"] == "gpt-3.5-turbo"
assert kwargs["max_tokens"] == 10
customHandler = MyCustomHandler()
def test_chat_completion_optional_params():
# [PROXY: PROD TEST] - DO NOT DELETE
# This tests if all the /chat/completion params are passed to litellm
try:
# Your test data
litellm.set_verbose=True
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "hi"
},
],
"max_tokens": 10,
"user": "proxy-user"
}
litellm.callbacks = [customHandler]
print("testing proxy server: optional params")
response = client.post("/v1/chat/completions", json=test_data)
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
# Run the test
# test_chat_completion_optional_params()
# Test Reading config.yaml file
from litellm.proxy.proxy_server import load_router_config
def test_load_router_config():
try:
print("testing reading config")
# this is a basic config.yaml with only a model
filepath = os.path.dirname(os.path.abspath(__file__))
result = load_router_config(router=None, config_file_path=f"{filepath}/example_config_yaml/simple_config.yaml")
print(result)
assert len(result[1]) == 1
# this is a load balancing config yaml
result = load_router_config(router=None, config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml")
print(result)
assert len(result[1]) == 2
# config with general settings - custom callbacks
result = load_router_config(router=None, config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml")
print(result)
assert len(result[1]) == 2
except Exception as e:
pytest.fail("Proxy: Got exception reading config", e)
# test_load_router_config()
|