test

# What does this PR do?


## Test Plan

tests/unit/providers/agents/meta_reference/test_openai_responses.py

@@ -146,14 +146,17 @@ async def test_create_openai_response_with_string_input(openai_responses_impl, m
     # For streaming response, collect all chunks
     chunks = [chunk async for chunk in result]
 
-    mock_inference_api.openai_chat_completion.assert_called_once_with(
-        model=model,
-        messages=[OpenAIUserMessageParam(role="user", content="What is the capital of Ireland?", name=None)],
-        response_format=None,
-        tools=None,
-        stream=True,
-        temperature=0.1,
-    )
+    # Verify the inference API was called with the correct params
+    call_args = mock_inference_api.openai_chat_completion.call_args
+    params = call_args.args[0]  # params is passed as first positional arg
+    assert params.model == model
+    assert params.messages == [
+        OpenAIUserMessageParam(role="user", content="What is the capital of Ireland?", name=None)
+    ]
+    assert params.response_format is None
+    assert params.tools is None
+    assert params.stream is True
+    assert params.temperature == 0.1
 
     # Should have content part events for text streaming
     # Expected: response.created, content_part.added, output_text.delta, content_part.done, response.completed
@@ -228,13 +231,15 @@ async def test_create_openai_response_with_string_input_with_tools(openai_respon
 
         # Verify
         first_call = mock_inference_api.openai_chat_completion.call_args_list[0]
-        assert first_call.kwargs["messages"][0].content == "What is the capital of Ireland?"
-        assert first_call.kwargs["tools"] is not None
-        assert first_call.kwargs["temperature"] == 0.1
+        first_params = first_call.args[0]
+        assert first_params.messages[0].content == "What is the capital of Ireland?"
+        assert first_params.tools is not None
+        assert first_params.temperature == 0.1
 
         second_call = mock_inference_api.openai_chat_completion.call_args_list[1]
-        assert second_call.kwargs["messages"][-1].content == "Dublin"
-        assert second_call.kwargs["temperature"] == 0.1
+        second_params = second_call.args[0]
+        assert second_params.messages[-1].content == "Dublin"
+        assert second_params.temperature == 0.1
 
         openai_responses_impl.tool_groups_api.get_tool.assert_called_once_with("web_search")
         openai_responses_impl.tool_runtime_api.invoke_tool.assert_called_once_with(
@@ -309,9 +314,10 @@ async def test_create_openai_response_with_tool_call_type_none(openai_responses_
 
     # Verify inference API was called correctly (after iterating over result)
     first_call = mock_inference_api.openai_chat_completion.call_args_list[0]
-    assert first_call.kwargs["messages"][0].content == input_text
-    assert first_call.kwargs["tools"] is not None
-    assert first_call.kwargs["temperature"] == 0.1
+    first_params = first_call.args[0]
+    assert first_params.messages[0].content == input_text
+    assert first_params.tools is not None
+    assert first_params.temperature == 0.1
 
     # Check response.created event (should have empty output)
     assert chunks[0].type == "response.created"
@@ -386,9 +392,10 @@ async def test_create_openai_response_with_tool_call_function_arguments_none(ope
 
     # Verify inference API was called correctly (after iterating over result)
     first_call = mock_inference_api.openai_chat_completion.call_args_list[0]
-    assert first_call.kwargs["messages"][0].content == input_text
-    assert first_call.kwargs["tools"] is not None
-    assert first_call.kwargs["temperature"] == 0.1
+    first_params = first_call.args[0]
+    assert first_params.messages[0].content == input_text
+    assert first_params.tools is not None
+    assert first_params.temperature == 0.1
 
     # Check response.created event (should have empty output)
     assert chunks[0].type == "response.created"
@@ -435,9 +442,10 @@ async def test_create_openai_response_with_tool_call_function_arguments_none(ope
 
     # Verify inference API was called correctly (after iterating over result)
     first_call = mock_inference_api.openai_chat_completion.call_args_list[0]
-    assert first_call.kwargs["messages"][0].content == input_text
-    assert first_call.kwargs["tools"] is not None
-    assert first_call.kwargs["temperature"] == 0.1
+    first_params = first_call.args[0]
+    assert first_params.messages[0].content == input_text
+    assert first_params.tools is not None
+    assert first_params.temperature == 0.1
 
     # Check response.created event (should have empty output)
     assert chunks[0].type == "response.created"
@@ -485,7 +493,9 @@ async def test_create_openai_response_with_multiple_messages(openai_responses_im
 
     # Verify the the correct messages were sent to the inference API i.e.
     # All of the responses message were convered to the chat completion message objects
-    inference_messages = mock_inference_api.openai_chat_completion.call_args_list[0].kwargs["messages"]
+    call_args = mock_inference_api.openai_chat_completion.call_args_list[0]
+    params = call_args.args[0]
+    inference_messages = params.messages
     for i, m in enumerate(input_messages):
         if isinstance(m.content, str):
             assert inference_messages[i].content == m.content
@@ -653,7 +663,8 @@ async def test_create_openai_response_with_instructions(openai_responses_impl, m
     # Verify
     mock_inference_api.openai_chat_completion.assert_called_once()
     call_args = mock_inference_api.openai_chat_completion.call_args
-    sent_messages = call_args.kwargs["messages"]
+    params = call_args.args[0]
+    sent_messages = params.messages
 
     # Check that instructions were prepended as a system message
     assert len(sent_messages) == 2
@@ -691,7 +702,8 @@ async def test_create_openai_response_with_instructions_and_multiple_messages(
     # Verify
     mock_inference_api.openai_chat_completion.assert_called_once()
     call_args = mock_inference_api.openai_chat_completion.call_args
-    sent_messages = call_args.kwargs["messages"]
+    params = call_args.args[0]
+    sent_messages = params.messages
 
     # Check that instructions were prepended as a system message
     assert len(sent_messages) == 4  # 1 system + 3 input messages
@@ -751,7 +763,8 @@ async def test_create_openai_response_with_instructions_and_previous_response(
     # Verify
     mock_inference_api.openai_chat_completion.assert_called_once()
     call_args = mock_inference_api.openai_chat_completion.call_args
-    sent_messages = call_args.kwargs["messages"]
+    params = call_args.args[0]
+    sent_messages = params.messages
 
     # Check that instructions were prepended as a system message
     assert len(sent_messages) == 4, sent_messages
@@ -987,8 +1000,9 @@ async def test_create_openai_response_with_text_format(
 
     # Verify
     first_call = mock_inference_api.openai_chat_completion.call_args_list[0]
-    assert first_call.kwargs["messages"][0].content == input_text
-    assert first_call.kwargs["response_format"] == response_format
+    first_params = first_call.args[0]
+    assert first_params.messages[0].content == input_text
+    assert first_params.response_format == response_format
 
 
 async def test_create_openai_response_with_invalid_text_format(openai_responses_impl, mock_inference_api):

tests/unit/providers/inference/test_remote_vllm.py

@@ -13,6 +13,7 @@ import pytest
 from llama_stack.apis.inference import (
     OpenAIAssistantMessageParam,
     OpenAIChatCompletion,
+    OpenaiChatCompletionRequest,
     OpenAIChoice,
     ToolChoice,
 )
@@ -56,13 +57,14 @@ async def test_old_vllm_tool_choice(vllm_inference_adapter):
         mock_client_property.return_value = mock_client
 
         # No tools but auto tool choice
-        await vllm_inference_adapter.openai_chat_completion(
-            "mock-model",
-            [],
+        params = OpenaiChatCompletionRequest(
+            model="mock-model",
+            messages=[{"role": "user", "content": "test"}],
             stream=False,
             tools=None,
             tool_choice=ToolChoice.auto.value,
         )
+        await vllm_inference_adapter.openai_chat_completion(params)
         mock_client.chat.completions.create.assert_called()
         call_args = mock_client.chat.completions.create.call_args
         # Ensure tool_choice gets converted to none for older vLLM versions
@@ -171,9 +173,12 @@ async def test_openai_chat_completion_is_async(vllm_inference_adapter):
         )
 
     async def do_inference():
-        await vllm_inference_adapter.openai_chat_completion(
-            "mock-model", messages=["one fish", "two fish"], stream=False
+        params = OpenaiChatCompletionRequest(
+            model="mock-model",
+            messages=[{"role": "user", "content": "one fish two fish"}],
+            stream=False,
         )
+        await vllm_inference_adapter.openai_chat_completion(params)
 
     with patch.object(VLLMInferenceAdapter, "client", new_callable=PropertyMock) as mock_create_client:
         mock_client = MagicMock()
@@ -186,3 +191,48 @@ async def test_openai_chat_completion_is_async(vllm_inference_adapter):
 
         assert mock_create_client.call_count == 4  # no cheating
         assert total_time < (sleep_time * 2), f"Total time taken: {total_time}s exceeded expected max"
+
+
+async def test_extra_body_forwarding(vllm_inference_adapter):
+    """
+    Test that extra_body parameters (e.g., chat_template_kwargs) are correctly
+    forwarded to the underlying OpenAI client.
+    """
+    mock_model = Model(identifier="mock-model", provider_resource_id="mock-model", provider_id="vllm-inference")
+    vllm_inference_adapter.model_store.get_model.return_value = mock_model
+
+    with patch.object(VLLMInferenceAdapter, "client", new_callable=PropertyMock) as mock_client_property:
+        mock_client = MagicMock()
+        mock_client.chat.completions.create = AsyncMock(
+            return_value=OpenAIChatCompletion(
+                id="chatcmpl-abc123",
+                created=1,
+                model="mock-model",
+                choices=[
+                    OpenAIChoice(
+                        message=OpenAIAssistantMessageParam(
+                            content="test response",
+                        ),
+                        finish_reason="stop",
+                        index=0,
+                    )
+                ],
+            )
+        )
+        mock_client_property.return_value = mock_client
+
+        # Test with chat_template_kwargs for Granite thinking mode
+        params = OpenaiChatCompletionRequest(
+            model="mock-model",
+            messages=[{"role": "user", "content": "test"}],
+            stream=False,
+            chat_template_kwargs={"thinking": True},
+        )
+        await vllm_inference_adapter.openai_chat_completion(params)
+
+        # Verify that the client was called with extra_body containing chat_template_kwargs
+        mock_client.chat.completions.create.assert_called_once()
+        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
+        assert "extra_body" in call_kwargs
+        assert "chat_template_kwargs" in call_kwargs["extra_body"]
+        assert call_kwargs["extra_body"]["chat_template_kwargs"] == {"thinking": True}

tests/unit/providers/utils/inference/test_openai_mixin.py

@@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, patch
 import pytest
 from pydantic import BaseModel, Field
 
-from llama_stack.apis.inference import Model, OpenAIUserMessageParam
+from llama_stack.apis.inference import Model, OpenaiChatCompletionRequest, OpenAIUserMessageParam
 from llama_stack.apis.models import ModelType
 from llama_stack.core.request_headers import request_provider_data_context
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
@@ -271,7 +271,8 @@ class TestOpenAIMixinImagePreprocessing:
             with patch("llama_stack.providers.utils.inference.openai_mixin.localize_image_content") as mock_localize:
                 mock_localize.return_value = (b"fake_image_data", "jpeg")
 
-                await mixin.openai_chat_completion(model="test-model", messages=[message])
+                params = OpenaiChatCompletionRequest(model="test-model", messages=[message])
+                await mixin.openai_chat_completion(params)
 
             mock_localize.assert_called_once_with("http://example.com/image.jpg")
 
@@ -303,7 +304,8 @@ class TestOpenAIMixinImagePreprocessing:
 
         with patch.object(type(mixin), "client", new_callable=PropertyMock, return_value=mock_client):
             with patch("llama_stack.providers.utils.inference.openai_mixin.localize_image_content") as mock_localize:
-                await mixin.openai_chat_completion(model="test-model", messages=[message])
+                params = OpenaiChatCompletionRequest(model="test-model", messages=[message])
+                await mixin.openai_chat_completion(params)
 
             mock_localize.assert_not_called()