incorporate my feature into the remove_code_blocks function. And add the test to test_vllm.py instead of a entire separate test.

Author: dog-last

mem0/memory/main.py

@@ -31,7 +31,6 @@
     parse_vision_messages,
     process_telemetry_filters,
     remove_code_blocks,
-    remove_thinking_tags,
 )
 from mem0.utils.factory import (
     EmbedderFactory,
@@ -363,8 +362,6 @@ def _add_to_vector_store(self, messages, metadata, filters, infer):
 
         try:
             response = remove_code_blocks(response)
-            if '</think>' in response:
-                response=remove_thinking_tags(response)
             new_retrieved_facts = json.loads(response)["facts"]
         except Exception as e:
             logger.error(f"Error in new_retrieved_facts: {e}")
@@ -419,8 +416,6 @@ def _add_to_vector_store(self, messages, metadata, filters, infer):
                     new_memories_with_actions = {}
                 else:
                     response = remove_code_blocks(response)
-                    if '</think>' in response:
-                        response=remove_thinking_tags(response)
                     new_memories_with_actions = json.loads(response)
             except Exception as e:
                 logger.error(f"Invalid JSON response: {e}")
@@ -894,8 +889,7 @@ def _create_procedural_memory(self, messages, metadata=None, prompt=None):
 
         try:
             procedural_memory = self.llm.generate_response(messages=parsed_messages)
-            if '</think>' in procedural_memory:
-                procedural_memory=remove_thinking_tags(procedural_memory)
+            procedural_memory = remove_code_blocks(procedural_memory)
         except Exception as e:
             logger.error(f"Error generating procedural memory summary: {e}")
             raise
@@ -1220,8 +1214,6 @@ async def _add_to_vector_store(
         )
         try:
             response = remove_code_blocks(response)
-            if '</think>' in response:
-                response=remove_thinking_tags(response)
             new_retrieved_facts = json.loads(response)["facts"]
         except Exception as e:
             logger.error(f"Error in new_retrieved_facts: {e}")
@@ -1270,8 +1262,6 @@ async def process_fact_for_search(new_mem_content):
                     messages=[{"role": "user", "content": function_calling_prompt}],
                     response_format={"type": "json_object"},
                 )
-                if '</think>' in response:
-                    response=remove_thinking_tags(response)
             except Exception as e:
                 logger.error(f"Error in new memory actions response: {e}")
                 response = ""
