diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 2c096cafced..865dc71939d 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -35,6 +35,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): "n", "presence_penalty", "response_format", + "reasoning_effort", ] def is_tool_choice_option(self, tool_choice: Optional[Union[str, dict]]) -> bool: @@ -124,16 +125,18 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): None if model.startswith("deployment/") else api_params["project_id"] ) return payload - + @staticmethod - def _apply_prompt_template_core(model: str, messages: List[Dict[str, str]], hf_template_fn) -> Optional[str]: + def _apply_prompt_template_core( + model: str, messages: List[Dict[str, str]], hf_template_fn + ) -> Optional[str]: """Core logic for applying prompt templates""" from litellm.litellm_core_utils.prompt_templates.factory import ( custom_prompt, ibm_granite_pt, mistral_instruct_pt, ) - + if WatsonXModelPattern.GRANITE_CHAT.value in model: return ibm_granite_pt(messages=messages) elif WatsonXModelPattern.IBM_MISTRAL.value in model: @@ -147,9 +150,18 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): elif WatsonXModelPattern.LLAMA3_INSTRUCT.value in model: return custom_prompt( role_dict={ - "system": {"pre_message": "<|start_header_id|>system<|end_header_id|>\n", "post_message": "<|eot_id|>"}, - "user": {"pre_message": "<|start_header_id|>user<|end_header_id|>\n", "post_message": "<|eot_id|>"}, - "assistant": {"pre_message": "<|start_header_id|>assistant<|end_header_id|>\n", "post_message": "<|eot_id|>"}, + "system": { + "pre_message": "<|start_header_id|>system<|end_header_id|>\n", + "post_message": "<|eot_id|>", + }, + "user": { + "pre_message": "<|start_header_id|>user<|end_header_id|>\n", + "post_message": "<|eot_id|>", + }, + "assistant": { + "pre_message": "<|start_header_id|>assistant<|end_header_id|>\n", + "post_message": "<|eot_id|>", + }, }, messages=messages, initial_prompt_value="<|begin_of_text|>", @@ -158,7 +170,9 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): return None @staticmethod - async def aapply_prompt_template(model: str, messages: List[Dict[str, str]]) -> Optional[str]: + async def aapply_prompt_template( + model: str, messages: List[Dict[str, str]] + ) -> Optional[str]: """Apply prompt template (async version)""" import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -204,9 +218,11 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): final_prompt_value="<|start_header_id|>assistant<|end_header_id|>\n", ) return None - + @staticmethod - def apply_prompt_template(model: str, messages: List[Dict[str, str]]) -> Optional[str]: + def apply_prompt_template( + model: str, messages: List[Dict[str, str]] + ) -> Optional[str]: """Apply prompt template (sync version)""" from litellm.litellm_core_utils.prompt_templates.factory import ( hf_chat_template, @@ -215,4 +231,3 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): return IBMWatsonXChatConfig._apply_prompt_template_core( model=model, messages=messages, hf_template_fn=hf_chat_template ) - diff --git a/tests/test_litellm/llms/watsonx/test_watsonx.py b/tests/test_litellm/llms/watsonx/test_watsonx.py index b76d7e8a2be..a41316bb47e 100644 --- a/tests/test_litellm/llms/watsonx/test_watsonx.py +++ b/tests/test_litellm/llms/watsonx/test_watsonx.py @@ -211,28 +211,28 @@ def test_watsonx_completion_regular_model_includes_model_id( async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): """ Test that gpt-oss-120b model transforms messages to proper format instead of simple concatenation. - + This test starts from litellm.acompletion and verifies what gets sent in the final POST request body. Input messages should be transformed using the HuggingFace chat template from openai/gpt-oss-120b, not just concatenated as "You are chatgpt Hi there". """ monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") - + # Test with gpt-oss model using watsonx_text provider (text generation endpoint) model = "watsonx_text/openai/gpt-oss-120b" - + # Input messages messages = [ {"role": "system", "content": "You are chatgpt"}, - {"role": "user", "content": "Hi there"} + {"role": "user", "content": "Hi there"}, ] - + # Mock the HTTP client from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - + client = AsyncHTTPHandler() - + # Mock the token call mock_token_response = Mock() mock_token_response.json.return_value = { @@ -240,25 +240,27 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): "expires_in": 3600, } mock_token_response.raise_for_status = Mock() - + # Mock the completion call mock_completion_response = Mock() mock_completion_response.status_code = 200 mock_completion_response.json.return_value = { - "results": [{ - "generated_text": "Hello! How can I help you?", - "generated_token_count": 10, - "input_token_count": 5 - }], - "model_id": "openai/gpt-oss-120b" + "results": [ + { + "generated_text": "Hello! How can I help you?", + "generated_token_count": 10, + "input_token_count": 5, + } + ], + "model_id": "openai/gpt-oss-120b", } - + with patch.object(client, "post") as mock_post, patch.object( litellm.module_level_client, "post", return_value=mock_token_response ): # Set the mock to return the completion response mock_post.return_value = mock_completion_response - + try: # Call acompletion with messages await litellm.acompletion( @@ -270,14 +272,16 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): except Exception as e: # May fail due to incomplete mocking, but we should have captured the request print(f"Exception (may be expected): {e}") - + # Verify the POST was called - assert mock_post.call_count >= 1, f"POST should have been called at least once, got {mock_post.call_count}" - + assert ( + mock_post.call_count >= 1 + ), f"POST should have been called at least once, got {mock_post.call_count}" + # Get the request body from the first call call_args = mock_post.call_args json_data = json.loads(call_args.kwargs["data"]) - + print(f"\n{'='*80}") print(f"Input messages to litellm.acompletion:") print(json.dumps(messages, indent=2)) @@ -285,14 +289,14 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): print(f"Final POST request body:") print(json.dumps(json_data, indent=2)) print(f"{'='*80}\n") - + # Verify the transformed input is in the request assert "input" in json_data, "Request should have 'input' field" transformed_prompt = json_data["input"] - + print(f"Transformed prompt: {repr(transformed_prompt)}") print(f"Prompt length: {len(transformed_prompt)}") - + # Verify it's NOT simple concatenation simple_concat = "You are chatgpt Hi there" assert transformed_prompt != simple_concat, ( @@ -300,13 +304,17 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch): f"Expected: Chat template with <|start|> tags\n" f"Got: {transformed_prompt}" ) - + # Verify it contains proper chat template formatting assert "<|start|>" in transformed_prompt, "Prompt should contain <|start|> tag" assert "<|message|>" in transformed_prompt, "Prompt should contain <|message|> tag" assert "<|end|>" in transformed_prompt, "Prompt should contain <|end|> tag" - assert "You are chatgpt" in transformed_prompt, "Prompt should contain system message content" - assert "Hi there" in transformed_prompt, "Prompt should contain user message content" + assert ( + "You are chatgpt" in transformed_prompt + ), "Prompt should contain system message content" + assert ( + "Hi there" in transformed_prompt + ), "Prompt should contain user message content" @pytest.mark.asyncio @@ -324,24 +332,85 @@ async def test_watsonx_gpt_oss_uses_async_http_handler(): mock_async_client = MagicMock() mock_get = AsyncMock() mock_async_client.get = mock_get - + # Create mock response for chat template file mock_response = MagicMock() mock_response.status_code = 200 mock_response.content = b"test template content" mock_get.return_value = mock_response - + # Test the async function directly - with patch("litellm.litellm_core_utils.prompt_templates.huggingface_template_handler.get_async_httpx_client", return_value=mock_async_client): + with patch( + "litellm.litellm_core_utils.prompt_templates.huggingface_template_handler.get_async_httpx_client", + return_value=mock_async_client, + ): result = await _aget_chat_template_file(hf_model_name="test/model") - + # Verify async HTTP client was called assert mock_get.called, "Async HTTP client's get method should be called" assert mock_get.await_count > 0, "Async HTTP client's get should be awaited" - + # Verify it was called with HuggingFace URL call_args = mock_get.call_args assert call_args is not None, "get should have been called with arguments" called_url = call_args.kwargs.get("url", "") - assert "huggingface.co/test/model" in called_url, f"Should call HuggingFace API for test/model, got: {called_url}" + assert ( + "huggingface.co/test/model" in called_url + ), f"Should call HuggingFace API for test/model, got: {called_url}" assert result["status"] == "success", "Should return success status" + + +def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch): + """ + Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload. + """ + monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id") + monkeypatch.setenv("WATSONX_API_BASE", "https://test-api.watsonx.ai") + + model = "watsonx/openai/gpt-oss-120b" + messages = [{"role": "user", "content": "Test message"}] + + client = HTTPHandler() + + # Mock the token generation call + mock_token_response = Mock() + mock_token_response.json.return_value = { + "access_token": "mock_access_token", + "expires_in": 3600, + } + mock_token_response.raise_for_status = Mock() + + # Call litellm.completion with the new parameter + with patch.object(client, "post") as mock_post, patch.object( + litellm.module_level_client, "post", return_value=mock_token_response + ): + try: + completion( + model=model, + messages=messages, + api_key="test_api_key", + client=client, + reasoning_effort="low", + ) + except Exception as e: + print(f"Caught expected exception: {e}") + + # Verify the parameter is in the final request payload + assert ( + mock_post.call_count == 1 + ), "The completion endpoint should have been called once." + + # Get the JSON data sent in the POST request + request_kwargs = mock_post.call_args.kwargs + json_data = json.loads(request_kwargs["data"]) + + print("\nRequest payload sent to WatsonX API:") + print(json.dumps(json_data, indent=2)) + + # Check for the parameter at the top level of the payload + assert ( + "reasoning_effort" in json_data + ), "'reasoning_effort' should be at the top level of the payload." + assert ( + json_data["reasoning_effort"] == "low" + ), "The value of 'reasoning_effort' should be 'low'."