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
This commit is contained in:
Jerry Wei 2026-04-06 17:50:15 -05:00
parent b8290cfeb5
commit 07d7e956f8
2 changed files with 9 additions and 3 deletions

View file

@ -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,

View file

@ -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