memory_ops_sdk.py•21.5 kB
#!/usr/bin/env python3
"""
MemOS Memory Ops SDK
轻量级记忆操作SDK，提供简单易用的一行代码接口
基于MVP管理器封装，便于其他脚本集成记忆功能
使用示例:
from memory_ops_sdk import MemoryOps
# 初始化
mem = MemoryOps()
# 添加记忆到不同类型
mem.add("今天学习了Python编程", memory_type='code_snippet_mem', tags=["学习", "Python"])
mem.add("遇到了导入错误", memory_type='error_log_mem', tags=["错误", "导入"])
# 查询记忆
results = mem.query("Python编程", memory_type='code_snippet_mem')
# 列出Memory类型
types = mem.list_memory_types()
"""
import os
import sys
from pathlib import Path
from typing import List, Dict, Any, Optional, Union
from datetime import datetime
# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent))
# 添加src目录到Python路径
sys.path.insert(0, str(Path(__file__).parent / "src"))
try:
# 导入bootstrap来解决依赖问题
import bootstrap
bootstrap.setup_mock_modules()
except ImportError:
pass # bootstrap是可选的
try:
from memos.configs.mem_cube import MultiMemoryMemCubeConfig
from memos.configs.memory import MemoryConfigFactory
from memos.mem_cube.multi import MultiMemoryMemCube
except ImportError as e:
raise ImportError(f"无法导入MultiMemoryMemCube: {e}. 请确保在MemOS项目目录中运行")
class MemoryOps:
"""
轻量级记忆操作SDK
提供简单易用的记忆管理接口，封装MVP管理器的复杂操作
"""
def __init__(self, data_dir: Optional[str] = None, auto_init: bool = True, verbose: bool = False):
"""
初始化Memory Ops SDK
Args:
data_dir: 数据目录路径，默认为当前目录下的memos_data
auto_init: 是否自动初始化，False时需要手动调用init()
verbose: 是否显示详细日志
"""
self.data_dir = data_dir or "./memos_data"
self.verbose = verbose
self._cube = None
self._initialized = False
# 配置多Memory模块
self.memory_configs = {
'general_mem': MemoryConfigFactory(backend='general_text'),
'code_snippet_mem': MemoryConfigFactory(backend='general_text'),
'error_log_mem': MemoryConfigFactory(backend='general_text')
}
if auto_init:
self.init()
def init(self) -> bool:
"""
初始化SDK
Returns:
bool: 初始化是否成功
"""
try:
if self.verbose:
print(f"🚀 初始化MemOS SDK...")
print(f" 数据目录: {self.data_dir}")
print(f" Memory模块: {list(self.memory_configs.keys())}")
# 创建MultiMemoryMemCube配置
config = MultiMemoryMemCubeConfig(
memory_configs=self.memory_configs
)
# 初始化MultiMemoryMemCube
self._cube = MultiMemoryMemCube(config)
self._initialized = True
if self.verbose:
print("✅ MemOS SDK初始化成功")
return True
except Exception as e:
if self.verbose:
print(f"❌ MemOS SDK初始化失败: {e}")
return False
def _ensure_initialized(self):
"""确保SDK已初始化"""
if not self._initialized or not self._cube:
raise RuntimeError("SDK未初始化，请先调用init()方法")
def add(self, text: str, memory_type: str = 'general_mem', tags: Optional[List[str]] = None,
metadata: Optional[Dict[str, Any]] = None) -> bool:
"""
添加记忆（一行代码接口）
Args:
text: 记忆内容
memory_type: Memory类型 ('general_mem', 'code_snippet_mem', 'error_log_mem')
tags: 标签列表
metadata: 元数据字典
Returns:
bool: 是否添加成功
Example:
mem.add("今天学习了Python编程", memory_type='code_snippet_mem', tags=["学习", "Python"])
"""
self._ensure_initialized()
try:
# 检查memory_type是否有效
if memory_type not in self.memory_configs:
if self.verbose:
print(f"❌ 无效的Memory类型: {memory_type}")
return False
# 自动添加时间戳到元数据
if metadata is None:
metadata = {}
if 'timestamp' not in metadata:
metadata['timestamp'] = datetime.now().isoformat()
# 添加标签到元数据
if tags:
metadata['tags'] = tags
# 获取指定的Memory模块
memory = self._cube.get_memory(memory_type)
if memory is None:
if self.verbose:
print(f"❌ Memory模块 {memory_type} 未初始化")
return False
# 构造记忆项
memory_item = {
"memory": text,
"metadata": metadata
}
# 添加记忆
memory.add([memory_item])
if self.verbose:
print(f"✅ 成功添加到 {memory_type}: {text[:50]}...")
return True
except Exception as e:
if self.verbose:
print(f"❌ 添加记忆失败: {e}")
return False
def query(self, q: str, memory_type: Optional[str] = None, limit: int = 5) -> List[Dict[str, Any]]:
"""
查询记忆（一行代码接口）
Args:
q: 查询文本
memory_type: Memory类型，None表示搜索所有类型
limit: 返回结果数量限制
Returns:
List[Dict]: 查询结果列表
Example:
results = mem.query("Python编程", memory_type='code_snippet_mem')
"""
self._ensure_initialized()
try:
all_results = []
# 确定要搜索的Memory类型
search_types = [memory_type] if memory_type else list(self.memory_configs.keys())
for mem_type in search_types:
if mem_type not in self.memory_configs:
continue
memory = self._cube.get_memory(mem_type)
if memory is None:
continue
try:
results = memory.search(q, top_k=limit)
# 转换为字典格式并添加类型信息
for item in results:
result_dict = {
"id": item.id,
"memory": item.memory,
"metadata": item.metadata.model_dump() if hasattr(item.metadata, 'model_dump') else item.metadata,
"memory_type": mem_type
}
all_results.append(result_dict)
except Exception as e:
if self.verbose:
print(f"❌ 在 {mem_type} 中搜索失败: {e}")
continue
# 按相关性排序并限制结果数量
all_results = all_results[:limit]
if self.verbose:
print(f"🔍 查询'{q}' 找到 {len(all_results)} 条结果")
return all_results
except Exception as e:
if self.verbose:
print(f"❌ 查询失败: {e}")
return []
def list_memory_types(self) -> List[str]:
"""
列出所有可用的Memory类型
Returns:
List[str]: Memory类型列表
Example:
types = mem.list_memory_types()
"""
return list(self.memory_configs.keys())
def stats(self) -> Dict[str, Any]:
"""
获取记忆统计信息
Returns:
Dict: 统计信息字典
Example:
info = mem.stats()
"""
self._ensure_initialized()
try:
stats = {}
for mem_type in self.memory_configs.keys():
memory = self._cube.get_memory(mem_type)
if memory is not None:
try:
all_memories = memory.get_all()
stats[mem_type] = {
"count": len(all_memories),
"status": "active",
"type": type(memory).__name__
}
except Exception as e:
stats[mem_type] = {
"count": 0,
"status": f"error: {e}",
"type": "unknown"
}
else:
stats[mem_type] = {
"count": 0,
"status": "uninitialized",
"type": "none"
}
stats["timestamp"] = datetime.now().isoformat()
if self.verbose:
print("📊 获取统计信息成功")
return stats
except Exception as e:
if self.verbose:
print(f"❌ 获取统计信息失败: {e}")
return {}
def dump(self, memory_type: Optional[str] = None, output_file: Optional[str] = None) -> List[Dict[str, Any]]:
"""
导出记忆
Args:
memory_type: Memory类型，None表示导出所有类型
output_file: 输出文件路径（可选）
Returns:
List[Dict]: 记忆列表
Example:
memories = mem.dump(memory_type='code_snippet_mem', output_file="backup.json")
"""
self._ensure_initialized()
try:
all_memories = []
# 确定要导出的Memory类型
export_types = [memory_type] if memory_type else list(self.memory_configs.keys())
for mem_type in export_types:
if mem_type not in self.memory_configs:
continue
memory = self._cube.get_memory(mem_type)
if memory is None:
continue
try:
memories = memory.get_all()
for item in memories:
memory_dict = {
"id": item.id,
"memory": item.memory,
"metadata": item.metadata.model_dump() if hasattr(item.metadata, 'model_dump') else item.metadata,
"memory_type": mem_type
}
all_memories.append(memory_dict)
except Exception as e:
if self.verbose:
print(f"❌ 导出 {mem_type} 失败: {e}")
continue
if output_file:
import json
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(all_memories, f, ensure_ascii=False, indent=2)
if self.verbose:
print(f"📁 已导出 {len(all_memories)} 条记忆到 {output_file}")
return all_memories
except Exception as e:
if self.verbose:
print(f"❌ 导出记忆失败: {e}")
return []
def close(self):
"""
关闭SDK，清理资源
Example:
mem.close()
"""
if self._cube:
# 清理资源（如果有的话）
self._cube = None
self._initialized = False
if self.verbose:
print("🔒 MemOS SDK已关闭")
def __enter__(self):
"""上下文管理器入口"""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""上下文管理器出口"""
self.close()
def __repr__(self):
"""字符串表示"""
status = "已初始化" if self._initialized else "未初始化"
memory_types = list(self.memory_configs.keys())
return f"MemoryOps(status={status}, memory_types={memory_types}, data_dir='{self.data_dir}')"
# 便捷函数，用于快速创建SDK实例
def create_memory_ops(data_dir: Optional[str] = None, **kwargs) -> MemoryOps:
"""
创建MemoryOps实例的便捷函数
Args:
data_dir: 数据目录
**kwargs: 其他参数传递给MemoryOps构造函数
Returns:
MemoryOps: SDK实例
Example:
mem = create_memory_ops("./my_memories")
"""
return MemoryOps(data_dir=data_dir, **kwargs)
# 全局单例实例（可选使用）
_global_instance = None
def get_global_memory_ops(**kwargs) -> MemoryOps:
"""
获取全局MemoryOps单例实例
Args:
**kwargs: 首次创建时的参数
Returns:
MemoryOps: 全局SDK实例
Example:
mem = get_global_memory_ops()
"""
global _global_instance
if _global_instance is None:
_global_instance = MemoryOps(**kwargs)
return _global_instance
class MultiMemoryOps:
"""
多Memory模块操作SDK
支持6类专用Memory模块的统一管理和操作
"""
def __init__(self, data_dir: Optional[str] = None, auto_init: bool = True, verbose: bool = False):
"""
初始化Multi Memory Ops SDK
Args:
data_dir: 数据目录路径，默认为当前目录下的memos_data
auto_init: 是否自动初始化，默认True
verbose: 是否显示详细日志，默认False
"""
# 导入bootstrap来解决依赖问题
try:
import bootstrap # noqa: F401
except ImportError:
pass # bootstrap是可选的
self.data_dir = data_dir or os.path.join(os.getcwd(), "memos_data")
self.verbose = verbose
self.memory_cube = None
# 6类专用Memory配置 - 修复为MemoryConfigFactory对象
self.memory_configs = {
'code_snippet_mem': {'backend': 'uninitialized', 'description': '代码片段记忆'},
'error_log_mem': {'backend': 'uninitialized', 'description': '错误日志记忆'},
'api_sample_mem': {'backend': 'uninitialized', 'description': 'API示例记忆'},
'file_chunk_mem': {'backend': 'uninitialized', 'description': '文件分块记忆'},
'git_commit_mem': {'backend': 'uninitialized', 'description': 'Git提交记忆'},
'product_graph_mem': {'backend': 'uninitialized', 'description': '商品图谱记忆'}
}
if auto_init:
self.initialize()
def initialize(self) -> bool:
"""
初始化多Memory系统
Returns:
bool: 初始化是否成功
"""
try:
from memos.configs.mem_cube import MultiMemoryMemCubeConfig
from memos.configs.memory import MemoryConfigFactory
from memos.mem_cube.multi import MultiMemoryMemCube
# 创建Memory配置 - 修复配置对象创建
mem_configs = {}
for name, config_info in self.memory_configs.items():
mem_configs[name] = MemoryConfigFactory(
backend=config_info['backend'],
config={}
)
# 创建MemCube配置
cube_config = MultiMemoryMemCubeConfig(
memory_configs=mem_configs
)
# 初始化MemCube
self.memory_cube = MultiMemoryMemCube(cube_config)
if self.verbose:
print(f"✅ 多Memory系统初始化成功")
print(f"📁 数据目录: {self.data_dir}")
print(f"🧠 Memory模块: {list(self.memory_configs.keys())}")
return True
except Exception as e:
if self.verbose:
print(f"❌ 多Memory系统初始化失败: {e}")
import traceback
traceback.print_exc()
return False
def add_to_memory(self, memory_name: str, content: str, metadata: Optional[Dict[str, Any]] = None) -> bool:
"""
向指定Memory添加内容
Args:
memory_name: Memory模块名称
content: 要添加的内容
metadata: 元数据
Returns:
bool: 添加是否成功
"""
if not self.memory_cube:
if self.verbose:
print("❌ Memory系统未初始化")
return False
if memory_name not in self.memory_configs:
if self.verbose:
print(f"❌ 未知的Memory模块: {memory_name}")
return False
try:
memory = self.memory_cube.get_memory(memory_name)
if memory is None:
if self.verbose:
print(f"❌ Memory模块 {memory_name} 未初始化")
return False
# 构造记忆项
memory_item = {
"memory": content,
"metadata": metadata or {}
}
memory.add([memory_item])
if self.verbose:
print(f"✅ 成功添加到 {memory_name}: {content[:50]}...")
return True
except Exception as e:
if self.verbose:
print(f"❌ 添加到 {memory_name} 失败: {e}")
return False
def search_memory(self, memory_name: str, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
"""
在指定Memory中搜索
Args:
memory_name: Memory模块名称
query: 搜索查询
top_k: 返回结果数量
Returns:
List[Dict]: 搜索结果列表
"""
if not self.memory_cube:
if self.verbose:
print("❌ Memory系统未初始化")
return []
if memory_name not in self.memory_configs:
if self.verbose:
print(f"❌ 未知的Memory模块: {memory_name}")
return []
try:
memory = self.memory_cube.get_memory(memory_name)
if memory is None:
if self.verbose:
print(f"❌ Memory模块 {memory_name} 未初始化")
return []
results = memory.search(query, top_k=top_k)
# 转换为字典格式
formatted_results = []
for item in results:
formatted_results.append({
"id": item.id,
"memory": item.memory,
"metadata": item.metadata.model_dump() if hasattr(item.metadata, 'model_dump') else item.metadata
})
if self.verbose:
print(f"🔍 在 {memory_name} 中找到 {len(formatted_results)} 条结果")
return formatted_results
except Exception as e:
if self.verbose:
print(f"❌ 在 {memory_name} 中搜索失败: {e}")
return []
def list_memory_modules(self) -> Dict[str, str]:
"""
列出所有Memory模块
Returns:
Dict[str, str]: Memory模块名称和描述的映射
"""
if not self.memory_cube:
return {}
return {name: config['description'] for name, config in self.memory_configs.items()}
def get_memory_stats(self) -> Dict[str, Dict[str, Any]]:
"""
获取所有Memory模块的统计信息
Returns:
Dict: 各Memory模块的统计信息
"""
if not self.memory_cube:
return {}
stats = {}
for name in self.memory_configs.keys():
memory = self.memory_cube.get_memory(name)
if memory is not None:
try:
all_memories = memory.get_all()
stats[name] = {
"count": len(all_memories),
"status": "active",
"type": type(memory).__name__
}
except Exception as e:
stats[name] = {
"count": 0,
"status": f"error: {e}",
"type": "unknown"
}
else:
stats[name] = {
"count": 0,
"status": "uninitialized",
"type": "none"
}
return stats