diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index d664acb5d58..0002a597d6a 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -11,6 +11,7 @@ A2A Streaming Events (in order): """ from collections.abc import AsyncIterator, Callable, Coroutine, Mapping +from types import MappingProxyType from typing import Any, Final import litellm @@ -21,6 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( A2AStreamingContext, ) from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager +from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS from litellm.interactions.agents.utils import merge_agent_headers from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse @@ -44,6 +46,32 @@ _AGENT_ONLY_PARAMS: Final = frozenset( ) +def bridge_model_name(litellm_params: Mapping[str, object]) -> str: + """The litellm model string a bridge agent completes with: ``/``. + + A provider bridged on chat names a real upstream endpoint with ``model``, so a missing one stays empty for that + provider to reject or ignore instead of becoming a made-up ``agent`` endpoint. + """ + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + default_model: Final = "" if custom_llm_provider in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS else "agent" + configured_model: Final = litellm_params.get("model") + model: Final = configured_model if isinstance(configured_model, str) else default_model + if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"): + return model + return f"{custom_llm_provider}/{model}" + + +def agent_completion_kwargs(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + """The agent's litellm_params that are valid litellm.acompletion kwargs (api_key, api_base, ...).""" + return MappingProxyType( + { + key: value + for key, value in litellm_params.items() + if key not in ("model", "custom_llm_provider") and key not in _AGENT_ONLY_PARAMS + } + ) + + class A2ACompletionBridgeHandler: """ Static methods for handling A2A requests via LiteLLM completion. @@ -64,16 +92,7 @@ class A2ACompletionBridgeHandler: # Transform A2A message to OpenAI format openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) - # Get completion params - custom_llm_provider: Final = litellm_params.get("custom_llm_provider") - model: Final[str] = litellm_params.get("model", "agent") - - # Build full model string if provider specified - # Skip prepending if model already starts with the provider prefix - if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"): - full_model = f"{custom_llm_provider}/{model}" - else: - full_model = model + full_model: Final = bridge_model_name(litellm_params) if stream: verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base) @@ -87,13 +106,7 @@ class A2ACompletionBridgeHandler: "api_base": api_base, "stream": stream, } - # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) - litellm_params_to_add: Final = { - k: v - for k, v in litellm_params.items() - if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS - } - completion_params.update(litellm_params_to_add) + completion_params.update(agent_completion_kwargs(litellm_params)) # Apply forward metadata AFTER the litellm_params merge so the helper # sees any agent-owner-configured ``extra_body.metadata`` and can keep # those keys authoritative over the client-supplied A2A metadata. diff --git a/litellm/constants.py b/litellm/constants.py index 847b8564507..584006be7f9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -17,7 +17,7 @@ OAUTH_TOKEN_EXCHANGE_CACHE_SAFETY_MARGIN_SECONDS: Final = 60 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_PROFILE: Final = "jwt_bearer_obo" MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_SCOPE: Final = "https://graph.microsoft.com/.default" # Replies depend on the caller's delegated identity, which response-cache keys do not include. -RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot"}) +RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot", "databricks_agent"}) MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC" GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth" GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type" @@ -845,6 +845,8 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset( } ) +A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"}) + LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [ "openai", "azure", diff --git a/litellm/llms/databricks/agent/__init__.py b/litellm/llms/databricks/agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/databricks/agent/responses_output.py b/litellm/llms/databricks/agent/responses_output.py new file mode 100644 index 00000000000..cbcb6913577 --- /dev/null +++ b/litellm/llms/databricks/agent/responses_output.py @@ -0,0 +1,102 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.databricks.common_utils import DatabricksException +from litellm.types.utils import GenericStreamingChunk, Usage + + +class _ContentPart(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + text: str | None = None + + +class OutputItem(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + id: str | None = None + content: tuple[_ContentPart, ...] | None = None + + +class AgentUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + input_tokens: int + output_tokens: int + total_tokens: int | None = None + + def to_usage(self) -> Usage: + return Usage( + prompt_tokens=self.input_tokens, + completion_tokens=self.output_tokens, + total_tokens=self.total_tokens if self.total_tokens is not None else self.input_tokens + self.output_tokens, + ) + + +class AgentResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + output: tuple[OutputItem, ...] + custom_outputs: object = None + usage: AgentUsage | None = None + + +class _StreamEvent(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "" + item_id: str | None = None + delta: str | None = None + item: OutputItem | None = None + message: object = None + + +_MESSAGE_SEPARATOR: Final = "\n\n" + + +def output_text(items: Sequence[OutputItem] | None) -> str: + texts: Final = (_message_text(item) for item in items or ()) + return _MESSAGE_SEPARATOR.join(text for text in texts if text) + + +def _message_text(item: OutputItem) -> str: + if item.type != "message": + return "" + return "".join(part.text or "" for part in item.content or () if part.type == "output_text") + + +def _text_chunk(text: str) -> GenericStreamingChunk: + return GenericStreamingChunk(text=text, is_finished=False, finish_reason="", usage=None, index=0, tool_use=None) + + +class DatabricksAgentResponsesIterator(BaseModelResponseIterator): + _streamed_item_ids: frozenset[str] = frozenset() + _emitted_text: bool = False + _last_text_item_id: str | None = None + + def chunk_parser(self, chunk: Mapping[str, object]) -> GenericStreamingChunk: + event: Final = _StreamEvent.model_validate(chunk) + if event.type == "response.output_text.delta": + if event.item_id is not None: + self._streamed_item_ids = self._streamed_item_ids | {event.item_id} + return self._item_text_chunk(event.delta or "", event.item_id) + if event.type == "response.output_item.done": + if event.item is None or event.item.id in self._streamed_item_ids: + return _text_chunk("") + return self._item_text_chunk(output_text((event.item,)), event.item.id) + if event.type == "error": + raise DatabricksException(status_code=500, message=f"Databricks agent stream error: {event.message}") + return _text_chunk("") + + def _item_text_chunk(self, text: str, item_id: str | None) -> GenericStreamingChunk: + if not text: + return _text_chunk("") + separator: Final = _MESSAGE_SEPARATOR if self._emitted_text and item_id != self._last_text_item_id else "" + self._emitted_text = True + self._last_text_item_id = item_id + return _text_chunk(f"{separator}{text}") diff --git a/litellm/llms/databricks/agent/transformation.py b/litellm/llms/databricks/agent/transformation.py new file mode 100644 index 00000000000..fb0fdd22132 --- /dev/null +++ b/litellm/llms/databricks/agent/transformation.py @@ -0,0 +1,467 @@ +"""Databricks ResponsesAgent endpoints. + +Model Serving: POST {workspace}/serving-endpoints/{endpoint}/invocations +Model Serving, unified: POST {workspace}/serving-endpoints/responses with the endpoint as the body "model" +Databricks Apps: POST {app}/responses +""" + +import hashlib +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import quote, urlsplit + +import httpx +from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter, ValidationError + +import litellm +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.databricks.agent.responses_output import ( + AgentResponse, + DatabricksAgentResponsesIterator, + output_text, +) +from litellm.llms.databricks.common_utils import DatabricksBase, DatabricksException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + +PROVIDER_NAME: Final = "databricks_agent" +_MODEL_PREFIX: Final = f"{PROVIDER_NAME}/" +_INVOCATIONS_SUFFIX: Final = "/invocations" +_APP_RESPONSES_SUFFIX: Final = "/responses" +_ENDPOINT_URL_SUFFIXES: Final = (_INVOCATIONS_SUFFIX, _APP_RESPONSES_SUFFIX) +_SERVING_ENDPOINTS_SEGMENT: Final = "/serving-endpoints/" +_UNIFIED_RESPONSES_SUFFIX: Final = "/serving-endpoints/responses" +_APP_HOST_SUFFIX: Final = ".databricksapps.com" +_TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 + + +class _TokenCache(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: str, *, ttl: int) -> None: ... + + +def _new_token_cache() -> _TokenCache: + return InMemoryCache(max_size_in_memory=64) + + +_ENV_M2M_TOKENS: Final = _new_token_cache() +_INPUT_ROLES: Final = frozenset({"system", "developer", "user", "assistant"}) +_BODY_PARAMS: Final = ( + "max_output_tokens", + "temperature", + "top_p", + "user", + "custom_inputs", + "context", + "databricks_options", +) +_SAMPLING_PARAMS: Final = ("temperature", "top_p", "user") +_OUTPUT_TOKEN_PARAMS: Final = ("max_completion_tokens", "max_tokens") +_STRING_METADATA: Final = TypeAdapter(Mapping[str, StrictStr]) +_METADATA: Final = TypeAdapter(Mapping[str, object]) +_PINNED_BODY_FIELDS: Final = ("model", "input") + + +def endpoint_name_from_model(model: str) -> str: + return model.removeprefix(_MODEL_PREFIX).strip() + + +def _normalized_base(api_base: str) -> str: + return api_base.strip().rstrip("/") + + +def _is_unified_responses_url(api_base: str) -> bool: + return _normalized_base(api_base).endswith(_UNIFIED_RESPONSES_SUFFIX) + + +def _is_databricks_app_host(api_base: str) -> bool: + return (urlsplit(_normalized_base(api_base)).hostname or "").endswith(_APP_HOST_SUFFIX) + + +def is_databricks_app_url(api_base: str) -> bool: + if _is_unified_responses_url(api_base): + return False + return _is_databricks_app_host(api_base) or _normalized_base(api_base).endswith(_APP_RESPONSES_SUFFIX) + + +def workspace_url_from_api_base(api_base: str | None) -> str | None: + if not api_base or not api_base.strip() or is_databricks_app_url(api_base): + return None + parts: Final = urlsplit(_normalized_base(api_base)) + if not parts.scheme or not parts.netloc: + return None + return f"{parts.scheme}://{parts.netloc}" + + +def _required_endpoint_name(model: str) -> str: + endpoint: Final = endpoint_name_from_model(model) + if not endpoint: + raise DatabricksException( + status_code=400, + message=( + "model must name the Databricks serving endpoint (model: databricks_agent/) " + "unless api_base is the endpoint's /invocations URL or a Databricks App /responses URL" + ), + ) + return endpoint + + +def _app_endpoint_url(app_base: str) -> str: + return f"{app_base}{_APP_RESPONSES_SUFFIX}" if urlsplit(app_base).path in ("", "/") else app_base + + +def resolve_endpoint_url(api_base: str, model: str) -> str: + base: Final = _normalized_base(api_base) + if _is_unified_responses_url(base): + _required_endpoint_name(model) + return base + if base.endswith(_ENDPOINT_URL_SUFFIXES): + return base + if _is_databricks_app_host(base): + return _app_endpoint_url(base) + if _SERVING_ENDPOINTS_SEGMENT in base: + return f"{base}{_INVOCATIONS_SUFFIX}" + workspace: Final = base.removesuffix("/serving-endpoints") + return f"{workspace}/serving-endpoints/{quote(_required_endpoint_name(model), safe='')}/invocations" + + +def _body_model_field(api_base: object, model: str) -> Mapping[str, str]: + if not isinstance(api_base, str) or not _is_unified_responses_url(api_base): + return {} + return {"model": _required_endpoint_name(model)} + + +class _InputContentPart(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str = "unknown" + + +class _InputMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + content: str | tuple[_InputContentPart, ...] | None = None + tool_calls: tuple[object, ...] | None = None + + +def _drop_params_enabled(litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or litellm_params.get("drop_params") is True + + +def _input_item(message: AllMessageValues, drop_unsupported: bool) -> Mapping[str, str] | None: + role: Final = message["role"] + if role not in _INPUT_ROLES: + if drop_unsupported: + return None + raise DatabricksException( + status_code=400, + message=( + f"message role '{role}' is not supported by a Databricks ResponsesAgent; " + "send system, developer, user, or assistant messages, or set drop_params: true to drop it" + ), + ) + unsupported: Final = _unsupported_message_parts(message) + if unsupported and not drop_unsupported: + raise DatabricksException( + status_code=400, + message=( + f"databricks_agent sends text content only; remove {', '.join(unsupported)} from the " + "message, or set drop_params: true to drop them" + ), + ) + content: Final = convert_content_list_to_str(message) + if unsupported and not content: + return None + return {"role": role, "content": content} + + +def _input_items(messages: Sequence[AllMessageValues], drop_unsupported: bool) -> Sequence[Mapping[str, str]]: + items: Final = (_input_item(message, drop_unsupported) for message in messages) + return [item for item in items if item is not None] + + +def _unsupported_message_parts(message: AllMessageValues) -> tuple[str, ...]: + parsed: Final = _InputMessage.model_validate(message) + content_parts: Final = parsed.content if isinstance(parsed.content, tuple) else () + content_types: Final[frozenset[str]] = frozenset(part.type for part in content_parts) - {"text"} + tool_calls: Final[frozenset[str]] = frozenset({"tool_calls"}) if parsed.tool_calls else frozenset() + return tuple(sorted(content_types | tool_calls)) + + +def _has_authorization(headers: Mapping[str, object]) -> bool: + return any(name.lower() == "authorization" for name in headers) + + +def _without_authorization(headers: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in headers.items() if name.lower() != "authorization"} + + +def _env_oauth_m2m_token(api_base: str | None) -> str | None: + client_id: Final = get_secret_str("DATABRICKS_CLIENT_ID") + client_secret: Final = get_secret_str("DATABRICKS_CLIENT_SECRET") + if not (client_id and client_secret): + return None + workspace: Final = get_secret_str("DATABRICKS_HOST") or workspace_url_from_api_base(api_base) + if workspace is None: + raise DatabricksException( + status_code=400, + message=( + "DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET are set but the OAuth token endpoint is unknown: " + "set DATABRICKS_HOST to the workspace URL (a Databricks App URL is not a workspace)" + ), + ) + return _cached_m2m_token(workspace, client_id, client_secret) + + +def _cached_m2m_token(workspace: str, client_id: str, client_secret: str) -> str: + secret_digest: Final = hashlib.sha256(client_secret.encode()).hexdigest()[:16] + cache_key: Final = f"{workspace}|{client_id}|{secret_digest}" + cached: Final[object] = _ENV_M2M_TOKENS.get_cache(cache_key) + if isinstance(cached, str): + return cached + token: Final = DatabricksBase().get_oauth_m2m_token_and_lifetime(workspace, client_id, client_secret) + ttl: Final = token.expires_in - _TOKEN_EXPIRY_BUFFER_SECONDS + if ttl > 0: + _ENV_M2M_TOKENS.set_cache(cache_key, token.access_token, ttl=ttl) + return token.access_token + + +def _resolved_env_reference(value: str | None) -> str | None: + if value is not None and value.startswith("os.environ/"): + return get_secret_str(value) + return value + + +def _max_output_tokens(non_default_params: Mapping[str, object]) -> Mapping[str, object]: + limits: Final = (non_default_params[key] for key in _OUTPUT_TOKEN_PARAMS if non_default_params.get(key) is not None) + limit: Final = next(limits, None) + return {} if limit is None else {"max_output_tokens": limit} + + +def _string_metadata_entries(metadata: object) -> Mapping[str, str]: + try: + entries: Final = _METADATA.validate_python(metadata) + except ValidationError: + return {} + return {key: value for key, value in entries.items() if isinstance(value, str)} + + +def _has_string_values_only(metadata: object) -> bool: + try: + _STRING_METADATA.validate_python(metadata) + except ValidationError: + return False + return True + + +def _without_pinned_fields(extra_body: Mapping[str, object], drop_params: bool) -> Mapping[str, object]: + pinned: Final = tuple(key for key in _PINNED_BODY_FIELDS if key in extra_body) + if not pinned: + return extra_body + if not drop_params: + raise DatabricksException( + status_code=400, + message=( + f"extra_body cannot set {', '.join(pinned)} on a Databricks ResponsesAgent request, " + "since the agent's endpoint and the messages set them; remove them, " + "or set drop_params: true to drop them" + ), + ) + return {key: value for key, value in extra_body.items() if key not in _PINNED_BODY_FIELDS} + + +def _with_string_metadata(extra_body: Mapping[str, object], drop_params: bool) -> Mapping[str, object]: + if "metadata" not in extra_body or _has_string_values_only(extra_body["metadata"]): + return extra_body + if not drop_params: + raise DatabricksException( + status_code=400, + message=( + "metadata sent to a Databricks ResponsesAgent must map names to string values; " + "send string values only, or set drop_params: true to drop the others" + ), + ) + kept: Final = _string_metadata_entries(extra_body["metadata"]) + rest: Final = {key: value for key, value in extra_body.items() if key != "metadata"} + return {**rest, **({"metadata": kept} if kept else {})} + + +class DatabricksAgentConfig(BaseConfig): + @staticmethod + def resolve_api_base_and_key(api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]: + return ( + _resolved_env_reference(api_base) or get_secret_str("DATABRICKS_API_BASE"), + _resolved_env_reference(api_key), + ) + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: mirrors override contract + return ["stream", *_OUTPUT_TOKEN_PARAMS, *_SAMPLING_PARAMS] + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + sampling: Final = {key: non_default_params[key] for key in _SAMPLING_PARAMS if key in non_default_params} + return {**optional_params, **sampling, **_max_output_tokens(non_default_params)} + + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + bearer: Final = api_key or _env_oauth_m2m_token(api_base) or get_secret_str("DATABRICKS_API_KEY") + if bearer: + return { + **_without_authorization(headers), + "Content-Type": "application/json", + "Authorization": f"Bearer {bearer}", + } + if _has_authorization(headers): + return {**headers, "Content-Type": "application/json"} + raise DatabricksException( + status_code=400, + message=( + "Missing Databricks credentials: set api_key to a personal access token, set the " + "DATABRICKS_API_KEY environment variable, set DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET " + "for OAuth M2M, or register the agent with OAuth client credentials" + ), + ) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + if not api_base: + raise DatabricksException( + status_code=400, + message=( + "api_base is required for databricks_agent: the workspace URL, the Databricks App URL, " + "or the full endpoint URL (set it or the DATABRICKS_API_BASE environment variable)" + ), + ) + return resolve_endpoint_url(api_base, model) + + def transform_request( + self, + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: mirrors override contract + body_params: Final = {key: optional_params[key] for key in _BODY_PARAMS if key in optional_params} + stream: Final = {"stream": True} if optional_params.get("stream") is True else {} + return { + **_body_model_field(litellm_params.get("api_base"), model), + "input": _input_items(messages, _drop_params_enabled(litellm_params)), + **stream, + **body_params, + } + + def transform_extra_body( + self, + extra_body: Mapping[str, object], + request: Mapping[str, object], + model: str, + litellm_params: Mapping[str, object], + ) -> Mapping[str, object]: + drop_params: Final = _drop_params_enabled(litellm_params) + return _with_string_metadata(_without_pinned_fields(extra_body, drop_params), drop_params) + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + logging_obj.post_call( + input=messages, + api_key="", + original_response=raw_response.text, + additional_args={"complete_input_dict": request_data}, + ) + try: + agent_response: Final = AgentResponse.model_validate_json(raw_response.content) + except ValidationError as e: + raise DatabricksException( + status_code=502, + message=f"Databricks agent returned a response that is not a ResponsesAgent response: {e}", + ) + + content: Final = output_text(agent_response.output) + message: Final = Message( + content=content, + role="assistant", + provider_specific_fields=( + {"custom_outputs": agent_response.custom_outputs} if agent_response.custom_outputs is not None else None + ), + ) + return ModelResponse( + id=model_response.id, + created=model_response.created, + model=model, + choices=[Choices(finish_reason="stop", index=0, message=message)], + usage=agent_response.usage.to_usage() + if agent_response.usage is not None + else self._estimated_usage(model, messages, content), + ) + + @staticmethod + def _estimated_usage(model: str, messages: Sequence[AllMessageValues], content: str) -> Usage: + from litellm.utils import token_counter + + prompt_tokens: Final = token_counter(model=model, messages=list(messages)) + completion_tokens: Final = token_counter(model=model, text=content, count_response_tokens=True) + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + def get_model_response_iterator( + self, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, + sync_stream: bool, + json_mode: bool | None = False, + ) -> DatabricksAgentResponsesIterator: + return DatabricksAgentResponsesIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: + return DatabricksException(status_code=status_code, message=error_message) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index f2f9df422e8..5e85c91a6b3 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -15,13 +15,24 @@ import re from typing import Final, Literal from urllib.parse import urlsplit, urlunsplit +from pydantic import BaseModel, ConfigDict + from litellm.llms.base_llm.chat.transformation import BaseLLMException +_DEFAULT_M2M_TOKEN_LIFETIME_SECONDS: Final = 3600 + class DatabricksException(BaseLLMException): pass +class OAuthM2MToken(BaseModel): + model_config = ConfigDict(frozen=True) + + access_token: str + expires_in: int = _DEFAULT_M2M_TOKEN_LIFETIME_SECONDS + + class DatabricksBase: """ Base class for Databricks integration with authentication, @@ -206,12 +217,20 @@ class DatabricksBase: return f"{normalized_api_base}/ai-gateway/mlflow/v1" return api_base - def _get_oauth_m2m_token( + def get_oauth_m2m_token( self, api_base: str, client_id: str, client_secret: str, ) -> str: + return self.get_oauth_m2m_token_and_lifetime(api_base, client_id, client_secret).access_token + + def get_oauth_m2m_token_and_lifetime( + self, + api_base: str, + client_id: str, + client_secret: str, + ) -> OAuthM2MToken: """ Obtain an OAuth M2M access token using client credentials flow. @@ -224,7 +243,7 @@ class DatabricksBase: client_secret: OAuth client secret Returns: - Access token string + The access token and its lifetime in seconds Raises: DatabricksException: If token request fails @@ -258,8 +277,7 @@ class DatabricksBase: message=f"OAuth M2M token request failed: {response.text}", ) - token_data: Final = response.json() - return token_data["access_token"] + return OAuthM2MToken.model_validate(response.json()) def _get_databricks_credentials( self, api_key: str | None, api_base: str | None, headers: dict | None @@ -346,7 +364,7 @@ class DatabricksBase: if client_id and client_secret and api_base: # Use OAuth M2M flow (preferred for production) verbose_logger.debug("Using OAuth M2M authentication for Databricks") - access_token: Final = self._get_oauth_m2m_token(api_base, client_id, client_secret) + access_token: Final = self.get_oauth_m2m_token(api_base, client_id, client_secret) headers = headers or {} headers["Authorization"] = f"Bearer {access_token}" headers["Content-Type"] = "application/json" diff --git a/litellm/main.py b/litellm/main.py index 4528d4914f1..f4c06565e85 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5143,6 +5143,33 @@ def _complete_langflow(ctx: CompletionDispatchContext) -> _CompletionDispatchRes ) +def _complete_databricks_agent(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: + from litellm.llms.databricks.agent.transformation import DatabricksAgentConfig + + api_base, api_key = DatabricksAgentConfig.resolve_api_base_and_key( + api_base=ctx.api_base or litellm.api_base, + api_key=ctx.api_key or litellm.api_key, + ) + return base_llm_http_handler.completion( + model=ctx.model, + stream=ctx.stream, + messages=ctx.messages, + acompletion=ctx.acompletion, + api_base=api_base, + model_response=ctx.model_response, + optional_params=ctx.optional_params, + litellm_params={**ctx.litellm_params, "api_base": api_base}, + shared_session=ctx.shared_session, + custom_llm_provider=ctx.custom_llm_provider, + timeout=ctx.timeout, + headers=ctx.headers or litellm.headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=ctx.logging, + client=_dispatch_client_http(ctx), + ) + + @tracer.wrap() @client def completion( @@ -6101,6 +6128,9 @@ def completion( # LangFlow - Visual AI Agent Platform response = _complete_langflow(_dispatch_ctx) + elif custom_llm_provider == "databricks_agent": + response = _complete_databricks_agent(_dispatch_ctx) # rebind-ok: dispatch chain binds response per branch + else: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) return response diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 304f282d1f8..5c8849e50ca 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -34,8 +34,9 @@ from litellm.proxy.a2a.version_convert import ( normalize_stream_event, ) from litellm.proxy.agent_endpoints.databricks_oauth import ( - DATABRICKS_OAUTH_PARAM, - resolve_databricks_app_auth_header, + AgentBackendAuth, + has_databricks_oauth, + resolve_databricks_backend_auth, ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -166,13 +167,17 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) -async def _resolve_backend_auth_header( +async def _resolve_backend_auth( litellm_params: dict[str, object], custom_llm_provider: object, -) -> Mapping[str, str] | None: - if litellm_params.get(DATABRICKS_OAUTH_PARAM): - return await resolve_databricks_app_auth_header(litellm_params) - return await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider) +) -> AgentBackendAuth: + databricks_auth: Final = await resolve_databricks_backend_auth(litellm_params) + if has_databricks_oauth(litellm_params): + return databricks_auth + return AgentBackendAuth( + litellm_params=databricks_auth.litellm_params, + header=await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider), + ) def _forwarding_headers( @@ -819,6 +824,7 @@ async def invoke_agent_a2a( if header_name: dynamic_headers[header_name] = val + backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider) agent_extra_headers: Final = _forwarding_headers( caller_identity=caller_identity, request_data=data, @@ -826,8 +832,9 @@ async def invoke_agent_a2a( dynamic_headers=dynamic_headers or None, static_headers=static_headers or None, ), - backend_auth_header=await _resolve_backend_auth_header(litellm_params, custom_llm_provider), + backend_auth_header=backend_auth.header, ) + forwarded_litellm_params: Final = backend_auth.litellm_params # Merge agent-level guardrails into data so post_call_success_hook and # _handle_stream_message both pick them up. A2A agents use model @@ -872,7 +879,7 @@ async def invoke_agent_a2a( model=f"a2a_agent/{agent_name}", request=a2a_request, api_base=agent_url, - litellm_params=litellm_params, + litellm_params=forwarded_litellm_params, agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), @@ -912,7 +919,7 @@ async def invoke_agent_a2a( api_base=agent_url, request_id=request_id if request_id is not None else "", params=params, - litellm_params=litellm_params, + litellm_params=forwarded_litellm_params, agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index c57315ebc21..9a6f674db2c 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -5,17 +5,115 @@ Handles routing for A2A agents (models with "a2a/" prefix). Looks up agents in the registry and injects their API base URL. """ +from collections.abc import Mapping from typing import Any, Final from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + agent_completion_kwargs, + bridge_model_name, +) +from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.databricks_oauth import ( + resolve_databricks_backend_auth, +) +from litellm.proxy.agent_endpoints.utils import merge_agent_headers +from litellm.types.agents import AgentResponse + +_HEADERS: Final = TypeAdapter(Mapping[str, str]) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_AGENT_PINNED_KWARGS: Final = ("api_base", "api_key") +_AGENT_LAYERED_KWARGS: Final = ("custom_inputs", "extra_body") +_CLIENT_TARGET_KWARGS: Final = ("api_base", "base_url") +_PROXY_HEADER_PREFIX: Final = "x-litellm-" + + +def _agent_extra_headers(agent_kwargs: Mapping[str, object]) -> Mapping[str, str] | None: + extra_headers: Final = agent_kwargs.get("extra_headers") + return None if extra_headers is None else _HEADERS.validate_python(extra_headers) + + +def _drop_params_enabled(data: Mapping[str, object], litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or data.get("drop_params") is True or litellm_params.get("drop_params") is True + + +def _client_extra_headers(data: Mapping[str, object], litellm_params: Mapping[str, object]) -> Mapping[str, str] | None: + extra_headers: Final = data.get("extra_headers") + if not extra_headers: + return None + try: + headers: Final = _HEADERS.validate_python(extra_headers) + except ValidationError as e: + if _drop_params_enabled(data, litellm_params): + return None + raise HTTPException( + status_code=400, + detail=( + "extra_headers must be an object mapping header names to string values, " + "or set drop_params: true to drop it" + ), + ) from e + return {name: value for name, value in headers.items() if not name.lower().startswith(_PROXY_HEADER_PREFIX)} + + +def _mappings(client_value: object, agent_value: object) -> tuple[Mapping[str, object], Mapping[str, object]] | None: + try: + return _OBJECT_MAPPING.validate_python(client_value), _OBJECT_MAPPING.validate_python(agent_value) + except ValidationError: + return None + + +def _agent_entries_over_client(client_value: object, agent_value: object) -> object: + mappings: Final = _mappings(client_value, agent_value) + return agent_value if mappings is None else {**mappings[0], **mappings[1]} + + +def _agent_keys_over_client(client_value: object, agent_value: object) -> object: + mappings: Final = _mappings(client_value, agent_value) + if mappings is None: + return agent_value + client, agent = mappings + return {**client, **{key: _agent_entries_over_client(client.get(key), value) for key, value in agent.items()}} + + +def _with_headers(data: Mapping[str, object], extra_headers: Mapping[str, str] | None) -> Mapping[str, object]: + request: Final = {key: value for key, value in data.items() if key != "extra_headers"} + return {**request, **({"extra_headers": extra_headers} if extra_headers else {})} + + +def _bridge_request_data( + data: Mapping[str, object], + litellm_params: Mapping[str, object], + static_headers: Mapping[str, str] | None, +) -> Mapping[str, object]: + agent_kwargs: Final = agent_completion_kwargs(litellm_params) + pinned: Final = {key: agent_kwargs[key] for key in _AGENT_PINNED_KWARGS if key in agent_kwargs} + agent_headers: Final = merge_agent_headers( + dynamic_headers=static_headers, + static_headers=_agent_extra_headers(agent_kwargs), + ) + extra_headers: Final = merge_agent_headers( + dynamic_headers=_client_extra_headers(data, litellm_params), + static_headers=agent_headers, + ) + layered: Final = { + key: _agent_keys_over_client(data[key], agent_kwargs[key]) + for key in _AGENT_LAYERED_KWARGS + if key in data and key in agent_kwargs + } + return _with_headers( + {**agent_kwargs, **data, **layered, **pinned, "model": bridge_model_name(litellm_params)}, + extra_headers, + ) async def route_a2a_agent_request( - data: dict, + data: dict[str, object], # mutable-ok: the URL agent path writes api_base into the caller's request route_type: str, user_api_key_dict: UserAPIKeyAuth | None = None, ) -> Any | None: @@ -68,6 +166,12 @@ async def route_a2a_agent_request( detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.", ) + litellm_params: Final = agent.litellm_params or {} + + if litellm_params.get("custom_llm_provider") in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: + bridge_auth: Final = await resolve_databricks_backend_auth(litellm_params) + return _call_route(route_type, _bridge_agent_request(data, agent, bridge_auth.litellm_params)) + # Get API base URL from agent config if not agent.agent_card_params or "url" not in agent.agent_card_params: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) @@ -78,4 +182,19 @@ async def route_a2a_agent_request( data["api_base"] = agent.agent_card_params["url"] verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"]) - return getattr(litellm, f"{route_type}")(**data) + return _call_route(route_type, data) + + +def _bridge_agent_request( + data: Mapping[str, object], + agent: AgentResponse, + litellm_params: Mapping[str, object], +) -> Mapping[str, object]: + card_url: Final[object] = (agent.agent_card_params or {}).get("url") + api_base: Final = {"api_base": card_url} if card_url else {} + request: Final = {key: value for key, value in data.items() if key not in _CLIENT_TARGET_KWARGS} + return _bridge_request_data({**request, **api_base}, litellm_params, agent.static_headers) + + +def _call_route(route_type: str, request_data: Mapping[str, object]) -> object: + return getattr(litellm, route_type)(**request_data) diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 38a76ea6890..7425c240d10 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -30,18 +30,24 @@ from dataclasses import dataclass from typing import Final import httpx +from pydantic import TypeAdapter, ValidationError +import litellm from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.databricks.agent.transformation import is_databricks_app_url, workspace_url_from_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider DATABRICKS_OAUTH_PARAM: Final = "databricks_oauth" +DATABRICKS_AGENT_PROVIDER: Final = "databricks_agent" +_FLAT_OAUTH_FIELDS: Final = ("client_id", "client_secret", "workspace_url", "scope") _DEFAULT_SCOPE: Final = "all-apis" _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60 _DEFAULT_TTL_SECONDS: Final = 3600 +_OAUTH_BLOCK: Final = TypeAdapter(Mapping[str, object]) def _resolve_secret(value: object) -> str | None: @@ -75,41 +81,118 @@ class DatabricksAppOAuthConfig: return f"{self.token_url}|{self.client_id}|{self.scope}|{secret_digest}" +def _workspace_origin(url: object) -> str | None: + return workspace_url_from_api_base(url) if isinstance(url, str) else None + + +@dataclass(frozen=True, slots=True) +class DatabricksOAuthConfigError: + message: str + + +def _config_error(message: str) -> litellm.BadRequestError: + return litellm.BadRequestError(message=message, model="a2a", llm_provider=DATABRICKS_AGENT_PROVIDER) + + +def _is_databricks_agent(litellm_params: Mapping[str, object]) -> bool: + return litellm_params.get("custom_llm_provider") == DATABRICKS_AGENT_PROVIDER + + +def _flat_databricks_agent_oauth(litellm_params: Mapping[str, object]) -> Mapping[str, object] | None: + """The Admin UI catalog cannot nest fields, so a ``databricks_agent`` spells its OAuth block as flat + ``client_id`` / ``client_secret`` / ``workspace_url`` / ``scope``. ``workspace_url`` falls back to the + ``api_base`` origin for Model Serving; a Databricks App host is never a token endpoint, so an App must set it.""" + if not _is_databricks_agent(litellm_params): + return None + if not litellm_params.get("client_id") and not litellm_params.get("client_secret"): + return None + return { + "client_id": litellm_params.get("client_id"), + "client_secret": litellm_params.get("client_secret"), + "workspace_url": litellm_params.get("workspace_url") + or _workspace_origin(_resolve_secret(litellm_params.get("api_base"))), + "scope": litellm_params.get("scope"), + } + + +def _raw_oauth_block(litellm_params: Mapping[str, object]) -> object | None: + nested: Final = litellm_params.get(DATABRICKS_OAUTH_PARAM) + if not _is_databricks_agent(litellm_params): + return nested or None + return nested if nested is not None else _flat_databricks_agent_oauth(litellm_params) + + +def has_databricks_oauth(litellm_params: Mapping[str, object]) -> bool: + return _raw_oauth_block(litellm_params) is not None + + +def without_databricks_oauth_params( + litellm_params: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: asend_message takes litellm_params as a dict + consumed: Final = ( + frozenset({DATABRICKS_OAUTH_PARAM, *_FLAT_OAUTH_FIELDS}) + if _is_databricks_agent(litellm_params) + else frozenset({DATABRICKS_OAUTH_PARAM}) + ) + return {key: value for key, value in litellm_params.items() if key not in consumed} + + def parse_databricks_oauth_config( litellm_params: Mapping[str, object] | None, ) -> DatabricksAppOAuthConfig | None: """Build a Databricks App OAuth config from an agent's ``litellm_params``. - Returns ``None`` when the agent has no ``databricks_oauth`` block. Raises - ``ValueError`` when the block is present but incomplete, so misconfiguration - surfaces loudly instead of silently sending an unauthenticated request. + Returns ``None`` when the agent has no ``databricks_oauth`` block. A block that is present but incomplete + raises instead of sending an unauthenticated request: a 400 ``litellm.BadRequestError`` for a + ``databricks_agent``, and the ``ValueError`` URL agents raised before ``databricks_agent`` existed. """ + config: Final = databricks_oauth_config(litellm_params) + if not isinstance(config, DatabricksOAuthConfigError): + return config + if litellm_params and _is_databricks_agent(litellm_params): + raise _config_error(config.message) + raise ValueError(config.message) + + +def databricks_oauth_config( + litellm_params: Mapping[str, object] | None, +) -> DatabricksAppOAuthConfig | DatabricksOAuthConfigError | None: if not litellm_params: return None - raw: Final = litellm_params.get(DATABRICKS_OAUTH_PARAM) + raw: Final = _raw_oauth_block(litellm_params) if raw is None: return None - if not isinstance(raw, dict): - raise ValueError(f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, got {type(raw).__name__}") - - client_id: Final = _resolve_secret(raw.get("client_id")) - client_secret: Final = _resolve_secret(raw.get("client_secret")) - workspace_url: Final = _resolve_secret(raw.get("workspace_url")) - - missing: Final = [ - name - for name, value in ( - ("client_id", client_id), - ("client_secret", client_secret), - ("workspace_url", workspace_url), + try: + block: Final = _OAUTH_BLOCK.validate_python(raw) + except ValidationError: + return DatabricksOAuthConfigError( + f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, got {type(raw).__name__}" ) - if not value - ] - if missing: - raise ValueError(f"Databricks App OAuth config is missing required field(s): {', '.join(missing)}") - scope: Final = _resolve_secret(raw.get("scope")) or _DEFAULT_SCOPE + client_id: Final = _resolve_secret(block.get("client_id")) + client_secret: Final = _resolve_secret(block.get("client_secret")) + workspace_url: Final = _resolve_secret(block.get("workspace_url")) + + if not (client_id and client_secret and workspace_url): + missing: Final = tuple( + name + for name, value in ( + ("client_id", client_id), + ("client_secret", client_secret), + ("workspace_url", workspace_url), + ) + if not value + ) + return DatabricksOAuthConfigError( + f"Databricks App OAuth config is missing required field(s): {', '.join(missing)}" + ) + if is_databricks_app_url(workspace_url): + return DatabricksOAuthConfigError( + "workspace_url must be the Databricks workspace URL, not a Databricks App URL" + ) + + scope: Final = _resolve_secret(block.get("scope")) or _DEFAULT_SCOPE return DatabricksAppOAuthConfig( client_id=client_id, @@ -215,6 +298,13 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): databricks_app_oauth_token_cache: Final = DatabricksAppOAuthTokenCache() +async def _resolve_token(litellm_params: Mapping[str, object] | None) -> str | None: + config: Final = parse_databricks_oauth_config(litellm_params) + if config is None: + return None + return await databricks_app_oauth_token_cache.async_get_token(config) + + async def resolve_databricks_app_auth_header( litellm_params: Mapping[str, object] | None, ) -> dict[str, str] | None: @@ -222,9 +312,29 @@ async def resolve_databricks_app_auth_header( Returns ``None`` when the agent is not configured for Databricks App OAuth. """ - config: Final = parse_databricks_oauth_config(litellm_params) - if config is None: - return None + token: Final = await _resolve_token(litellm_params) + return None if token is None else {"Authorization": f"Bearer {token}"} - token: Final = await databricks_app_oauth_token_cache.async_get_token(config) - return {"Authorization": f"Bearer {token}"} + +@dataclass(frozen=True, slots=True) +class AgentBackendAuth: + litellm_params: dict[str, object] # mutable-ok: asend_message takes litellm_params as a dict + header: Mapping[str, str] | None + + +async def resolve_databricks_backend_auth(litellm_params: Mapping[str, object]) -> AgentBackendAuth: + """A ``databricks_agent`` completes with the minted token as its ``api_key``, so the token outranks a personal + access token and a client ``Authorization`` header; any other agent sends it as the outbound header.""" + return await _backend_auth(litellm_params, parse_databricks_oauth_config(litellm_params)) + + +async def _backend_auth( + litellm_params: Mapping[str, object], config: DatabricksAppOAuthConfig | None +) -> AgentBackendAuth: + forwarded: Final = without_databricks_oauth_params(litellm_params) + token: Final = None if config is None else await databricks_app_oauth_token_cache.async_get_token(config) + if token is None: + return AgentBackendAuth(litellm_params=forwarded, header=None) + if _is_databricks_agent(litellm_params): + return AgentBackendAuth(litellm_params={**forwarded, "api_key": token}, header=None) + return AgentBackendAuth(litellm_params=forwarded, header={"Authorization": f"Bearer {token}"}) diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index cc2fc17d759..f289a6c1805 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -305,6 +305,77 @@ "litellm_params_template": { "custom_llm_provider": "watsonx_orchestrate" } + }, + { + "agent_type": "databricks_agent", + "agent_type_display_name": "Databricks Agent", + "description": "Connect to Databricks Agent Platform ResponsesAgents on Model Serving or Databricks Apps", + "logo_url": "/ui/assets/logos/databricks.svg", + "credential_fields": [ + { + "key": "api_base", + "label": "Databricks URL", + "placeholder": "https://.cloud.databricks.com", + "tooltip": "The workspace URL for a Model Serving endpoint, or a full URL: https:///serving-endpoints//invocations, https:///serving-endpoints/responses, or https://.databricksapps.com/responses", + "required": true, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "model", + "label": "Serving endpoint name", + "placeholder": "my-agent-endpoint", + "tooltip": "The Model Serving endpoint that hosts the agent. Leave empty when the Databricks URL is an endpoint /invocations URL or a Databricks App /responses URL", + "required": false, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "api_key", + "label": "Personal access token", + "placeholder": null, + "tooltip": "A Databricks PAT. Works for Model Serving only; Databricks Apps accept OAuth client credentials only", + "required": false, + "field_type": "password", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "client_id", + "label": "OAuth client ID", + "placeholder": null, + "tooltip": "Service principal application ID for OAuth M2M (client credentials). Required for Databricks Apps", + "required": false, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "client_secret", + "label": "OAuth client secret", + "placeholder": null, + "tooltip": "Service principal OAuth secret", + "required": false, + "field_type": "password", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "workspace_url", + "label": "Workspace URL (OAuth)", + "placeholder": "https://.cloud.databricks.com", + "tooltip": "The workspace that mints the OAuth token. Defaults to the Databricks URL origin; set it for a Databricks App, whose host is not the workspace", + "required": false, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + } + ], + "litellm_params_template": { + "custom_llm_provider": "databricks_agent" + } } ] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d83e874dee6..5f97c25d3d2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4305,6 +4305,7 @@ class LlmProviders(str, Enum): TRITON = "triton" PREDIBASE = "predibase" DATABRICKS = "databricks" + DATABRICKS_AGENT = "databricks_agent" EMPOWER = "empower" GITHUB = "github" RAGFLOW = "ragflow" diff --git a/litellm/utils.py b/litellm/utils.py index 22d52fe8d74..1bc3b55992c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8777,6 +8777,10 @@ class ProviderConfigManager: lambda: ProviderConfigManager._get_langflow_config(), False, ), + LlmProviders.DATABRICKS_AGENT: ( + lambda: ProviderConfigManager._get_databricks_agent_config(), + False, + ), LlmProviders.GDC: ( lambda: litellm.GDCGeminiConfig(), False, @@ -8879,6 +8883,12 @@ class ProviderConfigManager: return LangFlowConfig() + @staticmethod + def _get_databricks_agent_config() -> BaseConfig: + from litellm.llms.databricks.agent.transformation import DatabricksAgentConfig + + return DatabricksAgentConfig() + @staticmethod def get_provider_chat_config( model: str, diff --git a/tests/integration/providers/test_databricks_agent_wire.py b/tests/integration/providers/test_databricks_agent_wire.py new file mode 100644 index 00000000000..cd5fd168375 --- /dev/null +++ b/tests/integration/providers/test_databricks_agent_wire.py @@ -0,0 +1,984 @@ +import asyncio +import base64 +import json +import re +import threading +import uuid +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from urllib.parse import parse_qs + +import httpx +import openai +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_ENDPOINT: Final = "support-agent" +_INVOCATIONS: Final = f"/serving-endpoints/{_ENDPOINT}/invocations" +_UNIFIED: Final = "/serving-endpoints/responses" +_APP: Final = "/responses" +_TOKEN: Final = "/oidc/v1/token" +_PAT: Final = "synthetic-databricks-pat" +_CLIENT_ID: Final = "synthetic-databricks-sp" +_CLIENT_SECRET: Final = "synthetic-databricks-sp-secret" +_MINTED: Final = "synthetic-databricks-minted-token" +_REPLY: Final = "the export job is waiting on a lock" +_PROMPT: Final = "Why does the export job hang?" +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 11, "output_tokens": 3, "total_tokens": 14} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_BUDGET: Final = int(2 * graceful_stop_seconds() + 120) + + +def _basic(client_id: str, client_secret: str) -> str: + return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + + +def _agent_reply(text: str = _REPLY) -> bytes: + return json.dumps( + { + "output": [ + { + "type": "message", + "id": "msg-1", + "role": "assistant", + "content": [{"type": "output_text", "text": text}], + } + ], + "usage": _USAGE, + } + ).encode() + + +def _agent_stream(text: str = _REPLY) -> bytes: + head, tail = text[: len(text) // 2], text[len(text) // 2 :] + events: Final = ( + {"type": "response.output_text.delta", "item_id": "msg-1", "delta": head}, + {"type": "response.output_text.delta", "item_id": "msg-1", "delta": tail}, + { + "type": "response.output_item.done", + "item": {"type": "message", "id": "msg-1", "content": [{"type": "output_text", "text": text}]}, + }, + ) + return "".join(f"data: {json.dumps(event)}\n\n" for event in events).encode() + + +def _token_reply() -> Reply: + return Reply(body=json.dumps({"access_token": _MINTED, "token_type": "Bearer", "expires_in": 3600}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Databricks: + token_status: Callable[[], int] = lambda: 200 + endpoint_reply: Callable[[Request], Reply] | None = None + + def __call__(self, request: Request) -> Reply: + if request.target == _TOKEN: + status: Final = self.token_status() + if status != 200: + return Reply(status=status, body=b'{"error":"invalid_client"}') + return _token_reply() + if self.endpoint_reply is not None: + return self.endpoint_reply(request) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(body=_agent_stream(), content_type="text/event-stream") + return Reply(body=_agent_reply()) + + +def _calls_to(requests: Sequence[Request], target: str) -> tuple[Request, ...]: + return tuple(request for request in requests if request.target == target) + + +def _body(request: Request) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def _sdk_base_url(gateway: Gateway) -> str: + return f"{str(gateway.client.base_url).rstrip('/')}/v1" + + +def _register( + gateway: Gateway, + cleanups: Callable[[Callable[[], None]], object], + card_url: str, + static_headers: dict[str, JsonValue] | None = None, + **params: JsonValue, +) -> str: + name: Final = f"dbx-agent-{uuid.uuid4().hex[:12]}" + created: Final = gateway.request( + "POST", + "/v1/agents", + { + "agent_name": name, + "agent_card_params": { + "protocolVersion": "0.3", + "name": name, + "description": "Databricks Agent agent", + "url": card_url, + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {"streaming": True}, + "skills": [], + }, + "litellm_params": {"custom_llm_provider": "databricks_agent", **params}, + **({"static_headers": static_headers} if static_headers else {}), + }, + ) + assert created.status_code == 200, created.text + agent_id: Final = created.json()["agent_id"] + + def delete() -> None: + deleted: Final = gateway.request("DELETE", f"/v1/agents/{agent_id}") + assert deleted.status_code == 200, deleted.text + + cleanups(delete) + return name + + +@contextmanager +def _agent( + gateway: Gateway, + databricks: _Databricks, + *, + url_path: str = "", + static_headers: dict[str, JsonValue] | None = None, + **params: JsonValue, +) -> Iterator[tuple[Wire, str]]: + with wire_server(databricks) as wire, gateway.scenario() as scenario: + params_with_base: Final = {"api_base": wire.url + url_path, **params} + yield ( + wire, + _register(gateway, scenario.cleanups.callback, wire.url + url_path, static_headers, **params_with_base), + ) + + +def _chat(gateway: Gateway, name: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": f"a2a/{name}", "messages": [{"role": "user", "content": _PROMPT}], **extra}, + ) + + +def _a2a( + gateway: Gateway, name: str, method: str, marker: str, metadata: dict[str, JsonValue] | None = None +) -> httpx.Response: + message: Final[dict[str, JsonValue]] = { + "role": "user", + "messageId": marker, + "parts": [{"kind": "text", "text": _PROMPT}], + **({"metadata": metadata} if metadata is not None else {}), + } + return gateway.request( + "POST", + f"/a2a/{name}", + {"jsonrpc": "2.0", "id": marker, "method": method, "params": {"message": message}}, + ) + + +def _sse(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _spend_status(request_id: str) -> str: + rows: Final = eventually( + lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)), + lambda found: len(found) == 1, + seconds=70, + ) + return str(rows[0]["status"]) + + +def _assert_invocation(request: Request, bearer: str, *, stream: bool = False) -> None: + assert request.method == "POST" and request.target == _INVOCATIONS + assert request.headers["authorization"] == f"Bearer {bearer}" + assert _body(request) == { + "input": [{"role": "user", "content": _PROMPT}], + **({"stream": True} if stream else {}), + } + + +def _assert_token_request(request: Request) -> None: + assert request.method == "POST" + assert request.headers["authorization"] == _basic(_CLIENT_ID, _CLIENT_SECRET) + assert parse_qs(request.body.decode()) == {"grant_type": ["client_credentials"], "scope": ["all-apis"]} + + +def test_chat_on_a_pat_agent_reaches_the_serving_endpoint_and_bills_the_agent_usage(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name) + assert response.status_code == 200, response.text + payload: Final = response.json() + assert payload["choices"][0]["message"] == {"content": _REPLY, "role": "assistant"} + assert payload["choices"][0]["finish_reason"] == "stop" + assert payload["usage"]["prompt_tokens"] == 11 and payload["usage"]["completion_tokens"] == 3 + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + assert _spend_status(payload["id"]) == "success" + + +def test_async_chat_stream_on_a_pat_agent_sends_the_reply_once(gateway: Gateway) -> None: + async def stream(name: str) -> str: + client: Final = openai.AsyncOpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0) + chunks: Final = await client.chat.completions.create( + model=f"a2a/{name}", messages=[{"role": "user", "content": _PROMPT}], stream=True + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in chunks if chunk.choices]) + + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + assert asyncio.run(stream(name)) == _REPLY + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT, stream=True) + + +def test_sync_sdk_chat_on_a_pat_agent_returns_the_agent_reply(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + client: Final = openai.OpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0) + completion: Final = client.chat.completions.create( + model=f"a2a/{name}", messages=[{"role": "user", "content": _PROMPT}] + ) + assert completion.choices[0].message.content == _REPLY + assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 1 + + +def test_responses_on_a_pat_agent_returns_the_reply_in_output(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = gateway.request("POST", "/v1/responses", {"model": f"a2a/{name}", "input": _PROMPT}) + assert response.status_code == 200, response.text + texts: Final = tuple( + part["text"] + for item in response.json()["output"] + if item["type"] == "message" + for part in item["content"] + if part["type"] == "output_text" + ) + assert texts == (_REPLY,) + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + + +@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop-params")) +def test_messages_on_a_pat_agent_sends_max_tokens_as_max_output_tokens(gateway: Gateway, drop_params: bool) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": f"a2a/{name}", + "max_tokens": 64, + "messages": [{"role": "user", "content": _PROMPT}], + **({"drop_params": True} if drop_params else {}), + }, + ) + assert response.status_code == 200, response.text + assert response.json()["content"] == [{"type": "text", "text": _REPLY}] + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + assert _body(invocations[0]) == {"input": [{"role": "user", "content": _PROMPT}], "max_output_tokens": 64} + + +def test_chat_sampling_params_reach_the_agent_in_its_own_shape(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name, max_completion_tokens=32, temperature=0.2, top_p=0.9, user="u-1") + assert response.status_code == 200, response.text + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + assert _body(invocations[0]) == { + "input": [{"role": "user", "content": _PROMPT}], + "max_output_tokens": 32, + "temperature": 0.2, + "top_p": 0.9, + "user": "u-1", + } + + +def test_a_reply_with_several_message_items_reads_the_same_streamed_or_not(gateway: Gateway) -> None: + first, second = "Let me look at the export queue.", "The export job is waiting on a lock." + + def message(item_id: str, text: str) -> dict[str, JsonValue]: + return {"type": "message", "id": item_id, "content": [{"type": "output_text", "text": text}]} + + def reply(request: Request) -> Reply: + if _body(request).get("stream") is not True: + output: Final = (message("msg-1", first), {"type": "function_call", "id": "fc-1"}, message("msg-2", second)) + return Reply(body=json.dumps({"output": output, "usage": _USAGE}).encode()) + events: Final = ( + {"type": "response.output_text.delta", "item_id": "msg-1", "delta": first}, + {"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc-1"}}, + {"type": "response.output_text.delta", "item_id": "msg-2", "delta": second}, + {"type": "response.output_item.done", "item": message("msg-2", second)}, + ) + return Reply( + body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + content_type="text/event-stream", + ) + + with _agent(gateway, _Databricks(endpoint_reply=reply), model=_ENDPOINT, api_key=_PAT) as (wire, name): + client: Final = openai.OpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0) + messages: Final = [{"role": "user", "content": _PROMPT}] + whole: Final = client.chat.completions.create(model=f"a2a/{name}", messages=messages) + streamed: Final = client.chat.completions.create(model=f"a2a/{name}", messages=messages, stream=True) + expected: Final = f"{first}\n\n{second}" + assert whole.choices[0].message.content == expected + assert "".join(chunk.choices[0].delta.content or "" for chunk in streamed if chunk.choices) == expected + assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 2 + + +def test_identical_chat_requests_on_an_agent_each_reach_the_agent(gateway: Gateway) -> None: + replies: Final = iter(("first answer", "second answer")) + databricks: Final = _Databricks(endpoint_reply=lambda _: Reply(body=_agent_reply(next(replies)))) + with _agent(gateway, databricks, model=_ENDPOINT, api_key=_PAT) as (wire, name): + answers: Final = tuple(_chat(gateway, name) for _ in range(2)) + assert [answer.json()["choices"][0]["message"]["content"] for answer in answers] == [ + "first answer", + "second answer", + ], [answer.text for answer in answers] + assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 2 + + +def test_static_headers_sit_between_client_and_agent_extra_headers_on_chat(gateway: Gateway) -> None: + with _agent( + gateway, + _Databricks(), + static_headers={"X-Static": "static", "X-Shared": "from-static", "X-Both": "from-static"}, + model=_ENDPOINT, + api_key=_PAT, + extra_headers={"x-both": "from-agent"}, + ) as (wire, name): + response: Final = _chat(gateway, name, extra_headers={"x-shared": "from-client", "X-Trace": "trace-1"}) + assert response.status_code == 200, response.text + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + headers: Final = invocations[0].headers + assert (headers["x-static"], headers["x-shared"], headers["x-both"]) == ("static", "from-static", "from-agent") + assert headers["x-trace"] == "trace-1" + _assert_invocation(invocations[0], _PAT) + + +def test_chat_never_forwards_client_litellm_identity_headers_to_the_agent(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat( + gateway, + name, + extra_headers={"X-LiteLLM-User-Id": "spoofed-user", "x-litellm-team-id": "spoofed-team", "X-Trace": "t-1"}, + ) + assert response.status_code == 200, response.text + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + headers: Final = invocations[0].headers + assert "x-litellm-user-id" not in headers and "x-litellm-team-id" not in headers, headers + assert headers["x-trace"] == "t-1" + _assert_invocation(invocations[0], _PAT) + + +def test_the_agent_custom_inputs_and_extra_body_keys_win_over_the_callers_on_chat(gateway: Gateway) -> None: + with _agent( + gateway, + _Databricks(), + model=_ENDPOINT, + api_key=_PAT, + custom_inputs={"tenant": "owner"}, + extra_body={"metadata": {"tenant": "owner"}}, + ) as (wire, name): + response: Final = _chat( + gateway, + name, + custom_inputs={"tenant": "caller", "session": "s-1"}, + extra_body={"metadata": {"tenant": "caller", "turn": "2"}}, + ) + assert response.status_code == 200, response.text + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + assert _body(invocations[0]) == { + "input": [{"role": "user", "content": _PROMPT}], + "custom_inputs": {"tenant": "owner", "session": "s-1"}, + "metadata": {"tenant": "owner", "turn": "2"}, + } + + +@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop-params")) +def test_a2a_message_metadata_reaches_the_agent_only_as_string_values(gateway: Gateway, drop_params: bool) -> None: + params: Final[dict[str, JsonValue]] = {"drop_params": True} if drop_params else {} + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, **params) as (wire, name): + strings: Final = _a2a(gateway, name, "message/send", f"meta-{uuid.uuid4().hex}", {"session": "s-1"}) + assert strings.status_code == 200 and "error" not in strings.json(), strings.text + assert _body(_calls_to(wire.drain(), _INVOCATIONS)[0])["metadata"] == {"session": "s-1"} + mixed: Final = _a2a(gateway, name, "message/send", f"meta-{uuid.uuid4().hex}", {"session": "s-1", "turn": 1}) + sent: Final = _calls_to(wire.drain(), _INVOCATIONS) + if drop_params: + assert mixed.status_code == 200 and "error" not in mixed.json(), mixed.text + assert [_body(request)["metadata"] for request in sent] == [{"session": "s-1"}] + else: + assert "string values" in mixed.text, mixed.text + assert sent == () + + +def test_a2a_message_send_on_a_pat_agent_answers_with_the_agent_text(gateway: Gateway) -> None: + marker: Final = f"send-{uuid.uuid4().hex}" + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _a2a(gateway, name, "message/send", marker) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["id"] == marker and "error" not in body, body + assert body["result"]["parts"] == [{"kind": "text", "text": _REPLY}] + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + + +def test_a2a_message_stream_on_a_pat_agent_streams_the_text_and_completes(gateway: Gateway) -> None: + marker: Final = f"stream-{uuid.uuid4().hex}" + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _a2a(gateway, name, "message/stream", marker) + assert response.status_code == 200, response.text + events: Final = tuple(event["result"] for event in _sse(response.text)) + artifacts: Final = tuple( + part["text"] + for event in events + if event["kind"] == "artifact-update" + for part in event["artifact"]["parts"] + ) + assert artifacts == (_REPLY,) + assert events[-1]["kind"] == "status-update" and events[-1]["final"] is True + assert events[-1]["status"]["state"] == "completed" + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT, stream=True) + + +def test_the_unified_responses_url_names_the_endpoint_in_the_body(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), url_path=_UNIFIED, model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == _REPLY + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", _UNIFIED)] + assert received[0].headers["authorization"] == f"Bearer {_PAT}" + assert _body(received[0]) == {"model": _ENDPOINT, "input": [{"role": "user", "content": _PROMPT}]} + + +def test_an_app_agent_mints_its_token_at_the_workspace_and_calls_the_app(gateway: Gateway) -> None: + with ( + wire_server(_Databricks()) as workspace, + _agent( + gateway, + _Databricks(), + url_path=_APP, + client_id=_CLIENT_ID, + client_secret=_CLIENT_SECRET, + workspace_url=workspace.url, + ) as (app, name), + ): + response: Final = _chat(gateway, name) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == _REPLY + minted: Final = workspace.drain() + assert [request.target for request in minted] == [_TOKEN] + _assert_token_request(minted[0]) + called: Final = app.drain() + assert [(request.method, request.target) for request in called] == [("POST", _APP)] + assert called[0].headers["authorization"] == f"Bearer {_MINTED}" + assert _body(called[0]) == {"input": [{"role": "user", "content": _PROMPT}]} + + +def test_a_nested_oauth_block_mints_once_for_chat_and_a2a_calls_on_one_agent(gateway: Gateway) -> None: + def oauth(url: str) -> dict[str, JsonValue]: + return {"client_id": _CLIENT_ID, "client_secret": _CLIENT_SECRET, "workspace_url": url} + + with wire_server(_Databricks()) as wire, gateway.scenario() as scenario: + name: Final = _register( + gateway, + scenario.cleanups.callback, + wire.url, + model=_ENDPOINT, + api_base=wire.url, + databricks_oauth=oauth(wire.url), + ) + for _ in range(3): + response = _chat(gateway, name) + assert response.status_code == 200, response.text + sent: Final = _a2a(gateway, name, "message/send", f"send-{uuid.uuid4().hex}") + assert sent.status_code == 200 and "error" not in sent.json(), sent.text + received: Final = wire.drain() + tokens: Final = _calls_to(received, _TOKEN) + assert len(tokens) <= 2, [request.target for request in received] + _assert_token_request(tokens[0]) + invocations: Final = _calls_to(received, _INVOCATIONS) + assert len(invocations) == 4 + for invocation in invocations: + _assert_invocation(invocation, _MINTED) + + +def test_the_minted_token_outranks_the_agent_pat_and_a_client_authorization(gateway: Gateway) -> None: + with _agent( + gateway, + _Databricks(), + model=_ENDPOINT, + api_key=_PAT, + client_id=_CLIENT_ID, + client_secret=_CLIENT_SECRET, + ) as (wire, name): + response: Final = _chat(gateway, name, extra_headers={"Authorization": "Bearer client-supplied"}) + assert response.status_code == 200, response.text + received: Final = wire.drain() + assert [request.target for request in received] == [_TOKEN, _INVOCATIONS] + _assert_invocation(received[1], _MINTED) + + +def test_the_agent_pat_outranks_a_client_authorization(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name, extra_headers={"Authorization": "Bearer client-supplied"}) + assert response.status_code == 200, response.text + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + + +def test_empty_oauth_fields_beside_a_pat_keep_the_pat_path(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, client_id="", client_secret="") as (wire, name): + response: Final = _chat(gateway, name) + assert response.status_code == 200, response.text + received: Final = wire.drain() + assert [request.target for request in received] == [_INVOCATIONS] + _assert_invocation(received[0], _PAT) + + +def test_custom_inputs_reach_the_agent_beside_the_input(gateway: Gateway) -> None: + custom_inputs: Final[dict[str, JsonValue]] = {"session_id": "s-1", "tier": "gold"} + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name, custom_inputs=custom_inputs) + assert response.status_code == 200, response.text + received: Final = wire.drain() + assert len(received) == 1, received + assert _body(received[0]) == {"input": [{"role": "user", "content": _PROMPT}], "custom_inputs": custom_inputs} + + +def test_two_agents_on_one_endpoint_name_never_answer_from_each_other(gateway: Gateway) -> None: + def answering(text: str) -> _Databricks: + return _Databricks(endpoint_reply=lambda _: Reply(body=_agent_reply(text))) + + with ( + _agent(gateway, answering("from workspace one"), model=_ENDPOINT, api_key=_PAT) as (first_wire, first), + _agent(gateway, answering("from workspace two"), model=_ENDPOINT, api_key=_PAT) as (second_wire, second), + ): + responses: Final = tuple(_chat(gateway, name) for name in (first, second, first)) + assert all(response.status_code == 200 for response in responses), [r.text for r in responses] + answers: Final = tuple(response.json()["choices"][0]["message"]["content"] for response in responses) + assert answers == ("from workspace one", "from workspace two", "from workspace one") + assert len(_calls_to(first_wire.drain(), _INVOCATIONS)) == 2 + assert len(_calls_to(second_wire.drain(), _INVOCATIONS)) == 1 + + +def test_a_workspace_agent_with_no_endpoint_is_refused_before_any_call(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), api_key=_PAT) as (wire, name): + response: Final = _chat(gateway, name) + assert response.status_code == 400, response.text + assert "model must name the Databricks serving endpoint" in response.text + assert wire.drain() == () + + +def test_an_app_agent_without_a_workspace_url_is_refused_before_its_secret_leaves(gateway: Gateway) -> None: + with _agent(gateway, _Databricks(), url_path=_APP, client_id=_CLIENT_ID, client_secret=_CLIENT_SECRET) as ( + wire, + name, + ): + response: Final = _chat(gateway, name) + assert response.status_code == 400, response.text + assert "workspace_url" in response.text + assert wire.drain() == () + + +@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params")) +def test_malformed_client_extra_headers_follow_drop_params(gateway: Gateway, drop_params: bool, status: int) -> None: + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat( + gateway, name, extra_headers=["Authorization", "x"], **({"drop_params": True} if drop_params else {}) + ) + assert response.status_code == status, response.text + if drop_params: + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + else: + assert "extra_headers must be an object" in response.text + assert wire.drain() == () + + +@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params")) +def test_an_image_part_follows_drop_params(gateway: Gateway, drop_params: bool, status: int) -> None: + content: Final[list[JsonValue]] = [ + {"type": "text", "text": _PROMPT}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ] + with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": f"a2a/{name}", + "messages": [{"role": "user", "content": content}], + **({"drop_params": True} if drop_params else {}), + }, + ) + assert response.status_code == status, response.text + if drop_params: + received: Final = wire.drain() + assert len(received) == 1, received + _assert_invocation(received[0], _PAT) + else: + assert "image_url" in response.text + assert wire.drain() == () + + +@pytest.mark.parametrize( + ("reply", "status"), + ( + (Reply(status=401, body=b'{"error_code":"UNAUTHENTICATED","message":"Invalid token"}'), 401), + (Reply(status=500, body=b'{"error_code":"INTERNAL_ERROR","message":"agent crashed"}'), 500), + (Reply(body=b"{}"), 502), + ), + ids=("upstream-401", "upstream-500", "not-a-responses-agent"), +) +def test_a_failing_agent_keeps_its_status_and_unrelated_models_keep_serving( + gateway: Gateway, reply: Reply, status: int +) -> None: + with ( + _agent(gateway, _Databricks(endpoint_reply=lambda _: reply), model=_ENDPOINT, api_key=_PAT) as (wire, name), + gateway.scenario() as scenario, + ): + response: Final = _chat(gateway, name) + assert response.status_code == status, response.text + assert len(_calls_to(wire.drain(), _INVOCATIONS)) >= 1 + unrelated: Final = gateway.chat(scenario.model()) + assert unrelated["choices"], unrelated + + +def test_a_refused_token_request_answers_an_error_without_calling_the_agent(gateway: Gateway) -> None: + with _agent( + gateway, + _Databricks(token_status=lambda: 400), + model=_ENDPOINT, + client_id=_CLIENT_ID, + client_secret=_CLIENT_SECRET, + ) as (wire, name): + response: Final = _chat(gateway, name) + assert response.status_code >= 400, response.text + received: Final = wire.drain() + assert [request.target for request in received] == [_TOKEN] + + +def test_reading_an_agent_back_redacts_its_pat_and_client_secret(gateway: Gateway) -> None: + with _agent( + gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, client_id=_CLIENT_ID, client_secret=_CLIENT_SECRET + ) as (_, name): + listed: Final = gateway.request("GET", "/v1/agents") + assert listed.status_code == 200, listed.text + text: Final = listed.text + assert name in text + assert _PAT not in text and _CLIENT_SECRET not in text + + +@pytest.mark.parametrize( + "oauth_fields", + (None, (), ("client_id",), ("client_id", "client_secret", "workspace_url")), + ids=( + "plain", + "with-empty-databricks-oauth", + "with-partial-databricks-oauth", + "with-complete-databricks-oauth", + ), +) +def test_a_url_agent_keeps_its_a2a_call_to_the_card_url(gateway: Gateway, oauth_fields: tuple[str, ...] | None) -> None: + marker: Final = f"url-agent-{uuid.uuid4().hex}" + + def peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(status=404) + outbound: Final = _body(request) + return Reply( + body=json.dumps( + { + "jsonrpc": "2.0", + "id": outbound["id"], + "result": { + "kind": "message", + "role": "agent", + "messageId": marker, + "parts": [{"kind": "text", "text": _REPLY}], + }, + } + ).encode() + ) + + with wire_server(peer) as wire, gateway.scenario() as scenario: + complete_block: Final = {"client_id": _CLIENT_ID, "client_secret": _CLIENT_SECRET, "workspace_url": wire.url} + oauth_block: Final = None if oauth_fields is None else {field: complete_block[field] for field in oauth_fields} + created: Final = gateway.request( + "POST", + "/v1/agents", + { + "agent_name": marker, + "agent_card_params": { + "protocolVersion": "0.3", + "name": marker, + "description": "URL agent", + "url": wire.url + "/", + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {"streaming": False}, + "skills": [], + }, + **({"litellm_params": {"databricks_oauth": oauth_block}} if oauth_block is not None else {}), + }, + ) + assert created.status_code == 200, created.text + agent_id: Final = created.json()["agent_id"] + scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}")) + response: Final = _chat(gateway, marker, extra_headers={"Authorization": "Bearer client-token"}) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == _REPLY + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert [request.target for request in posts] == ["/"] + assert _body(posts[0])["method"] == "message/send" + assert posts[0].headers["authorization"] == "Bearer client-token" + + +def test_a_url_agent_with_an_incomplete_databricks_oauth_answers_a2a_with_an_internal_error(gateway: Gateway) -> None: + marker: Final = f"url-agent-{uuid.uuid4().hex}" + with wire_server(lambda _: Reply(status=404)) as wire, gateway.scenario() as scenario: + created: Final = gateway.request( + "POST", + "/v1/agents", + { + "agent_name": marker, + "agent_card_params": { + "protocolVersion": "0.3", + "name": marker, + "description": "URL agent", + "url": wire.url + "/", + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {"streaming": False}, + "skills": [], + }, + "litellm_params": {"databricks_oauth": {"client_id": _CLIENT_ID}}, + }, + ) + assert created.status_code == 200, created.text + agent_id: Final = created.json()["agent_id"] + scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}")) + response: Final = _a2a(gateway, marker, "message/send", marker) + assert response.status_code == 500, response.text + assert response.json()["error"]["code"] == -32603 + assert "missing required field" in response.json()["error"]["message"] + assert tuple(request for request in wire.drain() if request.method == "POST") == () + + +@pytest.mark.parametrize( + "pinned", + ({"model": "another-endpoint"}, {"input": [{"role": "user", "content": "injected"}]}), + ids=("model", "input"), +) +@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params")) +def test_extra_body_never_replaces_the_agent_endpoint_or_input( + gateway: Gateway, pinned: dict[str, JsonValue], drop_params: bool, status: int +) -> None: + custom_inputs: Final[dict[str, JsonValue]] = {"tenant": "t-1"} + with _agent(gateway, _Databricks(), url_path=_UNIFIED, model=_ENDPOINT, api_key=_PAT) as (wire, name): + response: Final = _chat( + gateway, + name, + extra_body={**pinned, "custom_inputs": custom_inputs}, + **({"drop_params": True} if drop_params else {}), + ) + assert response.status_code == status, response.text + received: Final = wire.drain() + if drop_params: + assert [(request.method, request.target) for request in received] == [("POST", _UNIFIED)] + assert _body(received[0]) == { + "model": _ENDPOINT, + "input": [{"role": "user", "content": _PROMPT}], + "custom_inputs": custom_inputs, + } + else: + assert f"cannot set {next(iter(pinned))}" in response.text + assert received == () + + +def test_a_token_outage_mid_burst_fails_its_calls_and_one_mint_serves_the_recovery(gateway: Gateway) -> None: + outage: Final = threading.Event() + outage.set() + + def token_status() -> int: + return 503 if outage.is_set() else 200 + + async def burst(name: str, size: int) -> tuple[int, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + responses: Final = await asyncio.gather( + *( + client.post( + "/v1/chat/completions", + json={"model": f"a2a/{name}", "messages": [{"role": "user", "content": _PROMPT}]}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + for _ in range(size) + ) + ) + return tuple(response.status_code for response in responses) + + with _agent( + gateway, + _Databricks(token_status=token_status), + model=_ENDPOINT, + client_id=_CLIENT_ID, + client_secret=_CLIENT_SECRET, + ) as (wire, name): + failed: Final = asyncio.run(burst(name, 12)) + assert all(status >= 400 for status in failed), failed + assert _calls_to(wire.drain(), _INVOCATIONS) == () + outage.clear() + recovered: Final = asyncio.run(burst(name, 12)) + assert recovered == (200,) * 12, recovered + received: Final = wire.drain() + assert len(_calls_to(received, _TOKEN)) <= 2 + for invocation in _calls_to(received, _INVOCATIONS): + _assert_invocation(invocation, _MINTED) + assert len(_calls_to(received, _INVOCATIONS)) == 12 + + +def _owned_config( + directory: Path, wire: Wire, api_key: str | None = _PAT, *, allow_client_side_credentials: bool = False +) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = _JSON_OBJECT.validate_python(base.get("general_settings") or {}) + name: Final = "dbx-owned-agent" + config: Final[dict[str, JsonValue]] = { + **base, + "general_settings": { + **general_settings, + **({"allow_client_side_credentials": True} if allow_client_side_credentials else {}), + }, + "agents": [ + { + "agent_name": name, + "agent_card_params": { + "protocolVersion": "0.3", + "name": name, + "description": "Databricks Agent agent", + "url": wire.url, + "version": "1.0.0", + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "capabilities": {"streaming": True}, + "skills": [], + }, + "litellm_params": { + "custom_llm_provider": "databricks_agent", + "model": _ENDPOINT, + "api_base": wire.url, + **({"api_key": api_key} if api_key is not None else {}), + }, + } + ], + } + path: Final = directory / "databricks-agent-owned.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.timeout(_OWNED_PROXY_BUDGET) +def test_a_client_api_base_or_base_url_never_takes_the_agent_pat_off_its_card_url( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_Databricks()) as wire, wire_server(lambda _: Reply(body=_agent_reply())) as collector: + config: Final = _owned_config(tmp_path, wire, allow_client_side_credentials=True) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + rejected: Final = _chat(owned.gateway, "dbx-owned-agent", api_base=collector.url) + served: Final = _chat(owned.gateway, "dbx-owned-agent", base_url=collector.url) + assert rejected.status_code == 400, rejected.text + assert served.status_code == 200, served.text + assert collector.drain() == () + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1 + _assert_invocation(invocations[0], _PAT) + + +@pytest.mark.timeout(_OWNED_PROXY_BUDGET) +def test_a_killed_worker_leaves_the_sibling_serving_the_databricks_agent(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_Databricks()) as wire: + with owned_proxy_process(gateway, tmp_path, {}, config=_owned_config(tmp_path, wire), workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds() + 60, + ) + first: Final = _chat(candidate, "dbx-owned-agent") + assert first.status_code == 200, first.text + replacement: Final = owned.kill_worker(workers[0]) + assert replacement not in workers + served: Final = tuple(_chat(candidate, "dbx-owned-agent") for _ in range(6)) + assert all(response.status_code == 200 for response in served), [r.text for r in served] + assert {response.json()["choices"][0]["message"]["content"] for response in served} == {_REPLY} + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 7 + for invocation in invocations: + _assert_invocation(invocation, _PAT) + + +@pytest.mark.timeout(_OWNED_PROXY_BUDGET) +def test_the_env_service_principal_mints_once_over_an_env_pat_and_env_references_resolve_on_an_api_created_agent( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_Databricks()) as wire: + environment: Final = { + "DATABRICKS_CLIENT_ID": _CLIENT_ID, + "DATABRICKS_CLIENT_SECRET": _CLIENT_SECRET, + "DATABRICKS_HOST": wire.url, + "DATABRICKS_API_KEY": "env-pat", + "DBX_AGENT_HOST": wire.url, + "DBX_AGENT_PAT": _PAT, + } + with owned_proxy_process(gateway, tmp_path, environment, config=_owned_config(tmp_path, wire, None)) as owned: + candidate: Final = owned.gateway + served: Final = tuple(_chat(candidate, "dbx-owned-agent") for _ in range(3)) + assert all(response.status_code == 200 for response in served), [r.text for r in served] + received: Final = wire.drain() + assert len(_calls_to(received, _TOKEN)) == 1, received + _assert_token_request(_calls_to(received, _TOKEN)[0]) + for invocation in _calls_to(received, _INVOCATIONS): + _assert_invocation(invocation, _MINTED) + assert len(_calls_to(received, _INVOCATIONS)) == 3 + + with candidate.scenario() as scenario: + name: Final = _register( + candidate, + scenario.cleanups.callback, + wire.url, + model=_ENDPOINT, + api_base="os.environ/DBX_AGENT_HOST", + api_key="os.environ/DBX_AGENT_PAT", + ) + response: Final = _chat(candidate, name) + assert response.status_code == 200, response.text + invocations: Final = _calls_to(wire.drain(), _INVOCATIONS) + assert len(invocations) == 1, invocations + _assert_invocation(invocations[0], _PAT) diff --git a/tests/unit/a2a_protocol/litellm_completion_bridge/__init__.py b/tests/unit/a2a_protocol/litellm_completion_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/a2a_protocol/litellm_completion_bridge/test_handler.py b/tests/unit/a2a_protocol/litellm_completion_bridge/test_handler.py new file mode 100644 index 00000000000..515105d6ddf --- /dev/null +++ b/tests/unit/a2a_protocol/litellm_completion_bridge/test_handler.py @@ -0,0 +1,151 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + A2ACompletionBridgeHandler, + agent_completion_kwargs, + bridge_model_name, +) + +WORKSPACE: Final = "https://adb-1.azuredatabricks.net" +INVOCATIONS_URL: Final = f"{WORKSPACE}/serving-endpoints/my-agent/invocations" +APP_URL: Final = "https://my-app-1.azure.databricksapps.com/responses" +MESSAGE_ITEM: Final = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong"}], +} + + +@pytest.fixture +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +def _a2a_params() -> dict[str, object]: + return {"message": {"role": "user", "parts": [{"kind": "text", "text": "ping"}], "messageId": "m-1"}} + + +@pytest.mark.parametrize( + ("litellm_params", "expected"), + [ + ({"custom_llm_provider": "databricks_agent", "model": "my-agent"}, "databricks_agent/my-agent"), + ( + {"custom_llm_provider": "databricks_agent", "model": "databricks_agent/my-agent"}, + "databricks_agent/my-agent", + ), + ({"custom_llm_provider": "databricks_agent"}, "databricks_agent/"), + ({"custom_llm_provider": "databricks_agent", "model": None}, "databricks_agent/"), + ({"custom_llm_provider": "databricks_agent", "model": ""}, "databricks_agent/"), + ({"custom_llm_provider": "langgraph"}, "langgraph/agent"), + ({"custom_llm_provider": "langgraph", "model": None}, "langgraph/agent"), + ({"model": "gpt-5"}, "gpt-5"), + ], +) +def test_bridge_model_name(litellm_params: dict[str, object], expected: str) -> None: + assert bridge_model_name(litellm_params) == expected + + +def test_agent_completion_kwargs_drops_model_provider_and_agent_only_keys() -> None: + kwargs = agent_completion_kwargs( + { + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + "is_public": True, + "agent_name": "dbx", + "agent_id": "id-1", + "agent_card_params": {"url": WORKSPACE}, + A2A_USER_API_KEY_HASH_PARAM: "hash", + } + ) + assert dict(kwargs) == {"api_base": WORKSPACE, "api_key": "pat-1"} + + +def test_agent_completion_kwargs_never_expands_proxy_env_secrets(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_ONLY_SECRET", "proxy-secret-value") + kwargs = agent_completion_kwargs( + { + "custom_llm_provider": "langgraph", + "api_base": "https://agent.example.com", + "api_key": "os.environ/PROXY_ONLY_SECRET", + } + ) + assert dict(kwargs) == {"api_base": "https://agent.example.com", "api_key": "os.environ/PROXY_ONLY_SECRET"} + + +@respx.mock +async def test_databricks_agent_message_send_through_the_bridge(httpx_transport: None) -> None: + route = respx.post(INVOCATIONS_URL).mock( + return_value=httpx.Response(200, json={"object": "response", "output": [MESSAGE_ITEM]}) + ) + response = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-1", + params=_a2a_params(), + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + }, + api_base=None, + ) + sent = route.calls.last.request + assert sent.headers["Authorization"] == "Bearer pat-1" + assert json.loads(sent.content) == {"input": [{"role": "user", "content": "ping"}]} + result = response["result"] + assert result["kind"] == "message" + assert result["parts"] == [{"kind": "text", "text": "pong"}] + + +@respx.mock +async def test_databricks_app_credential_wins_over_a_forwarded_client_authorization(httpx_transport: None) -> None: + route = respx.post(APP_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-1", + params=_a2a_params(), + litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "minted-oauth"}, + api_base=APP_URL, + agent_extra_headers={"authorization": "Bearer client-token", "X-LiteLLM-User-Id": "u-1"}, + ) + sent = route.calls.last.request + assert sent.headers.get_list("Authorization") == ["Bearer minted-oauth"] + assert sent.headers["X-LiteLLM-User-Id"] == "u-1" + + +@respx.mock +async def test_databricks_agent_message_stream_through_the_bridge(httpx_transport: None) -> None: + events = [ + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"}, + {"type": "response.output_item.done", "item": MESSAGE_ITEM}, + ] + body = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + respx.post(INVOCATIONS_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + ) + chunks = [ + chunk + async for chunk in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-1", + params=_a2a_params(), + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + }, + ) + ] + kinds = [chunk["result"]["kind"] for chunk in chunks] + assert kinds == ["task", "status-update", "artifact-update", "status-update"] + assert chunks[2]["result"]["artifact"]["parts"] == [{"kind": "text", "text": "pong"}] + assert chunks[3]["result"]["status"]["state"] == "completed" diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 1ed513d9db1..a9541f9fcb8 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -172,6 +172,50 @@ async def test_async_copilot_requests_bypass_response_cache(copilot_response_cac _assert_copilot_cache_empty(copilot_response_cache) +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_databricks_agent_replies_are_never_served_from_the_response_cache( + copilot_response_cache: Cache, stream: bool +) -> None: + workspace: Final = "https://adb-cache.azuredatabricks.net" + replies: Final = iter(("first agent answer", "second agent answer")) + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + text: Final = next(replies) + if json.loads(request.content).get("stream"): + frame: Final = {"type": "response.output_text.delta", "item_id": "m", "delta": text} + return httpx.Response(200, content=f"data: {json.dumps(frame)}\n\n".encode(), request=request) + message: Final = {"type": "message", "id": "m", "content": [{"type": "output_text", "text": text}]} + return httpx.Response(200, json={"output": [message]}, request=request) + + async def answer() -> str: + response: Final = await litellm.acompletion( + model="databricks_agent/support-agent", + messages=[{"role": "user", "content": "same question"}], + api_base=workspace, + api_key="dbx-pat", + client=client, + stream=stream, + caching=True, + ) + if isinstance(response, CustomStreamWrapper): + return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + return cast(str, response.choices[0].message.content) # cast-ok: a ModelResponse's message content + + client: Final = AsyncHTTPHandler(transport=_CopilotCacheAsyncTransport(respond)) + async with client.client: + first: Final = await answer() + await asyncio.gather(*_PENDING_CACHE_WRITES) + second: Final = await answer() + await asyncio.gather(*_PENDING_CACHE_WRITES) + + assert (first, second) == ("first agent answer", "second agent answer") + assert len(requests) == 2 + _assert_copilot_cache_empty(copilot_response_cache) + + @pytest.mark.asyncio async def test_existing_provider_response_cache_still_hits(copilot_response_cache: Cache) -> None: messages: Final = [{"role": "user", "content": "cached provider control"}] diff --git a/tests/unit/llms/databricks/agent/__init__.py b/tests/unit/llms/databricks/agent/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/agent/test_databricks_agent_transformation.py b/tests/unit/llms/databricks/agent/test_databricks_agent_transformation.py new file mode 100644 index 00000000000..a754cd23c01 --- /dev/null +++ b/tests/unit/llms/databricks/agent/test_databricks_agent_transformation.py @@ -0,0 +1,725 @@ +import json +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest +import responses +import respx + +import litellm +from litellm.llms.databricks.agent.responses_output import ( + AgentResponse, + DatabricksAgentResponsesIterator, + output_text, +) +from litellm.llms.databricks.agent.transformation import ( + DatabricksAgentConfig, + is_databricks_app_url, + resolve_endpoint_url, + workspace_url_from_api_base, +) +from litellm.llms.databricks.common_utils import DatabricksException +from litellm.types.utils import ModelResponse + +WORKSPACE: Final = "https://adb-1.azuredatabricks.net" +TOKEN_URL: Final = f"{WORKSPACE}/oidc/v1/token" +INVOCATIONS_URL: Final = f"{WORKSPACE}/serving-endpoints/my-agent/invocations" +UNIFIED_URL: Final = f"{WORKSPACE}/serving-endpoints/responses" +APP_HOST: Final = "https://my-app-1.azure.databricksapps.com" +APP_URL: Final = f"{APP_HOST}/responses" +MESSAGE_ITEM: Final = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], +} +_TOOL_CALL: Final = { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, +} + + +@pytest.fixture +def env_oauth_m2m(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATABRICKS_HOST", raising=False) + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "sp-id") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", f"sp-secret-{uuid4().hex}") + + +@pytest.fixture +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +def _sse(events: list[dict[str, object]]) -> str: + return "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + + +@pytest.mark.parametrize( + ("api_base", "model", "expected"), + [ + (WORKSPACE, "my-agent", INVOCATIONS_URL), + (f"{WORKSPACE}/", "databricks_agent/my-agent", INVOCATIONS_URL), + (f"{WORKSPACE}/serving-endpoints", "my-agent", INVOCATIONS_URL), + (INVOCATIONS_URL, "anything", INVOCATIONS_URL), + (f"{WORKSPACE}/serving-endpoints/my-agent", "", INVOCATIONS_URL), + (f"{APP_URL}/", "agent", APP_URL), + (APP_HOST, "", APP_URL), + (f"{APP_HOST}/", "agent", APP_URL), + (f"{APP_HOST}/api/agent", "", f"{APP_HOST}/api/agent"), + (f"{UNIFIED_URL}/", "databricks_agent/my-agent", UNIFIED_URL), + (WORKSPACE, "team/agent v2", f"{WORKSPACE}/serving-endpoints/team%2Fagent%20v2/invocations"), + ], +) +def test_resolve_endpoint_url(api_base: str, model: str, expected: str) -> None: + assert resolve_endpoint_url(api_base, model) == expected + + +@pytest.mark.parametrize("app_base", [APP_HOST, f"{APP_HOST}/", APP_URL, "https://other.example/responses"]) +def test_an_app_url_is_never_a_workspace(app_base: str) -> None: + assert is_databricks_app_url(app_base) + assert workspace_url_from_api_base(app_base) is None + + +@pytest.mark.parametrize("api_base", [WORKSPACE, f"{WORKSPACE}/serving-endpoints", INVOCATIONS_URL, UNIFIED_URL]) +def test_model_serving_urls_resolve_to_the_workspace(api_base: str) -> None: + assert not is_databricks_app_url(api_base) + assert workspace_url_from_api_base(api_base) == WORKSPACE + + +def test_resolve_endpoint_url_needs_an_endpoint_name_for_a_workspace_url() -> None: + with pytest.raises(DatabricksException, match="model must name the Databricks serving endpoint") as info: + resolve_endpoint_url(WORKSPACE, "databricks_agent/") + assert info.value.status_code == 400 + + +def test_unified_responses_url_needs_an_endpoint_name_for_the_body_model() -> None: + with pytest.raises(DatabricksException, match="model must name the Databricks serving endpoint") as info: + resolve_endpoint_url(UNIFIED_URL, "databricks_agent/") + assert info.value.status_code == 400 + + +def test_get_complete_url_requires_api_base() -> None: + with pytest.raises(DatabricksException, match="api_base is required for databricks_agent"): + DatabricksAgentConfig().get_complete_url( + api_base=None, api_key=None, model="my-agent", optional_params={}, litellm_params={} + ) + + +def test_transform_request_maps_messages_to_responses_input_with_passthrough_fields() -> None: + body = DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": [{"type": "text", "text": "ping"}, {"type": "text", "text": "pong?"}]}, + {"role": "assistant", "content": "pong"}, + ], + optional_params={"custom_inputs": {"tenant": "t-1"}, "context": {"conversation_id": "c-1"}, "ignored": 1}, + litellm_params={}, + headers={}, + ) + assert body == { + "input": [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "pingpong?"}, + {"role": "assistant", "content": "pong"}, + ], + "custom_inputs": {"tenant": "t-1"}, + "context": {"conversation_id": "c-1"}, + } + + +def test_transform_request_sets_stream_only_when_streaming() -> None: + config = DatabricksAgentConfig() + messages = [{"role": "user", "content": "ping"}] + assert "stream" not in config.transform_request("m", messages, {"stream": False}, {}, {}) + assert config.transform_request("m", messages, {"stream": True}, {}, {})["stream"] is True + + +def test_transform_request_rejects_roles_a_responses_agent_cannot_take(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(DatabricksException, match="role 'tool' is not supported") as info: + DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[{"role": "tool", "content": "42", "tool_call_id": "call_1"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert info.value.status_code == 400 + + +def test_transform_request_rejects_assistant_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(DatabricksException, match="remove tool_calls from the message") as info: + DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_TOOL_CALL]}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert info.value.status_code == 400 + + +def test_transform_request_drops_tool_turns_and_empty_messages_under_drop_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + body = DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_TOOL_CALL]}, + {"role": "tool", "content": "sunny", "tool_call_id": "call_1"}, + {"role": "assistant", "content": "Checking the weather", "tool_calls": [_TOOL_CALL]}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}]}, + {"role": "user", "content": "and tomorrow?"}, + ], + optional_params={}, + litellm_params={"drop_params": True}, + headers={}, + ) + assert body == { + "input": [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": "Checking the weather"}, + {"role": "user", "content": "and tomorrow?"}, + ] + } + + +@pytest.mark.parametrize( + "part", + [ + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + {"type": "file", "file": {"file_id": "file-1"}}, + ], +) +def test_transform_request_rejects_content_a_responses_agent_cannot_take( + part: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(DatabricksException, match=str(part["type"])) as info: + DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[{"role": "user", "content": [{"type": "text", "text": "describe this"}, part]}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert info.value.status_code == 400 + + +@pytest.mark.parametrize( + ("global_drop_params", "litellm_params"), + [(True, {}), (False, {"drop_params": True})], +) +def test_transform_request_drops_non_text_content_under_drop_params( + global_drop_params: bool, litellm_params: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + body = DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + } + ], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + assert body == {"input": [{"role": "user", "content": "describe this"}]} + + +def test_transform_request_joins_text_content_parts() -> None: + body = DatabricksAgentConfig().transform_request( + model="my-agent", + messages=[{"role": "user", "content": [{"type": "text", "text": "pi"}, {"type": "text", "text": "ng"}]}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body == {"input": [{"role": "user", "content": "ping"}]} + + +def test_validate_environment_lets_the_api_key_replace_a_header_authorization() -> None: + headers = DatabricksAgentConfig().validate_environment( + headers={"authorization": "Bearer client-token", "X-LiteLLM-User-Id": "u1"}, + model="my-agent", + messages=[], + optional_params={}, + litellm_params={}, + api_key="pat-1", + ) + assert headers == {"X-LiteLLM-User-Id": "u1", "Content-Type": "application/json", "Authorization": "Bearer pat-1"} + + +def test_validate_environment_uses_a_header_authorization_when_no_api_key_is_set() -> None: + headers = DatabricksAgentConfig().validate_environment( + headers={"authorization": "Bearer client-token"}, + model="my-agent", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + assert headers == {"authorization": "Bearer client-token", "Content-Type": "application/json"} + + +def test_validate_environment_uses_the_pat_as_bearer() -> None: + headers = DatabricksAgentConfig().validate_environment( + headers={}, model="my-agent", messages=[], optional_params={}, litellm_params={}, api_key="pat-1" + ) + assert headers == {"Content-Type": "application/json", "Authorization": "Bearer pat-1"} + + +def test_validate_environment_without_any_credential_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + with pytest.raises(DatabricksException, match="Missing Databricks credentials") as info: + DatabricksAgentConfig().validate_environment( + headers={}, model="my-agent", messages=[], optional_params={}, litellm_params={}, api_key=None + ) + assert info.value.status_code == 400 + + +def _validate_environment(api_key: str | None, api_base: str | None) -> dict[str, object]: + return DatabricksAgentConfig().validate_environment( + headers={"authorization": "Bearer client-token"}, + model="my-agent", + messages=[], + optional_params={}, + litellm_params={}, + api_key=api_key, + api_base=api_base, + ) + + +@responses.activate +@pytest.mark.parametrize("api_base", [WORKSPACE, f"{WORKSPACE}/serving-endpoints", INVOCATIONS_URL]) +def test_validate_environment_mints_an_oauth_m2m_token_from_the_env_service_principal( + env_oauth_m2m: None, api_base: str +) -> None: + token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 3600}) + headers = _validate_environment(api_key=None, api_base=api_base) + assert headers == {"Content-Type": "application/json", "Authorization": "Bearer m2m-token"} + assert token.calls[0].request.headers["Authorization"].startswith("Basic ") + assert "grant_type=client_credentials" in str(token.calls[0].request.body) + + +@responses.activate +def test_validate_environment_mints_the_env_token_from_databricks_host_for_an_app( + env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABRICKS_HOST", f"{WORKSPACE}/") + responses.post(TOKEN_URL, json={"access_token": "m2m-token"}) + assert _validate_environment(api_key=None, api_base=APP_HOST)["Authorization"] == "Bearer m2m-token" + + +@responses.activate +def test_validate_environment_never_posts_the_env_secret_to_an_app_host(env_oauth_m2m: None) -> None: + with pytest.raises(DatabricksException, match="set DATABRICKS_HOST") as info: + _validate_environment(api_key=None, api_base=APP_URL) + assert info.value.status_code == 400 + assert not responses.calls + + +@responses.activate +def test_validate_environment_prefers_the_api_key_over_the_env_service_principal(env_oauth_m2m: None) -> None: + assert _validate_environment(api_key="pat-1", api_base=WORKSPACE)["Authorization"] == "Bearer pat-1" + assert not responses.calls + + +def test_output_text_joins_message_items_and_skips_tool_trace_items() -> None: + output = [ + {"type": "function_call", "name": "lookup", "arguments": "{}", "call_id": "c1"}, + {"type": "function_call_output", "call_id": "c1", "output": "secret trace"}, + { + "type": "message", + "content": [{"type": "output_text", "text": "Let me check."}, {"type": "refusal", "refusal": "no"}], + }, + {"type": "message", "content": [{"type": "output_text", "text": "It is sunny."}]}, + ] + assert output_text(AgentResponse.model_validate({"output": output}).output) == "Let me check.\n\nIt is sunny." + assert output_text(AgentResponse.model_validate({"output": output[:2]}).output) == "" + assert output_text(None) == "" + + +def test_transform_response_reads_usage_and_custom_outputs() -> None: + raw = httpx.Response( + 200, + json={ + "object": "response", + "output": [MESSAGE_ITEM], + "usage": {"input_tokens": 7, "output_tokens": 3}, + "custom_outputs": {"trace_id": "tr-1"}, + }, + request=httpx.Request("POST", INVOCATIONS_URL), + ) + logging_obj = MagicMock() + response = DatabricksAgentConfig().transform_response( + model="my-agent", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "ping"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert logging_obj.post_call.call_args.kwargs["original_response"] == raw.text + assert response.choices[0].message.content == "pong" + assert response.choices[0].message.provider_specific_fields == {"custom_outputs": {"trace_id": "tr-1"}} + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 3, 10) + + +def test_transform_response_estimates_usage_when_the_agent_reports_none() -> None: + raw = httpx.Response( + 200, json={"object": "response", "output": [MESSAGE_ITEM]}, request=httpx.Request("POST", INVOCATIONS_URL) + ) + response = DatabricksAgentConfig().transform_response( + model="my-agent", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "ping"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.choices[0].message.provider_specific_fields is None + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +@pytest.mark.parametrize("body", ["login", "[]", "{}", '{"output": "pong"}']) +def test_transform_response_rejects_a_body_that_is_not_a_responses_agent_response(body: str) -> None: + raw = httpx.Response(200, text=body, request=httpx.Request("POST", APP_URL)) + with pytest.raises(DatabricksException, match="not a ResponsesAgent response") as info: + DatabricksAgentConfig().transform_response( + model="my-agent", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert info.value.status_code == 502 + + +def test_stream_iterator_emits_deltas_once_and_whole_items_only_when_nothing_streamed() -> None: + iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True) + chunks = [ + iterator.chunk_parser({"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"}), + iterator.chunk_parser({"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"}), + iterator.chunk_parser({"type": "response.output_item.done", "item": MESSAGE_ITEM}), + iterator.chunk_parser({"type": "response.output_item.done", "item": {**MESSAGE_ITEM, "id": "msg_2"}}), + iterator.chunk_parser({"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc_1"}}), + ] + assert [chunk["text"] for chunk in chunks] == ["po", "ng", "", "\n\npong", ""] + assert all(chunk["is_finished"] is False for chunk in chunks) + + +def test_stream_iterator_separates_consecutive_message_items_like_the_non_streaming_reply() -> None: + iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True) + events: list[dict[str, object]] = [ + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": ""}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Let me check."}, + {"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc_1"}}, + {"type": "response.output_text.delta", "item_id": "msg_2", "delta": "It is"}, + {"type": "response.output_text.delta", "item_id": "msg_2", "delta": " sunny."}, + ] + streamed = "".join(iterator.chunk_parser(event)["text"] for event in events) + whole = output_text( + AgentResponse.model_validate( + { + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Let me check."}]}, + {"type": "function_call", "id": "fc_1"}, + {"type": "message", "content": [{"type": "output_text", "text": "It is sunny."}]}, + ] + } + ).output + ) + assert streamed == whole == "Let me check.\n\nIt is sunny." + + +def test_stream_iterator_surfaces_agent_errors() -> None: + iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True) + with pytest.raises(DatabricksException, match="tool failed"): + iterator.chunk_parser({"type": "error", "message": "tool failed"}) + + +@respx.mock +async def test_acompletion_posts_responses_input_to_model_serving(httpx_transport: None) -> None: + route = respx.post(INVOCATIONS_URL).mock( + return_value=httpx.Response(200, json={"object": "response", "output": [MESSAGE_ITEM]}) + ) + response = await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + api_base=WORKSPACE, + api_key="pat-1", + custom_inputs={"tenant": "t-1"}, + ) + sent = route.calls.last.request + assert sent.headers["Authorization"] == "Bearer pat-1" + assert json.loads(sent.content) == { + "input": [{"role": "user", "content": "ping"}], + "custom_inputs": {"tenant": "t-1"}, + } + assert response.choices[0].message.content == "pong" + assert response.model == "my-agent" + + +@respx.mock +async def test_acompletion_names_the_endpoint_in_the_body_for_the_unified_responses_url( + httpx_transport: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABRICKS_API_BASE", UNIFIED_URL) + monkeypatch.setenv("DATABRICKS_API_KEY", "pat-env") + route = respx.post(UNIFIED_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + response = await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + ) + sent = route.calls.last.request + assert sent.headers["Authorization"] == "Bearer pat-env" + assert json.loads(sent.content) == {"model": "my-agent", "input": [{"role": "user", "content": "ping"}]} + assert response.choices[0].message.content == "pong" + + +@respx.mock +async def test_acompletion_forwards_extra_headers_and_keeps_the_minted_token(httpx_transport: None) -> None: + route = respx.post(APP_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + await litellm.acompletion( + model="databricks_agent/agent", + messages=[{"role": "user", "content": "ping"}], + api_base=APP_URL, + extra_headers={"Authorization": "Bearer minted-oauth", "X-LiteLLM-Trace-Id": "tr-1"}, + ) + sent = route.calls.last.request + assert sent.headers["Authorization"] == "Bearer minted-oauth" + assert sent.headers["X-LiteLLM-Trace-Id"] == "tr-1" + + +@respx.mock +async def test_acompletion_streams_output_text_deltas(httpx_transport: None) -> None: + events: list[dict[str, object]] = [ + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"}, + {"type": "response.output_item.done", "item": MESSAGE_ITEM}, + ] + route = respx.post(INVOCATIONS_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_sse(events)) + ) + stream = await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + api_base=WORKSPACE, + api_key="pat-1", + stream=True, + ) + deltas = [chunk.choices[0].delta.content async for chunk in stream if chunk.choices[0].delta.content] + assert json.loads(route.calls.last.request.content)["stream"] is True + assert "".join(deltas) == "pong" + + +@pytest.mark.parametrize("env_pat", (None, "pat-env"), ids=("alone", "beside-an-env-pat")) +@responses.activate +@respx.mock +async def test_acompletion_authenticates_with_the_env_service_principal( + httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch, env_pat: str | None +) -> None: + if env_pat is None: + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + else: + monkeypatch.setenv("DATABRICKS_API_KEY", env_pat) + responses.post(TOKEN_URL, json={"access_token": "m2m-token"}) + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + response = await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + api_base=WORKSPACE, + ) + assert route.calls.last.request.headers["Authorization"] == "Bearer m2m-token" + assert response.choices[0].message.content == "pong" + + +async def test_unsupported_openai_params_follow_drop_params(httpx_transport: None) -> None: + messages = [{"role": "user", "content": "ping"}] + with pytest.raises(litellm.UnsupportedParamsError): + await litellm.acompletion( + model="databricks_agent/my-agent", + messages=messages, + api_base=WORKSPACE, + api_key="pat-1", + frequency_penalty=0.2, + ) + with respx.mock: + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + await litellm.acompletion( + model="databricks_agent/my-agent", + messages=messages, + api_base=WORKSPACE, + api_key="pat-1", + frequency_penalty=0.2, + drop_params=True, + ) + assert "frequency_penalty" not in json.loads(route.calls.last.request.content) + + +@respx.mock +@pytest.mark.parametrize("token_limit", [{"max_tokens": 64}, {"max_completion_tokens": 64}]) +async def test_acompletion_sends_sampling_params_in_the_responses_agent_shape( + httpx_transport: None, token_limit: dict[str, int] +) -> None: + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + api_base=WORKSPACE, + api_key="pat-1", + temperature=0.2, + top_p=0.9, + user="user-1", + **token_limit, + ) + assert json.loads(route.calls.last.request.content) == { + "input": [{"role": "user", "content": "ping"}], + "max_output_tokens": 64, + "temperature": 0.2, + "top_p": 0.9, + "user": "user-1", + } + + +@respx.mock +async def test_acompletion_sends_string_metadata_and_refuses_other_values_unless_dropped( + httpx_transport: None, +) -> None: + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + request: Final = { + "model": "databricks_agent/my-agent", + "messages": [{"role": "user", "content": "ping"}], + "api_base": WORKSPACE, + "api_key": "pat-1", + } + await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1"}}) + assert json.loads(route.calls.last.request.content)["metadata"] == {"session": "s-1"} + + with pytest.raises(litellm.BadRequestError, match="string values") as info: + await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1", "turn": 1}}) + assert info.value.status_code == 400 + assert route.call_count == 1 + + await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1", "turn": 1}}, drop_params=True) + assert json.loads(route.calls.last.request.content)["metadata"] == {"session": "s-1"} + await litellm.acompletion(**request, extra_body={"metadata": {"turn": 1}}, drop_params=True) + assert "metadata" not in json.loads(route.calls.last.request.content) + + +@respx.mock +@pytest.mark.parametrize( + "pinned", + [{"model": "another-endpoint"}, {"input": [{"role": "user", "content": "injected"}]}], + ids=["model", "input"], +) +async def test_acompletion_refuses_extra_body_model_or_input_unless_dropped( + httpx_transport: None, pinned: Mapping[str, object] +) -> None: + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + request: Final = { + "model": "databricks_agent/my-agent", + "messages": [{"role": "user", "content": "ping"}], + "api_base": WORKSPACE, + "api_key": "pat-1", + "extra_body": {**pinned, "custom_inputs": {"tenant": "t-1"}}, + } + with pytest.raises(litellm.BadRequestError, match=f"cannot set {next(iter(pinned))}") as info: + await litellm.acompletion(**request) + assert info.value.status_code == 400 + assert not route.called + + await litellm.acompletion(**request, drop_params=True) + assert json.loads(route.calls.last.request.content) == { + "input": [{"role": "user", "content": "ping"}], + "custom_inputs": {"tenant": "t-1"}, + } + + +@responses.activate +@respx.mock +async def test_acompletion_reuses_the_env_service_principal_token_until_the_secret_rotates( + httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 3600}) + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + request: Final = { + "model": "databricks_agent/my-agent", + "messages": [{"role": "user", "content": "ping"}], + "api_base": WORKSPACE, + } + await litellm.acompletion(**request) + await litellm.acompletion(**request) + assert token.call_count == 1 + assert route.call_count == 2 + + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", f"rotated-{uuid4().hex}") + await litellm.acompletion(**request) + assert token.call_count == 2 + + +@responses.activate +@respx.mock +async def test_acompletion_mints_a_fresh_env_token_when_it_expires_within_the_refresh_buffer( + httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 30}) + respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + request: Final = { + "model": "databricks_agent/my-agent", + "messages": [{"role": "user", "content": "ping"}], + "api_base": WORKSPACE, + } + await litellm.acompletion(**request) + await litellm.acompletion(**request) + assert token.call_count == 2 + + +@respx.mock +async def test_acompletion_resolves_os_environ_api_base_and_api_key( + httpx_transport: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("MY_DBX_HOST", WORKSPACE) + monkeypatch.setenv("MY_DBX_PAT", "pat-from-env") + route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]})) + response = await litellm.acompletion( + model="databricks_agent/my-agent", + messages=[{"role": "user", "content": "ping"}], + api_base="os.environ/MY_DBX_HOST", + api_key="os.environ/MY_DBX_PAT", + ) + assert route.calls.last.request.headers["Authorization"] == "Bearer pat-from-env" + assert response.choices[0].message.content == "pong" diff --git a/tests/unit/llms/databricks/test_databricks_partner_integration.py b/tests/unit/llms/databricks/test_databricks_partner_integration.py index c6ad78366f7..8c12f1aa1b4 100644 --- a/tests/unit/llms/databricks/test_databricks_partner_integration.py +++ b/tests/unit/llms/databricks/test_databricks_partner_integration.py @@ -197,7 +197,7 @@ class TestOAuthM2M: mock_response.json.return_value = {"access_token": "test-access-token"} with patch("requests.post", return_value=mock_response) as mock_post: - token = databricks_base._get_oauth_m2m_token( + token = databricks_base.get_oauth_m2m_token( api_base="https://adb-123.azuredatabricks.net/serving-endpoints", client_id="test-client-id", client_secret="test-client-secret", @@ -219,7 +219,7 @@ class TestOAuthM2M: with patch("requests.post", return_value=mock_response): with pytest.raises(DatabricksException) as exc_info: - databricks_base._get_oauth_m2m_token( + databricks_base.get_oauth_m2m_token( api_base="https://adb-123.azuredatabricks.net", client_id="bad-client-id", client_secret="bad-secret", @@ -235,7 +235,7 @@ class TestOAuthM2M: mock_response.json.return_value = {"access_token": "token"} with patch("requests.post", return_value=mock_response) as mock_post: - databricks_base._get_oauth_m2m_token( + databricks_base.get_oauth_m2m_token( api_base="https://adb-123.azuredatabricks.net/serving-endpoints", client_id="id", client_secret="secret", @@ -254,7 +254,7 @@ class TestOAuthM2M: mock_response.json.return_value = {"access_token": "token"} with patch("requests.post", return_value=mock_response) as mock_post: - databricks_base._get_oauth_m2m_token( + databricks_base.get_oauth_m2m_token( api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1", client_id="id", client_secret="secret", @@ -278,7 +278,7 @@ class TestValidateEnvironmentWithOAuth: databricks_base = DatabricksBase() with patch.object( - databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + databricks_base, "get_oauth_m2m_token", return_value="oauth-token" ) as mock_oauth: api_base, headers = databricks_base.databricks_validate_environment( api_key=None, @@ -546,7 +546,7 @@ class TestAuthenticationPriority: databricks_base = DatabricksBase() with patch.object( - databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + databricks_base, "get_oauth_m2m_token", return_value="oauth-token" ) as mock_oauth: api_base, headers = databricks_base.databricks_validate_environment( api_key=None, # No PAT provided - OAuth should be used @@ -570,7 +570,7 @@ class TestAuthenticationPriority: # Mock the OAuth call - it will be attempted but PAT should override with patch.object( - databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + databricks_base, "get_oauth_m2m_token", return_value="oauth-token" ): api_base, headers = databricks_base.databricks_validate_environment( api_key="dapi-explicit-pat", diff --git a/tests/unit/proxy/agent_endpoints/test_agent_headers.py b/tests/unit/proxy/agent_endpoints/test_agent_headers.py index e894f4ad69a..6a24939f41b 100644 --- a/tests/unit/proxy/agent_endpoints/test_agent_headers.py +++ b/tests/unit/proxy/agent_endpoints/test_agent_headers.py @@ -9,10 +9,17 @@ Tests cover: - merge_agent_headers utility """ +import base64 +import json import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx + +import litellm # --------------------------------------------------------------------------- # Helper: build a minimal mock agent @@ -378,6 +385,72 @@ async def test_databricks_oauth_overrides_static_authorization(): assert headers.get("Authorization") == "Bearer oauth-wins" +@respx.mock +@pytest.mark.asyncio +async def test_databricks_agent_flat_oauth_send_reaches_the_endpoint_with_the_minted_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + from litellm.proxy.agent_endpoints.databricks_oauth import databricks_app_oauth_token_cache + from litellm.types.agents import AgentResponse + + workspace: Final = "https://adb-1.azuredatabricks.net" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + databricks_app_oauth_token_cache.flush_cache() + registry: Final = AgentRegistry() + registry.register_agent( + AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": workspace}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": workspace, + "client_id": "cid", + "client_secret": "secret", + }, + ) + ) + token: Final = respx.post(f"{workspace}/oidc/v1/token").mock( + return_value=httpx.Response(200, json={"access_token": "flat-minted", "expires_in": 3600}) + ) + endpoint: Final = respx.post(f"{workspace}/serving-endpoints/my-agent/invocations").mock( + return_value=httpx.Response( + 200, + json={ + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong"}], + } + ] + }, + ) + ) + + with ( + patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry), + ): + response = await invoke_agent_a2a( + agent_id="dbx-agent", + request=_make_mock_request(), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + basic: Final = base64.b64encode(b"cid:secret").decode() + assert token.calls.last.request.headers["Authorization"] == f"Basic {basic}" + sent: Final = endpoint.calls.last.request + assert sent.headers["Authorization"] == "Bearer flat-minted" + assert json.loads(sent.content) == {"input": [{"role": "user", "content": "Hello"}]} + assert json.loads(response.body)["result"]["parts"][0]["text"] == "pong" + + @pytest.mark.asyncio async def test_non_databricks_agent_skips_oauth_resolution(): """Agents without a databricks_oauth block never enter the OAuth path.""" @@ -385,13 +458,10 @@ async def test_non_databricks_agent_skips_oauth_resolution(): mock_agent.litellm_params = {"require_trace_id_on_calls_to_agent": False} mock_request = _make_mock_request() - with patch( - "litellm.proxy.agent_endpoints.a2a_endpoints.resolve_databricks_app_auth_header", - new_callable=AsyncMock, - ) as mock_resolve: + with patch("litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client") as mock_token_client: mock_asend = await _invoke(mock_agent, mock_request, None) - mock_resolve.assert_not_called() + mock_token_client.assert_not_called() headers = mock_asend.call_args.kwargs.get("agent_extra_headers") assert headers == {"x-custom": "v", "X-LiteLLM-User-Id": "u1"} assert "Authorization" not in headers diff --git a/tests/unit/proxy/agent_endpoints/test_databricks_oauth.py b/tests/unit/proxy/agent_endpoints/test_databricks_oauth.py index 39f51ee0401..cd02ca4314e 100644 --- a/tests/unit/proxy/agent_endpoints/test_databricks_oauth.py +++ b/tests/unit/proxy/agent_endpoints/test_databricks_oauth.py @@ -8,17 +8,22 @@ 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, ) @@ -46,9 +51,7 @@ def _mock_http_handler(access_token="tok-abc", expires_in=3600, post_error=None) handler.post.side_effect = post_error else: response = MagicMock() - response.json = MagicMock( - return_value={"access_token": access_token, "expires_in": expires_in} - ) + response.json = MagicMock(return_value={"access_token": access_token, "expires_in": expires_in}) handler.post.return_value = response return handler @@ -111,23 +114,49 @@ def test_parse_custom_scope(): assert config.scope == "custom-scope" -@pytest.mark.parametrize( - "missing_field", ["client_id", "client_secret", "workspace_url"] +AGENT_KINDS: Final = pytest.mark.parametrize( + ("agent_params", "error"), + [({}, ValueError), ({"custom_llm_provider": "databricks_agent"}, litellm.BadRequestError)], + ids=["url-agent", "databricks-agent"], ) -def test_parse_raises_on_missing_field(missing_field): + + +@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(ValueError, match=missing_field): - parse_databricks_oauth_config({"databricks_oauth": block}) + with pytest.raises(error, match=missing_field): + parse_databricks_oauth_config({**agent_params, "databricks_oauth": block}) -def test_parse_raises_on_non_mapping_block(): - with pytest.raises(ValueError, match="mapping"): - parse_databricks_oauth_config({"databricks_oauth": "not-a-dict"}) +@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): @@ -179,9 +208,7 @@ async def test_fetch_token_posts_client_credentials_with_basic_auth(): } # 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 call.kwargs["headers"]["Authorization"] == _expected_basic_auth("cid", "secret") assert "auth" not in call.kwargs @@ -304,9 +331,7 @@ async def test_http_status_error_raises_value_error(): 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 - ) + post_error=httpx.HTTPStatusError("unauthorized", request=request, response=error_response) ) with patch( @@ -494,3 +519,169 @@ async def test_resolve_returns_bearer_header(): 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", + } diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index b03dbede869..9e8910e44de 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -507,6 +507,7 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "charity_engine", "chutes", "darkbloom", + "databricks_agent", "gdc", "helicone", "inception", diff --git a/tests/unit/proxy/test_route_a2a_models.py b/tests/unit/proxy/test_route_a2a_models.py index 0429dd97a1c..b54f07f3afb 100644 --- a/tests/unit/proxy/test_route_a2a_models.py +++ b/tests/unit/proxy/test_route_a2a_models.py @@ -4,14 +4,22 @@ Test A2A model routing in proxy. Maps to: litellm/proxy/agent_endpoints/a2a_routing.py """ - - +import base64 +import json +from collections.abc import Iterator from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest +import respx +from fastapi import HTTPException +import litellm from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.databricks_oauth import databricks_app_oauth_token_cache from litellm.proxy.route_llm_request import route_request +from litellm.types.agents import AgentResponse @pytest.mark.asyncio @@ -183,3 +191,477 @@ async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_re assert call_kwargs["model"] == f"a2a/{agent_name}" assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com" prisma_client.db.litellm_agentstable.find_unique.assert_awaited() + + +WORKSPACE = "https://adb-1.azuredatabricks.net" +APP_URL = "https://my-app-1.azure.databricksapps.com/responses" +AGENT_REPLY = { + "output": [ + {"type": "message", "id": "msg_1", "role": "assistant", "content": [{"type": "output_text", "text": "pong"}]} + ] +} +A2A_REPLY = { + "jsonrpc": "2.0", + "id": "1", + "result": {"kind": "message", "role": "agent", "messageId": "m-1", "parts": [{"kind": "text", "text": "pong"}]}, +} + + +@pytest.fixture +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + databricks_app_oauth_token_cache.flush_cache() + yield + databricks_app_oauth_token_cache.flush_cache() + + +def _registry_with(agent: AgentResponse) -> AgentRegistry: + registry = AgentRegistry() + registry.register_agent(agent) + return registry + + +def _token_route(workspace: str) -> respx.Route: + return respx.post(f"{workspace}/oidc/v1/token").mock( + return_value=httpx.Response(200, json={"access_token": "minted-oauth", "expires_in": 3600}) + ) + + +@respx.mock +async def test_route_a2a_bridge_agent_calls_the_endpoint_with_the_agent_params(httpx_transport: None) -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + "custom_inputs": {"tenant": "t-1"}, + "is_public": True, + }, + ) + endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock( + return_value=httpx.Response(200, json=AGENT_REPLY) + ) + data = {"model": "a2a/dbx-agent", "messages": [{"role": "user", "content": "ping"}]} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + sent = endpoint.calls.last.request + assert sent.headers["Authorization"] == "Bearer pat-1" + assert json.loads(sent.content) == { + "input": [{"role": "user", "content": "ping"}], + "custom_inputs": {"tenant": "t-1"}, + } + assert response.choices[0].message.content == "pong" + + +@respx.mock +async def test_route_a2a_bridge_agent_merges_client_static_and_agent_extra_headers(httpx_transport: None) -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + static_headers={"X-Static": "static", "x-shared": "from-static", "X-Both": "from-static"}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + "extra_headers": {"X-Agent": "agent", "x-both": "from-agent"}, + }, + ) + endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock( + return_value=httpx.Response(200, json=AGENT_REPLY) + ) + data = { + "model": "a2a/dbx-agent", + "messages": [{"role": "user", "content": "ping"}], + "extra_headers": {"X-Trace": "trace-1", "X-Shared": "from-client", "X-LiteLLM-User-Id": "spoofed-user"}, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + sent = endpoint.calls.last.request + assert "X-LiteLLM-User-Id" not in sent.headers + assert sent.headers["X-Trace"] == "trace-1" + assert sent.headers["X-Static"] == "static" + assert sent.headers["X-Agent"] == "agent" + assert sent.headers.get_list("X-Shared") == ["from-static"] + assert sent.headers.get_list("X-Both") == ["from-agent"] + assert sent.headers["Authorization"] == "Bearer pat-1" + + +async def test_route_a2a_bridge_agent_lets_the_request_win_except_for_the_agent_credentials() -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + "metadata": {"tags": ["agent-tag"]}, + "custom_inputs": {"tenant": "t-1"}, + "timeout": 30, + }, + ) + proxy_metadata = {"user_api_key_hash": "key-hash", "user_api_key_team_id": "team-1"} + data = { + "model": "a2a/dbx-agent", + "messages": [{"role": "user", "content": "ping"}], + "metadata": proxy_metadata, + "api_base": "https://elsewhere.example.com", + "api_key": "client-key", + } + completion = Mock(return_value="reply") + + with ( + patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)), + patch.object(litellm, "acompletion", completion), + ): + await route_a2a_agent_request(data=data, route_type="acompletion") + + sent = completion.call_args.kwargs + assert sent["metadata"] == proxy_metadata + assert (sent["api_base"], sent["api_key"]) == (WORKSPACE, "pat-1") + assert (sent["custom_inputs"], sent["timeout"]) == ({"tenant": "t-1"}, 30) + assert sent["model"] == "databricks_agent/my-agent" + + +@respx.mock +async def test_route_a2a_bridge_agent_keeps_its_custom_inputs_and_extra_body_keys_over_the_callers( + httpx_transport: None, +) -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "model": "my-agent", + "api_base": WORKSPACE, + "api_key": "pat-1", + "custom_inputs": {"tenant": "owner"}, + "extra_body": {"metadata": {"tenant": "owner"}}, + }, + ) + endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock( + return_value=httpx.Response(200, json=AGENT_REPLY) + ) + data = { + "model": "a2a/dbx-agent", + "messages": [{"role": "user", "content": "ping"}], + "custom_inputs": {"tenant": "caller", "session": "s-1"}, + "extra_body": {"metadata": {"tenant": "caller", "turn": "2"}}, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert json.loads(endpoint.calls.last.request.content) == { + "input": [{"role": "user", "content": "ping"}], + "custom_inputs": {"tenant": "owner", "session": "s-1"}, + "metadata": {"tenant": "owner", "turn": "2"}, + } + + +@respx.mock +async def test_route_a2a_bridge_agent_with_an_empty_oauth_block_fails_before_calling_the_app( + httpx_transport: None, +) -> None: + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", "databricks_oauth": {}}, + ) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}]} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + with pytest.raises(litellm.BadRequestError, match="missing required field") as info: + await route_a2a_agent_request(data=data, route_type="acompletion") + + assert info.value.status_code == 400 + assert not app.called + + +@respx.mock +@pytest.mark.parametrize( + "oauth_params", + [ + {"databricks_oauth": {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE}}, + {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE, "scope": "all-apis"}, + ], +) +async def test_route_a2a_bridge_agent_with_databricks_oauth_sends_the_minted_token( + httpx_transport: None, oauth_params: dict[str, object] +) -> None: + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", **oauth_params}, + ) + token = _token_route(WORKSPACE) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = { + "model": "a2a/dbx-app", + "messages": [{"role": "user", "content": "ping"}], + "extra_headers": {"X-Tenant": "t"}, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert token.calls.last.request.headers["Authorization"] == f"Basic {base64.b64encode(b'sp:s').decode()}" + sent = app.calls.last.request + assert sent.headers["Authorization"] == "Bearer minted-oauth" + assert sent.headers["X-Tenant"] == "t" + assert json.loads(sent.content) == {"input": [{"role": "user", "content": "ping"}]} + assert response.choices[0].message.content == "pong" + + +@respx.mock +@pytest.mark.parametrize( + "oauth_block", + [ + {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE}, + {"client_id": "sp"}, + "not-a-mapping", + {"client_id": "sp", "client_secret": "s", "workspace_url": "https://my-app-1.azure.databricksapps.com"}, + ], + ids=["complete", "partial", "not-a-mapping", "app-host-as-workspace"], +) +async def test_route_a2a_url_agent_chat_keeps_the_client_authorization_over_its_databricks_oauth( + httpx_transport: None, oauth_block: object +) -> None: + agent_url = "https://my-a2a-app.azure.databricksapps.com" + agent = AgentResponse( + agent_id="a2a-app-id", + agent_name="a2a-app", + agent_card_params={"url": agent_url}, + litellm_params={"databricks_oauth": oauth_block}, + ) + token = _token_route(WORKSPACE) + app = respx.post(agent_url).mock(return_value=httpx.Response(200, json=A2A_REPLY)) + data = { + "model": "a2a/a2a-app", + "messages": [{"role": "user", "content": "ping"}], + "extra_headers": {"Authorization": "Bearer client-token"}, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + sent = app.calls.last.request + assert sent.headers["Authorization"] == "Bearer client-token" + assert json.loads(sent.content)["method"] == "message/send" + assert not token.called + assert response.choices[0].message.content == "pong" + + +@respx.mock +async def test_route_a2a_agent_with_another_provider_keeps_calling_its_card_url_over_a2a( + httpx_transport: None, +) -> None: + agent_url = "https://pydantic-agent.example.com" + agent = AgentResponse( + agent_id="pydantic-agent-id", + agent_name="pydantic-agent", + agent_card_params={"url": agent_url}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents", "model": "pydantic-agent"}, + ) + agent_endpoint = respx.post(agent_url).mock(return_value=httpx.Response(200, json=A2A_REPLY)) + data = {"model": "a2a/pydantic-agent", "messages": [{"role": "user", "content": "ping"}]} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert json.loads(agent_endpoint.calls.last.request.content)["method"] == "message/send" + assert response.choices[0].message.content == "pong" + + +@respx.mock +@pytest.mark.parametrize("client_target", ["api_base", "base_url"]) +async def test_route_a2a_bridge_agent_ignores_a_client_target_over_its_card_url( + httpx_transport: None, client_target: str +) -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + litellm_params={"custom_llm_provider": "databricks_agent", "model": "my-agent", "api_key": "pat-1"}, + ) + endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock( + return_value=httpx.Response(200, json=AGENT_REPLY) + ) + data = { + "model": "a2a/dbx-agent", + client_target: "https://collector.example.com", + "messages": [{"role": "user", "content": "ping"}], + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert endpoint.call_count == 1 + assert endpoint.calls.last.request.headers["Authorization"] == "Bearer pat-1" + assert response.choices[0].message.content == "pong" + + +@respx.mock +@pytest.mark.parametrize("endpoint_name", [{"model": ""}, {"model": None}, {}]) +async def test_route_a2a_bridge_agent_without_an_endpoint_name_answers_400_before_calling_databricks( + httpx_transport: None, endpoint_name: dict[str, object] +) -> None: + agent = AgentResponse( + agent_id="dbx-agent-id", + agent_name="dbx-agent", + agent_card_params={"url": WORKSPACE}, + litellm_params={ + "custom_llm_provider": "databricks_agent", + "api_base": WORKSPACE, + "api_key": "p", + **endpoint_name, + }, + ) + workspace = respx.route(host="adb-1.azuredatabricks.net").mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = {"model": "a2a/dbx-agent", "messages": [{"role": "user", "content": "ping"}]} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + with pytest.raises(litellm.BadRequestError, match="must name the Databricks serving endpoint"): + await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert not workspace.called + + +@respx.mock +async def test_route_a2a_bridge_app_agent_without_an_endpoint_name_calls_the_app(httpx_transport: None) -> None: + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"}, + ) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}]} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert app.called + assert response.choices[0].message.content == "pong" + + +@respx.mock +@pytest.mark.parametrize( + ("credential", "env_api_key", "expected_authorization"), + [ + ({"api_key": "pat-1"}, None, "Bearer pat-1"), + ({}, "env-pat", "Bearer env-pat"), + ({"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE}, "env-pat", "Bearer minted-oauth"), + ( + {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE, "api_key": "pat-1"}, + None, + "Bearer minted-oauth", + ), + ({}, None, "Bearer client-token"), + ], +) +async def test_route_a2a_bridge_agent_credential_wins_over_a_client_authorization( + httpx_transport: None, + monkeypatch: pytest.MonkeyPatch, + credential: dict[str, object], + env_api_key: str | None, + expected_authorization: str, +) -> None: + if env_api_key is None: + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + else: + monkeypatch.setenv("DATABRICKS_API_KEY", env_api_key) + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", **credential}, + ) + _token_route(WORKSPACE) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = { + "model": "a2a/dbx-app", + "messages": [{"role": "user", "content": "ping"}], + "extra_headers": {"authorization": "Bearer client-token", "X-Tenant": "t"}, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + sent = app.calls.last.request + assert sent.headers.get_list("Authorization") == [expected_authorization] + assert sent.headers["X-Tenant"] == "t" + + +@respx.mock +@pytest.mark.parametrize("extra_headers", [5, ["X-A"], {"X-N": 5}, "x"]) +async def test_route_a2a_bridge_agent_answers_400_for_malformed_client_extra_headers( + httpx_transport: None, extra_headers: object +) -> None: + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"}, + ) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}], "extra_headers": extra_headers} + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + with pytest.raises(HTTPException) as info: + await route_a2a_agent_request(data=data, route_type="acompletion") + + assert info.value.status_code == 400 + assert "extra_headers" in str(info.value.detail) + assert not app.called + + +@respx.mock +@pytest.mark.parametrize( + ("extra_headers", "drop_params"), + [ + (5, {"drop_params": True}), + ({"X-N": 5}, {"drop_params": True}), + ("", {}), + ({}, {}), + (None, {}), + ], +) +async def test_route_a2a_bridge_agent_calls_without_client_extra_headers_it_drops_or_finds_empty( + httpx_transport: None, extra_headers: object, drop_params: dict[str, object] +) -> None: + agent = AgentResponse( + agent_id="dbx-app-id", + agent_name="dbx-app", + agent_card_params={"url": APP_URL}, + litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"}, + ) + app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY)) + data = { + "model": "a2a/dbx-app", + "messages": [{"role": "user", "content": "ping"}], + "extra_headers": extra_headers, + **drop_params, + } + + with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)): + response = await (await route_a2a_agent_request(data=data, route_type="acompletion")) + + assert app.calls.last.request.headers["Authorization"] == "Bearer p" + assert "X-N" not in app.calls.last.request.headers + assert response.choices[0].message.content == "pong" diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 200d9a0000c..8d3e341096c 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -54,6 +54,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/llms/azure/common_utils.py", "litellm/llms/bedrock/base_aws_llm.py", "litellm/llms/custom_httpx/http_handler.py", + "litellm/llms/databricks/agent/transformation.py", "litellm/llms/gigachat/authenticator.py", "litellm/llms/github_copilot/per_user_auth.py", "litellm/llms/litellm_proxy/skills/handler.py", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts index 89528d91b8d..a1edf0a9829 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts @@ -96,6 +96,21 @@ describe("detectAgentType", () => { expect(detectAgentType(agent)).toBe("bedrock_agentcore"); }); + + it("detects databricks_agent agents from custom_llm_provider so the edit form shows the Databricks fields", () => { + const agent = { + agent_id: "agent-2", + agent_name: "dbx-agent", + agent_card_params: { url: "https://adb-1.azuredatabricks.net" }, + litellm_params: { + custom_llm_provider: "databricks_agent", + model: "my-agent-endpoint", + api_base: "https://adb-1.azuredatabricks.net", + }, + } as unknown as Agent; + + expect(detectAgentType(agent)).toBe("databricks_agent"); + }); }); describe("API agent metadata validation", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts index 5e8fb0478f7..b79889c0216 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts @@ -12,6 +12,7 @@ export const detectAgentType = (agent: Agent): string => { // Check by custom_llm_provider first if (customProvider === "langflow") return "langflow"; + if (customProvider === "databricks_agent") return "databricks_agent"; if (customProvider === "langgraph") return "langgraph"; if (customProvider === "azure_ai") return "azure_ai_foundry"; if (customProvider === "bedrock") return "bedrock_agentcore";