From 3ffea554bc67d26efc6d96094005e58ac74d3d1f Mon Sep 17 00:00:00 2001 From: Rakhee Singh Date: Wed, 8 Apr 2026 19:22:02 +0530 Subject: [PATCH] fix(azure_openai): forward response_format to Azure OpenAI API (#4689) --- mem0/llms/azure_openai.py | 2 ++ tests/llms/test_azure_openai.py | 44 +++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/mem0/llms/azure_openai.py b/mem0/llms/azure_openai.py index 3abd7de49..7a088a17c 100644 --- a/mem0/llms/azure_openai.py +++ b/mem0/llms/azure_openai.py @@ -134,6 +134,8 @@ class AzureOpenAILLM(LLMBase): "messages": messages, }) + if response_format: + params["response_format"] = response_format if tools: params["tools"] = tools params["tool_choice"] = tool_choice diff --git a/tests/llms/test_azure_openai.py b/tests/llms/test_azure_openai.py index c10dc1635..ee6c01d7c 100644 --- a/tests/llms/test_azure_openai.py +++ b/tests/llms/test_azure_openai.py @@ -91,6 +91,50 @@ def test_generate_response_with_tools(mock_openai_client): assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."} +def test_generate_response_with_response_format(mock_openai_client): + config = AzureOpenAIConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P) + llm = AzureOpenAILLM(config) + messages = [ + {"role": "system", "content": "You are a memory extraction assistant."}, + {"role": "user", "content": "I like hiking on weekends."}, + ] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content='{"facts": ["User likes hiking on weekends"]}'))] + mock_openai_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages, response_format={"type": "json_object"}) + + mock_openai_client.chat.completions.create.assert_called_once_with( + model=MODEL, + messages=messages, + temperature=TEMPERATURE, + max_tokens=MAX_TOKENS, + top_p=TOP_P, + response_format={"type": "json_object"}, + ) + assert response == '{"facts": ["User likes hiking on weekends"]}' + + +def test_generate_response_without_response_format(mock_openai_client): + config = AzureOpenAIConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P) + llm = AzureOpenAILLM(config) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell me a joke."}, + ] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="Why did the chicken cross the road?"))] + mock_openai_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages) + + call_kwargs = mock_openai_client.chat.completions.create.call_args[1] + assert "response_format" not in call_kwargs + assert response == "Why did the chicken cross the road?" + + def test_reasoning_model_with_reasoning_effort(mock_openai_client): """Test that reasoning_effort is passed to the API for Azure reasoning models.""" config = AzureOpenAIConfig(model="o3-mini", reasoning_effort="low")