setup_memory.py

"""
Memory 资源设置脚本

此脚本用于创建 AgentCore Memory Service 资源，包括：
1. 短期记忆（STM）：存储会话内的原始对话消息
2. 长期记忆（LTM）：跨会话提取和存储用户偏好及语义记忆

运行此脚本一次即可创建 Memory 资源，然后使用返回的 Memory ID 配置 Agent。

使用方法：
    python setup_memory.py
    
输出：
    - STM Memory ID：用于会话内记忆
    - LTM Memory ID：用于跨会话记忆
    
环境变量：
    AWS_REGION: AWS 区域（默认 us-west-2）
"""

import os
import sys
import uuid
import logging
from bedrock_agentcore.memory import MemoryClient


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def create_stm_memory(client: MemoryClient, name_suffix: str = None) -> dict:
    """
    创建短期记忆（STM）资源
    
    短期记忆只存储原始对话消息，不进行智能提取。
    适用于会话内的上下文理解。
    
    参数:
        client: MemoryClient 实例
        name_suffix: 名称后缀（可选，默认使用随机 UUID）
        
    返回:
        创建的 Memory 资源信息
    """
    if name_suffix is None:
        name_suffix = uuid.uuid4().hex[:8]
    
    memory_name = f"NL2SQL_STM_{name_suffix}"
    
    logger.info(f"正在创建短期记忆（STM）资源: {memory_name}")
    
    try:
        # 创建 STM，不使用任何提取策略
        stm = client.create_memory_and_wait(
            name=memory_name,
            strategies=[],  # 空列表表示不使用提取策略
            event_expiry_days=7  # 保留 7 天的对话历史
        )
        
        logger.info(f"✅ 短期记忆（STM）创建成功")
        logger.info(f"   Memory ID: {stm['id']}")
        logger.info(f"   名称: {stm['name']}")
        logger.info(f"   功能特点:")
        logger.info(f"   - 存储原始对话消息")
        logger.info(f"   - 仅在同一会话内有效")
        logger.info(f"   - 即时检索，无需处理时间")
        logger.info(f"   - 对话历史保留 7 天")
        
        return stm
        
    except Exception as e:
        logger.error(f"❌ 创建短期记忆（STM）失败: {str(e)}")
        raise


def create_ltm_memory(client: MemoryClient, name_suffix: str = None) -> dict:
    """
    创建长期记忆（LTM）资源
    
    长期记忆使用智能提取策略，自动从对话中提取用户偏好和语义信息。
    适用于跨会话的个性化服务。
    
    参数:
        client: MemoryClient 实例
        name_suffix: 名称后缀（可选，默认使用随机 UUID）
        
    返回:
        创建的 Memory 资源信息
    """
    if name_suffix is None:
        name_suffix = uuid.uuid4().hex[:8]
    
    memory_name = f"NL2SQL_LTM_{name_suffix}"
    
    logger.info(f"正在创建长期记忆（LTM）资源: {memory_name}")
    
    try:
        # 创建 LTM，使用用户偏好和语义记忆策略
        ltm = client.create_memory_and_wait(
            name=memory_name,
            strategies=[
                # 用户偏好记忆策略
                # 自动提取用户的偏好信息，如"我喜欢查看销售数据"
                {
                    "userPreferenceMemoryStrategy": {
                        "name": "user_preferences",
                        "namespaces": ["/user/preferences", "/user/query_patterns"]
                    }
                },
                # 语义记忆策略
                # 自动提取和存储重要的事实和知识，如"我们公司的财年从4月开始"
                {
                    "semanticMemoryStrategy": {
                        "name": "semantic_facts",
                        "namespaces": ["/user/facts", "/business/context"]
                    }
                }
            ],
            event_expiry_days=30  # 保留 30 天的对话历史
        )
        
        logger.info(f"✅ 长期记忆（LTM）创建成功")
        logger.info(f"   Memory ID: {ltm['id']}")
        logger.info(f"   名称: {ltm['name']}")
        logger.info(f"   功能特点:")
        logger.info(f"   - 包含 STM 的所有功能")
        logger.info(f"   - 自动提取用户偏好（如常用查询模式）")
        logger.info(f"   - 自动提取语义事实（如业务上下文）")
        logger.info(f"   - 跨会话记忆，不同会话也能访问")
        logger.info(f"   - 提取处理需要 5-10 秒")
        logger.info(f"   - 对话历史保留 30 天")
        
        return ltm
        
    except Exception as e:
        logger.error(f"❌ 创建长期记忆（LTM）失败: {str(e)}")
        raise