@@ -1795,8 +1785,8 @@ async def _create_procedural_memory(self, messages, metadata=None, llm=None, pro
                 procedural_memory = response.content
             else:
                 procedural_memory = await asyncio.to_thread(self.llm.generate_response, messages=parsed_messages)
-                if '</think>' in procedural_memory:
-                    procedural_memory=remove_thinking_tags
+                procedural_memory = remove_code_blocks(procedural_memory)
+        
         except Exception as e:
             logger.error(f"Error generating procedural memory summary: {e}")
             raise

mem0/memory/utils.py

@@ -43,7 +43,9 @@ def remove_code_blocks(content: str) -> str:
     """
     pattern = r"^```[a-zA-Z0-9]*\n([\s\S]*?)\n```$"
     match = re.match(pattern, content.strip())
-    return match.group(1).strip() if match else content.strip()
+    match_res=match.group(1).strip() if match else content.strip()
+    return re.sub(r"<think>.*?</think>", "", match_res, flags=re.DOTALL).strip()
+
 
 
 def extract_json(text):
@@ -183,15 +185,3 @@ def sanitize_relationship_for_cypher(relationship) -> str:
 
     return re.sub(r"_+", "_", sanitized).strip("_")
 
-
-def remove_thinking_tags(text: str) -> str:
-    """
-    Removes <think>...</think> tags from the input text.
-
-    Args:
-        text (str): The input text potentially containing <think>...</think> tags.
-
-    Returns:
-        str: The text with all <think></think>tags and their content removed.
-    """
-    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

tests/llms/test_vllm.py

@@ -1,9 +1,11 @@
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 
+from mem0 import AsyncMemory, Memory
 from mem0.configs.llms.base import BaseLlmConfig
 from mem0.llms.vllm import VllmLLM
+from mem0.memory.utils import remove_code_blocks
 
 
 @pytest.fixture
@@ -84,3 +86,132 @@ def test_generate_response_with_tools(mock_vllm_client):
     assert len(response["tool_calls"]) == 1
     assert response["tool_calls"][0]["name"] == "add_memory"
     assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
+
+
+
+def test_remove_thinking_tags():
+    # with thinking tag
+    assert remove_code_blocks('<think>abc</think>{"facts":["test fact"]}') == '{"facts":["test fact"]}'
+    
+    # thinking content with multiple lines
+    assert remove_code_blocks('<think>\nabc\n</think>\n{"facts":["test fact"]}') == '{"facts":["test fact"]}'
+    
+    # more than one thinking tag(rare in practice)
+    assert remove_code_blocks('<think>A</think><think>B</think>{"facts":["test fact"]}') == '{"facts":["test fact"]}'
+    
+    # no tag
+    assert remove_code_blocks('{"facts":["test fact"]}') == '{"facts":["test fact"]}'
+
+
+
+
+def create_mocked_memory():
+    """Create a fully mocked Memory instance for testing."""
+    with patch('mem0.utils.factory.LlmFactory.create') as mock_llm_factory, \
+         patch('mem0.utils.factory.EmbedderFactory.create') as mock_embedder_factory, \
+         patch('mem0.utils.factory.VectorStoreFactory.create') as mock_vector_factory, \
+         patch('mem0.memory.storage.SQLiteManager') as mock_sqlite:
+
+        mock_llm = MagicMock()
+        mock_llm_factory.return_value = mock_llm
+
+        mock_embedder = MagicMock()
+        mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
+        mock_embedder_factory.return_value = mock_embedder
+
+        mock_vector_store = MagicMock()
+        mock_vector_store.search.return_value = []
+        mock_vector_store.add.return_value = None
+        mock_vector_factory.return_value = mock_vector_store
+
+        mock_sqlite.return_value = MagicMock()
+
+        memory = Memory()
+        memory.api_version = "v1.0"
+        return memory, mock_llm, mock_vector_store
+
+
+def create_mocked_async_memory():
+    """Create a fully mocked AsyncMemory instance for testing."""
+    with patch('mem0.utils.factory.LlmFactory.create') as mock_llm_factory, \
+         patch('mem0.utils.factory.EmbedderFactory.create') as mock_embedder_factory, \
+         patch('mem0.utils.factory.VectorStoreFactory.create') as mock_vector_factory, \
+         patch('mem0.memory.storage.SQLiteManager') as mock_sqlite:
+
+        mock_llm = MagicMock()
+        mock_llm_factory.return_value = mock_llm
+
+        mock_embedder = MagicMock()
+        mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
+        mock_embedder_factory.return_value = mock_embedder
+
+        mock_vector_store = MagicMock()
+        mock_vector_store.search.return_value = []
+        mock_vector_store.add.return_value = None
+        mock_vector_factory.return_value = mock_vector_store
+
+        mock_sqlite.return_value = MagicMock()
+
+        memory = AsyncMemory()
+        memory.api_version = "v1.0"
+        return memory, mock_llm, mock_vector_store
+
+
+def test_thinking_tags_in_add_to_vector_store():
+    """Test thinking tags handling in Memory._add_to_vector_store (sync)."""
+    memory, mock_llm, mock_vector_store = create_mocked_memory()
+    
+    # Mock LLM responses for both phases
+    mock_llm.generate_response.side_effect = [
+        '        <think>Sync fact extraction</think>  \n{"facts": ["User loves sci-fi"]}',
+        '        <think>Sync memory actions</think>  \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
+    ]
+    
+    mock_vector_store.search.return_value = []
+    
+    result = memory._add_to_vector_store(
+        messages=[{"role": "user", "content": "I love sci-fi movies"}],
+        metadata={}, 
+        filters={}, 
+        infer=True
+    )
+    
+    assert len(result) == 1
+    assert result[0]["memory"] == "Loves sci-fi"
+    assert result[0]["event"] == "ADD"
+
+
+
+@pytest.mark.asyncio
+async def test_async_thinking_tags_in_add_to_vector_store():
+    """Test thinking tags handling in AsyncMemory._add_to_vector_store."""
+    memory, mock_llm, mock_vector_store = create_mocked_async_memory()
+    
+    # Directly mock llm.generate_response instead of via asyncio.to_thread
+    mock_llm.generate_response.side_effect = [
+        '        <think>Async fact extraction</think>  \n{"facts": ["User loves sci-fi"]}',
+        '        <think>Async memory actions</think>  \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
+    ]
+    
+    # Mock asyncio.to_thread to call the function directly (bypass threading)
+    async def mock_to_thread(func, *args, **kwargs):
+        if func == mock_llm.generate_response:
+            return func(*args, **kwargs)
+        elif hasattr(func, '__name__') and 'embed' in func.__name__:
+            return [0.1, 0.2, 0.3]
+        elif hasattr(func, '__name__') and 'search' in func.__name__:
+            return []
+        else:
+            return func(*args, **kwargs)
+    
+    with patch('mem0.memory.main.asyncio.to_thread', side_effect=mock_to_thread):
+        result = await memory._add_to_vector_store(
+            messages=[{"role": "user", "content": "I love sci-fi movies"}],
+            metadata={}, 
+            effective_filters={}, 
+            infer=True
+        )
+    
+    assert len(result) == 1
+    assert result[0]["memory"] == "Loves sci-fi"
+    assert result[0]["event"] == "ADD"