From 51ce6f13477ad08b59701c21f963b7009f5b5a2a Mon Sep 17 00:00:00 2001
From: dog-last <2384399200@qq.com>
Date: Fri, 3 Oct 2025 21:43:08 +0800
Subject: [PATCH] Bug fix of thinking llm in vllm (#3510)
---
mem0/memory/main.py | 3 ++
mem0/memory/utils.py | 5 +-
tests/llms/test_vllm.py | 116 +++++++++++++++++++++++++++++++++++++++-
3 files changed, 122 insertions(+), 2 deletions(-)
diff --git a/mem0/memory/main.py b/mem0/memory/main.py
index efcbc52d6..79719566a 100644
--- a/mem0/memory/main.py
+++ b/mem0/memory/main.py
@@ -889,6 +889,7 @@ class Memory(MemoryBase):
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 @@ class AsyncMemory(MemoryBase):
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
diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py
index 466abdbe3..0872465da 100644
--- a/mem0/memory/utils.py
+++ b/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".*?", "", 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("_")
+
diff --git a/tests/llms/test_vllm.py b/tests/llms/test_vllm.py
index 7b2bfcac5..b0acf4631 100644
--- a/tests/llms/test_vllm.py
+++ b/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 = [
+ ' Sync fact extraction \n{"facts": ["User loves sci-fi"]}',
+ ' Sync memory actions \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 = [
+ ' Async fact extraction \n{"facts": ["User loves sci-fi"]}',
+ ' Async memory actions \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"
\ No newline at end of file