def main():
    """
    主函数：创建 STM 和 LTM Memory 资源
    """
    print("=" * 70)
    print("AgentCore Memory 资源创建工具")
    print("=" * 70)
    print()
    
    # 获取 AWS 区域
    region = os.getenv('AWS_REGION', 'us-west-2')
    logger.info(f"使用 AWS 区域: {region}")
    
    try:
        # 初始化 Memory Client
        logger.info("正在连接到 AgentCore Memory Service...")
        client = MemoryClient(region_name=region)
        logger.info("✅ 成功连接到 Memory Service")
        print()
        
        # 生成唯一的名称后缀
        name_suffix = uuid.uuid4().hex[:8]
        
        # 创建短期记忆（STM）
        print("-" * 70)
        print("1. 创建短期记忆（STM）")
        print("-" * 70)
        stm = create_stm_memory(client, name_suffix)
        print()
        
        # 创建长期记忆（LTM）
        print("-" * 70)
        print("2. 创建长期记忆（LTM）")
        print("-" * 70)
        ltm = create_ltm_memory(client, name_suffix)
        print()
        
        # 输出使用说明
        print("=" * 70)
        print("✅ Memory 资源创建完成！")
        print("=" * 70)
        print()
        print("请选择要使用的 Memory 类型：")
        print()
        print("【选项 1】使用短期记忆（STM）- 适合快速原型和简单场景")
        print(f"  export MEMORY_ID={stm['id']}")
        print()
        print("【选项 2】使用长期记忆（LTM）- 适合生产环境和个性化服务")
        print(f"  export MEMORY_ID={ltm['id']}")
        print()
        print("=" * 70)
        print()
        print("使用示例：")
        print()
        print("# 1. 设置环境变量")
        print(f"export MEMORY_ID={ltm['id']}")
        print()
        print("# 2. 在代码中使用")
        print("from bedrock_agentcore.memory import MemoryClient")
        print("from nl2sql_agent import MemoryHook")
        print()
        print("memory_client = MemoryClient(region_name='us-west-2')")
        print("memory_hook = MemoryHook(")
        print("    memory_client=memory_client,")
        print(f"    memory_id='{ltm['id']}'")
        print(")")
        print()
        print("# 3. 部署到 AgentCore Runtime")
        print(f"export MEMORY_ID={ltm['id']}")
        print("agentcore configure -e agent.py")
        print("agentcore launch")
        print()
        print("=" * 70)
        
        # 保存 Memory ID 到文件（可选）
        config_file = "memory_config.txt"
        try:
            with open(config_file, 'w', encoding='utf-8') as f:
                f.write(f"# AgentCore Memory 配置\n")
                f.write(f"# 创建时间: {uuid.uuid4()}\n\n")
                f.write(f"# 短期记忆（STM）\n")
                f.write(f"STM_MEMORY_ID={stm['id']}\n")
                f.write(f"STM_MEMORY_NAME={stm['name']}\n\n")
                f.write(f"# 长期记忆（LTM）\n")
                f.write(f"LTM_MEMORY_ID={ltm['id']}\n")
                f.write(f"LTM_MEMORY_NAME={ltm['name']}\n\n")
                f.write(f"# 推荐使用 LTM\n")
                f.write(f"MEMORY_ID={ltm['id']}\n")
            
            logger.info(f"Memory 配置已保存到文件: {config_file}")
        except Exception as e:
            logger.warning(f"保存配置文件失败: {str(e)}")
        
        return 0
        
    except Exception as e:
        logger.error(f"❌ 创建 Memory 资源失败: {str(e)}")
        logger.error("请检查：")
        logger.error("  1. AWS 凭证是否正确配置")
        logger.error("  2. 是否有权限访问 AgentCore Memory Service")
        logger.error("  3. AWS 区域是否正确")
        return 1


if __name__ == "__main__":
    sys.exit(main())