From 962fd672279fa2c0220309cc1296bd7951724c23 Mon Sep 17 00:00:00 2001 From: Juan Cruz-Benito Date: Mon, 23 Jun 2025 19:14:10 +0200 Subject: [PATCH] Fixing watsonx error: 'model_id' or 'model' cannot be specified in the request body for models in a deployment space (#11854) * Fixing watsonx error: 'model_id' or 'model' cannot be specified in the request body for models in a deployment space * Revert "Fixing watsonx error: 'model_id' or 'model' cannot be specified in the request body for models in a deployment space" This reverts commit 9d16a3000b24d255f70ab74c01eb6267354df151. * Implementing feedback from code review --- litellm/llms/watsonx/chat/handler.py | 4 +- litellm/llms/watsonx/chat/transformation.py | 14 +- .../test_litellm/llms/watsonx/test_watsonx.py | 205 ++++++++++++++++++ 3 files changed, 220 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/llms/watsonx/test_watsonx.py diff --git a/litellm/llms/watsonx/chat/handler.py b/litellm/llms/watsonx/chat/handler.py index 45378c55292..5c19757fecb 100644 --- a/litellm/llms/watsonx/chat/handler.py +++ b/litellm/llms/watsonx/chat/handler.py @@ -52,7 +52,7 @@ class WatsonXChatHandler(OpenAILikeChatHandler): litellm_params=litellm_params, ) - ## UPDATE PAYLOAD (optional params) + ## UPDATE PAYLOAD (optional params and special cases for models deployed in spaces) watsonx_auth_payload = watsonx_chat_transformation._prepare_payload( model=model, api_params=api_params, @@ -70,7 +70,7 @@ class WatsonXChatHandler(OpenAILikeChatHandler): ) return super().completion( - model=model, + model=watsonx_auth_payload.get("model_id", None), messages=messages, api_base=api_base, custom_llm_provider=custom_llm_provider, diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 3c2d1c6f0bf..71d8bba4ef4 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -7,7 +7,7 @@ Docs: https://cloud.ibm.com/apidocs/watsonx-ai#text-chat from typing import List, Optional, Tuple, Union from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.watsonx import WatsonXAIEndpoint +from litellm.types.llms.watsonx import WatsonXAIEndpoint, WatsonXAPIParams from ....utils import _remove_additional_properties, _remove_strict_from_schema from ...openai.chat.gpt_transformation import OpenAIGPTConfig @@ -108,3 +108,15 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): url=url, api_version=optional_params.pop("api_version", None) ) return url + + def _prepare_payload(self, model: str, api_params: WatsonXAPIParams) -> dict: + """ + Prepare payload for deployment models. + Deployment models cannot have 'model_id' or 'model' in the request body. + """ + payload: dict = {} + payload["model_id"] = None if model.startswith("deployment/") else model + payload["project_id"] = ( + None if model.startswith("deployment/") else api_params["project_id"] + ) + return payload diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py new file mode 100644 index 00000000000..13ca6ad8e6c --- /dev/null +++ b/tests/test_litellm/llms/watsonx/test_watsonx.py @@ -0,0 +1,205 @@ +import json +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm import completion +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from unittest.mock import patch, Mock +import pytest +from typing import Optional + + +@pytest.fixture +def watsonx_chat_completion_call(): + def _call( + model="watsonx/my-test-model", + messages=None, + api_key="test_api_key", + space_id: Optional[str] = None, + headers=None, + client=None, + patch_token_call=True, + ): + if messages is None: + messages = [{"role": "user", "content": "Hello, how are you?"}] + if client is None: + client = HTTPHandler() + + if patch_token_call: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "mock_access_token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() # No-op to simulate no exception + + with patch.object(client, "post") as mock_post, patch.object( + litellm.module_level_client, "post", return_value=mock_response + ) as mock_get: + try: + completion( + model=model, + messages=messages, + api_key=api_key, + headers=headers or {}, + client=client, + space_id=space_id, + ) + except Exception as e: + print(e) + + return mock_post, mock_get + else: + with patch.object(client, "post") as mock_post: + try: + completion( + model=model, + messages=messages, + api_key=api_key, + headers=headers or {}, + client=client, + space_id=space_id, + ) + except Exception as e: + print(e) + return mock_post, None + + return _call + + +def test_watsonx_deployment_model_id_not_in_payload( + monkeypatch, watsonx_chat_completion_call +): + """Test that deployment models do not include 'model_id' in the request payload""" + monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") + monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") + model = "watsonx/deployment/test-deployment-id" + messages = [{"role": "user", "content": "Test message"}] + + mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) + + assert mock_post.call_count == 1 + json_data = json.loads(mock_post.call_args.kwargs["data"]) + # Ensure model_id is not in the payload for deployment models + assert "model_id" not in json_data or json_data["model_id"] is None + # Ensure project_id is also not in the payload for deployment models + assert "project_id" not in json_data or json_data["project_id"] is None + + +def test_watsonx_regular_model_includes_model_id( + monkeypatch, watsonx_chat_completion_call +): + """Test that regular models include 'model_id' in the request payload""" + monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") + monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") + model = "watsonx/regular-model" + messages = [{"role": "user", "content": "Test message"}] + + mock_post, _ = watsonx_chat_completion_call(model=model, messages=messages) + + assert mock_post.call_count == 1 + json_data = json.loads(mock_post.call_args.kwargs["data"]) + # Ensure model_id is included in the payload for regular models + assert "model_id" in json_data + assert json_data["model_id"] == "regular-model" # Provider prefix is stripped + # Ensure project_id is also included for regular models + assert "project_id" in json_data + + +@pytest.fixture +def watsonx_completion_call(): + def _call( + model="watsonx_text/my-test-model", + prompt="Hello, how are you?", + api_key="test_api_key", + space_id: Optional[str] = None, + headers=None, + client=None, + patch_token_call=True, + ): + if client is None: + client = HTTPHandler() + + if patch_token_call: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "mock_access_token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + + with patch.object(client, "post") as mock_post, patch.object( + litellm.module_level_client, "post", return_value=mock_response + ) as mock_get: + try: + litellm.text_completion( + model=model, + prompt=prompt, + api_key=api_key, + headers=headers or {}, + client=client, + space_id=space_id, + ) + except Exception as e: + print(e) + + return mock_post, mock_get + else: + with patch.object(client, "post") as mock_post: + try: + litellm.text_completion( + model=model, + prompt=prompt, + api_key=api_key, + headers=headers or {}, + client=client, + space_id=space_id, + ) + except Exception as e: + print(e) + return mock_post, None + + return _call + + +def test_watsonx_completion_deployment_model_id_not_in_payload( + monkeypatch, watsonx_completion_call +): + """Test that deployment models do not include 'model_id' in completion request payload""" + monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") + monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") + model = "watsonx_text/deployment/test-deployment-id" + prompt = "Test prompt" + + mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) + + assert mock_post.call_count == 1 + json_data = json.loads(mock_post.call_args.kwargs["data"]) + # Ensure model_id is not in the payload for deployment models + assert "model_id" not in json_data + # Ensure project_id is also not in the payload for deployment models + assert "project_id" not in json_data + + +def test_watsonx_completion_regular_model_includes_model_id( + monkeypatch, watsonx_completion_call +): + """Test that regular models include 'model_id' in completion request payload""" + monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") + monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") + model = "watsonx_text/regular-model" + prompt = "Test prompt" + + mock_post, _ = watsonx_completion_call(model=model, prompt=prompt) + + assert mock_post.call_count == 1 + json_data = json.loads(mock_post.call_args.kwargs["data"]) + # Ensure model_id is included in the payload for regular models + assert "model_id" in json_data + assert json_data["model_id"] == "regular-model" # Provider prefix is stripped + # Ensure project_id is also included for regular models + assert "project_id" in json_data