feat(models): support Databricks-hosted models

[naorpeled]

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_BASE settings are exported to the DATABRICKS_API_KEY / DATABRICKS_API_BASE environment 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 global api_base used by other providers.

Usage

[config] # in configuration.toml
model = "databricks/databricks-claude-sonnet-4"
fallback_models = ["databricks/databricks-claude-sonnet-4"]
[databricks] # in .secrets.toml
api_key = "..." # Databricks personal access token (PAT)
api_base = "https://adb-xxxx.azuredatabricks.net/serving-endpoints"

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] section
• docs/.../changing_a_model.md: add a Databricks section
• Unit tests covering the provider wiring (key+base set, base optional, nothing exported when unset)

Notes

• Scoped to PAT/key auth (the method requested in the issue). Databricks OAuth/service-principal auth is out of scope for this PR.

[qodo-free-for-open-source-projects]

PR Summary by Qodo

Add Databricks-hosted model support via LiteLLM env-var credentials
✨ Enhancement 🧪 Tests 📝 Documentation ⚙️ Configuration changes 🕐 10-20 Minutes

Description

• Export Databricks PAT and optional API base to LiteLLM’s Databricks provider env vars.
• Document Databricks model configuration and add secrets template entries.
• Add unit tests to ensure env-var wiring is correct and non-leaky.

Diagram

graph TD
  A["PR-Agent settings"] --> B["get_settings()"] --> C["LiteLLMAIHandler.__init__"] --> D["Env: DATABRICKS_*"] --> E{{"LiteLLM Databricks"}} --> F{{"Databricks endpoint"}}

Loading

High-Level Assessment

The following are alternative approaches to this PR:

1. Gate env-var export by selected model prefix

• ➕ Avoids setting Databricks env vars in processes that never use Databricks models
• ➕ Reduces chance of unintended interaction if multiple providers are configured
• ➖ Slightly more logic in init (must parse model/fallbacks)
• ➖ Less consistent with current pattern that exports provider vars when configured

2. Centralize provider env-var exports in a mapping table

• ➕ Reduces repetitive boilerplate across providers (DeepSeek/Mistral/Codestral/Databricks)
• ➕ Makes adding future providers less error-prone
• ➖ Bigger refactor than needed for a single provider
• ➖ Harder to review/validate behavior equivalence across all existing providers

Recommendation: Current approach (exporting Databricks-specific env vars when settings are present) is the best fit for a minimal, low-risk enhancement and matches existing provider patterns. Consider gating by model prefix or a centralized mapping only if more providers are added or init logic grows further.

Files changed (4) +116 / -0

Enhancement (1) +8 / -0

litellm_ai_handler.pyExport Databricks PAT and API base to LiteLLM env vars +8/-0

Export Databricks PAT and API base to LiteLLM env vars

• Extends LiteLLMAIHandler.__init__ to export DATABRICKS.API_KEY and DATABRICKS.API_BASE settings into DATABRICKS_API_KEY and DATABRICKS_API_BASE environment variables. This enables LiteLLM’s Databricks provider to authenticate using PAT/key credentials.

pr_agent/algo/ai_handlers/litellm_ai_handler.py

Tests (1) +89 / -0

test_litellm_databricks_provider.pyAdd unit tests for Databricks env-var wiring +89/-0

Add unit tests for Databricks env-var wiring

• Adds tests verifying that Databricks settings are exported to LiteLLM’s expected env vars, that api_base is optional, and that nothing is exported when settings are absent. Uses environment isolation to prevent leakage between tests.

tests/unittest/test_litellm_databricks_provider.py

Documentation (1) +15 / -0

changing_a_model.mdDocument Databricks model configuration and endpoint naming +15/-0

Document Databricks model configuration and endpoint naming

• Adds a Databricks section describing how to configure PR-Agent to use a Databricks-hosted model via the databricks/ prefix. Includes a TOML example and links to LiteLLM’s Databricks provider documentation.

docs/docs/usage-guide/changing_a_model.md

Other (1) +4 / -0

.secrets_template.tomlAdd Databricks secrets template block +4/-0

Add Databricks secrets template block

• Introduces a new [databricks] section with api_key (PAT) and api_base placeholders, guiding users to provide the required credentials.

pr_agent/settings/.secrets_template.toml

