diff --git a/litellm/__init__.py b/litellm/__init__.py index 6b824067818..aa5c547bdd1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -140,7 +140,8 @@ open_ai_embedding_models = [ ] from .timeout import timeout -from .utils import client, logging, exception_type, get_optional_params, modify_integration, token_counter, cost_per_token, completion_cost, load_test_model, get_litellm_params +from .testing import * +from .utils import client, logging, exception_type, get_optional_params, modify_integration, token_counter, cost_per_token, completion_cost, get_litellm_params from .main import * # Import all the symbols from main.py from .integrations import * from openai.error import AuthenticationError, InvalidRequestError, RateLimitError, ServiceUnavailableError, OpenAIError \ No newline at end of file diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index 3a9b2741d5d..07167d14986 100644 Binary files a/litellm/__pycache__/__init__.cpython-311.pyc and b/litellm/__pycache__/__init__.cpython-311.pyc differ diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index afc8469fd96..d2da312cb36 100644 Binary files a/litellm/__pycache__/main.cpython-311.pyc and b/litellm/__pycache__/main.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index 92c3af615df..c8658d0785a 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/testing.py b/litellm/testing.py new file mode 100644 index 00000000000..9b13f1f714d --- /dev/null +++ b/litellm/testing.py @@ -0,0 +1,57 @@ +import litellm +import time +from concurrent.futures import ThreadPoolExecutor + +def batch_completion(*args, **kwargs): + batch_messages = args[1] if len(args) > 1 else kwargs.get("messages") + results = [] + completions = [] + exceptions = [] + times = [] + with ThreadPoolExecutor() as executor: + for message_list in batch_messages: + if len(args) > 1: + args_modified = list(args) + args_modified[1] = message_list + future = executor.submit(litellm.completion, *args_modified) + else: + kwargs_modified = dict(kwargs) + kwargs_modified["messages"] = message_list + future = executor.submit(litellm.completion, *args, **kwargs_modified) + completions.append(future) + + # Retrieve the results and calculate elapsed time for each completion call + for future in completions: + start_time = time.time() + try: + result = future.result() + end_time = time.time() + elapsed_time = end_time - start_time + result_dict = {"status": "succeeded", "response": future.result(), "response_time": elapsed_time} + results.append(result_dict) + except Exception as e: + end_time = time.time() + elapsed_time = end_time - start_time + result_dict = {"status": "succeeded", "response": e, "response_time": elapsed_time} + results.append(result_dict) + + return results + +def load_test_model(model: str, custom_llm_provider: str = None, custom_api_base: str = None, prompt: str = None, num_calls: int = None, force_timeout: int = None): + test_prompt = "Hey, how's it going" + test_calls = 100 + if prompt: + test_prompt = prompt + if num_calls: + test_calls = num_calls + messages = [[{"role": "user", "content": test_prompt}] for _ in range(test_calls)] + start_time = time.time() + try: + results = batch_completion(model=model, messages=messages, custom_llm_provider=custom_llm_provider, custom_api_base = custom_api_base, force_timeout=force_timeout) + end_time = time.time() + response_time = end_time - start_time + return {"total_response_time": response_time, "calls_made": test_calls, "prompt": test_prompt, "results": results} + except Exception as e: + end_time = time.time() + response_time = end_time - start_time + return {"total_response_time": response_time, "calls_made": test_calls, "prompt": test_prompt, "exception": e} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 6d676f2a94a..9edb2a2abbe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.400" +version = "0.1.401" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"