AI项目实战（7）数据分析 Agent

dangerous = ["DROP", "DELETE", "INSERT", "UPDATE", "ALTER", "CREATE"]
for kw in dangerous:
    if kw in sql_upper:
        return f"错误：不允许执行包含 {kw} 的 SQL 语句。"

forbidden = ["import os", "import subprocess", "eval(", "exec(", "open("]
for fb in forbidden:
    if fb in code:
        return f"错误：代码包含禁止的操作（{fb}），已被沙箱拦截。"

import os
os.system("rm -rf /")  # 一行代码，天台见

"""P04 数据分析 Agent - 领域模型层"""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class QueryStatus(str, Enum):
    """查询执行状态。"""
    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"
    FORBIDDEN = "forbidden"


@dataclass
class ColumnInfo:
    """数据库列信息。"""
    name: str
    type: str
    nullable: bool = True
    description: str = ""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class TableSchema:
    """数据库表结构信息。"""
    table_name: str
    columns: list[ColumnInfo] = field(default_factory=list)
    row_count: int = 0
    sample_rows: list[dict] = field(default_factory=list)

    def get_column_names(self) -> list[str]:
        return [c.name for c in self.columns]

    def to_markdown(self) -> str:
        """转换为 Markdown 表格格式。"""
        lines = [f"### 表 {self.table_name}", ""]
        lines.append(f"**总行数**: {self.row_count}")
        lines.append("")
        lines.append("| 列名 | 类型 | 说明 |")
        lines.append("|------|------|------|")
        for col in self.columns:
            lines.append(f"| {col.name} | {col.type} | {col.description or '-'} |")
        return "\n".join(lines)


@dataclass
class QueryResult:
    """SQL 查询结果。"""
    sql: str
    status: QueryStatus
    columns: list[str] = field(default_factory=list)
    rows: list[list[Any]] = field(default_factory=list)
    row_count: int = 0
    error_message: str = ""
    execution_time: float = 0.0
    executed_at: str = field(default_factory=lambda: datetime.now().isoformat())

    @property
    def success(self) -> bool:
        return self.status == QueryStatus.SUCCESS

    def to_markdown(self, max_rows: int = 20) -> str:
        if not self.success:
            return f"**查询失败**: {self.error_message}"
        if not self.rows:
            return "查询结果为空。"
        lines = [f"**查询结果**（共 {self.row_count} 行）:", ""]
        lines.append("| " + " | ".join(self.columns) + " |")
        lines.append("|" + "|".join(["------"] * len(self.columns)) + "|")
        for row in self.rows[:max_rows]:
            values = [str(v) if v is not None else "NULL" for v in row]
            lines.append("| " + " | ".join(values) + " |")
        if len(self.rows) > max_rows:
            lines.append(f"\n... (显示前 {max_rows} 行)")
        return "\n".join(lines)

"""P04 数据分析 Agent - MCP 数据访问层"""
from __future__ import annotations

import sqlite3
import time
from contextlib import contextmanager
from pathlib import Path

from core.config import DATA_DIR
from core.logging_conf import get_logger

from .models import ColumnInfo, QueryResult, QueryStatus, TableSchema

logger = get_logger("p04.data_analysis.mcp")
DB_PATH = DATA_DIR / "analysis.db"

FORBIDDEN_KEYWORDS = {"INSERT", "UPDATE", "DELETE", "DROP", "ALTER",
                      "CREATE", "TRUNCATE", "GRANT", "REVOKE"}


@contextmanager
def get_db_connection() -> sqlite3.Connection:
    """获取数据库连接（上下文管理器，自动关闭）。"""
    conn = sqlite3.connect(str(DB_PATH), timeout=10.0)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()


def validate_sql(sql: str) -> tuple[bool, str]:
    """验证 SQL 安全性。"""
    sql_upper = sql.strip().upper()

    # 必须以 SELECT 开头
    if not sql_upper.startswith("SELECT"):
        return False, "只允许执行 SELECT 查询"

    # 检查禁止的关键字
    for keyword in FORBIDDEN_KEYWORDS:
        if keyword in sql_upper:
            return False, f"不允许执行包含 {keyword} 的 SQL 语句"

    return True, ""


def execute_query(sql: str, max_rows: int = 500) -> QueryResult:
    """安全执行 SQL 查询。"""
    start_time = time.time()
    logger.info("执行 SQL 查询: %s", sql[:200])

    is_safe, error_msg = validate_sql(sql)
    if not is_safe:
        return QueryResult(
            sql=sql,
            status=QueryStatus.FORBIDDEN,
            error_message=error_msg,
        )

    try:
        with get_db_connection() as conn:
            cur = conn.execute(sql)
            columns = [desc[0] for desc in cur.description]
            rows = [list(row) for row in cur.fetchmany(max_rows)]

            return QueryResult(
                sql=sql,
                status=QueryStatus.SUCCESS,
                columns=columns,
                rows=rows,
                row_count=len(rows),
                execution_time=time.time() - start_time,
            )
    except sqlite3.Error as e:
        return QueryResult(
            sql=sql,
            status=QueryStatus.ERROR,
            error_message=f"SQL 执行错误: {e}",
            execution_time=time.time() - start_time,
        )


def get_table_schema(table_name: str = "sales") -> TableSchema:
    """获取数据库表结构信息。"""
    with get_db_connection() as conn:
        pragma = conn.execute(f"PRAGMA table_info({table_name})").fetchall()
        columns = [ColumnInfo(name=col["name"], type=col["type"]) for col in pragma]
        count = conn.execute(f"SELECT COUNT(*) FROM {table_name}").fetchone()[0]

        return TableSchema(
            table_name=table_name,
            columns=columns,
            row_count=count,
        )

