mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
b9f01c9f5b
* fix(databricks/common_utils.py): fix custom endpoint check Fixes https://github.com/BerriAI/litellm/issues/9915 * fix(common_utils.py): add unit test to ensure custom_endpoint=False is handled correctly Fixes https://github.com/BerriAI/litellm/issues/9915
108 lines
4.2 KiB
Python
108 lines
4.2 KiB
Python
from typing import Literal, Optional, Tuple
|
|
|
|
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|
|
|
|
|
class DatabricksException(BaseLLMException):
|
|
pass
|
|
|
|
|
|
class DatabricksBase:
|
|
def _get_api_base(self, api_base: Optional[str]) -> str:
|
|
if api_base is None:
|
|
try:
|
|
from databricks.sdk import WorkspaceClient
|
|
|
|
databricks_client = WorkspaceClient()
|
|
|
|
api_base = (
|
|
api_base or f"{databricks_client.config.host}/serving-endpoints"
|
|
)
|
|
|
|
return api_base
|
|
except ImportError:
|
|
raise DatabricksException(
|
|
status_code=400,
|
|
message=(
|
|
"Either set the DATABRICKS_API_BASE and DATABRICKS_API_KEY environment variables, "
|
|
"or install the databricks-sdk Python library."
|
|
),
|
|
)
|
|
return api_base
|
|
|
|
def _get_databricks_credentials(
|
|
self, api_key: Optional[str], api_base: Optional[str], headers: Optional[dict]
|
|
) -> Tuple[str, dict]:
|
|
headers = headers or {"Content-Type": "application/json"}
|
|
try:
|
|
from databricks.sdk import WorkspaceClient
|
|
|
|
databricks_client = WorkspaceClient()
|
|
|
|
api_base = api_base or f"{databricks_client.config.host}/serving-endpoints"
|
|
|
|
if api_key is None:
|
|
databricks_auth_headers: dict[
|
|
str, str
|
|
] = databricks_client.config.authenticate()
|
|
headers = {**databricks_auth_headers, **headers}
|
|
|
|
return api_base, headers
|
|
except ImportError:
|
|
raise DatabricksException(
|
|
status_code=400,
|
|
message=(
|
|
"If the Databricks base URL and API key are not set, the databricks-sdk "
|
|
"Python library must be installed. Please install the databricks-sdk, set "
|
|
"{LLM_PROVIDER}_API_BASE and {LLM_PROVIDER}_API_KEY environment variables, "
|
|
"or provide the base URL and API key as arguments."
|
|
),
|
|
)
|
|
|
|
def databricks_validate_environment(
|
|
self,
|
|
api_key: Optional[str],
|
|
api_base: Optional[str],
|
|
endpoint_type: Literal["chat_completions", "embeddings"],
|
|
custom_endpoint: Optional[bool],
|
|
headers: Optional[dict],
|
|
) -> Tuple[str, dict]:
|
|
if api_key is None and not headers: # handle empty headers
|
|
if custom_endpoint is True:
|
|
raise DatabricksException(
|
|
status_code=400,
|
|
message="Missing API Key - A call is being made to LLM Provider but no key is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params",
|
|
)
|
|
else:
|
|
api_base, headers = self._get_databricks_credentials(
|
|
api_base=api_base, api_key=api_key, headers=headers
|
|
)
|
|
|
|
if api_base is None:
|
|
if custom_endpoint:
|
|
raise DatabricksException(
|
|
status_code=400,
|
|
message="Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params",
|
|
)
|
|
else:
|
|
api_base, headers = self._get_databricks_credentials(
|
|
api_base=api_base, api_key=api_key, headers=headers
|
|
)
|
|
|
|
if headers is None:
|
|
headers = {
|
|
"Authorization": "Bearer {}".format(api_key),
|
|
"Content-Type": "application/json",
|
|
}
|
|
else:
|
|
if api_key is not None:
|
|
headers.update({"Authorization": "Bearer {}".format(api_key)})
|
|
|
|
if api_key is not None:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
if endpoint_type == "chat_completions" and custom_endpoint is not True:
|
|
api_base = "{}/chat/completions".format(api_base)
|
|
elif endpoint_type == "embeddings" and custom_endpoint is not True:
|
|
api_base = "{}/embeddings".format(api_base)
|
|
return api_base, headers
|