diff --git a/litellm/__init__.py b/litellm/__init__.py index 25cae832826..0b2ec8820b6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -483,7 +483,12 @@ openai_compatible_providers: List = [ "azure_ai", "github", ] - +openai_text_completion_compatible_providers: List = ( + [ # providers that support `/v1/completions` + "together_ai", + "fireworks_ai", + ] +) # well supported replicate llms replicate_models: List = [ diff --git a/litellm/main.py b/litellm/main.py index cb355561916..2095322328f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1209,6 +1209,9 @@ def completion( custom_llm_provider == "text-completion-openai" or "ft:babbage-002" in model or "ft:davinci-002" in model # support for finetuned completion models + or custom_llm_provider + in litellm.openai_text_completion_compatible_providers + and kwargs.get("text_completion") is True ): openai.api_type = "openai" @@ -4099,8 +4102,8 @@ def text_completion( kwargs.pop("prompt", None) - if ( - _model is not None and custom_llm_provider == "openai" + if _model is not None and ( + custom_llm_provider == "openai" ): # for openai compatible endpoints - e.g. vllm, call the native /v1/completions endpoint for text completion calls if _model not in litellm.open_ai_chat_completion_models: model = "text-completion-openai/" + _model diff --git a/litellm/tests/test_text_completion.py b/litellm/tests/test_text_completion.py index 70325c44e22..e6a4a0499cc 100644 --- a/litellm/tests/test_text_completion.py +++ b/litellm/tests/test_text_completion.py @@ -4239,3 +4239,14 @@ def test_completion_vllm(): mock_call.assert_called_once() assert "hello" in mock_call.call_args.kwargs["extra_body"] + + +def test_completion_fireworks_ai_multiple_choices(): + litellm.set_verbose = True + response = litellm.text_completion( + model="fireworks_ai/llama-v3p1-8b-instruct", + prompt=["halo", "hi", "halo", "hi"], + ) + print(response.choices) + + assert len(response.choices) == 4