From b8290cfeb5c4fdd5c5ad1821048804f9fe014fcd Mon Sep 17 00:00:00 2001
From: Jerry Wei <1506599306@qq.com>
Date: Mon, 6 Apr 2026 17:03:35 -0500
Subject: [PATCH] refactor(predibase): migrate transform_request and
transform_response to transformation.py
---
litellm/llms/predibase/chat/handler.py | 267 +++---------------
litellm/llms/predibase/chat/transformation.py | 200 ++++++++++++-
.../llms/test_predibase_transformation.py | 178 ++++++++++++
3 files changed, 411 insertions(+), 234 deletions(-)
create mode 100644 tests/test_litellm/llms/test_predibase_transformation.py
diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py
index 79936764acd..2f8d10de307 100644
--- a/litellm/llms/predibase/chat/handler.py
+++ b/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|>",
- "",
- "",
- ]
- 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,9 +144,9 @@ 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,
@@ -363,17 +176,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,
)
@@ -417,17 +229,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 litellm.PredibaseConfig().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,
)
diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py
index 9fbb9d6c9e2..ea8f7e7d9d7 100644
--- a/litellm/llms/predibase/chat/transformation.py
+++ b/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.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,
@@ -135,9 +143,117 @@ class PredibaseConfig(BaseConfig):
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
+ )
+
+ 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)
+
+ prompt_tokens = 0
+ try:
+ prompt_tokens = litellm.token_counter(messages=messages)
+ except Exception:
+ 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:
+ 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 +263,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|>",
+ "",
+ "",
+ ]
+ 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=)`) 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]
diff --git a/tests/test_litellm/llms/test_predibase_transformation.py b/tests/test_litellm/llms/test_predibase_transformation.py
new file mode 100644
index 00000000000..9285c786ec2
--- /dev/null
+++ b/tests/test_litellm/llms/test_predibase_transformation.py
@@ -0,0 +1,178 @@
+from unittest.mock import Mock
+
+import httpx
+import pytest
+
+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_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",
+ "details": {
+ "finish_reason": "eos_token",
+ "tokens": [{"logprob": -0.2}, {"logprob": None}],
+ "best_of_sequences": [
+ {
+ "generated_text": "secondary-output",
+ "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