Bug fix of thinking llm in vllm (#3510)

Author: dog-last

mem0/memory/main.py

@@ -889,6 +889,7 @@ def _create_procedural_memory(self, messages, metadata=None, prompt=None):
 
         try:
             procedural_memory = self.llm.generate_response(messages=parsed_messages)
+            procedural_memory = remove_code_blocks(procedural_memory)
         except Exception as e:
             logger.error(f"Error generating procedural memory summary: {e}")
             raise
@@ -1784,6 +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)
+                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):
@@ -182,3 +184,4 @@ def sanitize_relationship_for_cypher(relationship) -> str:
         sanitized = sanitized.replace(old, new)
 
     return re.sub(r"_+", "_", sanitized).strip("_")
+

tests/llms/test_vllm.py

@@ -1,7 +1,8 @@
-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
 
@@ -84,3 +85,116 @@ 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 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_sync():
+    """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_async():
+    """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"