Skip to content

Commit ffdf167

Browse files
author
久氢
committed
feat(memory_collection): support AsyncMemory and add vector dimension configuration
Change-Id: I715f9c34f82bd49bd56611ca16d21b9dd0eee953 Co-developed-by: Cursor <noreply@cursor.com> Signed-off-by: 久氢 <mapenghui.mph@alibaba-inc.com>
1 parent 5110e13 commit ffdf167

File tree

6 files changed

+237
-32
lines changed

6 files changed

+237
-32
lines changed

‎agentrun/memory_collection/__memory_collection_async_template.py‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -228,15 +228,15 @@ async def to_mem0_memory_async(
228228
config: Optional[Config] = None,
229229
history_db_path: Optional[str] = None,
230230
):
231-
"""将 MemoryCollection 转换为 agentrun-mem0ai Memory 客户端（异步）
231+
"""将 MemoryCollection 转换为 agentrun-mem0ai AsyncMemory 客户端（异步）
232232
233233
Args:
234234
memory_collection_name: 记忆集合名称
235235
config: AgentRun 配置
236236
history_db_path: mem0 历史数据库路径（可选）
237237
238238
Returns:
239-
Memory: agentrun-mem0ai Memory 客户端实例
239+
AsyncMemory: agentrun-mem0ai AsyncMemory 客户端实例
240240
241241
Raises:
242242
ImportError: 如果未安装 agentrun-mem0ai 包
@@ -247,10 +247,10 @@ async def to_mem0_memory_async(
247247
... "memoryCollection010901",
248248
... config=config
249249
... )
250-
>>> memory.add("用户喜欢吃苹果", user_id="user123")
250+
>>> await memory.add("用户喜欢吃苹果", user_id="user123")
251251
"""
252252
try:
253-
from agentrun_mem0 import Memory
253+
from agentrun_mem0 import AsyncMemory
254254
except ImportError as e:
255255
raise ImportError(
256256
"agentrun-mem0ai package is required. Install it with: pip"
@@ -267,8 +267,8 @@ async def to_mem0_memory_async(
267267
memory_collection, config, history_db_path
268268
)
269269

270-
# 创建并返回 Memory 实例
271-
return Memory.from_config(mem0_config)
270+
# 创建并返回 AsyncMemory 实例
271+
return await AsyncMemory.from_config(mem0_config)
272272

273273
@staticmethod
274274
def _convert_vpc_endpoint_to_public(endpoint: str) -> str:
@@ -392,13 +392,25 @@ async def _build_mem0_config_async(
392392
)
393393
)
394394

395+
embedder_config_dict = {
396+
"model": embedder_config.config.model,
397+
"openai_base_url": base_url,
398+
"api_key": api_key,
399+
}
400+
401+
# 从 vector_store_config 中获取向量维度
402+
if (
403+
memory_collection.vector_store_config
404+
and memory_collection.vector_store_config.config
405+
and memory_collection.vector_store_config.config.vector_dimension
406+
):
407+
embedder_config_dict["embedding_dims"] = (
408+
memory_collection.vector_store_config.config.vector_dimension
409+
)
410+
395411
mem0_config["embedder"] = {
396412
"provider": "openai", # mem0 使用 openai 兼容接口
397-
"config": {
398-
"model": embedder_config.config.model,
399-
"openai_base_url": base_url,
400-
"api_key": api_key,
401-
},
413+
"config": embedder_config_dict,
402414
}
403415

404416
# 添加历史数据库路径

‎agentrun/memory_collection/memory_collection.py‎

Lines changed: 40 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -407,15 +407,15 @@ async def to_mem0_memory_async(
407407
config: Optional[Config] = None,
408408
history_db_path: Optional[str] = None,
409409
):
410-
"""将 MemoryCollection 转换为 agentrun-mem0ai Memory 客户端（异步）
410+
"""将 MemoryCollection 转换为 agentrun-mem0ai AsyncMemory 客户端（异步）
411411
412412
Args:
413413
memory_collection_name: 记忆集合名称
414414
config: AgentRun 配置
415415
history_db_path: mem0 历史数据库路径（可选）
416416
417417
Returns:
418-
Memory: agentrun-mem0ai Memory 客户端实例
418+
AsyncMemory: agentrun-mem0ai AsyncMemory 客户端实例
419419
420420
Raises:
421421
ImportError: 如果未安装 agentrun-mem0ai 包
@@ -426,10 +426,10 @@ async def to_mem0_memory_async(
426426
... "memoryCollection010901",
427427
... config=config
428428
... )
429-
>>> memory.add("用户喜欢吃苹果", user_id="user123")
429+
>>> await memory.add("用户喜欢吃苹果", user_id="user123")
430430
"""
431431
try:
432-
from agentrun_mem0 import Memory
432+
from agentrun_mem0 import AsyncMemory
433433
except ImportError as e:
434434
raise ImportError(
435435
"agentrun-mem0ai package is required. Install it with: pip"
@@ -446,8 +446,8 @@ async def to_mem0_memory_async(
446446
memory_collection, config, history_db_path
447447
)
448448

449-
# 创建并返回 Memory 实例
450-
return Memory.from_config(mem0_config)
449+
# 创建并返回 AsyncMemory 实例
450+
return await AsyncMemory.from_config(mem0_config)
451451

452452
@classmethod
453453
def to_mem0_memory(
@@ -620,13 +620,25 @@ async def _build_mem0_config_async(
620620
)
621621
)
622622

623+
embedder_config_dict = {
624+
"model": embedder_config.config.model,
625+
"openai_base_url": base_url,
626+
"api_key": api_key,
627+
}
628+
629+
# 从 vector_store_config 中获取向量维度
630+
if (
631+
memory_collection.vector_store_config
632+
and memory_collection.vector_store_config.config
633+
and memory_collection.vector_store_config.config.vector_dimension
634+
):
635+
embedder_config_dict["embedding_dims"] = (
636+
memory_collection.vector_store_config.config.vector_dimension
637+
)
638+
623639
mem0_config["embedder"] = {
624640
"provider": "openai", # mem0 使用 openai 兼容接口
625-
"config": {
626-
"model": embedder_config.config.model,
627-
"openai_base_url": base_url,
628-
"api_key": api_key,
629-
},
641+
"config": embedder_config_dict,
630642
}
631643

632644
# 添加历史数据库路径
@@ -732,13 +744,25 @@ def _build_mem0_config(
732744
model_service_name, config
733745
)
734746

747+
embedder_config_dict = {
748+
"model": embedder_config.config.model,
749+
"openai_base_url": base_url,
750+
"api_key": api_key,
751+
}
752+
753+
# 从 vector_store_config 中获取向量维度
754+
if (
755+
memory_collection.vector_store_config
756+
and memory_collection.vector_store_config.config
757+
and memory_collection.vector_store_config.config.vector_dimension
758+
):
759+
embedder_config_dict["embedding_dims"] = (
760+
memory_collection.vector_store_config.config.vector_dimension
761+
)
762+
735763
mem0_config["embedder"] = {
736764
"provider": "openai", # mem0 使用 openai 兼容接口
737-
"config": {
738-
"model": embedder_config.config.model,
739-
"openai_base_url": base_url,
740-
"api_key": api_key,
741-
},
765+
"config": embedder_config_dict,
742766
}
743767

744768
# 添加历史数据库路径

‎codegen/codegen.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,7 @@ def _generate_sync_code_for_file(async_file):
206206
content = (
207207
line.replace("AsyncClient", "Client")
208208
.replace("AsyncOpenAI", "OpenAI")
209+
.replace("AsyncMemory", "Memory")
209210
.replace("async_playwright", "sync_playwright")
210211
.replace("asyncio.gather(*", "(")
211212
.replace("_async", "")

‎examples/memory_collection_example.py‎

Lines changed: 69 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@
2323
from agentrun.utils.config import Config
2424

2525

26-
async def main():
27-
"""主函数 / Main function"""
26+
async def memory_collection_basic_example():
27+
"""MemoryCollection 基础操作示例 / MemoryCollection Basic Operations Example"""
2828

2929
# 创建配置
3030
# Create configuration
@@ -155,8 +155,73 @@ async def main():
155155

156156
traceback.print_exc()
157157

158-
print("\n✅ 示例完成")
158+
159+
async def async_mem0_memory_example():
160+
"""异步 AsyncMemory 转换方法使用示例"""
161+
print("\n=== 异步转换为 mem0ai AsyncMemory 客户端 ===")
162+
163+
try:
164+
# 使用高层 API 的 to_mem0_memory_async 方法
165+
# Use high-level API's to_mem0_memory_async method
166+
memory = await MemoryCollection.to_mem0_memory_async(
167+
"memoryCollection010901"
168+
)
169+
print(f"✅ 成功创建 mem0ai AsyncMemory 客户端")
170+
print(f" 类型: {type(memory)}")
171+
172+
# 使用 mem0ai AsyncMemory 客户端进行操作
173+
# Use mem0ai AsyncMemory client for operations
174+
user_id = "user456"
175+
176+
# 添加记忆
177+
# Add memory
178+
result = await memory.add(
179+
"我喜欢喝咖啡和茶",
180+
user_id=user_id,
181+
metadata={"category": "beverage"},
182+
)
183+
print(f"\n✅ 添加记忆成功:")
184+
for idx, res in enumerate(result.get("results", []), 1):
185+
print(f" {idx}. ID: {res.get('id')}, 事件: {res.get('event')}")
186+
187+
# 搜索记忆
188+
# Search memory
189+
search_results = await memory.search(
190+
"用户喜欢喝什么？", user_id=user_id
191+
)
192+
print(f"\n✅ 搜索记忆结果:")
193+
for idx, result in enumerate(search_results.get("results", []), 1):
194+
print(
195+
f" {idx}. 内容: {result.get('memory')}, 相似度:"
196+
f" {result.get('score', 0):.4f}"
197+
)
198+
199+
# 获取所有记忆
200+
# Get all memories
201+
all_memories = await memory.get_all(user_id=user_id)
202+
print(f"\n✅ 获取所有记忆:")
203+
for idx, mem in enumerate(all_memories.get("results", []), 1):
204+
print(f" {idx}. {mem.get('memory')}")
205+
206+
except ImportError as e:
207+
print(f"⚠️ mem0ai 未安装: {e}")
208+
print(" 安装方法: pip install agentrun-sdk[mem0]")
209+
except Exception as e:
210+
print(f"❌ AsyncMemory 操作失败: {e}")
211+
import traceback
212+
213+
traceback.print_exc()
159214

160215

161216
if __name__ == "__main__":
162-
asyncio.run(main())
217+
import asyncio
218+
219+
# 运行基础操作示例
220+
# Run basic operations example
221+
asyncio.run(memory_collection_basic_example())
222+
223+
# 运行异步 AsyncMemory 示例
224+
# Run async AsyncMemory example
225+
asyncio.run(async_mem0_memory_example())
226+
227+
print("\n✅ 示例完成")

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ dependencies = [
1616
"alibabacloud-agentrun20250910>=5.2.0",
1717
"alibabacloud_tea_openapi>=0.4.2",
1818
"alibabacloud_bailian20231229>=2.6.2",
19-
"agentrun-mem0ai>=0.0.6",
19+
"agentrun-mem0ai>=0.0.10",
2020
]
2121

2222
[project.optional-dependencies]

‎tests/unittests/memory_collection/test_memory_collection.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,3 +364,106 @@ def test_from_inner_object_with_extra(self):
364364
assert (
365365
memory_collection.memory_collection_name == "test-memory-collection"
366366
)
367+
368+
369+
class TestMemoryCollectionToMem0Memory:
370+
"""测试 MemoryCollection.to_mem0_memory 和 to_mem0_memory_async 方法"""
371+
372+
@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
373+
def test_to_mem0_memory_sync(self, mock_control_api_class):
374+
"""测试同步转换为 mem0 Memory 客户端"""
375+
# Mock MemoryCollection 数据
376+
mock_control_api = MagicMock()
377+
mock_data = MockMemoryCollectionData()
378+
mock_control_api.get_memory_collection.return_value = mock_data
379+
mock_control_api_class.return_value = mock_control_api
380+
381+
# Mock Memory.from_config
382+
with patch("agentrun_mem0.Memory") as mock_memory_class:
383+
mock_memory_instance = MagicMock()
384+
mock_memory_class.from_config.return_value = mock_memory_instance
385+
386+
# 调用 to_mem0_memory
387+
result = MemoryCollection.to_mem0_memory(
388+
"test-memory-collection",
389+
config=Config(
390+
access_key_id="test-key",
391+
access_key_secret="test-secret",
392+
region_id="cn-hangzhou",
393+
),
394+
)
395+
396+
# 验证返回的是 Memory 实例
397+
assert result == mock_memory_instance
398+
mock_memory_class.from_config.assert_called_once()
399+
400+
@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
401+
@pytest.mark.asyncio
402+
async def test_to_mem0_memory_async(self, mock_control_api_class):
403+
"""测试异步转换为 mem0 AsyncMemory 客户端"""
404+
# Mock MemoryCollection 数据
405+
mock_control_api = MagicMock()
406+
mock_data = MockMemoryCollectionData()
407+
mock_control_api.get_memory_collection_async = AsyncMock(
408+
return_value=mock_data
409+
)
410+
mock_control_api_class.return_value = mock_control_api
411+
412+
# Mock AsyncMemory.from_config
413+
with patch("agentrun_mem0.AsyncMemory") as mock_async_memory_class:
414+
mock_async_memory_instance = MagicMock()
415+
# from_config 是异步方法，需要返回 AsyncMock
416+
mock_async_memory_class.from_config = AsyncMock(
417+
return_value=mock_async_memory_instance
418+
)
419+
420+
# 调用 to_mem0_memory_async
421+
result = await MemoryCollection.to_mem0_memory_async(
422+
"test-memory-collection",
423+
config=Config(
424+
access_key_id="test-key",
425+
access_key_secret="test-secret",
426+
region_id="cn-hangzhou",
427+
),
428+
)
429+
430+
# 验证返回的是 AsyncMemory 实例
431+
assert result == mock_async_memory_instance
432+
mock_async_memory_class.from_config.assert_called_once()
433+
434+
@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
435+
def test_to_mem0_memory_import_error(self, mock_control_api_class):
436+
"""测试 mem0 包未安装时的错误处理"""
437+
mock_control_api = MagicMock()
438+
mock_control_api.get_memory_collection.return_value = (
439+
MockMemoryCollectionData()
440+
)
441+
mock_control_api_class.return_value = mock_control_api
442+
443+
# Mock import 失败
444+
with patch("builtins.__import__", side_effect=ImportError("No module")):
445+
with pytest.raises(ImportError) as exc_info:
446+
MemoryCollection.to_mem0_memory("test-memory-collection")
447+
448+
assert "agentrun-mem0ai package is required" in str(exc_info.value)
449+
450+
@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
451+
@pytest.mark.asyncio
452+
async def test_to_mem0_memory_async_import_error(
453+
self, mock_control_api_class
454+
):
455+
"""测试异步方法 mem0 包未安装时的错误处理"""
456+
mock_control_api = MagicMock()
457+
mock_control_api.get_memory_collection_async = AsyncMock(
458+
return_value=MockMemoryCollectionData()
459+
)
460+
mock_control_api_class.return_value = mock_control_api
461+
462+
# Mock import 失败
463+
with patch("builtins.__import__", side_effect=ImportError("No module")):
464+
with pytest.raises(ImportError) as exc_info:
465+
await MemoryCollection.to_mem0_memory_async(
466+
"test-memory-collection"
467+
)
468+
469+
assert "agentrun-mem0ai package is required" in str(exc_info.value)

0 commit comments

Comments
 (0)