scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

add reasoning effort and test (#16261)

Browse source
This commit is contained in:
Alan Ponnachan 2025-11-05 08:23:02 +05:30 • committed by GitHub
parent 4a83ae0695
commit b564ed81f5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 126 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/llms/watsonx/chat/transformation.py
View file

@ -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
)

133
tests/test_litellm/llms/watsonx/test_watsonx.py
View file

@ -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'."