example_agent_usage.py

"""
NL2SQL Agent 使用示例

演示如何使用 NL2SQLAgent 处理自然语言查询。
"""

import logging
import os
from nl2sql_agent import NL2SQLAgent, DatabaseConfig

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

logger = logging.getLogger(__name__)


def example_basic_usage():
    """
    示例 1: 基本使用（不使用 Memory）
    """
    logger.info("\n" + "=" * 60)
    logger.info("示例 1: 基本使用")
    logger.info("=" * 60)
    
    try:
        # 1. 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 2. 初始化 Agent（使用 Mock LLM 进行演示）
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True  # 使用 Mock LLM 避免实际调用 Bedrock
        )
        
        # 3. 处理查询
        response = agent.process_query(
            user_query="查询所有客户的数量",
            session_id="demo-session-1"
        )
        
        # 4. 查看结果
        logger.info("\n查询结果:")
        logger.info(f"成功: {response.success}")
        logger.info(f"分析报告:\n{response.analysis}")
        logger.info(f"执行的 SQL: {response.sql_statement}")
        logger.info(f"执行时间: {response.execution_time:.2f}秒")
        logger.info(f"重试次数: {response.retry_count}")
        
        # 5. 清理资源
        agent.close()
        
    except Exception as e:
        logger.error(f"示例执行失败: {str(e)}")


def example_with_memory():
    """
    示例 2: 使用 Memory Service（支持多轮对话）
    """
    logger.info("\n" + "=" * 60)
    logger.info("示例 2: 使用 Memory Service")
    logger.info("=" * 60)
    
    try:
        # 1. 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 2. 初始化 Agent（提供 memory_id）
        memory_id = os.getenv("MEMORY_ID", "demo-memory-id")
        
        agent = NL2SQLAgent(
            db_config=db_config,
            memory_id=memory_id,
            use_mock_llm=True
        )
        
        # 3. 第一轮对话
        logger.info("\n第一轮对话:")
        response1 = agent.process_query(
            user_query="查询订单总金额超过 1000 的客户",
            session_id="demo-session-2"
        )
        logger.info(f"分析报告:\n{response1.analysis}")
        
        # 4. 第二轮对话（基于上下文）
        logger.info("\n第二轮对话:")
        response2 = agent.process_query(
            user_query="这些客户都来自哪些城市？",  # Agent 会理解"这些客户"
            session_id="demo-session-2"  # 使用相同的 session_id
        )
        logger.info(f"分析报告:\n{response2.analysis}")
        
        # 5. 清理资源
        agent.close()
        
    except Exception as e:
        logger.error(f"示例执行失败: {str(e)}")
        logger.info("注意: 如果没有配置 AWS 凭证或 Memory ID，此示例可能会失败")


def example_with_env_vars():
    """
    示例 3: 使用环境变量配置
    """
    logger.info("\n" + "=" * 60)
    logger.info("示例 3: 使用环境变量配置")
    logger.info("=" * 60)
    
    try:
        # 检查环境变量是否已设置
        required_vars = ['DB_ENDPOINT', 'DB_NAME', 'DB_USER', 'DB_PASSWORD']
        missing_vars = [var for var in required_vars if not os.getenv(var)]
        
        if missing_vars:
            logger.warning(f"缺少环境变量: {', '.join(missing_vars)}")
            logger.info("跳过此示例，请设置环境变量后重试")
            return
        
        # 从环境变量加载配置
        db_config = DatabaseConfig.from_env()
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            memory_id=os.getenv('MEMORY_ID'),
            use_mock_llm=True
        )
        
        # 处理查询
        response = agent.process_query(
            user_query="查询所有客户的数量",
            session_id="demo-session-3"
        )
        
        logger.info(f"\n分析报告:\n{response.analysis}")
        
        # 清理资源
        agent.close()
        
    except Exception as e:
        logger.error(f"示例执行失败: {str(e)}")


def example_multi_turn_conversation():
    """
    示例 4: 多轮对话场景
    """
    logger.info("\n" + "=" * 60)
    logger.info("示例 4: 多轮对话场景")
    logger.info("=" * 60)
    
    try:
        # 初始化 Agent
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        session_id = "demo-session-4"
        
        # 第一轮：查询概览
        logger.info("\n第一轮：查询概览")
        response1 = agent.process_query(
            user_query="查询每个城市的客户数量",
            session_id=session_id
        )
        logger.info(f"分析报告:\n{response1.analysis}")
        
        # 第二轮：基于前一次结果的追问
        logger.info("\n第二轮：基于前一次结果的追问")
        response2 = agent.process_query(
            user_query="他们的平均订单金额是多少？",
            session_id=session_id
        )
        logger.info(f"分析报告:\n{response2.analysis}")
        
        # 第三轮：换个角度分析
        logger.info("\n第三轮：换个角度分析")
        response3 = agent.process_query(
            user_query="换个角度，按订单数量排序",
            session_id=session_id
        )
        logger.info(f"分析报告:\n{response3.analysis}")
        
        # 清理资源
        agent.close()
        
    except Exception as e:
        logger.error(f"示例执行失败: {str(e)}")


def main():
    """
    运行所有示例
    """
    logger.info("\n" + "=" * 60)
    logger.info("NL2SQL Agent 使用示例")
    logger.info("=" * 60)
    
    examples = [
        ("基本使用", example_basic_usage),
        ("使用 Memory Service", example_with_memory),
        ("使用环境变量配置", example_with_env_vars),
        ("多轮对话场景", example_multi_turn_conversation),
    ]
    
    for example_name, example_func in examples:
        try:
            logger.info(f"\n运行示例: {example_name}")
            example_func()
        except Exception as e:
            logger.error(f"示例 '{example_name}' 执行失败: {str(e)}")
        
        logger.info("\n" + "-" * 60)
    
    logger.info("\n所有示例执行完成")


if __name__ == "__main__":
    main()