diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 342945a6a6a..8723778779a 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -189,7 +189,15 @@ class PredibaseConfig(BaseConfig): sum_logprob, # [TODO] move this to using the actual logprobs ) - if "best_of" in optional_params and optional_params["best_of"] > 1: + effective_best_of = optional_params.get("best_of") + if effective_best_of is None: + effective_best_of = request_data.get("parameters", {}).get("best_of", 0) + try: + best_of_value = int(effective_best_of) + except (TypeError, ValueError): + best_of_value = 0 + + if best_of_value > 1: if ( "details" in completion_response and "best_of_sequences" in completion_response["details"] diff --git a/tests/test_litellm/llms/test_predibase_transformation.py b/tests/test_litellm/llms/test_predibase_transformation.py index db1fd479150..be13dbb577e 100644 --- a/tests/test_litellm/llms/test_predibase_transformation.py +++ b/tests/test_litellm/llms/test_predibase_transformation.py @@ -290,6 +290,48 @@ def test_predibase_transform_response_best_of_with_empty_generated_text(monkeypa assert result.choices[1].message.content is None +def test_predibase_transform_response_best_of_from_request_data(monkeypatch): + config = PredibaseConfig() + logging_obj = Mock() + encoding = Mock() + encoding.encode.return_value = [1] + monkeypatch.setattr("litellm.token_counter", lambda messages: 1) + + raw_response = httpx.Response( + status_code=200, + json={ + "generated_text": "primary-output", + "details": { + "finish_reason": "stop", + "tokens": [], + "best_of_sequences": [ + { + "generated_text": "secondary-output", + "finish_reason": "length", + "tokens": [], + } + ], + }, + }, + ) + + result = config.transform_response( + model="predibase-model", + raw_response=raw_response, + model_response=_build_model_response(), + logging_obj=logging_obj, + request_data={"inputs": "hello", "parameters": {"best_of": 2}}, + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + litellm_params={}, + encoding=encoding, + api_key="test-key", + ) + + assert len(result.choices) == 2 + assert result.choices[1].message.content == "secondary-output" + + def test_predibase_transform_response_empty_output_sets_completion_tokens_zero(monkeypatch): config = PredibaseConfig() logging_obj = Mock()