feat(models): support Databricks-hosted models

docs/docs/usage-guide/changing_a_model.md

@@ -386,6 +386,21 @@ key = "..." # your Codestral api key
 
 (you can obtain a Codestral key from [here](https://console.mistral.ai/codestral))
 
+### Databricks
+
+To use a model hosted on Databricks (e.g. an Azure Databricks serving endpoint), set:
+
+```toml
+[config] # in configuration.toml
+model = "databricks/databricks-claude-sonnet-4"
+fallback_models = ["databricks/databricks-claude-sonnet-4"]
+[databricks] # in .secrets.toml
+api_key = "..." # your Databricks personal access token (PAT)
+api_base = "https://adb-xxxx.azuredatabricks.net/serving-endpoints" # your workspace serving-endpoints URL
+```
+
+The model name after the `databricks/` prefix is the name of your serving endpoint. See LiteLLM's [Databricks provider docs](https://docs.litellm.ai/docs/providers/databricks) for details.
+
 ### Openrouter
 
 To use model from Openrouter, for example, set:

pr_agent/algo/ai_handlers/litellm_ai_handler.py

@@ -200,6 +200,14 @@ def __init__(self):
         if get_settings().get("CODESTRAL.KEY", None):
             os.environ["CODESTRAL_API_KEY"] = get_settings().get("CODESTRAL.KEY")
 
+        # 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")
+
         # Check for Azure AD configuration
         if get_settings().get("AZURE_AD.CLIENT_ID", None):
             self.azure = True

pr_agent/settings/.secrets_template.toml

@@ -121,6 +121,10 @@ key = ""
 [deepinfra]
 key = ""
 
+[databricks]
+api_key = ""  # Databricks personal access token (PAT)
+api_base = ""  # e.g. https://adb-xxxx.azuredatabricks.net/serving-endpoints
+
 [azure_ad]
 # Azure AD authentication for OpenAI services
 client_id = ""  # Your Azure AD application client ID

tests/unittest/test_litellm_databricks_provider.py

@@ -0,0 +1,89 @@
+"""
+Tests for Databricks provider wiring in LiteLLMAIHandler.__init__.
+
+Verifies that DATABRICKS.API_KEY / DATABRICKS.API_BASE settings are exported to
+the env vars LiteLLM's Databricks provider reads (DATABRICKS_API_KEY /
+DATABRICKS_API_BASE), and that nothing is exported when they are unset.
+"""
+import os
+
+import pytest
+
+import pr_agent.algo.ai_handlers.litellm_ai_handler as litellm_handler
+from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAIHandler
+
+# Env vars LiteLLMAIHandler.__init__ branches on — clear them so the handler
+# under test isn't influenced by (or leaking into) the runner environment.
+_ISOLATED_ENV = (
+    "DATABRICKS_API_KEY",
+    "DATABRICKS_API_BASE",
+    "OPENAI_API_KEY",
+    "AWS_USE_IMDS",
+    "AWS_ACCESS_KEY_ID",
+    "AWS_SECRET_ACCESS_KEY",
+    "AWS_SESSION_TOKEN",
+    "AWS_REGION_NAME",
+)
+
+
+def _make_settings(overrides):
+    """Minimal settings whose top-level .get() returns the provided overrides."""
+    return type("Settings", (), {
+        "config": type("Config", (), {
+            "reasoning_effort": None,
+            "ai_timeout": 30,
+            "custom_reasoning_model": False,
+            "max_model_tokens": 32000,
+            "verbosity_level": 0,
+            "seed": -1,
+            "get": lambda self, key, default=None: default,
+        })(),
+        "litellm": type("LiteLLM", (), {
+            "get": lambda self, key, default=None: default,
+        })(),
+        "get": lambda self, key, default=None: overrides.get(key, default),
+    })()
+
+
+@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)
+
+
+def test_databricks_env_vars_exported_from_settings(monkeypatch):
+    overrides = {
+        "DATABRICKS.API_KEY": "dapi-test-123",
+        "DATABRICKS.API_BASE": "https://adb-1234.azuredatabricks.net/serving-endpoints",
+    }
+    monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides))
+
+    LiteLLMAIHandler()
+
+    assert os.environ["DATABRICKS_API_KEY"] == "dapi-test-123"
+    assert os.environ["DATABRICKS_API_BASE"] == "https://adb-1234.azuredatabricks.net/serving-endpoints"
+
+
+def test_databricks_env_vars_absent_when_unset(monkeypatch):
+    monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings({}))
+
+    LiteLLMAIHandler()
+
+    assert "DATABRICKS_API_KEY" not in os.environ
+    assert "DATABRICKS_API_BASE" not in os.environ
+
+
+def test_databricks_api_base_optional(monkeypatch):
+    """API base is optional (a workspace default may be configured elsewhere)."""
+    overrides = {"DATABRICKS.API_KEY": "dapi-only-key"}
+    monkeypatch.setattr(litellm_handler, "get_settings", lambda: _make_settings(overrides))
+
+    LiteLLMAIHandler()
+
+    assert os.environ["DATABRICKS_API_KEY"] == "dapi-only-key"
+    assert "DATABRICKS_API_BASE" not in os.environ