refactor(predibase): migrate transform_request and transform_response… - #25249
Conversation
… to transformation.py
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
Greptile SummaryThis PR completes the migration of Predibase's request-building, response-parsing, and URL-construction logic from Confidence Score: 5/5Safe to merge — clean refactoring with all prior P1 concerns resolved and comprehensive mock test coverage. All previously raised P1 findings (encoding type, config-instance reuse, best_of class-default visibility, missing tenant_id test) have been addressed. The only remaining finding is a P2 style issue (unused No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/llms/predibase/chat/transformation.py | Receives full transform_request, transform_response, get_complete_url, and output_parser implementations; encoding annotation corrected to Any, best_of fallback now reads from request_data["parameters"] to capture class-level config defaults. Minor: unused api_key param in get_complete_url. |
| litellm/llms/predibase/chat/handler.py | Slimmed to delegation only; process_response and output_parser removed; predibase_config instance correctly forwarded to async_completion with a None fallback guard; streaming paths unchanged. |
| tests/test_litellm/llms/test_predibase_transformation.py | New mock-only test file covering request transform, custom prompt, stream/non-stream URL resolution, missing tenant_id ValueError, best_of from optional_params and from request_data, invalid JSON, error-field, non-dict payload, empty generated_text, invalid best_of fallback, usage fallbacks, and async config reuse. No real network calls. |
Sequence Diagram
sequenceDiagram
participant Caller
participant Handler as PredibaseChatCompletion<br/>(handler.py)
participant Config as PredibaseConfig<br/>(transformation.py)
participant HTTP as HTTP Client
Caller->>Handler: completion(model, messages, ...)
Handler->>Config: validate_environment(...)
Config-->>Handler: headers
Handler->>Config: get_complete_url(api_base, model, tenant_id, stream)
Config-->>Handler: completion_url
Handler->>Config: transform_request(model, messages, optional_params)
Config-->>Handler: data {inputs, parameters}
Handler->>HTTP: POST completion_url, data
HTTP-->>Handler: raw_response
Handler->>Config: transform_response(model, raw_response, ...)
Config-->>Handler: ModelResponse
Handler-->>Caller: ModelResponse
Reviews (7): Last reviewed commit: "test(predibase): cover invalid best_of f..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -135,9 +143,117 @@ def transform_response( | |||
| api_key: Optional[str] = None, | |||
There was a problem hiding this comment.
encoding typed as str but used as tokenizer object
The encoding parameter is annotated as str, but lines 228–231 call .encode() on it and take the length of the result — the calling convention of a tiktoken or similar tokenizer object, not a Python string. The annotation is misleading and will produce MyPy errors on every .encode() invocation.
| api_key: Optional[str] = None, | |
| encoding: Any, |
There was a problem hiding this comment.
The fix is already in this PR! Line 142 of the diff shows encoding: Any instead of encoding: str:
def transform_response( # noqa: PLR0915
...
encoding: Any,This correctly reflects that encoding is a tokenizer object (likely tiktoken's Encoding type) that has an .encode() method, not a string.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules and .greptile/config.
| 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, | ||
| ) |
There was a problem hiding this comment.
async_completion creates a redundant fresh PredibaseConfig() instance
async_completion calls litellm.PredibaseConfig().transform_response(...) — constructing a brand-new instance — while the parent completion() method already holds a predibase_config instance created at line 72. That instance is never forwarded to async_completion, so the async path silently discards it and duplicates work. If instance-level state is added to PredibaseConfig in the future, the async path will silently ignore it.
Consider adding a predibase_config parameter to async_completion and passing the existing instance from completion():
# In completion(), pass the existing instance:
return self.async_completion(
...,
predibase_config=predibase_config,
)
# In async_completion(), accept and use it:
async def async_completion(self, ..., predibase_config: PredibaseConfig, ...) -> ModelResponse:
...
return predibase_config.transform_response(...)There was a problem hiding this comment.
Confirmed! The fix has been properly applied:
- Parameter added (line 210):
async_completionnow acceptspredibase_configparameter - Instance passed (line 152):
completion()passes its existingpredibase_configinstance toasync_completion - Fallback handled (lines 211-212): Creates new instance only if
Noneis passed - Used correctly (line 236): Calls
transform_response()on the passed/fallback instance
The async path now reuses the config instance from the parent method instead of silently discarding it. ✓
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ) | ||
| from litellm.types.utils import LiteLLMLoggingBaseClass | ||
| from litellm.utils import Choices, CustomStreamWrapper, Message, ModelResponse, Usage | ||
| from litellm.utils import CustomStreamWrapper, ModelResponse |
There was a problem hiding this comment.
@github-advanced-security I have fixed it.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
- 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
|
Predibase-targeted tests pass; proxy/prompts failures appear unrelated (MagicMock awaited) |
Add focused Predibase tests to cover remaining transformation and handler branches highlighted by Codecov, including env URL fallback and async/sync delegation paths. Made-with: Cursor
| 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, | ||
| ) |
There was a problem hiding this comment.
best_of default invisible to transform_response
transform_request merges class-level config defaults (e.g. PredibaseConfig.best_of = 2) into the outgoing request body on its own internal copy, but transform_response only sees the raw per-request request_optional_params (no config defaults merged). If best_of is non-None at the class level, the API will return best_of_sequences, but the "best_of" in optional_params guard in transform_response will be False, silently dropping all but the first choice.
The original handler.py avoided this by mutating optional_params in-place with config defaults before passing it to process_response. Now those two paths diverge.
One clean fix: check the already-merged request body instead of optional_params in transform_response:
# In transform_response, replace:
if "best_of" in optional_params and optional_params["best_of"] > 1:
# With:
effective_best_of = optional_params.get("best_of") or request_data.get("parameters", {}).get("best_of", 0)
if effective_best_of > 1:This makes transform_response self-consistent regardless of whether best_of came from a per-request param or a class-level default.
|
@Jerry-SDE can you include a screenshot of this working as expected? |
|
d4c0d55
into
BerriAI:litellm_oss_branch
4 participants
… to transformation.py
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test
Changes
litellm/llms/predibase/chat/handler.pyintolitellm/llms/predibase/chat/transformation.pyby implementingPredibaseConfig.transform_request().handler.pyintoPredibaseConfig.transform_response()(including generated text parsing, finish reason/logprob extraction, best-of handling, usage calculation, andx-*header forwarding).PredibaseConfig.get_complete_url()to centralize Predibase URL construction and/generatevs/generate_streamselection.output_parser()logic totransformation.pyand removed duplicate logic from the handler.PredibaseChatCompletioninhandler.pyto delegate transformation concerns toPredibaseConfigwhile preserving existing sync/async + streaming execution paths.tests/test_litellm/llms/test_predibase_transformation.pyfor:invalid json,errorfield)Relevant issues
Fixes #25205
Pre-Submission checklis
tests/test_litellm/directorymake test-unit(local environment blocked by optional enterprise dependency import in full suite; targeted tests for this change pass)
@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
🧹 Refactoring
✅ Test