refactor(predibase): migrate transform_request and transform_response… (#25249)

litellm/llms/predibase/chat/handler.py

@@ -2,27 +2,17 @@
 ## Controller file for Predibase Integration - https://predibase.com/
 
 import json
-import os
-import time
 from functools import partial
 from typing import Callable, Optional, Union
 
 import httpx  # type: ignore
 
 import litellm
-import litellm.litellm_core_utils
-import litellm.litellm_core_utils.litellm_logging
-from litellm.litellm_core_utils.core_helpers import map_finish_reason
-from litellm.litellm_core_utils.prompt_templates.factory import (
-    custom_prompt,
-    prompt_factory,
-)
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     get_async_httpx_client,
 )
-from litellm.types.utils import LiteLLMLoggingBaseClass
-from litellm.utils import Choices, CustomStreamWrapper, Message, ModelResponse, Usage
+from litellm.utils import CustomStreamWrapper, ModelResponse
 
 from ..common_utils import PredibaseError
 
@@ -60,162 +50,6 @@ class PredibaseChatCompletion:
     def __init__(self) -> None:
         super().__init__()
 
-    def output_parser(self, generated_text: str):
-        """
-        Parse the output text to remove any special characters. In our current approach we just check for ChatML tokens.
-
-        Initial issue that prompted this - https://github.com/BerriAI/litellm/issues/763
-        """
-        chat_template_tokens = [
-            "<|assistant|>",
-            "<|system|>",
-            "<|user|>",
-            "<s>",
-            "</s>",
-        ]
-        for token in chat_template_tokens:
-            if generated_text.strip().startswith(token):
-                generated_text = generated_text.replace(token, "", 1)
-            if generated_text.endswith(token):
-                generated_text = generated_text[::-1].replace(token[::-1], "", 1)[::-1]
-        return generated_text
-
-    def process_response(  # noqa: PLR0915
-        self,
-        model: str,
-        response: httpx.Response,
-        model_response: ModelResponse,
-        stream: bool,
-        logging_obj: LiteLLMLoggingBaseClass,
-        optional_params: dict,
-        api_key: str,
-        data: Union[dict, str],
-        messages: list,
-        print_verbose,
-        encoding,
-    ) -> ModelResponse:
-        ## LOGGING
-        logging_obj.post_call(
-            input=messages,
-            api_key=api_key,
-            original_response=response.text,
-            additional_args={"complete_input_dict": data},
-        )
-        print_verbose(f"raw model_response: {response.text}")
-        ## RESPONSE OBJECT
-        try:
-            completion_response = response.json()
-        except Exception:
-            raise PredibaseError(message=response.text, status_code=422)
-        if "error" in completion_response:
-            raise PredibaseError(
-                message=str(completion_response["error"]),
-                status_code=response.status_code,
-            )
-        else:
-            if not isinstance(completion_response, dict):
-                raise PredibaseError(
-                    status_code=422,
-                    message=f"'completion_response' is not a dictionary - {completion_response}",
-                )
-            elif "generated_text" not in completion_response:
-                raise PredibaseError(
-                    status_code=422,
-                    message=f"'generated_text' is not a key response dictionary - {completion_response}",
-                )
-            if len(completion_response["generated_text"]) > 0:
-                model_response.choices[0].message.content = self.output_parser(  # type: ignore
-                    completion_response["generated_text"]
-                )
-            ## GETTING LOGPROBS + FINISH REASON
-            if (
-                "details" in completion_response
-                and "tokens" in completion_response["details"]
-            ):
-                model_response.choices[0].finish_reason = map_finish_reason(
-                    completion_response["details"]["finish_reason"]
-                )
-                sum_logprob = 0
-                for token in completion_response["details"]["tokens"]:
-                    if token["logprob"] is not None:
-                        sum_logprob += token["logprob"]
-                setattr(
-                    model_response.choices[0].message,  # type: ignore
-                    "_logprob",
-                    sum_logprob,  # [TODO] move this to using the actual logprobs
-                )
-            if "best_of" in optional_params and optional_params["best_of"] > 1:
-                if (
-                    "details" in completion_response
-                    and "best_of_sequences" in completion_response["details"]
-                ):
-                    choices_list = []
-                    for idx, item in enumerate(
-                        completion_response["details"]["best_of_sequences"]
-                    ):
-                        sum_logprob = 0
-                        for token in item["tokens"]:
-                            if token["logprob"] is not None:
-                                sum_logprob += token["logprob"]
-                        if len(item["generated_text"]) > 0:
-                            message_obj = Message(
-                                content=self.output_parser(item["generated_text"]),
-                                logprobs=sum_logprob,
-                            )
-                        else:
-                            message_obj = Message(content=None)
-                        choice_obj = Choices(
-                            finish_reason=map_finish_reason(item["finish_reason"]),
-                            index=idx + 1,
-                            message=message_obj,
-                        )
-                        choices_list.append(choice_obj)
-                    model_response.choices.extend(choices_list)
-
-        ## CALCULATING USAGE
-        prompt_tokens = 0
-        try:
-            prompt_tokens = litellm.token_counter(messages=messages)
-        except Exception:
-            # this should remain non blocking we should not block a response returning if calculating usage fails
-            pass
-        output_text = model_response["choices"][0]["message"].get("content", "")
-        if output_text is not None and len(output_text) > 0:
-            completion_tokens = 0
-            try:
-                completion_tokens = len(
-                    encoding.encode(
-                        model_response["choices"][0]["message"].get("content", "")
-                    )
-                )  ##[TODO] use a model-specific tokenizer
-            except Exception:
-                # this should remain non blocking we should not block a response returning if calculating usage fails
-                pass
-        else:
-            completion_tokens = 0
-
-        total_tokens = prompt_tokens + completion_tokens
-
-        model_response.created = int(time.time())
-        model_response.model = model
-        usage = Usage(
-            prompt_tokens=prompt_tokens,
-            completion_tokens=completion_tokens,
-            total_tokens=total_tokens,
-        )
-        model_response.usage = usage  # type: ignore
-
-        ## RESPONSE HEADERS
-        predibase_headers = response.headers
-        response_headers = {}
-        for k, v in predibase_headers.items():
-            if k.startswith("x-"):
-                response_headers["llm_provider-{}".format(k)] = v
-
-        model_response._hidden_params["additional_headers"] = response_headers
-
-        return model_response
-
     def completion(
         self,
         model: str,
@@ -235,7 +69,8 @@ class PredibaseChatCompletion:
         logger_fn=None,
         headers: dict = {},
     ) -> Union[ModelResponse, CustomStreamWrapper]:
