diff --git a/mem0/llms/azure_openai.py b/mem0/llms/azure_openai.py index fbb69e48a..933cb9c22 100644 --- a/mem0/llms/azure_openai.py +++ b/mem0/llms/azure_openai.py @@ -1,3 +1,4 @@ +import copy import json import os from typing import Dict, List, Optional, Union @@ -69,6 +70,26 @@ class AzureOpenAILLM(LLMBase): default_headers=default_headers, ) + @staticmethod + def _rewrite_assistant_keyword(messages): + """ + Return a copy of ``messages`` with the word "assistant" replaced by "ai" + in the last message's textual content. + + Azure's content management policy can flag the literal word "assistant", + which makes ``add`` fail (see issue #2636). The rewrite targets that + trigger without mutating the caller's messages and without assuming the + content is a string, so multimodal (list) content passes through untouched. + """ + if not messages: + return messages + + messages = copy.deepcopy(messages) + last_content = messages[-1].get("content") + if isinstance(last_content, str): + messages[-1]["content"] = last_content.replace("assistant", "ai") + return messages + def _parse_response(self, response, tools): """ Process the response based on whether tools are used or not. @@ -121,14 +142,14 @@ class AzureOpenAILLM(LLMBase): str: The generated response. """ - user_prompt = messages[-1]["content"] - - user_prompt = user_prompt.replace("assistant", "ai") - - messages[-1]["content"] = user_prompt + # Azure's "Indirect Attacks" content filter can flag the literal word + # "assistant" in the prompt, so it is rewritten to "ai" before the request. + # Work on a copy so the caller's messages are left untouched and string-only + # content is handled without breaking multimodal (list) content. + messages = self._rewrite_assistant_keyword(messages) params = self._get_supported_params(messages=messages, **kwargs) - + # Add model and messages params.update({ "model": self.config.model, diff --git a/mem0/llms/azure_openai_structured.py b/mem0/llms/azure_openai_structured.py index e9eaf3471..66addde29 100644 --- a/mem0/llms/azure_openai_structured.py +++ b/mem0/llms/azure_openai_structured.py @@ -1,3 +1,4 @@ +import copy import json import os from typing import Dict, List, Optional @@ -66,11 +67,11 @@ class AzureOpenAIStructuredLLM(LLMBase): str: The generated response. """ - user_prompt = messages[-1]["content"] - - user_prompt = user_prompt.replace("assistant", "ai") - - messages[-1]["content"] = user_prompt + # Azure's "Indirect Attacks" content filter can flag the literal word + # "assistant" in the prompt, so it is rewritten to "ai" before the request. + # Work on a copy so the caller's messages are left untouched and string-only + # content is handled without breaking multimodal (list) content. + messages = self._rewrite_assistant_keyword(messages) is_reasoning = self._is_reasoning_model(self.config.model) params = { @@ -101,6 +102,26 @@ class AzureOpenAIStructuredLLM(LLMBase): response = self.client.chat.completions.create(**params) return self._parse_response(response, tools) + @staticmethod + def _rewrite_assistant_keyword(messages): + """ + Return a copy of ``messages`` with the word "assistant" replaced by "ai" + in the last message's textual content. + + Azure's content management policy can flag the literal word "assistant", + which makes ``add`` fail (see issue #2636). The rewrite targets that + trigger without mutating the caller's messages and without assuming the + content is a string, so multimodal (list) content passes through untouched. + """ + if not messages: + return messages + + messages = copy.deepcopy(messages) + last_content = messages[-1].get("content") + if isinstance(last_content, str): + messages[-1]["content"] = last_content.replace("assistant", "ai") + return messages + def _parse_response(self, response, tools): """ Process the response based on whether tools are used or not. diff --git a/tests/llms/test_azure_openai.py b/tests/llms/test_azure_openai.py index 37df2cd60..afb77d038 100644 --- a/tests/llms/test_azure_openai.py +++ b/tests/llms/test_azure_openai.py @@ -135,6 +135,52 @@ def test_generate_response_without_response_format(mock_openai_client): assert response == "Why did the chicken cross the road?" +def test_generate_response_does_not_mutate_caller_messages(mock_openai_client): + config = AzureOpenAIConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P) + llm = AzureOpenAILLM(config) + messages = [{"role": "user", "content": "my assistant helps me schedule meetings"}] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="ok"))] + mock_openai_client.chat.completions.create.return_value = mock_response + + llm.generate_response(messages) + + assert messages[-1]["content"] == "my assistant helps me schedule meetings" + + +def test_generate_response_rewrites_assistant_keyword_for_model_only(mock_openai_client): + config = AzureOpenAIConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P) + llm = AzureOpenAILLM(config) + messages = [{"role": "user", "content": "my assistant helps me"}] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="ok"))] + mock_openai_client.chat.completions.create.return_value = mock_response + + llm.generate_response(messages) + + sent_messages = mock_openai_client.chat.completions.create.call_args[1]["messages"] + assert sent_messages[-1]["content"] == "my ai helps me" + assert messages[-1]["content"] == "my assistant helps me" + + +def test_generate_response_handles_multimodal_content(mock_openai_client): + config = AzureOpenAIConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P) + llm = AzureOpenAILLM(config) + messages = [{"role": "user", "content": [{"type": "text", "text": "describe my assistant"}]}] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="ok"))] + mock_openai_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages) + + assert response == "ok" + sent_messages = mock_openai_client.chat.completions.create.call_args[1]["messages"] + assert sent_messages[-1]["content"] == [{"type": "text", "text": "describe my assistant"}] + + 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") diff --git a/tests/llms/test_azure_openai_structured.py b/tests/llms/test_azure_openai_structured.py index e910ff073..8af8384f7 100644 --- a/tests/llms/test_azure_openai_structured.py +++ b/tests/llms/test_azure_openai_structured.py @@ -264,3 +264,61 @@ def test_regular_model_sends_sampling_params(mock_azure_openai): assert "max_tokens" in call_kwargs # standard sampling params still forwarded assert "top_p" in call_kwargs assert call_kwargs["model"] == "gpt-4o" + + +@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI") +def test_generate_response_does_not_mutate_caller_messages(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="ok"))] + mock_client.chat.completions.create.return_value = mock_response + + messages = [{"role": "user", "content": "my assistant helps me schedule meetings"}] + llm.generate_response(messages) + + assert messages[-1]["content"] == "my assistant helps me schedule meetings" + + +@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI") +def test_generate_response_rewrites_assistant_keyword_for_model_only(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="ok"))] + mock_client.chat.completions.create.return_value = mock_response + + messages = [{"role": "user", "content": "my assistant helps me"}] + llm.generate_response(messages) + + sent_messages = mock_client.chat.completions.create.call_args[1]["messages"] + assert sent_messages[-1]["content"] == "my ai helps me" + assert messages[-1]["content"] == "my assistant helps me" + + +@mock.patch("mem0.llms.azure_openai_structured.AzureOpenAI") +def test_generate_response_handles_multimodal_content(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="ok"))] + mock_client.chat.completions.create.return_value = mock_response + + messages = [{"role": "user", "content": [{"type": "text", "text": "describe my assistant"}]}] + response = llm.generate_response(messages) + + assert response == "ok" + sent_messages = mock_client.chat.completions.create.call_args[1]["messages"] + assert sent_messages[-1]["content"] == [{"type": "text", "text": "describe my assistant"}]