test: fix linting testing

This commit is contained in:
Krrish Dholakia 2023-10-19 18:40:33 -07:00
parent 4198901a2d
commit 04f8840a92
3 changed files with 6 additions and 6 deletions

View file

@ -130,12 +130,12 @@ def litellm_completion(data: Dict,
if user_headers:
data["headers"] = user_headers
if type == "completion":
if data["model"] in model_router.get_model_names():
if model_router and data["model"] in model_router.get_model_names():
model_router.text_completion(**data)
else:
response = litellm.text_completion(**data)
elif type == "chat_completion":
if data["model"] in model_router.get_model_names():
if model_router and data["model"] in model_router.get_model_names():
model_router.completion(**data)
else:
response = litellm.completion(**data)

View file

@ -221,5 +221,5 @@ class Router:
# ------------
# Update usage
# ------------
self.cache.increment(tpm_key, total_tokens)
self.cache.increment(rpm_key, 1)
self.increment(tpm_key, total_tokens)
self.increment(rpm_key, 1)

View file

@ -43,7 +43,7 @@ model_list = [{ # list of model deployments
"rpm": 9000
}]
router = Router(model_list=model_list, redis_host=os.getenv("REDIS_HOST"), redis_password=os.getenv("REDIS_PASSWORD"), redis_port=os.getenv("REDIS_PORT"))
router = Router(model_list=model_list, redis_host=os.getenv("REDIS_HOST"), redis_password=os.getenv("REDIS_PASSWORD"), redis_port=int(os.getenv("REDIS_PORT"))) # type: ignore
completions = []
with ThreadPoolExecutor(max_workers=100) as executor:
@ -52,7 +52,7 @@ with ThreadPoolExecutor(max_workers=100) as executor:
"messages": [{"role": "user", "content": "Hey, how's it going?"}]
}
for _ in range(20):
future = executor.submit(router.completion, **kwargs)
future = executor.submit(router.completion, **kwargs) # type: ignore
completions.append(future)
# Retrieve the results from the futures