From 04f8840a922f66950e1e845987c73e6958a3fcfe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 19 Oct 2023 18:40:33 -0700 Subject: [PATCH] test: fix linting testing --- litellm/proxy/llm.py | 4 ++-- litellm/router.py | 4 ++-- litellm/tests/test_router.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/llm.py b/litellm/proxy/llm.py index 581c0890d8d..1d51f1666b2 100644 --- a/litellm/proxy/llm.py +++ b/litellm/proxy/llm.py @@ -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) diff --git a/litellm/router.py b/litellm/router.py index ded21e98b97..c800db39b22 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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) diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py index 8382ebabc9c..1a309c39f7b 100644 --- a/litellm/tests/test_router.py +++ b/litellm/tests/test_router.py @@ -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