fix: add missing _parse_response to AzureOpenAIStructuredLLM (#4434)

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mem0/llms/azure_openai_structured.py

@@ -1,3 +1,4 @@
+import json
 import os
 from typing import Dict, List, Optional
 
@@ -6,6 +7,7 @@ from openai import AzureOpenAI
 
 from mem0.configs.llms.base import BaseLlmConfig
 from mem0.llms.base import LLMBase
+from mem0.memory.utils import extract_json
 
 SCOPE = "https://cognitiveservices.azure.com/.default"
 
@@ -83,9 +85,35 @@ class AzureOpenAIStructuredLLM(LLMBase):
             params["tools"] = tools
             params["tool_choice"] = tool_choice
 
-        if tools:
-            params["tools"] = tools
-            params["tool_choice"] = tool_choice
-
         response = self.client.chat.completions.create(**params)
         return self._parse_response(response, tools)
+
+    def _parse_response(self, response, tools):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: The raw response from API.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if tools:
+            processed_response = {
+                "content": response.choices[0].message.content,
+                "tool_calls": [],
+            }
+
+            if response.choices[0].message.tool_calls:
+                for tool_call in response.choices[0].message.tool_calls:
+                    processed_response["tool_calls"].append(
+                        {
+                            "name": tool_call.function.name,
+                            "arguments": json.loads(extract_json(tool_call.function.arguments)),
+                        }
+                    )
+
+            return processed_response
+        else:
+            return response.choices[0].message.content

tests/llms/test_azure_openai_structured.py

@@ -1,4 +1,5 @@
 from unittest import mock
+from unittest.mock import Mock
 
 from mem0.llms.azure_openai_structured import SCOPE, AzureOpenAIStructuredLLM
 
@@ -98,3 +99,75 @@ def test_init_with_placeholder_api_key_uses_default_credential(
         args, kwargs = mock_azure_openai.call_args
         assert kwargs["api_key"] is None
         assert kwargs["azure_ad_token_provider"] == "token-provider"
+
+
+@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
+def test_generate_response_without_tools(mock_azure_openai):
+    mock_client = Mock()
+    mock_azure_openai.return_value = mock_client
+
+    config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
+    llm = AzureOpenAIStructuredLLM(config)
+
+    mock_response = Mock()
+    mock_response.choices = [Mock(message=Mock(content="Hello there!"))]
+    mock_client.chat.completions.create.return_value = mock_response
+
+    messages = [{"role": "user", "content": "Hi"}]
+    response = llm.generate_response(messages)
+
+    assert response == "Hello there!"
+
+
+@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
+def test_generate_response_with_tools(mock_azure_openai):
+    mock_client = Mock()
+    mock_azure_openai.return_value = mock_client
+
+    config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
+    llm = AzureOpenAIStructuredLLM(config)
+
+    mock_tool_call = Mock()
+    mock_tool_call.function.name = "add_memory"
+    mock_tool_call.function.arguments = '{"data": "sunny day"}'
+
+    mock_message = Mock()
+    mock_message.content = "I've added the memory."
+    mock_message.tool_calls = [mock_tool_call]
+
+    mock_response = Mock()
+    mock_response.choices = [Mock(message=mock_message)]
+    mock_client.chat.completions.create.return_value = mock_response
+
+    tools = [{"type": "function", "function": {"name": "add_memory"}}]
+    messages = [{"role": "user", "content": "Remember sunny day"}]
+    response = llm.generate_response(messages, tools=tools)
+
+    assert response["content"] == "I've added the memory."
+    assert len(response["tool_calls"]) == 1
+    assert response["tool_calls"][0]["name"] == "add_memory"
+    assert response["tool_calls"][0]["arguments"] == {"data": "sunny day"}
+
+
+@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI")
+def test_generate_response_with_tools_no_tool_calls(mock_azure_openai):
+    mock_client = Mock()
+    mock_azure_openai.return_value = mock_client
+
+    config = DummyConfig(model="test-model", azure_kwargs=DummyAzureKwargs(api_key="real-key"))
+    llm = AzureOpenAIStructuredLLM(config)
+
+    mock_message = Mock()
+    mock_message.content = "No tools needed."
+    mock_message.tool_calls = None
+
+    mock_response = Mock()
+    mock_response.choices = [Mock(message=mock_message)]
+    mock_client.chat.completions.create.return_value = mock_response
+
+    tools = [{"type": "function", "function": {"name": "add_memory"}}]
+    messages = [{"role": "user", "content": "Hello"}]
+    response = llm.generate_response(messages, tools=tools)
+
+    assert response["content"] == "No tools needed."
+    assert response["tool_calls"] == []