diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index b0d06ef62..5a0b450c7 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -137,13 +137,14 @@ def get_ollama_response( additional_args={"api_base": url, "complete_input_dict": data, "headers": {}, "acompletion": acompletion,}, ) if acompletion is True: - if optional_params.get("stream", False): + if optional_params.get("stream", False) == True: response = ollama_async_streaming(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) else: response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) return response - elif optional_params.get("stream", False): + elif optional_params.get("stream", False) == True: return ollama_completion_stream(url=url, data=data, logging_obj=logging_obj) + response = requests.post( url=f"{url}", json=data, diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 8cecd270a..b46fb0099 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -571,6 +571,22 @@ def test_completion_openai_litellm_key(): # test_completion_openai_litellm_key() +def test_completion_ollama_hosted(): + try: + litellm.set_verbose = True + response = completion( + model="ollama/phi", + messages=messages, + max_tokens=10, + api_base="https://test-ollama-endpoint.onrender.com" + ) + # Add any assertions here to check the response + print(response) + except Exception as e: + pytest.fail(f"Error occurred: {e}") + +# test_completion_ollama_hosted() + def test_completion_openrouter1(): try: response = completion( @@ -626,7 +642,7 @@ def test_completion_anyscale_with_functions(): print(response) except Exception as e: pytest.fail(f"Error occurred: {e}") -test_completion_anyscale_with_functions() +# test_completion_anyscale_with_functions() def test_completion_azure_key_completion_arg(): # this tests if we can pass api_key to completion, when it's not in the env diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 8e6bd346f..2dfd5778a 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -276,6 +276,33 @@ def test_completion_azure_function_calling_stream(): # test_completion_azure_function_calling_stream() +def test_completion_ollama_hosted_stream(): + try: + litellm.set_verbose = True + response = completion( + model="ollama/phi", + messages=messages, + max_tokens=10, + api_base="https://test-ollama-endpoint.onrender.com", + stream=True + ) + # Add any assertions here to check the response + complete_response = "" + # Add any assertions here to check the response + for idx, init_chunk in enumerate(response): + chunk, finished = streaming_format_tests(idx, init_chunk) + complete_response += chunk + if finished: + assert isinstance(init_chunk.choices[0], litellm.utils.StreamingChoices) + break + if complete_response.strip() == "": + raise Exception("Empty response received") + print(f"complete_response: {complete_response}") + except Exception as e: + pytest.fail(f"Error occurred: {e}") + +# test_completion_ollama_hosted_stream() + def test_completion_claude_stream(): try: messages = [