[qodo-free-for-open-source-projects]

Code Review by Qodo

🐞 Bugs (0) 📘 Rule violations (0) 📎 Requirement gaps (0) 🎨 UX issues (0) 🔗 Cross-repo conflicts (0) 📜 Skill insights (0)

1. Azure prefix bypasses Databricks ✓ Resolved 🐞 Bug ≡ Correctness

Description

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).

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[475]

+                api_base = os.environ.get("DATABRICKS_API_BASE") if model.startswith("databricks/") else self.api_base

Evidence

The code prefixes all models with azure/ when Azure is enabled, but the new Databricks
routing/auth guards depend on the model still starting with databricks/; after prefixing, those
checks won’t match and the Databricks-specific behavior will not apply.

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]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

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

---

2. Databricks api_base overridden ✓ Resolved 🐞 Bug ≡ Correctness

Description

LiteLLMAIHandler.chat_completion() always includes api_base=self.api_base in the LiteLLM call,
so if another provider configured self.api_base during __init__, databricks/* calls can be
sent to the wrong base URL despite DATABRICKS_API_BASE being set. The PR only guards api_key
forwarding for Databricks, leaving api_base unguarded and able to override Databricks endpoint
selection in multi-provider configs.

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[R551-555]

+                # Databricks authenticates via the DATABRICKS_API_KEY/DATABRICKS_API_BASE env vars,
+                # so don't override it with another provider's key in multi-provider configs.
+                if (litellm.api_key and litellm.api_key != DUMMY_LITELLM_API_KEY
+                        and not model.startswith("databricks/")):
             kwargs["api_key"] = litellm.api_key

Evidence

chat_completion() builds kwargs with api_base for every model, and other providers can set
self.api_base during initialization (e.g., OpenRouter). The PR change only prevents forwarding
api_key for databricks/*, so api_base can still be wrong for Databricks requests in
mixed-provider configurations.

pr_agent/algo/ai_handlers/litellm_ai_handler.py[472-478]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[224-235]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[549-556]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Databricks is configured via `DATABRICKS_API_KEY`/`DATABRICKS_API_BASE`, but `chat_completion()` always forwards `api_base=self.api_base`. When `self.api_base` was set by another provider (e.g., OpenRouter/Azure AD/Ollama), Databricks calls may be routed to the wrong host.
### Issue Context
The PR added a Databricks-specific guard for `api_key`, but `api_base` is still unconditionally included in the request kwargs.
### Fix
For `model.startswith("databricks/")`, do not pass a foreign `api_base`:
- either omit `api_base` entirely from `kwargs` for Databricks models, or
- set it explicitly from `DATABRICKS_API_BASE` (if present) and otherwise remove it.
Add a unit test mirroring the existing `api_key` guard test to assert `api_base` is not forwarded for `databricks/*` when `self.api_base` is set to another provider.
### Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[472-478]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[549-556]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[224-235]
- tests/unittest/test_litellm_api_key_guard.py[323-348]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

3. Wrong key forwarded ✓ Resolved 🐞 Bug ≡ Correctness

Description

LiteLLMAIHandler.chat_completion forwards litellm.api_key for all models when it is set, so
configuring another provider (e.g., GROQ/OPENROUTER) can cause that key to be passed for
databricks/* calls and override the intended DATABRICKS_API_KEY env-var auth. This can make
Databricks-hosted models fail authentication in multi-provider configurations.

Code

pr_agent/algo/ai_handlers/litellm_ai_handler.py[R203-209]

+        # Support Databricks-hosted models (e.g. Azure Databricks serving endpoints).
+        # Uses PAT/key authentication via LiteLLM's env vars.
+        # SEE https://docs.litellm.ai/docs/providers/databricks
+        if get_settings().get("DATABRICKS.API_KEY", None):
+            os.environ["DATABRICKS_API_KEY"] = get_settings().get("DATABRICKS.API_KEY")
+        if get_settings().get("DATABRICKS.API_BASE", None):
+            os.environ["DATABRICKS_API_BASE"] = get_settings().get("DATABRICKS.API_BASE")

Evidence

The handler stores other providers’ credentials in the global litellm.api_key (e.g.,
Groq/SambaNova/XAI), while Databricks credentials are only exported to env vars. Later,
chat_completion injects litellm.api_key into every request when it’s set (non-dummy), which can
override the Databricks env-var auth path for databricks/* models.

pr_agent/algo/ai_handlers/litellm_ai_handler.py[157-165]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[203-209]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[549-553]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Databricks support exports credentials via `DATABRICKS_API_KEY`/`DATABRICKS_API_BASE` env vars, but the request path may still pass `kwargs["api_key"] = litellm.api_key` whenever any provider populated `litellm.api_key`. This can accidentally send (and prefer) a non-Databricks key for `databricks/*` models.
## Issue Context
- `__init__` sets `litellm.api_key` for some providers (Groq/SambaNova/XAI/OpenRouter/etc.) without regard to which model is being called.
- `chat_completion` injects `api_key` into `acompletion()` kwargs based solely on whether `litellm.api_key` is set.
- Databricks auth is intended to come from `DATABRICKS_API_KEY` env var, so passing an explicit `api_key` from another provider can break Databricks auth.
## Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[549-553]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[157-165]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[203-209]
## Implementation notes
Update the `api_key` injection guard to be provider-aware (e.g., only inject for providers that require `api_key` in kwargs), and explicitly avoid injecting for `databricks/*` models so LiteLLM can use `DATABRICKS_API_KEY`/`DATABRICKS_API_BASE`.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

4. Test state leakage ✓ Resolved 🐞 Bug ☼ Reliability

Description

The new Databricks provider unit tests instantiate LiteLLMAIHandler(), which mutates global
litellm.api_key (dummy fallback) but the autouse fixture only cleans environment variables. This
can leak global state into later tests and make the suite order-dependent/flaky.

Code

tests/unittest/test_litellm_databricks_provider.py[R48-57]

+@pytest.fixture(autouse=True)
+def _isolate_env(monkeypatch):
+    for var in _ISOLATED_ENV:
+        monkeypatch.delenv(var, raising=False)
+    yield
+    # Drop anything the handler wrote so it can't leak into other tests;
+    # monkeypatch then restores any pre-existing originals.
+    for var in ("DATABRICKS_API_KEY", "DATABRICKS_API_BASE"):
+        os.environ.pop(var, None)
+

Evidence

The tests’ autouse fixture removes env vars but doesn’t restore litellm.api_key, while
LiteLLMAIHandler.__init__ sets litellm.api_key to a dummy value when OPENAI_API_KEY is not
present. Since the tests call LiteLLMAIHandler() multiple times, global litellm state can
persist beyond this test file.

tests/unittest/test_litellm_databricks_provider.py[48-57]
tests/unittest/test_litellm_databricks_provider.py[59-89]
pr_agent/algo/ai_handlers/litellm_ai_handler.py[53-57]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`LiteLLMAIHandler.__init__` mutates global `litellm` module state (notably `litellm.api_key` when `OPENAI_API_KEY` is absent). The new test file only isolates/cleans `os.environ`, so `litellm.api_key` can remain changed after the test and affect other tests.
## Issue Context
This repo already has tests that explicitly reset `litellm.api_key` to avoid cross-test pollution; this new test should do the same.
## Fix Focus Areas
- tests/unittest/test_litellm_databricks_provider.py[48-57]
- tests/unittest/test_litellm_databricks_provider.py[59-89]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[53-57]
## Implementation notes
In the autouse fixture, snapshot and restore `litellm.api_key` (and any other mutated `litellm` globals you touch/observe) or explicitly set them to a known value before/after `LiteLLMAIHandler()`.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

---

• Author self-review: I have reviewed the code review findings, and addressed the relevant ones.

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit 51be212

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit 173df5a

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit b74e88f

[naorpeled]

Test state leakage — Databricks provider tests mutate global litellm.api_key (dummy fallback) but the autouse fixture only cleans env vars.

Fixed in d590363. The autouse _isolate_env fixture now snapshots litellm.api_key before yielding and restores it on teardown, so the dummy fallback that LiteLLMAIHandler.__init__ writes (when OPENAI_API_KEY is absent) can no longer leak into order-dependent later tests — matching the cross-test-pollution handling the api_key guard tests already do. Verified by running the file in-process after seeding a sentinel litellm.api_key: the sentinel is intact afterward (it became "dummy_key" without the fix).

[qodo-free-for-open-source-projects]

Code review by qodo was updated up to the latest commit d590363