fix(llms): honor reasoning-model params in AzureOpenAIStructuredLLM (#5548)

Co-authored-by: ly-wang19 <ly-wang19@users.noreply.github.com>
This commit is contained in:
ly-wang19
2026-06-15 16:07:19 +05:30
committed by GitHub
co-authored by ly-wang19
parent 09a9e34382
commit a8eeddde64
2 changed files with 68 additions and 3 deletions
+13 -3
View File
@@ -72,16 +72,26 @@ class AzureOpenAIStructuredLLM(LLMBase):
messages[-1]["content"] = user_prompt messages[-1]["content"] = user_prompt
is_reasoning = self._is_reasoning_model(self.config.model)
params = { params = {
"model": self.config.model, "model": self.config.model,
"messages": messages, "messages": messages,
"temperature": self.config.temperature,
"top_p": self.config.top_p,
} }
if self._uses_max_completion_tokens(self.config.model): # Reasoning models (o1/o3/GPT-5 series) reject temperature/top_p; only
# forward the sampling params for non-reasoning models. Mirrors the
# reasoning-aware handling of OpenAIStructuredLLM (#5458).
if not is_reasoning:
params["temperature"] = self.config.temperature
params["top_p"] = self.config.top_p
# Reasoning models require max_completion_tokens rather than max_tokens.
if is_reasoning or self._uses_max_completion_tokens(self.config.model):
params["max_completion_tokens"] = self.config.max_tokens params["max_completion_tokens"] = self.config.max_tokens
else: else:
params["max_tokens"] = self.config.max_tokens params["max_tokens"] = self.config.max_tokens
if is_reasoning:
reasoning_effort = getattr(self.config, "reasoning_effort", None)
if reasoning_effort:
params["reasoning_effort"] = reasoning_effort
if response_format: if response_format:
params["response_format"] = response_format params["response_format"] = response_format
if tools: if tools:
@@ -29,6 +29,7 @@ class DummyConfig:
max_tokens=256, max_tokens=256,
top_p=1.0, top_p=1.0,
http_client=None, http_client=None,
reasoning_effort=None,
): ):
self.model = model self.model = model
self.azure_kwargs = azure_kwargs or DummyAzureKwargs() self.azure_kwargs = azure_kwargs or DummyAzureKwargs()
@@ -36,6 +37,7 @@ class DummyConfig:
self.max_tokens = max_tokens self.max_tokens = max_tokens
self.top_p = top_p self.top_p = top_p
self.http_client = http_client self.http_client = http_client
self.reasoning_effort = reasoning_effort
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI") @mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
@@ -209,3 +211,56 @@ def test_generate_response_with_tools_no_tool_calls(mock_azure_openai):
assert response["content"] == "No tools needed." assert response["content"] == "No tools needed."
assert response["tool_calls"] == [] assert response["tool_calls"] == []
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_reasoning_model_drops_sampling_params(mock_azure_openai):
"""Reasoning models (o1/o3/GPT-5) reject temperature/max_tokens/top_p."""
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(
model="o3-mini",
azure_kwargs=DummyAzureKwargs(api_key="real-key"),
reasoning_effort="low",
)
llm = AzureOpenAIStructuredLLM(config)
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="ok"))]
mock_client.chat.completions.create.return_value = mock_response
llm.generate_response([{"role": "user", "content": "Hi"}])
call_kwargs = mock_client.chat.completions.create.call_args[1]
assert "temperature" not in call_kwargs # reasoning models reject these
assert "max_tokens" not in call_kwargs
assert "top_p" not in call_kwargs
assert call_kwargs["reasoning_effort"] == "low"
assert call_kwargs["model"] == "o3-mini"
@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
def test_regular_model_sends_sampling_params(mock_azure_openai):
"""Regular models still receive the standard sampling params."""
mock_client = Mock()
mock_azure_openai.return_value = mock_client
config = DummyConfig(
model="gpt-4o",
azure_kwargs=DummyAzureKwargs(api_key="real-key"),
temperature=0.3,
)
llm = AzureOpenAIStructuredLLM(config)
mock_response = Mock()
mock_response.choices = [Mock(message=Mock(content="ok"))]
mock_client.chat.completions.create.return_value = mock_response
llm.generate_response([{"role": "user", "content": "Hi"}])
call_kwargs = mock_client.chat.completions.create.call_args[1]
assert call_kwargs["temperature"] == 0.3
assert "max_tokens" in call_kwargs # standard sampling params still forwarded
assert "top_p" in call_kwargs
assert call_kwargs["model"] == "gpt-4o"