agentlearning_repo/01_dataanalyst/agents/sql_agent.py

"""
NL2SQL Agent：将自然语言转换为 SQL 查询，自动执行并返回分析结果。
"""

import os
import sys

# 确保父目录可导入（支持直接运行 python agents/sql_agent.py）
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import create_agent
from config import llm, db

# ---- 创建 SQL 工具包 ----
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

# ---- System Prompt ----
SQL_AGENT_PROMPT = """你是一名专业的 SQL 数据分析师。

## 工作流程
1. 先用 sql_db_list_tables 查看数据库中有哪些表
2. 用 sql_db_schema 获取相关表的字段结构和类型
3. 生成 SQL 之前，用 sql_db_query_checker 检查语法
4. 确认无误后，用 sql_db_query 执行查询
5. 用中文总结查询结果，给出简洁的业务洞察

## 约束
- 只使用数据库中实际存在的表和字段，不要凭空编造
- 单次查询结果限制在 50 条以内
- 如果查询出错，分析错误原因后重新生成 SQL
- 回答要简洁专业，不要啰嗦
"""

# ---- 创建 Agent ----
sql_agent = create_agent(
    model=llm,
    tools=tools,
    system_prompt=SQL_AGENT_PROMPT,
)


if __name__ == "__main__":
    print(f"📊 数据库：{db}")
    print(f"   可用表：{db.get_usable_table_names()}")
    print(f"\nSQL 工具包已加载（{len(tools)} 个工具）：")
    for t in tools:
        print(f"   - {t.name}")
    print("\n✅ NL2SQL Agent 创建完成，可以开始提问了！")