"""P04 数据分析 Agent - 工具层"""
from __future__ import annotations

import io
import sys
import time

from langchain.tools import tool

from core.logging_conf import get_logger

from .mcp import execute_query, get_table_schema

logger = get_logger("p04.data_analysis.tools")
FORBIDDEN_IMPORTS = ["os", "subprocess", "sys", "shutil", "socket", "requests"]
FORBIDDEN_FUNCTIONS = ["eval(", "exec(", "open(", "compile(", "__import__"]


@tool
def query_database(sql: str) -> str:
    """对分析数据库执行 SQL 查询。

    注意：只允许 SELECT 查询，禁止任何修改数据的操作。
    """
    result = execute_query(sql)
    return result.to_markdown()


@tool
def get_table_schema_info(table_name: str = "sales") -> str:
    """获取数据库表结构信息（列名、类型、示例数据）。

    在编写 SQL 之前，**必须先调用此工具**了解表结构！
    """
    try:
        schema = get_table_schema(table_name)
        return schema.to_markdown()
    except Exception as e:
        return f"获取表结构失败: {e}"


@tool
def execute_python(code: str) -> str:
    """在安全沙箱中执行 Python 代码，用于数据分析。"""
    # 沙箱检查
    for forbidden in FORBIDDEN_IMPORTS:
        if f"import {forbidden}" in code or f"from {forbidden}" in code:
            return f"错误：禁止导入 {forbidden} 模块"

    for forbidden in FORBIDDEN_FUNCTIONS:
        if forbidden in code:
            return f"错误：禁止使用 {forbidden.replace('(', '')} 函数"

    # 重定向 stdout
    old_stdout = sys.stdout
    sys.stdout = captured_output = io.StringIO()

    try:
        # 预导入安全的常用库
        exec_globals = {"__builtins__": {
            "print": print, "len": len, "range": range,
            "list": list, "dict": dict, "str": str,
            "int": int, "float": float, "sum": sum,
            "max": max, "min": min, "sorted": sorted,
        }}
        exec("import math\nimport statistics\nimport json", exec_globals)
        exec(code, exec_globals)

        output = captured_output.getvalue()
        if len(output) > 5000:
            output = output[:5000] + "\n... (输出过长，已截断)"
        return output if output else "代码执行成功，无输出。"
    except Exception as e:
        return f"代码执行出错: {type(e).__name__}: {e}"
    finally:
        sys.stdout = old_stdout

"""P04 数据分析 Agent - 服务层"""
from __future__ import annotations

import hashlib
import time

from core.logging_conf import get_logger

from .mcp import execute_query, get_table_schema, list_tables
from .models import AnalysisReport, QueryResult

logger = get_logger("p04.data_analysis.service")
CACHE_TTL = 300
_query_cache: dict[str, tuple[QueryResult, float]] = {}


class DataAnalysisService:
    """数据分析服务。"""

    def query_with_cache(self, sql: str, use_cache: bool = True) -> QueryResult:
        """带缓存的查询。"""
        cache_key = hashlib.md5(sql.encode()).hexdigest()

        if use_cache and cache_key in _query_cache:
            result, timestamp = _query_cache[cache_key]
            if time.time() - timestamp <= CACHE_TTL:
                return result

        result = execute_query(sql)
        if use_cache and result.success:
            _query_cache[cache_key] = (result, time.time())
        return result

    def get_total_revenue(self, start_date: str = None, end_date: str = None) -> float:
        """快捷方法：查询总销售额。"""
        sql = "SELECT SUM(quantity * unit_price) FROM sales"
        conditions = []
        if start_date:
            conditions.append(f"date >= '{start_date}'")
        if end_date:
            conditions.append(f"date <= '{end_date}'")
        if conditions:
            sql += " WHERE " + " AND ".join(conditions)

        result = self.query_with_cache(sql)
        if result.success and result.rows:
            return float(result.rows[0][0] or 0)
        return 0.0

    def get_revenue_by_category(self) -> dict[str, float]:
        """快捷方法：按品类统计销售额。"""
        sql = """SELECT category, SUM(quantity * unit_price)
                 FROM sales GROUP BY category ORDER BY 2 DESC"""
        result = self.query_with_cache(sql)
        if result.success:
            return {row[0]: float(row[1] or 0) for row in result.rows}
        return {}

    def get_schema_summary(self) -> str:
        """获取数据库结构摘要（用于提示词）。"""
        tables = list_tables()
        summaries = []
        for table in tables:
            schema = get_table_schema(table)
            summaries.append(f"- **{table}**: {', '.join(schema.get_column_names())}")
        return "\n".join(summaries)

"""P04 数据分析 Agent - 项目入口层"""
from __future__ import annotations

from langchain.agents import create_agent

from core import BaseProject, build_chat_model, registry
from core.logging_conf import get_logger

from .prompts import build_system_prompt
from .tools import get_available_tools

logger = get_logger("p04.data_analysis.project")


class DataAnalysisProject(BaseProject):
    """数据分析 Agent 项目。"""

    id = "p04_data_analysis"
    name = "数据分析 Agent"
    description = "自然语言查询数据库，自动 SQL + Python 分析。"
    capabilities = ["MCP", "SQL", "Python 沙箱", "数据分析"]

    def build_agent(self) -> any:
        """构建数据分析 Agent。"""
        model = build_chat_model()
        tools = get_available_tools()
        system_prompt = build_system_prompt()
        return create_agent(model=model, tools=tools, system_prompt=system_prompt)


registry.register(DataAnalysisProject())