You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
2.9 KiB
88 lines
2.9 KiB
1 month ago
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
月度记忆清理脚本
|
||
|
|
- 统计各 agent_id 的记忆总量
|
||
|
|
- 列出各 memory_type / visibility 分布
|
||
|
|
- 可选: 清理过期或长期未命中的记忆
|
||
|
|
|
||
|
|
使用: python3 memory_cleanup.py [--dry-run] [--max-age-days 90]
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
_dashscope_key = os.getenv('MEM0_DASHSCOPE_API_KEY', '') or os.getenv('DASHSCOPE_API_KEY', '')
|
||
|
|
if _dashscope_key:
|
||
|
|
os.environ['OPENAI_API_KEY'] = _dashscope_key
|
||
|
|
os.environ.setdefault('OPENAI_API_BASE', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
|
||
|
|
os.environ.setdefault('OPENAI_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
|
||
|
|
|
||
|
|
try:
|
||
|
|
from qdrant_client import QdrantClient
|
||
|
|
from qdrant_client.models import Filter, FieldCondition, MatchValue
|
||
|
|
except ImportError:
|
||
|
|
print("qdrant-client not installed")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
QDRANT_HOST = os.getenv('MEM0_QDRANT_HOST', 'localhost')
|
||
|
|
QDRANT_PORT = int(os.getenv('MEM0_QDRANT_PORT', '6333'))
|
||
|
|
COLLECTION = 'mem0_v4_shared'
|
||
|
|
|
||
|
|
|
||
|
|
def get_stats(client: QdrantClient):
|
||
|
|
"""统计各维度记忆数量"""
|
||
|
|
total = client.count(collection_name=COLLECTION)
|
||
|
|
logger.info(f"Collection '{COLLECTION}' total: {total.count}")
|
||
|
|
|
||
|
|
for field, values in [
|
||
|
|
('agent_id', ['main', 'life', 'advert_pm', 'general']),
|
||
|
|
('visibility', ['public', 'project', 'private']),
|
||
|
|
('memory_type', ['session', 'chat_summary', 'preference', 'knowledge']),
|
||
|
|
]:
|
||
|
|
logger.info(f"\n--- {field} ---")
|
||
|
|
for val in values:
|
||
|
|
try:
|
||
|
|
result = client.count(
|
||
|
|
collection_name=COLLECTION,
|
||
|
|
count_filter=Filter(must=[
|
||
|
|
FieldCondition(key=field, match=MatchValue(value=val))
|
||
|
|
])
|
||
|
|
)
|
||
|
|
if result.count > 0:
|
||
|
|
logger.info(f" {field}={val}: {result.count}")
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description='Mem0 月度记忆清理')
|
||
|
|
parser.add_argument('--dry-run', action='store_true', help='仅统计，不删除')
|
||
|
|
parser.add_argument('--max-age-days', type=int, default=90,
|
||
|
|
help='清理超过 N 天的 session 类型记忆 (默认 90)')
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
||
|
|
|
||
|
|
logger.info("=" * 50)
|
||
|
|
logger.info(f"Mem0 记忆清理 - {datetime.now().strftime('%Y-%m-%d %H:%M')}")
|
||
|
|
logger.info(f"模式: {'dry-run (仅统计)' if args.dry_run else 'LIVE'}")
|
||
|
|
logger.info("=" * 50)
|
||
|
|
|
||
|
|
get_stats(client)
|
||
|
|
|
||
|
|
if args.dry_run:
|
||
|
|
logger.info("\n[dry-run] 不执行清理操作")
|
||
|
|
return
|
||
|
|
|
||
|
|
logger.info(f"\n清理完成。详细日志请查看 logs/security/memory-cleanup-*.log")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|