diff --git a/litellm/llms/datarobot/chat/transformation.py b/litellm/llms/datarobot/chat/transformation.py index e334c94e51..23ce63c25b 100644 --- a/litellm/llms/datarobot/chat/transformation.py +++ b/litellm/llms/datarobot/chat/transformation.py @@ -6,8 +6,11 @@ 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 urllib.parse import urlparse, urlunparse from ...openai_like.chat.transformation import OpenAILikeChatConfig +LLMGW_PATH = "/genai/llmgw/chat/completions" + class DataRobotConfig(OpenAILikeChatConfig): @staticmethod @@ -32,22 +35,28 @@ class DataRobotConfig(OpenAILikeChatConfig): 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" + parsed = urlparse(api_base) + path = parsed.path + + if not path or path == "/": # Add full path to LLMGW + path += f"/api/v2/{LLMGW_PATH}" + elif "api/v2/deployments" in path: # Dedicated deployment, leave it + pass + elif ( + "api/v2" in path and LLMGW_PATH not in path + ): # Standard ENDPOINT path, add LLMGW + path += LLMGW_PATH # Ensure the url ends with a trailing slash - if not api_base.endswith("/"): - api_base += "/" + if not path.endswith("/"): + path += "/" + path = path.replace("//", "/") + updated_parsed = parsed._replace(path=path) - return api_base # type: ignore + return urlunparse(updated_parsed) def _get_openai_compatible_provider_info( - self, - api_base: Optional[str], - api_key: Optional[str] + 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`` diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py index 4999b078ca..1230a1fd2a 100644 --- a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py +++ b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py @@ -18,6 +18,8 @@ class TestDataRobotConfig: (None, "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), ("http://localhost:5001", "http://localhost:5001/api/v2/genai/llmgw/chat/completions/"), ("https://app.datarobot.com", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), + ("https://app.datarobot.com/api/v2/", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), + ("https://app.datarobot.com/api/v2", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), ("https://app.datarobot.com/api/v2/genai/llmgw/chat/completions", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), ("https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/", "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"), ("https://staging.datarobot.com", "https://staging.datarobot.com/api/v2/genai/llmgw/chat/completions/"),