forked from BerriAI/proxy_load_tester
-
Notifications
You must be signed in to change notification settings - Fork 2
/
no_cache_hits.py
25 lines (21 loc) · 927 Bytes
/
no_cache_hits.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import os
import uuid
from locust import HttpUser, task, between
class MyUser(HttpUser):
wait_time = between(0.5, 1) # Random wait time between requests
@task
def litellm_completion(self):
# no cache hits with this
payload = {
"model": "fake-openai-endpoint",
"messages": [{"role": "user", "content": f"{uuid.uuid4()} This is a test there will be no cache hits and we'll fill up the context" * 150}],
"user": "my-new-end-user-1"
}
response = self.client.post("chat/completions", json=payload)
if response.status_code != 200:
# log the errors in error.txt
with open("error.txt", "a") as error_log:
error_log.write(response.text + "\n")
def on_start(self):
self.api_key = os.getenv('API_KEY', 'sk-1234')
self.client.headers.update({'Authorization': f'Bearer {self.api_key}'})