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.
87 lines
2.9 KiB
87 lines
2.9 KiB
#!/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()
|
|
|