Skip to content

feat(memory_collection): support AsyncMemory and add vector dimension… #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
OhYee merged 1 commit into from
Jan 20, 2026
Merged

feat(memory_collection): support AsyncMemory and add vector dimension… #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 23 additions & 11 deletions agentrun/memory_collection/__memory_collection_async_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,15 +228,15 @@ async def to_mem0_memory_async(
config: Optional[Config] = None,
history_db_path: Optional[str] = None,
):
"""将 MemoryCollection 转换为 agentrun-mem0ai Memory 客户端（异步）
"""将 MemoryCollection 转换为 agentrun-mem0ai AsyncMemory 客户端（异步）

Args:
memory_collection_name: 记忆集合名称
config: AgentRun 配置
history_db_path: mem0 历史数据库路径（可选）

Returns:
Memory: agentrun-mem0ai Memory 客户端实例
AsyncMemory: agentrun-mem0ai AsyncMemory 客户端实例

Raises:
ImportError: 如果未安装 agentrun-mem0ai 包
Expand All @@ -247,10 +247,10 @@ async def to_mem0_memory_async(
... "memoryCollection010901",
... config=config
... )
>>> memory.add("用户喜欢吃苹果", user_id="user123")
>>> await memory.add("用户喜欢吃苹果", user_id="user123")
"""
try:
from agentrun_mem0 import Memory
from agentrun_mem0 import AsyncMemory
except ImportError as e:
raise ImportError(
"agentrun-mem0ai package is required. Install it with: pip"
Expand All @@ -267,8 +267,8 @@ async def to_mem0_memory_async(
memory_collection, config, history_db_path
)

# 创建并返回 Memory 实例
return Memory.from_config(mem0_config)
# 创建并返回 AsyncMemory 实例
return await AsyncMemory.from_config(mem0_config)

@staticmethod
def _convert_vpc_endpoint_to_public(endpoint: str) -> str:
Expand Down Expand Up @@ -392,13 +392,25 @@ async def _build_mem0_config_async(
)
)

embedder_config_dict = {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
}

# 从 vector_store_config 中获取向量维度
if (
memory_collection.vector_store_config
and memory_collection.vector_store_config.config
and memory_collection.vector_store_config.config.vector_dimension
):
embedder_config_dict["embedding_dims"] = (
memory_collection.vector_store_config.config.vector_dimension
)

mem0_config["embedder"] = {
"provider": "openai", # mem0 使用 openai 兼容接口
"config": {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
},
"config": embedder_config_dict,
}

# 添加历史数据库路径
Expand Down
56 changes: 40 additions & 16 deletions agentrun/memory_collection/memory_collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,15 +407,15 @@ async def to_mem0_memory_async(
config: Optional[Config] = None,
history_db_path: Optional[str] = None,
):
"""将 MemoryCollection 转换为 agentrun-mem0ai Memory 客户端（异步）
"""将 MemoryCollection 转换为 agentrun-mem0ai AsyncMemory 客户端（异步）

Args:
memory_collection_name: 记忆集合名称
config: AgentRun 配置
history_db_path: mem0 历史数据库路径（可选）

Returns:
Memory: agentrun-mem0ai Memory 客户端实例
AsyncMemory: agentrun-mem0ai AsyncMemory 客户端实例

Raises:
ImportError: 如果未安装 agentrun-mem0ai 包
Expand All @@ -426,10 +426,10 @@ async def to_mem0_memory_async(
... "memoryCollection010901",
... config=config
... )
>>> memory.add("用户喜欢吃苹果", user_id="user123")
>>> await memory.add("用户喜欢吃苹果", user_id="user123")
"""
try:
from agentrun_mem0 import Memory
from agentrun_mem0 import AsyncMemory
except ImportError as e:
raise ImportError(
"agentrun-mem0ai package is required. Install it with: pip"
Expand All @@ -446,8 +446,8 @@ async def to_mem0_memory_async(
memory_collection, config, history_db_path
)

# 创建并返回 Memory 实例
return Memory.from_config(mem0_config)
# 创建并返回 AsyncMemory 实例
return await AsyncMemory.from_config(mem0_config)

@classmethod
def to_mem0_memory(
Expand Down Expand Up @@ -620,13 +620,25 @@ async def _build_mem0_config_async(
)
)

embedder_config_dict = {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
}

# 从 vector_store_config 中获取向量维度
if (
memory_collection.vector_store_config
and memory_collection.vector_store_config.config
and memory_collection.vector_store_config.config.vector_dimension
):
embedder_config_dict["embedding_dims"] = (
memory_collection.vector_store_config.config.vector_dimension
)
Comment on lines +623 to +637
Copy link

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new vector dimension configuration logic (lines 623-637 and 747-761) that extracts vector_dimension from vector_store_config and adds it as embedding_dims to the embedder configuration is not covered by tests. Consider adding test cases to verify that this configuration is correctly passed to the mem0 config when vector_dimension is present.

Copilot uses AI. Check for mistakes.

mem0_config["embedder"] = {
"provider": "openai", # mem0 使用 openai 兼容接口
"config": {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
},
"config": embedder_config_dict,
}

# 添加历史数据库路径
Expand Down Expand Up @@ -732,13 +744,25 @@ def _build_mem0_config(
model_service_name, config
)

embedder_config_dict = {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
}

# 从 vector_store_config 中获取向量维度
if (
memory_collection.vector_store_config
and memory_collection.vector_store_config.config
and memory_collection.vector_store_config.config.vector_dimension
):
embedder_config_dict["embedding_dims"] = (
memory_collection.vector_store_config.config.vector_dimension
)

mem0_config["embedder"] = {
"provider": "openai", # mem0 使用 openai 兼容接口
"config": {
"model": embedder_config.config.model,
"openai_base_url": base_url,
"api_key": api_key,
},
"config": embedder_config_dict,
}

# 添加历史数据库路径
Expand Down
1 change: 1 addition & 0 deletions codegen/codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ def _generate_sync_code_for_file(async_file):
content = (
line.replace("AsyncClient", "Client")
.replace("AsyncOpenAI", "OpenAI")
.replace("AsyncMemory", "Memory")
.replace("async_playwright", "sync_playwright")
.replace("asyncio.gather(*", "(")
.replace("_async", "")
Expand Down
73 changes: 69 additions & 4 deletions examples/memory_collection_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@
from agentrun.utils.config import Config


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

# 创建配置
# Create configuration
Expand Down Expand Up @@ -155,8 +155,73 @@ async def main():

traceback.print_exc()

print("\n✅ 示例完成")

async def async_mem0_memory_example():
"""异步 AsyncMemory 转换方法使用示例"""
print("\n=== 异步转换为 mem0ai AsyncMemory 客户端 ===")

try:
# 使用高层 API 的 to_mem0_memory_async 方法
# Use high-level API's to_mem0_memory_async method
memory = await MemoryCollection.to_mem0_memory_async(
"memoryCollection010901"
)
print(f"✅ 成功创建 mem0ai AsyncMemory 客户端")
print(f" 类型: {type(memory)}")

# 使用 mem0ai AsyncMemory 客户端进行操作
# Use mem0ai AsyncMemory client for operations
user_id = "user456"

# 添加记忆
# Add memory
result = await memory.add(
"我喜欢喝咖啡和茶",
user_id=user_id,
metadata={"category": "beverage"},
)
print(f"\n✅ 添加记忆成功:")
for idx, res in enumerate(result.get("results", []), 1):
print(f" {idx}. ID: {res.get('id')}, 事件: {res.get('event')}")

# 搜索记忆
# Search memory
search_results = await memory.search(
"用户喜欢喝什么？", user_id=user_id
)
print(f"\n✅ 搜索记忆结果:")
for idx, result in enumerate(search_results.get("results", []), 1):
print(
f" {idx}. 内容: {result.get('memory')}, 相似度:"
f" {result.get('score', 0):.4f}"
)

# 获取所有记忆
# Get all memories
all_memories = await memory.get_all(user_id=user_id)
print(f"\n✅ 获取所有记忆:")
for idx, mem in enumerate(all_memories.get("results", []), 1):
print(f" {idx}. {mem.get('memory')}")

except ImportError as e:
print(f"⚠️ mem0ai 未安装: {e}")
print(" 安装方法: pip install agentrun-sdk[mem0]")
except Exception as e:
print(f"❌ AsyncMemory 操作失败: {e}")
import traceback

traceback.print_exc()


if __name__ == "__main__":
asyncio.run(main())
import asyncio
Copy link

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The module asyncio is already imported at the top of the file (line 8), so this re-import is redundant and should be removed.

Suggested change
import asyncio

Copilot uses AI. Check for mistakes.

# 运行基础操作示例
# Run basic operations example
asyncio.run(memory_collection_basic_example())

# 运行异步 AsyncMemory 示例
# Run async AsyncMemory example
asyncio.run(async_mem0_memory_example())

print("\n✅ 示例完成")
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ dependencies = [
"alibabacloud-agentrun20250910>=5.2.0",
"alibabacloud_tea_openapi>=0.4.2",
"alibabacloud_bailian20231229>=2.6.2",
"agentrun-mem0ai>=0.0.6",
"agentrun-mem0ai>=0.0.10",
]

[project.optional-dependencies]
Expand Down
103 changes: 103 additions & 0 deletions tests/unittests/memory_collection/test_memory_collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,3 +364,106 @@ def test_from_inner_object_with_extra(self):
assert (
memory_collection.memory_collection_name == "test-memory-collection"
)


class TestMemoryCollectionToMem0Memory:
"""测试 MemoryCollection.to_mem0_memory 和 to_mem0_memory_async 方法"""

@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
def test_to_mem0_memory_sync(self, mock_control_api_class):
"""测试同步转换为 mem0 Memory 客户端"""
# Mock MemoryCollection 数据
mock_control_api = MagicMock()
mock_data = MockMemoryCollectionData()
mock_control_api.get_memory_collection.return_value = mock_data
mock_control_api_class.return_value = mock_control_api

# Mock Memory.from_config
with patch("agentrun_mem0.Memory") as mock_memory_class:
mock_memory_instance = MagicMock()
mock_memory_class.from_config.return_value = mock_memory_instance

# 调用 to_mem0_memory
result = MemoryCollection.to_mem0_memory(
"test-memory-collection",
config=Config(
access_key_id="test-key",
access_key_secret="test-secret",
region_id="cn-hangzhou",
),
)

# 验证返回的是 Memory 实例
assert result == mock_memory_instance
mock_memory_class.from_config.assert_called_once()

@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
@pytest.mark.asyncio
async def test_to_mem0_memory_async(self, mock_control_api_class):
"""测试异步转换为 mem0 AsyncMemory 客户端"""
# Mock MemoryCollection 数据
mock_control_api = MagicMock()
mock_data = MockMemoryCollectionData()
mock_control_api.get_memory_collection_async = AsyncMock(
return_value=mock_data
)
mock_control_api_class.return_value = mock_control_api

# Mock AsyncMemory.from_config
with patch("agentrun_mem0.AsyncMemory") as mock_async_memory_class:
mock_async_memory_instance = MagicMock()
# from_config 是异步方法，需要返回 AsyncMock
mock_async_memory_class.from_config = AsyncMock(
return_value=mock_async_memory_instance
)

# 调用 to_mem0_memory_async
result = await MemoryCollection.to_mem0_memory_async(
"test-memory-collection",
config=Config(
access_key_id="test-key",
access_key_secret="test-secret",
region_id="cn-hangzhou",
),
)

# 验证返回的是 AsyncMemory 实例
assert result == mock_async_memory_instance
mock_async_memory_class.from_config.assert_called_once()

@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
def test_to_mem0_memory_import_error(self, mock_control_api_class):
"""测试 mem0 包未安装时的错误处理"""
mock_control_api = MagicMock()
mock_control_api.get_memory_collection.return_value = (
MockMemoryCollectionData()
)
mock_control_api_class.return_value = mock_control_api

# Mock import 失败
with patch("builtins.__import__", side_effect=ImportError("No module")):
Copy link

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test is using patch("builtins.__import__") to mock import failures, but the actual code uses a try-except block with from agentrun_mem0 import Memory. This approach won't work because the from X import Y statement is not the same as calling __import__. The test should instead patch the module agentrun_mem0.Memory or use patch.dict('sys.modules') to simulate the module not being available.

Copilot uses AI. Check for mistakes.
with pytest.raises(ImportError) as exc_info:
MemoryCollection.to_mem0_memory("test-memory-collection")

assert "agentrun-mem0ai package is required" in str(exc_info.value)

@patch("agentrun.memory_collection.client.MemoryCollectionControlAPI")
@pytest.mark.asyncio
async def test_to_mem0_memory_async_import_error(
self, mock_control_api_class
):
"""测试异步方法 mem0 包未安装时的错误处理"""
mock_control_api = MagicMock()
mock_control_api.get_memory_collection_async = AsyncMock(
return_value=MockMemoryCollectionData()
)
mock_control_api_class.return_value = mock_control_api

# Mock import 失败
with patch("builtins.__import__", side_effect=ImportError("No module")):
Copy link

Copilot AI Jan 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test is using patch("builtins.__import__") to mock import failures, but the actual code uses a try-except block with from agentrun_mem0 import AsyncMemory. This approach won't work because the from X import Y statement is not the same as calling __import__. The test should instead patch the module agentrun_mem0.AsyncMemory or use patch.dict('sys.modules') to simulate the module not being available.

Copilot uses AI. Check for mistakes.
with pytest.raises(ImportError) as exc_info:
await MemoryCollection.to_mem0_memory_async(
"test-memory-collection"
)

assert "agentrun-mem0ai package is required" in str(exc_info.value)
Loading