mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
* feat(agents): support Databricks Agent Platform agents in the agent gateway Add a databricks_agent provider that speaks the MLflow ResponsesAgent wire (Responses-shaped input, output_text items, output_text.delta SSE) to Model Serving endpoints (/invocations or the unified /serving-endpoints/responses) and Databricks Apps (/responses), with a PAT or OAuth M2M client credentials Register it through the agents config block or the Admin UI agent catalog and call it through /a2a/<agent> (message/send, message/stream) or model: a2a/<agent> chat completions. The chat route now runs every completion-bridge agent through its own provider, and attaches the minted Databricks OAuth token for both bridge and URL agents. Consumed OAuth settings never reach the provider call * fix(agents): harden databricks_agent input, output, and OAuth handling Non-text content parts answer 400 unless drop_params is set, a reply without an output list answers 502, flat OAuth fields resolve an os.environ api_base, the bridge debug log no longer echoes the request model, routing tests stub only the HTTP boundary, and the catalog description fits the dropdown * fix(agents): merge agent and client extra_headers on the a2a chat route, and reject an empty databricks_oauth block * fix(agents): bridge only databricks_agent agents on the a2a chat route, keep the card url over a client api_base, and require workspace_url for a Databricks App * refactor(agents): read the a2a chat bridge providers from a constants set instead of branching on a provider name * fix(databricks_agent): let the configured credential win, reject malformed client extra_headers, and keep a missing endpoint name empty * test(agents): drive the flat OAuth a2a send through the real bridge * fix(agents): keep an empty databricks_oauth block ignored on url agents, as before this change * fix(agents): accept bare App hosts, env service principals, and os.environ refs on Databricks agents * fix(databricks_agent): address review caveats on chat bridging, params, and auth Map max_tokens and max_completion_tokens to max_output_tokens and pass temperature, top_p, and user, so /v1/messages works without drop_params. Send A2A message metadata only with string values, answering 400 or dropping the rest under drop_params. Join message items with a blank line in both stream and non-stream replies. Cache the env service principal token until shortly before it expires. Keep databricks_agent replies out of the response cache. On the chat bridge, let request fields win over the agent's stored params while pinning the agent's api_base and api_key, and apply the agent's static_headers between client and agent extra_headers, as /a2a does. Resolve os.environ/ only for databricks_agent api_base and api_key instead of every bridge provider. Ignore an unusable databricks_oauth block on URL agents on chat, as at the merge base, instead of answering 400. * test(internal_context): list the databricks agent token cache as process-local * fix(databricks_agent): keep the agent's endpoint, input, and target out of client control extra_body model and input answer 400, or are dropped under drop_params. A client api_base or base_url never replaces the agent's card URL on the bridge. Chat to a URL agent keeps its merge-base call, and an incomplete databricks_oauth block on a URL agent keeps its merge-base ValueError. * test(databricks): expect the merge-base 400 for a client api_base on agent chat * fix(databricks): keep agent-owned inputs and identity headers out of a chat caller's reach * fix(databricks): bound the agent-over-caller merge to two levels so it needs no recursion --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
687 lines
23 KiB
Python
687 lines
23 KiB
Python
"""
|
|
Unit tests for Databricks App OAuth M2M support for A2A agents.
|
|
|
|
Covers config parsing (including os.environ/ resolution and validation),
|
|
workspace token-URL construction, client_credentials token fetching, caching
|
|
with expiry buffering, and the public ``resolve_databricks_app_auth_header``
|
|
helper.
|
|
"""
|
|
|
|
import base64
|
|
from typing import Final
|
|
from unittest.mock import MagicMock, create_autospec, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
from litellm.proxy.agent_endpoints.databricks_oauth import (
|
|
DatabricksAppOAuthConfig,
|
|
DatabricksAppOAuthTokenCache,
|
|
has_databricks_oauth,
|
|
parse_databricks_oauth_config,
|
|
resolve_databricks_app_auth_header,
|
|
resolve_databricks_backend_auth,
|
|
without_databricks_oauth_params,
|
|
)
|
|
|
|
|
|
def _expected_basic_auth(client_id: str, client_secret: str) -> str:
|
|
token = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
|
|
return f"Basic {token}"
|
|
|
|
|
|
def _mock_http_handler(access_token="tok-abc", expires_in=3600, post_error=None):
|
|
"""Return a mock that mirrors litellm's ``AsyncHTTPHandler`` contract.
|
|
|
|
Two properties of the real handler matter for these tests and were the
|
|
source of a runtime bug the original suite missed:
|
|
|
|
1. ``post`` does not accept an ``auth`` kwarg. ``create_autospec`` enforces
|
|
the real signature, so reintroducing HTTP Basic via ``auth=`` fails with
|
|
``TypeError`` instead of silently passing.
|
|
2. ``post`` calls ``raise_for_status`` internally and raises
|
|
``httpx.HTTPStatusError`` itself on non-2xx; callers never inspect the
|
|
returned response's status. Error-path tests therefore raise from
|
|
``post`` rather than from ``response.raise_for_status``.
|
|
"""
|
|
handler = create_autospec(AsyncHTTPHandler, instance=True)
|
|
if post_error is not None:
|
|
handler.post.side_effect = post_error
|
|
else:
|
|
response = MagicMock()
|
|
response.json = MagicMock(return_value={"access_token": access_token, "expires_in": expires_in})
|
|
handler.post.return_value = response
|
|
return handler
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config parsing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_returns_none_without_block():
|
|
assert parse_databricks_oauth_config(None) is None
|
|
assert parse_databricks_oauth_config({}) is None
|
|
assert parse_databricks_oauth_config({"other": "value"}) is None
|
|
|
|
|
|
def test_parse_builds_config_and_token_url():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"databricks_oauth": {
|
|
"client_id": "cid",
|
|
"client_secret": "secret",
|
|
"workspace_url": "https://dbc-abc.cloud.databricks.com",
|
|
}
|
|
}
|
|
)
|
|
assert config == DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc-abc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
|
|
|
|
def test_parse_strips_serving_endpoints_and_trailing_slash():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"databricks_oauth": {
|
|
"client_id": "cid",
|
|
"client_secret": "secret",
|
|
"workspace_url": "https://dbc-abc.cloud.databricks.com/serving-endpoints/",
|
|
}
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert config.token_url == "https://dbc-abc.cloud.databricks.com/oidc/v1/token"
|
|
|
|
|
|
def test_parse_custom_scope():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"databricks_oauth": {
|
|
"client_id": "cid",
|
|
"client_secret": "secret",
|
|
"workspace_url": "https://dbc-abc.cloud.databricks.com",
|
|
"scope": "custom-scope",
|
|
}
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert config.scope == "custom-scope"
|
|
|
|
|
|
AGENT_KINDS: Final = pytest.mark.parametrize(
|
|
("agent_params", "error"),
|
|
[({}, ValueError), ({"custom_llm_provider": "databricks_agent"}, litellm.BadRequestError)],
|
|
ids=["url-agent", "databricks-agent"],
|
|
)
|
|
|
|
|
|
@AGENT_KINDS
|
|
@pytest.mark.parametrize("missing_field", ["client_id", "client_secret", "workspace_url"])
|
|
def test_parse_raises_on_missing_field(missing_field, agent_params, error):
|
|
block = {
|
|
"client_id": "cid",
|
|
"client_secret": "secret",
|
|
"workspace_url": "https://dbc-abc.cloud.databricks.com",
|
|
}
|
|
block.pop(missing_field)
|
|
with pytest.raises(error, match=missing_field):
|
|
parse_databricks_oauth_config({**agent_params, "databricks_oauth": block})
|
|
|
|
|
|
@AGENT_KINDS
|
|
def test_parse_raises_on_non_mapping_block(agent_params, error):
|
|
with pytest.raises(error, match="mapping"):
|
|
parse_databricks_oauth_config({**agent_params, "databricks_oauth": "not-a-dict"})
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("block", "error"),
|
|
[({}, "missing required field"), ("", "mapping")],
|
|
)
|
|
def test_parse_raises_on_an_empty_nested_block(block, error):
|
|
params = {"custom_llm_provider": "databricks_agent", "databricks_oauth": block}
|
|
assert has_databricks_oauth(params)
|
|
with pytest.raises(litellm.BadRequestError, match=error):
|
|
parse_databricks_oauth_config(params)
|
|
|
|
|
|
@pytest.mark.parametrize("block", [{}, ""])
|
|
@pytest.mark.asyncio
|
|
async def test_an_empty_nested_block_on_a_url_agent_is_ignored_as_before(block):
|
|
auth: Final = await resolve_databricks_backend_auth({"databricks_oauth": block, "api_key": "hop-key"})
|
|
assert auth.header is None
|
|
assert auth.litellm_params == {"api_key": "hop-key"}
|
|
|
|
|
|
def test_parse_resolves_os_environ_references(monkeypatch):
|
|
monkeypatch.setenv("MY_DBX_CLIENT_ID", "env-cid")
|
|
monkeypatch.setenv("MY_DBX_SECRET", "env-secret")
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"databricks_oauth": {
|
|
"client_id": "os.environ/MY_DBX_CLIENT_ID",
|
|
"client_secret": "os.environ/MY_DBX_SECRET",
|
|
"workspace_url": "https://dbc-abc.cloud.databricks.com",
|
|
}
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert config.client_id == "env-cid"
|
|
assert config.client_secret == "env-secret"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token fetching + caching
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_token_posts_client_credentials_with_basic_auth():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
client = _mock_http_handler(access_token="tok-1")
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
token = await cache.async_get_token(config)
|
|
|
|
assert token == "tok-1"
|
|
client.post.assert_awaited_once()
|
|
call = client.post.call_args
|
|
assert call.args[0] == config.token_url
|
|
assert call.kwargs["data"] == {
|
|
"grant_type": "client_credentials",
|
|
"scope": "all-apis",
|
|
}
|
|
# Databricks authenticates the client with HTTP Basic; it must be sent as a
|
|
# header because litellm's AsyncHTTPHandler.post has no ``auth`` parameter.
|
|
assert call.kwargs["headers"]["Authorization"] == _expected_basic_auth("cid", "secret")
|
|
assert "auth" not in call.kwargs
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_is_cached_across_calls():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
client = _mock_http_handler(access_token="tok-cached")
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
first = await cache.async_get_token(config)
|
|
second = await cache.async_get_token(config)
|
|
|
|
assert first == second == "tok-cached"
|
|
client.post.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_distinct_clients_do_not_share_token():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config_a = DatabricksAppOAuthConfig(
|
|
client_id="cid-a",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
config_b = DatabricksAppOAuthConfig(
|
|
client_id="cid-b",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
|
|
clients = [_mock_http_handler("tok-a"), _mock_http_handler("tok-b")]
|
|
|
|
def _next_client(*args, **kwargs):
|
|
return clients.pop(0)
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
side_effect=_next_client,
|
|
):
|
|
token_a = await cache.async_get_token(config_a)
|
|
token_b = await cache.async_get_token(config_b)
|
|
|
|
assert token_a == "tok-a"
|
|
assert token_b == "tok-b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ttl_applies_expiry_buffer():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
client = _mock_http_handler(access_token="tok", expires_in=600)
|
|
|
|
captured = {}
|
|
real_set = cache.set_cache
|
|
|
|
def _spy_set(key, value, **kwargs):
|
|
captured["ttl"] = kwargs.get("ttl")
|
|
return real_set(key, value, **kwargs)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
),
|
|
patch.object(cache, "set_cache", side_effect=_spy_set),
|
|
):
|
|
await cache.async_get_token(config)
|
|
|
|
assert captured["ttl"] == 600 - 60
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_access_token_raises():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
client = _mock_http_handler()
|
|
client.post.return_value.json.return_value = {"not_a_token": "x"}
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
with pytest.raises(ValueError, match="access_token"):
|
|
await cache.async_get_token(config)
|
|
|
|
|
|
def _config():
|
|
return DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_http_status_error_raises_value_error():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
request = httpx.Request("POST", _config().token_url)
|
|
error_response = httpx.Response(status_code=401, request=request)
|
|
client = _mock_http_handler(
|
|
post_error=httpx.HTTPStatusError("unauthorized", request=request, response=error_response)
|
|
)
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
with pytest.raises(ValueError, match="status 401"):
|
|
await cache.async_get_token(_config())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transport_error_raises_value_error():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
client = _mock_http_handler(post_error=httpx.ConnectError("boom"))
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
with pytest.raises(ValueError, match="token request failed"):
|
|
await cache.async_get_token(_config())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_object_json_body_raises():
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
client = _mock_http_handler()
|
|
client.post.return_value.json.return_value = ["not", "an", "object"]
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
with pytest.raises(ValueError, match="non-object JSON"):
|
|
await cache.async_get_token(_config())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("expires_in", [None, "not-a-number"])
|
|
async def test_invalid_expires_in_falls_back_to_default_ttl(expires_in):
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
client = _mock_http_handler(expires_in=expires_in)
|
|
|
|
captured = {}
|
|
real_set = cache.set_cache
|
|
|
|
def _spy_set(key, value, **kwargs):
|
|
captured["ttl"] = kwargs.get("ttl")
|
|
return real_set(key, value, **kwargs)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
),
|
|
patch.object(cache, "set_cache", side_effect=_spy_set),
|
|
):
|
|
await cache.async_get_token(_config())
|
|
|
|
# default TTL (3600) minus the 60s expiry buffer
|
|
assert captured["ttl"] == 3600 - 60
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_short_lived_token_not_cached():
|
|
"""A token whose lifetime is below the refresh buffer is never cached and
|
|
leaves no per-key lock behind."""
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = _config()
|
|
client = _mock_http_handler(access_token="short", expires_in=30)
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
await cache.async_get_token(config)
|
|
await cache.async_get_token(config)
|
|
|
|
assert cache.get_cache(config.cache_key) is None
|
|
assert config.cache_key not in cache._locks
|
|
assert client.post.await_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rotated_secret_forces_new_token():
|
|
"""Rotating client_secret changes the cache key so a fresh token is minted."""
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
old = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="old-secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
rotated = DatabricksAppOAuthConfig(
|
|
client_id="cid",
|
|
client_secret="new-secret",
|
|
token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
assert old.cache_key != rotated.cache_key
|
|
|
|
clients = [_mock_http_handler("old-token"), _mock_http_handler("new-token")]
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
side_effect=lambda *a, **k: clients.pop(0),
|
|
):
|
|
assert await cache.async_get_token(old) == "old-token"
|
|
assert await cache.async_get_token(rotated) == "new-token"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lock_pruned_when_token_evicted():
|
|
"""The per-key lock is removed when its cached token is deleted/evicted."""
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = _config()
|
|
client = _mock_http_handler("tok")
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
await cache.async_get_token(config)
|
|
|
|
assert config.cache_key in cache._locks
|
|
|
|
cache.delete_cache(config.cache_key)
|
|
|
|
assert config.cache_key not in cache._locks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_cache_clears_locks():
|
|
"""flush_cache drops the per-key locks alongside the cached tokens."""
|
|
cache = DatabricksAppOAuthTokenCache()
|
|
config = _config()
|
|
client = _mock_http_handler("tok")
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
await cache.async_get_token(config)
|
|
|
|
assert config.cache_key in cache._locks
|
|
|
|
cache.flush_cache()
|
|
|
|
assert cache._locks == {}
|
|
assert cache.get_cache(config.cache_key) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_returns_none_when_not_configured():
|
|
assert await resolve_databricks_app_auth_header(None) is None
|
|
assert await resolve_databricks_app_auth_header({"foo": "bar"}) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_returns_bearer_header():
|
|
from litellm.proxy.agent_endpoints.databricks_oauth import (
|
|
databricks_app_oauth_token_cache,
|
|
)
|
|
|
|
databricks_app_oauth_token_cache.flush_cache()
|
|
|
|
litellm_params = {
|
|
"databricks_oauth": {
|
|
"client_id": "resolve-cid",
|
|
"client_secret": "secret",
|
|
"workspace_url": "https://resolve.cloud.databricks.com",
|
|
}
|
|
}
|
|
client = _mock_http_handler(access_token="resolved-token")
|
|
|
|
with patch(
|
|
"litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
|
|
return_value=client,
|
|
):
|
|
header = await resolve_databricks_app_auth_header(litellm_params)
|
|
|
|
assert header == {"Authorization": "Bearer resolved-token"}
|
|
|
|
|
|
def test_parse_builds_config_from_flat_databricks_agent_fields():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": "https://adb-1.azuredatabricks.net/serving-endpoints",
|
|
"model": "my-agent",
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
}
|
|
)
|
|
assert config == DatabricksAppOAuthConfig(
|
|
client_id="sp-id",
|
|
client_secret="sp-secret",
|
|
token_url="https://adb-1.azuredatabricks.net/oidc/v1/token",
|
|
scope="all-apis",
|
|
)
|
|
|
|
|
|
def test_parse_flat_fields_take_the_workspace_from_an_os_environ_api_base(monkeypatch):
|
|
monkeypatch.setenv("MY_DBX_HOST", "https://adb-3.azuredatabricks.net/serving-endpoints")
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": "os.environ/MY_DBX_HOST",
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert config.token_url == "https://adb-3.azuredatabricks.net/oidc/v1/token"
|
|
|
|
|
|
def test_parse_flat_fields_prefer_an_explicit_workspace_url_for_an_app():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": "https://my-app-1.azure.databricksapps.com/responses",
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
"workspace_url": "https://adb-1.azuredatabricks.net",
|
|
"scope": "custom-scope",
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert config.token_url == "https://adb-1.azuredatabricks.net/oidc/v1/token"
|
|
assert config.scope == "custom-scope"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"app_base",
|
|
["https://my-app-1.azure.databricksapps.com/responses", "https://my-app-1.azure.databricksapps.com"],
|
|
)
|
|
def test_parse_flat_fields_of_an_app_without_a_workspace_url_raise_instead_of_using_the_app_host(app_base):
|
|
with pytest.raises(litellm.BadRequestError, match="workspace_url") as info:
|
|
parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": app_base,
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
}
|
|
)
|
|
assert info.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"workspace_url",
|
|
["https://my-app-1.azure.databricksapps.com", "https://my-app-1.azure.databricksapps.com/responses"],
|
|
)
|
|
def test_parse_rejects_an_app_url_typed_as_the_workspace_url(workspace_url):
|
|
with pytest.raises(litellm.BadRequestError, match="not a Databricks App URL") as info:
|
|
parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": "https://my-app-1.azure.databricksapps.com/responses",
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
"workspace_url": workspace_url,
|
|
}
|
|
)
|
|
assert info.value.status_code == 400
|
|
|
|
|
|
def test_parse_flat_fields_with_only_a_client_id_raise_on_the_missing_secret():
|
|
with pytest.raises(litellm.BadRequestError, match="client_secret"):
|
|
parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"api_base": "https://adb-1.azuredatabricks.net",
|
|
"client_id": "sp",
|
|
}
|
|
)
|
|
|
|
|
|
def test_parse_ignores_flat_client_fields_of_other_providers():
|
|
assert (
|
|
parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "azure_ai",
|
|
"client_id": "entra-id",
|
|
"client_secret": "entra-secret",
|
|
"tenant_id": "t",
|
|
}
|
|
)
|
|
is None
|
|
)
|
|
assert parse_databricks_oauth_config({"custom_llm_provider": "databricks_agent", "api_key": "pat"}) is None
|
|
|
|
|
|
def test_parse_nested_block_wins_over_flat_fields():
|
|
config = parse_databricks_oauth_config(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"client_id": "flat-id",
|
|
"client_secret": "flat-secret",
|
|
"databricks_oauth": {
|
|
"client_id": "nested-id",
|
|
"client_secret": "nested-secret",
|
|
"workspace_url": "https://adb-2.azuredatabricks.net",
|
|
},
|
|
}
|
|
)
|
|
assert config is not None
|
|
assert (config.client_id, config.token_url) == ("nested-id", "https://adb-2.azuredatabricks.net/oidc/v1/token")
|
|
|
|
|
|
def test_without_oauth_params_drops_every_consumed_key_of_a_databricks_agent():
|
|
forwarded = without_databricks_oauth_params(
|
|
{
|
|
"custom_llm_provider": "databricks_agent",
|
|
"model": "my-agent",
|
|
"api_base": "https://adb-1.azuredatabricks.net",
|
|
"client_id": "sp-id",
|
|
"client_secret": "sp-secret",
|
|
"workspace_url": "https://adb-1.azuredatabricks.net",
|
|
"scope": "all-apis",
|
|
"databricks_oauth": {"client_id": "nested-id", "client_secret": "nested-secret"},
|
|
"custom_inputs": {"tenant": "t-1"},
|
|
}
|
|
)
|
|
assert forwarded == {
|
|
"custom_llm_provider": "databricks_agent",
|
|
"model": "my-agent",
|
|
"api_base": "https://adb-1.azuredatabricks.net",
|
|
"custom_inputs": {"tenant": "t-1"},
|
|
}
|
|
|
|
|
|
def test_without_oauth_params_keeps_other_providers_client_credentials():
|
|
forwarded = without_databricks_oauth_params(
|
|
{
|
|
"custom_llm_provider": "azure_ai",
|
|
"client_id": "entra-id",
|
|
"client_secret": "entra-secret",
|
|
"tenant_id": "t",
|
|
"databricks_oauth": {"client_id": "sp"},
|
|
}
|
|
)
|
|
assert forwarded == {
|
|
"custom_llm_provider": "azure_ai",
|
|
"client_id": "entra-id",
|
|
"client_secret": "entra-secret",
|
|
"tenant_id": "t",
|
|
}
|