forked from phoenix/litellm-mirror
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
# What this tests ?
|
|
## Tests /chat/completions by generating a key and then making a chat completions request
|
|
import pytest
|
|
import asyncio
|
|
import aiohttp
|
|
|
|
|
|
async def generate_key(session):
|
|
url = "http://0.0.0.0:4000/key/generate"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data = {
|
|
"models": ["gpt-4"],
|
|
"duration": None,
|
|
}
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
print(response_text)
|
|
print()
|
|
|
|
if status != 200:
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
return await response.json()
|
|
|
|
|
|
async def chat_completion(session, key):
|
|
url = "http://0.0.0.0:4000/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"model": "gpt-4",
|
|
"messages": [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Hello!"},
|
|
],
|
|
}
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
print(response_text)
|
|
print()
|
|
|
|
if status != 200:
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_key_gen():
|
|
async with aiohttp.ClientSession() as session:
|
|
key_gen = await generate_key(session=session)
|
|
key = key_gen["key"]
|
|
await chat_completion(session=session, key=key)
|