From 07d7e956f86de2d618e0ccd4eb367ea79fdf1ccf Mon Sep 17 00:00:00 2001
From: Jerry Wei <1506599306@qq.com>
Date: Mon, 6 Apr 2026 17:50:15 -0500
Subject: [PATCH] refactor(predibase): address Greptile review and CodeQL
 warnings

- Fix type annotation for 'encoding' (str -> Any) to match usage
- Optimize async_completion by passing PredibaseConfig instance
- Resolve CodeQL circular import by using litellm.types.utils
- Add comments to empty except blocks for better clarity
- Ensure 100% test parity for transformation logic
---
 litellm/llms/predibase/chat/handler.py        | 6 +++++-
 litellm/llms/predibase/chat/transformation.py | 6 ++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py
index 2f8d10de307..07f2738aa96 100644
--- a/litellm/llms/predibase/chat/handler.py
+++ b/litellm/llms/predibase/chat/handler.py
@@ -150,6 +150,7 @@ class PredibaseChatCompletion:
                     logger_fn=logger_fn,
                     headers=headers,
                     timeout=timeout,
+                    predibase_config=predibase_config,
                 )  # type: ignore
 
         ### SYNC STREAMING
@@ -206,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},
@@ -229,7 +233,7 @@ class PredibaseChatCompletion:
             raise PredibaseError(
                 status_code=500, message="{}".format(str(e))
             )  # don't use verbose_logger.exception, if exception is raised
-        return litellm.PredibaseConfig().transform_response(
+        return predibase_config.transform_response(
             model=model,
             raw_response=response,
             model_response=model_response,
diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py
index ea8f7e7d9d7..342945a6a6a 100644
--- a/litellm/llms/predibase/chat/transformation.py
+++ b/litellm/llms/predibase/chat/transformation.py
@@ -13,7 +13,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
 )
 from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
 from litellm.types.llms.openai import AllMessageValues
-from litellm.utils import Choices, Message, ModelResponse, Usage
+from litellm.types.utils import Choices, Message, ModelResponse, Usage
 
 from ..common_utils import PredibaseError
 
@@ -139,7 +139,7 @@ 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:
@@ -219,6 +219,7 @@ class PredibaseConfig(BaseConfig):
         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:
@@ -230,6 +231,7 @@ class PredibaseConfig(BaseConfig):
                     )
                 )
             except Exception:
+                # Keep usage calculation non-blocking if encoding fails.
                 pass
         else:
             completion_tokens = 0