feat(models): support Databricks-hosted models#2464
Open
naorpeled wants to merge 3 commits into
Open
Conversation
Wire Databricks PAT/key authentication into the LiteLLM handler so models served from Databricks (e.g. Azure Databricks serving endpoints) can be used as the PR-Agent model. DATABRICKS.API_KEY / DATABRICKS.API_BASE settings are exported to the DATABRICKS_API_KEY / DATABRICKS_API_BASE env vars that LiteLLM's Databricks provider reads. - Add provider block to LiteLLMAIHandler.__init__ - Add [databricks] section to .secrets_template.toml - Document usage in changing_a_model.md - Add unit tests for the provider wiring Closes #2246
Contributor
PR Summary by QodoAdd Databricks-hosted model support via LiteLLM env-var credentials Description
Diagram
High-Level Assessment
Files changed (4)
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
Code Review by Qodo
Context used✅ Tickets:
🎫 Add support for Databricks hosted models 1. Azure prefix bypasses Databricks
|
![qodo-free-for-open-source-projects[bot]](https://avatars.githubusercontent.com/in/1068000?s=60&v=4){kind=small}
This comment was marked as spam.
This comment was marked as spam.
Address AI review on #2464: - chat_completion no longer injects litellm.api_key for databricks/* models, so another provider's key can't override DATABRICKS_API_KEY env-var auth in multi-provider configs. - consolidate test_litellm_databricks_provider.py to a single import style (CodeQL: module imported with both 'import' and 'import from'). - add regression test covering the databricks key guard.
Contributor
|
Code review by qodo was updated up to the latest commit 51be212 |
Address AI review on #2464: chat_completion always forwarded api_base=self.api_base, which another provider (OpenRouter/Ollama/Azure AD/OpenAI) may have set during __init__. For databricks/* models this could route requests to the wrong host and override DATABRICKS_API_BASE. Now uses the Databricks endpoint (env var, or None so LiteLLM reads it) for those models. Adds a regression test.
| # Databricks selects its endpoint via the DATABRICKS_API_BASE env var; don't let an | ||
| # api_base configured by another provider (OpenRouter/Ollama/Azure AD/OpenAI) during | ||
| # __init__ override it in multi-provider configs. None lets LiteLLM read the env var. | ||
| api_base = os.environ.get("DATABRICKS_API_BASE") if model.startswith("databricks/") else self.api_base |
Contributor
There was a problem hiding this comment.
1. Azure prefix bypasses databricks 🐞 Bug ≡ Correctness
When Azure mode is enabled, chat_completion() rewrites the model string to start with "azure/", so the new databricks/* checks never trigger and Databricks calls may be routed/authenticated using the wrong provider settings. This breaks Databricks usage in multi-provider configs that also enable Azure (OPENAI.API_TYPE=azure or AZURE_AD).
Agent Prompt
## Issue description
`LiteLLMAIHandler.chat_completion()` prepends `azure/` to *all* model names when `self.azure` is true. The new Databricks logic keys off `model.startswith("databricks/")` to (a) select `api_base` from `DATABRICKS_API_BASE` and (b) avoid forwarding `litellm.api_key`. If the model is rewritten first (e.g. `databricks/foo` -> `azure/databricks/foo`), both guards are bypassed.
## Issue Context
This shows up specifically in multi-provider configurations where Azure is enabled but the request model is a Databricks model.
## Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[418-421]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[471-482]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[553-559]
### Suggested direction
Compute `is_databricks = model.startswith("databricks/")` (or extract provider prefix) *before* any Azure rewriting, then:
- Only apply `model = "azure/" + model` when not `is_databricks` (and ideally when the model isn’t already provider-qualified).
- Use `is_databricks` for both the `api_base` selection and the `api_key` forwarding guard, rather than re-checking `startswith()` on a potentially mutated `model`.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Contributor
|
Code review by qodo was updated up to the latest commit 173df5a |
|
Persistent review updated to latest commit 173df5a |
3 similar comments
|
Persistent review updated to latest commit 173df5a |
|
Persistent review updated to latest commit 173df5a |
|
Persistent review updated to latest commit 173df5a |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds support for using models hosted on Databricks (including Azure Databricks serving endpoints) as the PR-Agent model, via LiteLLM's Databricks provider with PAT/key authentication.
Closes #2246.
Why
LiteLLM already supports Databricks with key auth, but PR-Agent never wired the credentials through, so users with a self-managed (Azure) Databricks instance couldn't point PR-Agent at their existing endpoint.
How
DATABRICKS.API_KEY/DATABRICKS.API_BASEsettings are exported to theDATABRICKS_API_KEY/DATABRICKS_API_BASEenvironment variables that LiteLLM's Databricks provider reads. This mirrors the existing env-var provider pattern (DeepSeek, Mistral, Codestral) and is provider-scoped, so it doesn't affect the globalapi_baseused by other providers.Usage
The name after
databricks/is your serving endpoint name.Changes
LiteLLMAIHandler.__init__: export Databricks credentials to LiteLLM's env vars.secrets_template.toml: add[databricks]sectiondocs/.../changing_a_model.md: add a Databricks sectionNotes