Spaces:
Configuration error
Configuration error
""" | |
Support for OpenAI's `/v1/chat/completions` endpoint. | |
Calls done in OpenAI/openai.py as DataRobot is openai-compatible. | |
""" | |
from typing import Optional, Tuple | |
from litellm.secret_managers.main import get_secret_str | |
from ...openai_like.chat.transformation import OpenAILikeChatConfig | |
class DataRobotConfig(OpenAILikeChatConfig): | |
def _resolve_api_key(api_key: Optional[str] = None) -> str: | |
"""Attempt to ensure that the API key is set, preferring the user-provided key | |
over the secret manager key (``DATAROBOT_API_TOKEN``). | |
If both are None, a fake API key is returned for testing. | |
""" | |
return api_key or get_secret_str("DATAROBOT_API_TOKEN") or "fake-api-key" | |
def _resolve_api_base(api_base: Optional[str] = None) -> Optional[str]: | |
"""Attempt to ensure that the API base is set, preferring the user-provided key | |
over the secret manager key (``DATAROBOT_ENDPOINT``). | |
If both are None, a default Llamafile server URL is returned. | |
See: https://github.com/Mozilla-Ocho/llamafile/blob/bd1bbe9aabb1ee12dbdcafa8936db443c571eb9d/README.md#L61 | |
""" | |
api_base = api_base or get_secret_str("DATAROBOT_ENDPOINT") | |
if api_base is None: | |
api_base = "https://app.datarobot.com" | |
# If the api_base is a deployment URL, we do not append the chat completions path | |
if "api/v2/deployments" not in api_base: | |
# If the api_base is not a deployment URL, we need to append the chat completions path | |
if "api/v2/genai/llmgw/chat/completions" not in api_base: | |
api_base += "/api/v2/genai/llmgw/chat/completions" | |
# Ensure the url ends with a trailing slash | |
if not api_base.endswith("/"): | |
api_base += "/" | |
return api_base # type: ignore | |
def _get_openai_compatible_provider_info( | |
self, | |
api_base: Optional[str], | |
api_key: Optional[str] | |
) -> Tuple[Optional[str], Optional[str]]: | |
"""Attempts to ensure that the API base and key are set, preferring user-provided values, | |
before falling back to secret manager values (``DATAROBOT_ENDPOINT`` and ``DATAROBOT_API_TOKEN`` | |
respectively). | |
If an API key cannot be resolved via either method, a fake key is returned. | |
""" | |
api_base = DataRobotConfig._resolve_api_base(api_base) | |
dynamic_api_key = DataRobotConfig._resolve_api_key(api_key) | |
return api_base, dynamic_api_key | |
def get_complete_url( | |
self, | |
api_base: Optional[str], | |
api_key: Optional[str], | |
model: str, | |
optional_params: dict, | |
litellm_params: dict, | |
stream: Optional[bool] = None, | |
) -> str: | |
""" | |
Get the complete URL for the API call. Datarobot's API base is set to | |
the complete value, so it does not need to be updated to additionally add | |
chat completions. | |
Returns: | |
str: The complete URL for the API call. | |
""" | |
return str(api_base) # type: ignore | |