From 67518387f187efd3f982514a01a5ef51be4669ab Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 14 Dec 2023 19:48:18 +0530 Subject: [PATCH] (test) ollama test --- litellm/tests/test_ollama.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/litellm/tests/test_ollama.py b/litellm/tests/test_ollama.py index fa000b4ce0f..4a602fe64b6 100644 --- a/litellm/tests/test_ollama.py +++ b/litellm/tests/test_ollama.py @@ -16,15 +16,22 @@ import litellm from litellm.utils import get_optional_params, get_llm_provider def test_get_ollama_params(): - converted_params = get_optional_params(custom_llm_provider="ollama", model="llama2", max_tokens=20, temperature=0.5, stream=True) - print("Converted params", converted_params) - assert converted_params == {'num_predict': 20, 'stream': True, 'temperature': 0.5}, f"{converted_params} != {'num_predict': 20, 'stream': True, 'temperature': 0.5}" + try: + converted_params = get_optional_params(custom_llm_provider="ollama", model="llama2", max_tokens=20, temperature=0.5, stream=True) + print("Converted params", converted_params) + assert converted_params == {'num_predict': 20, 'stream': True, 'temperature': 0.5}, f"{converted_params} != {'num_predict': 20, 'stream': True, 'temperature': 0.5}" + except Exception as e: + pytest.fail(f"Error occurred: {e}") # test_get_ollama_params() def test_get_ollama_model(): - model, custom_llm_provider, _, _ = get_llm_provider("ollama/code-llama-22") - print("Model", "custom_llm_provider", model, custom_llm_provider) - assert custom_llm_provider == "ollama", f"{custom_llm_provider} != ollama" - assert model == "code-llama-22", f"{model} != code-llama-22" + try: + model, custom_llm_provider, _, _ = get_llm_provider("ollama/code-llama-22") + print("Model", "custom_llm_provider", model, custom_llm_provider) + assert custom_llm_provider == "ollama", f"{custom_llm_provider} != ollama" + assert model == "code-llama-22", f"{model} != code-llama-22" + except Exception as e: + pytest.fail(f"Error occurred: {e}") + # test_get_ollama_model() \ No newline at end of file