-        headers = litellm.PredibaseConfig().validate_environment(
+        predibase_config = litellm.PredibaseConfig()
+        headers = predibase_config.validate_environment(
             api_key=api_key,
             headers=headers,
             messages=messages,
@@ -243,54 +78,32 @@ class PredibaseChatCompletion:
             model=model,
             litellm_params=litellm_params,
         )
-        completion_url = ""
-        input_text = ""
-        base_url = "https://serving.app.predibase.com"
-
-        if "https" in model:
-            completion_url = model
-        elif api_base:
-            base_url = api_base
-        elif "PREDIBASE_API_BASE" in os.environ:
-            base_url = os.getenv("PREDIBASE_API_BASE", "")
-
-        completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}"
-
-        if optional_params.get("stream", False) is True:
-            completion_url += "/generate_stream"
-        else:
-            completion_url += "/generate"
-
-        if model in custom_prompt_dict:
-            # check if the model has a registered custom prompt
-            model_prompt_details = custom_prompt_dict[model]
-            prompt = custom_prompt(
-                role_dict=model_prompt_details["roles"],
-                initial_prompt_value=model_prompt_details["initial_prompt_value"],
-                final_prompt_value=model_prompt_details["final_prompt_value"],
-                messages=messages,
-            )
-        else:
-            prompt = prompt_factory(model=model, messages=messages)
-
-        ## Load Config
-        config = litellm.PredibaseConfig.get_config()
-        for k, v in config.items():
-            if (
-                k not in optional_params
-            ):  # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
-                optional_params[k] = v
-
-        stream = optional_params.pop("stream", False)
-
-        data = {
-            "inputs": prompt,
-            "parameters": optional_params,
+        request_optional_params = {**optional_params}
+        stream = request_optional_params.get("stream", False)
+        request_litellm_params = {
+            **litellm_params,
+            "custom_prompt_dict": custom_prompt_dict,
+            "predibase_tenant_id": tenant_id,
         }
-        input_text = prompt
+        completion_url = predibase_config.get_complete_url(
+            api_base=api_base,
+            api_key=api_key,
+            model=model,
+            optional_params=request_optional_params,
+            litellm_params=request_litellm_params,
+            stream=stream,
+        )
+        data = predibase_config.transform_request(
+            model=model,
+            messages=messages,
+            optional_params=request_optional_params,
+            litellm_params=request_litellm_params,
+            headers=headers,
+        )
+
         ## LOGGING
         logging_obj.pre_call(
-            input=input_text,
+            input=data.get("inputs", ""),
             api_key=api_key,
             additional_args={
                 "complete_input_dict": data,
@@ -313,8 +126,8 @@ class PredibaseChatCompletion:
                     encoding=encoding,
                     api_key=api_key,
                     logging_obj=logging_obj,
-                    optional_params=optional_params,
-                    litellm_params=litellm_params,
+                    optional_params=request_optional_params,
+                    litellm_params=request_litellm_params,
                     logger_fn=logger_fn,
                     headers=headers,
                     timeout=timeout,
@@ -331,12 +144,13 @@ class PredibaseChatCompletion:
                     encoding=encoding,
                     api_key=api_key,
                     logging_obj=logging_obj,
-                    optional_params=optional_params,
+                    optional_params=request_optional_params,
                     stream=False,
-                    litellm_params=litellm_params,
+                    litellm_params=request_litellm_params,
                     logger_fn=logger_fn,
                     headers=headers,
                     timeout=timeout,
+                    predibase_config=predibase_config,
                 )  # type: ignore
 
         ### SYNC STREAMING
@@ -363,17 +177,16 @@ class PredibaseChatCompletion:
                 data=json.dumps(data),
                 timeout=timeout,  # type: ignore
             )
-        return self.process_response(
+        return predibase_config.transform_response(
             model=model,
-            response=response,
+            raw_response=response,
             model_response=model_response,
-            stream=optional_params.get("stream", False),
             logging_obj=logging_obj,  # type: ignore
-            optional_params=optional_params,
+            optional_params=request_optional_params,
             api_key=api_key,
-            data=data,
+            request_data=data,
             messages=messages,
-            print_verbose=print_verbose,
+            litellm_params=request_litellm_params,
             encoding=encoding,
         )
 
@@ -394,7 +207,10 @@ class PredibaseChatCompletion:
         litellm_params=None,
         logger_fn=None,
         headers={},
+        predibase_config=None,
     ) -> ModelResponse:
+        if predibase_config is None:
+            predibase_config = litellm.PredibaseConfig()
         async_handler = get_async_httpx_client(
             llm_provider=litellm.LlmProviders.PREDIBASE,
             params={"timeout": timeout},
@@ -417,17 +233,16 @@ class PredibaseChatCompletion:
             raise PredibaseError(
                 status_code=500, message="{}".format(str(e))
             )  # don't use verbose_logger.exception, if exception is raised
-        return self.process_response(
+        return predibase_config.transform_response(
             model=model,
-            response=response,
+            raw_response=response,
             model_response=model_response,
-            stream=stream,
             logging_obj=logging_obj,
             api_key=api_key,
-            data=data,
+            request_data=data,
             messages=messages,
-            print_verbose=print_verbose,
             optional_params=optional_params,
+            litellm_params=litellm_params or {},
             encoding=encoding,
         )
 

litellm/llms/predibase/chat/transformation.py

@@ -1,11 +1,19 @@
+import os
+import time
 from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union
 
 from httpx import Headers, Response
 
+import litellm
 from litellm.constants import DEFAULT_MAX_TOKENS
+from litellm.litellm_core_utils.core_helpers import map_finish_reason
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    custom_prompt,
+    prompt_factory,
+)
 from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import Choices, Message, ModelResponse, Usage
 
 from ..common_utils import PredibaseError
 
@@ -121,7 +129,7 @@ class PredibaseConfig(BaseConfig):
                 optional_params["response_format"] = value
         return optional_params
 
-    def transform_response(
+    def transform_response(  # noqa: PLR0915
         self,
         model: str,
         raw_response: Response,
@@ -131,13 +139,131 @@ class PredibaseConfig(BaseConfig):
         messages: List[AllMessageValues],
         optional_params: dict,
         litellm_params: dict,
-        encoding: str,
+        encoding: Any,
         api_key: Optional[str] = None,
         json_mode: Optional[bool] = None,
     ) -> ModelResponse:
-        raise NotImplementedError(
-            "Predibase transformation currently done in handler.py. Need to migrate to this file."
+        logging_obj.post_call(
+            input=messages,
+            api_key=api_key or "",
+            original_response=raw_response.text,
+            additional_args={"complete_input_dict": request_data},
         )
+        try:
+            completion_response = raw_response.json()
+        except Exception:
+            raise PredibaseError(message=raw_response.text, status_code=422)
+
+        if "error" in completion_response:
+            raise PredibaseError(
+                message=str(completion_response["error"]),
+                status_code=raw_response.status_code,
+            )
+        elif not isinstance(completion_response, dict):
+            raise PredibaseError(
+                status_code=422,
+                message=f"'completion_response' is not a dictionary - {completion_response}",
+            )
+        elif "generated_text" not in completion_response:
+            raise PredibaseError(
+                status_code=422,
+                message=f"'generated_text' is not a key response dictionary - {completion_response}",
+            )
+
+        if len(completion_response["generated_text"]) > 0:
+            model_response.choices[0].message.content = self.output_parser(  # type: ignore
+                completion_response["generated_text"]
+            )
+
+        if "details" in completion_response and "tokens" in completion_response["details"]:
+            model_response.choices[0].finish_reason = map_finish_reason(
+                completion_response["details"]["finish_reason"]
+            )
+            sum_logprob = 0
+            for token in completion_response["details"]["tokens"]:
+                if token["logprob"] is not None:
+                    sum_logprob += token["logprob"]
+            setattr(
+                model_response.choices[0].message,  # type: ignore
+                "_logprob",
+                sum_logprob,  # [TODO] move this to using the actual logprobs
+            )
+
+        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"]
+            ):
+                choices_list = []
+                for idx, item in enumerate(completion_response["details"]["best_of_sequences"]):
+                    sum_logprob = 0
+                    for token in item["tokens"]:
+                        if token["logprob"] is not None:
+                            sum_logprob += token["logprob"]
+                    if len(item["generated_text"]) > 0:
+                        message_obj = Message(
+                            content=self.output_parser(item["generated_text"]),
+                            logprobs=sum_logprob,
+                        )
+                    else:
+                        message_obj = Message(content=None)
+                    choice_obj = Choices(
+                        finish_reason=map_finish_reason(item["finish_reason"]),
+                        index=idx + 1,
+                        message=message_obj,
+                    )
+                    choices_list.append(choice_obj)
+                model_response.choices.extend(choices_list)
+
+        prompt_tokens = 0
+        try:
+            prompt_tokens = litellm.token_counter(messages=messages)
+        except Exception:
+            # Keep usage calculation non-blocking if token counting fails.
+            pass
+        output_text = model_response["choices"][0]["message"].get("content", "")
+        if output_text is not None and len(output_text) > 0:
+            completion_tokens = 0
+            try:
+                completion_tokens = len(
+                    encoding.encode(
+                        model_response["choices"][0]["message"].get("content", "")
+                    )
+                )
+            except Exception:
+                # Keep usage calculation non-blocking if encoding fails.
+                pass
+        else:
+            completion_tokens = 0
+
+        total_tokens = prompt_tokens + completion_tokens
+
+        model_response.created = int(time.time())
+        model_response.model = model
+        usage = Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+        )
+        model_response.usage = usage  # type: ignore
+
+        predibase_headers = raw_response.headers
+        response_headers = {}
+        for k, v in predibase_headers.items():
+            if k.startswith("x-"):
+                response_headers[f"llm_provider-{k}"] = v
+
+        model_response._hidden_params["additional_headers"] = response_headers
+
+        return model_response
 
     def transform_request(
         self,
@@ -147,9 +273,81 @@ class PredibaseConfig(BaseConfig):
         litellm_params: dict,
         headers: dict,
     ) -> dict:
-        raise NotImplementedError(
-            "Predibase transformation currently done in handler.py. Need to migrate to this file."
+        custom_prompt_dict = litellm_params.get("custom_prompt_dict", {})
+        if model in custom_prompt_dict:
+            model_prompt_details = custom_prompt_dict[model]
+            prompt = custom_prompt(
+                role_dict=model_prompt_details["roles"],
+                initial_prompt_value=model_prompt_details["initial_prompt_value"],
+                final_prompt_value=model_prompt_details["final_prompt_value"],
+                messages=messages,
+            )
+        else:
+            prompt = prompt_factory(model=model, messages=messages)
+
+        request_optional_params = {**optional_params}
+        config = self.get_config()
+        for k, v in config.items():
+            if k not in request_optional_params:
+                request_optional_params[k] = v
+
+        request_optional_params.pop("stream", None)
+        return {
+            "inputs": prompt,
+            "parameters": request_optional_params,
+        }
+
+    @staticmethod
+    def output_parser(generated_text: str) -> str:
+        """
+        Parse the output text to remove any special characters.
+
+        Initial issue that prompted this - https://github.com/BerriAI/litellm/issues/763
+        """
+        chat_template_tokens = [
+            "<|assistant|>",
+            "<|system|>",
+            "<|user|>",
+            "<s>",
+            "</s>",
+        ]
+        for token in chat_template_tokens:
+            if generated_text.strip().startswith(token):
+                generated_text = generated_text.replace(token, "", 1)
+            if generated_text.endswith(token):
+                generated_text = generated_text[::-1].replace(token[::-1], "", 1)[::-1]
+        return generated_text
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        tenant_id = litellm_params.get("predibase_tenant_id") or litellm_params.get(
+            "tenant_id"
         )
+        if tenant_id is None:
+            raise ValueError(
+                "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=<MY-ID>)`) or in env - `PREDIBASE_TENANT_ID`."
+            )
+
+        base_url = "https://serving.app.predibase.com"
+        if api_base:
+            base_url = api_base
+        elif "PREDIBASE_API_BASE" in os.environ:
+            base_url = os.getenv("PREDIBASE_API_BASE", "")
+
+        completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}"
+        should_stream = stream if stream is not None else optional_params.get("stream", False)
+        if should_stream is True:
+            completion_url += "/generate_stream"
+        else:
+            completion_url += "/generate"
+        return completion_url
 
     def get_error_class(
         self, error_message: str, status_code: int, headers: Union[dict, Headers]

tests/test_litellm/llms/test_predibase_transformation.py

@@ -0,0 +1,612 @@
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+
+from litellm.llms.predibase.chat.handler import PredibaseChatCompletion
+from litellm.llms.predibase.chat.transformation import PredibaseConfig
+from litellm.llms.predibase.common_utils import PredibaseError
+from litellm.utils import Choices, Message, ModelResponse
+
+
+def _build_model_response() -> ModelResponse:
+    return ModelResponse(
+        choices=[
+            Choices(
+                finish_reason=None,
+                index=0,
+                message=Message(role="assistant", content=""),
+            )
+        ]
+    )
+
+
+def test_predibase_transform_request_non_stream():
+    config = PredibaseConfig()
+    request_data = config.transform_request(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={"temperature": 0.2},
+        litellm_params={},
+        headers={},
+    )
+
+    assert request_data["inputs"]
+    assert request_data["parameters"]["temperature"] == 0.2
+    assert request_data["parameters"]["details"] is True
+    assert "stream" not in request_data["parameters"]
+
+
+def test_predibase_transform_request_custom_prompt(monkeypatch):
+    config = PredibaseConfig()
+
+    monkeypatch.setattr(
+        "litellm.llms.predibase.chat.transformation.custom_prompt",
+        lambda **kwargs: "custom-prompt",
+    )
+
+    request_data = config.transform_request(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={},
+        litellm_params={
+            "custom_prompt_dict": {
+                "predibase-model": {
+                    "roles": {},
+                    "initial_prompt_value": "",
+                    "final_prompt_value": "",
+                }
+            }
+        },
+        headers={},
+    )
+
+    assert request_data["inputs"] == "custom-prompt"
+
+
+def test_predibase_get_complete_url_stream_and_non_stream():
+    config = PredibaseConfig()
+    litellm_params = {"predibase_tenant_id": "tenant-123"}
+
+    non_stream_url = config.get_complete_url(
+        api_base="https://serving.example.com",
+        api_key="test-key",
+        model="predibase-model",
+        optional_params={"stream": False},
+        litellm_params=litellm_params,
+    )
+    stream_url = config.get_complete_url(
+        api_base="https://serving.example.com",
+        api_key="test-key",
+        model="predibase-model",
+        optional_params={"stream": True},
+        litellm_params=litellm_params,
+    )
+
+    assert non_stream_url.endswith("/generate")
+    assert stream_url.endswith("/generate_stream")
+
+
+def test_predibase_get_complete_url_missing_tenant_id():
+    config = PredibaseConfig()
+
+    with pytest.raises(ValueError, match="Missing Predibase Tenant ID"):
+        config.get_complete_url(
+            api_base=None,
+            api_key="test-key",
+            model="predibase-model",
+            optional_params={},
+            litellm_params={},
+        )
+
+
+def test_predibase_get_complete_url_with_tenant_id_key():
+    config = PredibaseConfig()
+
+    url = config.get_complete_url(
+        api_base="https://serving.example.com",
+        api_key="test-key",
+        model="predibase-model",
+        optional_params={},
+        litellm_params={"tenant_id": "tenant-xyz"},
+    )
+
+    assert "tenant-xyz" in url
+    assert url.endswith("/generate")
+
+
+def test_predibase_transform_response_success_best_of(monkeypatch):
+    config = PredibaseConfig()
+    logging_obj = Mock()
+    encoding = Mock()
+    encoding.encode.return_value = [1, 2, 3]
+    monkeypatch.setattr("litellm.token_counter", lambda messages: 5)
+
+    raw_response = httpx.Response(
+        status_code=200,
+        json={
+            "generated_text": "<|assistant|>primary-output</s>",
+            "details": {
+                "finish_reason": "eos_token",
+                "tokens": [{"logprob": -0.2}, {"logprob": None}],
+                "best_of_sequences": [
+                    {
+                        "generated_text": "<s>secondary-output</s>",
+                        "finish_reason": "length",
+                        "tokens": [{"logprob": -0.5}],
+                    }
+                ],
+            },
+        },
+        headers={"x-request-id": "req-123"},
+    )
+
+    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": {}},
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={"best_of": 2},
+        litellm_params={},
+        encoding=encoding,
+        api_key="test-key",
+    )
+
+    assert result.choices[0].message.content == "primary-output"
+    assert len(result.choices) == 2
+    assert result.choices[1].message.content == "secondary-output"
+    assert result.usage.prompt_tokens == 5
+    assert result.usage.completion_tokens == 3
+    assert (
+        result._hidden_params["additional_headers"]["llm_provider-x-request-id"]
+        == "req-123"
+    )
+
+
+def test_predibase_transform_response_invalid_json():
+    config = PredibaseConfig()
+
+    with pytest.raises(PredibaseError) as exc:
+        config.transform_response(
+            model="predibase-model",
+            raw_response=httpx.Response(status_code=200, content=b"not-json"),
+            model_response=_build_model_response(),
+            logging_obj=Mock(),
+            request_data={},
+            messages=[{"role": "user", "content": "hello"}],
+            optional_params={},
+            litellm_params={},
+            encoding=Mock(),
+            api_key="test-key",
+        )
+
+    assert exc.value.status_code == 422
+
+
+def test_predibase_transform_response_error_field():
+    config = PredibaseConfig()
+
+    with pytest.raises(PredibaseError) as exc:
+        config.transform_response(
+            model="predibase-model",
+            raw_response=httpx.Response(
+                status_code=400, json={"error": "invalid request"}
+            ),
+            model_response=_build_model_response(),
+            logging_obj=Mock(),
+            request_data={},
+            messages=[{"role": "user", "content": "hello"}],
+            optional_params={},
+            litellm_params={},
+            encoding=Mock(),
+            api_key="test-key",
+        )
+
+    assert exc.value.status_code == 400
+
+
+def test_predibase_transform_response_missing_generated_text():
+    config = PredibaseConfig()
+
+    with pytest.raises(PredibaseError, match="'generated_text' is not a key"):
+        config.transform_response(
+            model="predibase-model",
+            raw_response=httpx.Response(status_code=200, json={"details": {}}),
+            model_response=_build_model_response(),
+            logging_obj=Mock(),
+            request_data={},
+            messages=[{"role": "user", "content": "hello"}],
+            optional_params={},
+            litellm_params={},
+            encoding=Mock(),
+            api_key="test-key",
+        )
+
+
+def test_predibase_transform_response_non_dict_payload():
+    config = PredibaseConfig()
+    raw_response = Mock()
+    raw_response.text = "[]"
+    raw_response.status_code = 200
+    raw_response.headers = {}
+    raw_response.json.return_value = []
+
+    with pytest.raises(PredibaseError, match="'completion_response' is not a dictionary"):
+        config.transform_response(
+            model="predibase-model",
+            raw_response=raw_response,
+            model_response=_build_model_response(),
+            logging_obj=Mock(),
+            request_data={},
+            messages=[{"role": "user", "content": "hello"}],
+            optional_params={},
+            litellm_params={},
+            encoding=Mock(),
+            api_key="test-key",
+        )
+
+
+def test_predibase_transform_response_best_of_with_empty_generated_text(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": "",
+                        "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": {}},
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={"best_of": 2},
+        litellm_params={},
+        encoding=encoding,
+        api_key="test-key",
+    )
+
+    assert len(result.choices) == 2
+    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_best_of_invalid_value_falls_back(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": {}},
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={"best_of": "invalid-int"},
+        litellm_params={},
+        encoding=encoding,
+        api_key="test-key",
+    )
+
+    # Invalid best_of should safely fall back to 0 and not append extra choices.
+    assert len(result.choices) == 1
+    assert result.choices[0].message.content == "primary-output"
+
+
+def test_predibase_transform_response_empty_output_sets_completion_tokens_zero(monkeypatch):
+    config = PredibaseConfig()
+    logging_obj = Mock()
+    encoding = Mock()
+    monkeypatch.setattr("litellm.token_counter", lambda messages: 3)
+
+    raw_response = httpx.Response(
+        status_code=200,
+        json={"generated_text": "", "details": {"tokens": [], "finish_reason": "stop"}},
+    )
+
+    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": {}},
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={},
+        litellm_params={},
+        encoding=encoding,
+        api_key="test-key",
+    )
+
+    assert result.usage.prompt_tokens == 3
+    assert result.usage.completion_tokens == 0
+
+
+def test_predibase_get_complete_url_uses_env_base_url(monkeypatch):
+    config = PredibaseConfig()
+    monkeypatch.setenv("PREDIBASE_API_BASE", "https://env.predibase.com")
+
+    url = config.get_complete_url(
+        api_base=None,
+        api_key="test-key",
+        model="predibase-model",
+        optional_params={},
+        litellm_params={"predibase_tenant_id": "tenant-123"},
+    )
+
+    assert url.startswith("https://env.predibase.com/tenant-123/")
+
+
+def test_predibase_transform_response_usage_fallbacks(monkeypatch):
+    config = PredibaseConfig()
+    logging_obj = Mock()
+    encoding = Mock()
+    encoding.encode.side_effect = RuntimeError("encoding failure")
+    monkeypatch.setattr(
+        "litellm.token_counter", lambda messages: (_ for _ in ()).throw(RuntimeError())
+    )
+
+    raw_response = httpx.Response(
+        status_code=200,
+        json={"generated_text": "ok", "details": {"tokens": [], "finish_reason": "stop"}},
+    )
+
+    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": {}},
+        messages=[{"role": "user", "content": "hello"}],
+        optional_params={},
+        litellm_params={},
+        encoding=encoding,
+        api_key="test-key",
+    )
+
+    assert result.usage.prompt_tokens == 0
+    assert result.usage.completion_tokens == 0
+
+
+@pytest.mark.asyncio
+async def test_predibase_async_completion_uses_default_config_when_none(monkeypatch):
+    handler = PredibaseChatCompletion()
+    mock_response = httpx.Response(status_code=200, json={"generated_text": "ok"})
+
+    async_handler = Mock()
+    async_handler.post = AsyncMock(return_value=mock_response)
+    monkeypatch.setattr(
+        "litellm.llms.predibase.chat.handler.get_async_httpx_client",
+        lambda **kwargs: async_handler,
+    )
+
+    default_config = Mock()
+    default_config.transform_response.return_value = _build_model_response()
+    monkeypatch.setattr("litellm.PredibaseConfig", lambda: default_config)
+
+    result = await handler.async_completion(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        api_base="https://serving.example.com/x/generate",
+        model_response=_build_model_response(),
+        print_verbose=Mock(),
+        encoding=Mock(),
+        api_key="test-key",
+        logging_obj=Mock(),
+        stream=False,
+        data={"inputs": "hello", "parameters": {}},
+        optional_params={},
+        timeout=10,
+        litellm_params={},
+        headers={"Authorization": "Bearer test"},
+    )
+
+    assert result is default_config.transform_response.return_value
+    default_config.transform_response.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_predibase_async_completion_uses_passed_config(monkeypatch):
+    handler = PredibaseChatCompletion()
+    mock_response = httpx.Response(status_code=200, json={"generated_text": "ok"})
+
+    async_handler = Mock()
+    async_handler.post = AsyncMock(return_value=mock_response)
+    monkeypatch.setattr(
+        "litellm.llms.predibase.chat.handler.get_async_httpx_client",
+        lambda **kwargs: async_handler,
+    )
+
+    passed_config = Mock()
+    passed_config.transform_response.return_value = _build_model_response()
+
+    result = await handler.async_completion(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        api_base="https://serving.example.com/x/generate",
+        model_response=_build_model_response(),
+        print_verbose=Mock(),
+        encoding=Mock(),
+        api_key="test-key",
+        logging_obj=Mock(),
+        stream=False,
+        data={"inputs": "hello", "parameters": {}},
+        optional_params={},
+        timeout=10,
+        litellm_params={},
+        headers={"Authorization": "Bearer test"},
+        predibase_config=passed_config,
+    )
+
+    assert result is passed_config.transform_response.return_value
+    passed_config.transform_response.assert_called_once()
+
+
+def test_predibase_completion_sync_returns_transform_response(monkeypatch):
+    handler = PredibaseChatCompletion()
+    expected = _build_model_response()
+
+    def fake_validate_environment(self, **kwargs):
+        return {"Authorization": "Bearer test"}
+
+    def fake_get_complete_url(self, **kwargs):
+        return "https://serving.example.com/tenant/deployments/v2/llms/model/generate"
+
+    def fake_transform_request(self, **kwargs):
+        return {"inputs": "hello", "parameters": {}}
+
+    def fake_transform_response(self, **kwargs):
+        return expected
+
+    monkeypatch.setattr(PredibaseConfig, "validate_environment", fake_validate_environment)
+    monkeypatch.setattr(PredibaseConfig, "get_complete_url", fake_get_complete_url)
+    monkeypatch.setattr(PredibaseConfig, "transform_request", fake_transform_request)
+    monkeypatch.setattr(PredibaseConfig, "transform_response", fake_transform_response)
+    monkeypatch.setattr(
+        "litellm.module_level_client.post",
+        lambda *args, **kwargs: httpx.Response(status_code=200, json={"generated_text": "ok"}),
+    )
+
+    result = handler.completion(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        api_base="https://serving.example.com",
+        custom_prompt_dict={},
+        model_response=_build_model_response(),
+        print_verbose=Mock(),
+        encoding=Mock(),
+        api_key="test-key",
+        logging_obj=Mock(),
+        optional_params={},
+        litellm_params={},
+        tenant_id="tenant-123",
+        timeout=10,
+        acompletion=False,
+    )
+
+    assert result is expected
+
+
+def test_predibase_completion_passes_existing_config_to_async_completion(monkeypatch):
+    handler = PredibaseChatCompletion()
+    captured = {}
+
+    def fake_validate_environment(self, **kwargs):
+        captured["config_instance"] = self
+        return {"Authorization": "Bearer test"}
+
+    def fake_get_complete_url(self, **kwargs):
+        return "https://serving.example.com/tenant/deployments/v2/llms/model/generate"
+
+    def fake_transform_request(self, **kwargs):
+        return {"inputs": "hello", "parameters": {}}
+
+    def fake_async_completion(**kwargs):
+        captured["async_kwargs"] = kwargs
+        return "async-result"
+
+    monkeypatch.setattr(PredibaseConfig, "validate_environment", fake_validate_environment)
+    monkeypatch.setattr(PredibaseConfig, "get_complete_url", fake_get_complete_url)
+    monkeypatch.setattr(PredibaseConfig, "transform_request", fake_transform_request)
+    monkeypatch.setattr(handler, "async_completion", fake_async_completion)
+
+    result = handler.completion(
+        model="predibase-model",
+        messages=[{"role": "user", "content": "hello"}],
+        api_base="https://serving.example.com",
+        custom_prompt_dict={},
+        model_response=_build_model_response(),
+        print_verbose=Mock(),
+        encoding=Mock(),
+        api_key="test-key",
+        logging_obj=Mock(),
+        optional_params={},
+        litellm_params={},
+        tenant_id="tenant-123",
+        timeout=10,
+        acompletion=True,
+    )
+
+    assert result == "async-result"
+    assert captured["async_kwargs"]["predibase_config"] is captured["config_instance"]