From 166990bc81836f30f32cb3fe7462e332e8f07d15 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 16 Sep 2025 15:23:18 +0800 Subject: [PATCH] =?UTF-8?q?feat(MySQL):=20=E9=9B=86=E6=88=90=20MySQL=20?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E5=BA=93=E6=9F=A5=E8=AF=A2=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 添加 MySQL 连接管理器，支持连接和查询数据库 - 实现获取表名、描述表结构和执行 SQL 查询的工具 - 增强安全性，防止 SQL 注入和限制查询结果大小 - 更新 README.md，提供 MySQL 配置示例和使用说明 - 添加 MySQL 连接测试脚本，验证连接和工具功能 --- README.md | 53 ++++ pyproject.toml | 1 + src/.env.template | 9 + src/agents/chatbot/tools.py | 10 +- src/agents/common/base.py | 3 +- src/agents/common/mcp.py | 4 + src/agents/common/toolkits/mysql/__init__.py | 3 + .../common/toolkits/mysql/connection.py | 200 +++++++++++++ .../common/toolkits/mysql/exceptions.py | 34 +++ src/agents/common/toolkits/mysql/security.py | 87 ++++++ src/agents/common/toolkits/mysql/tools.py | 280 ++++++++++++++++++ src/agents/common/tools.py | 4 + test/test_mysql_connection.py | 206 +++++++++++++ test/test_mysql_import.py | 24 ++ .../ToolCallingResult/ToolResultRenderer.vue | 30 ++ 15 files changed, 942 insertions(+), 6 deletions(-) create mode 100644 src/agents/common/toolkits/mysql/__init__.py create mode 100644 src/agents/common/toolkits/mysql/connection.py create mode 100644 src/agents/common/toolkits/mysql/exceptions.py create mode 100644 src/agents/common/toolkits/mysql/security.py create mode 100644 src/agents/common/toolkits/mysql/tools.py create mode 100644 test/test_mysql_connection.py create mode 100644 test/test_mysql_import.py diff --git a/README.md b/README.md index a8caadd3..2cc8ad38 100644 --- a/README.md +++ b/README.md @@ -261,6 +261,59 @@ agent_manager.register_agent(ChatbotAgent) agent_manager.init_all_agents() ``` +### MySQL 数据库查询集成（Beta） + +项目目前已经支持智能体去查询 MySQL 数据库如果想要接入数据库，则可以在环境变量中配置如下信息： + +```sh +# 基础配置示例 +MYSQL_HOST=192.168.1.100 +MYSQL_USER=username +MYSQL_PASSWORD=your_secure_password +MYSQL_DATABASE=database_name +MYSQL_PORT=3306 +MYSQL_CHARSET=utf8mb4 +``` + +然后在智能体配置的工具中，勾选上 `mysql` 开头的几个工具即可。 + +- **mysql_list_tables**: 获取数据库中的所有表名 +- **mysql_describe_table**: 获取指定表的详细结构信息 +- **mysql_query**: 执行只读的 SQL 查询语句 + +**安全特性** + +- ✅ 只允许 SELECT、SHOW、DESCRIBE、EXPLAIN 操作 +- ✅ 表名参数验证，严格的 SQL 注入防护 +- ✅ 查询超时控制（默认10秒，最大60秒） +- ✅ 结果大小、行数限制（默认10000字符，100行，最大1000行） + +**注意事项** + +1. 确保数据库用户只有只读权限 +2. 大表查询建议使用 LIMIT 子句 +3. 复杂查询可能需要调整超时时间 +4. 查询结果过大会被自动截断并提示 + +### 图表可视化绘制 - MCP（Beta） + +这个是基于 @antvis 团队开发的 [可视化图表-MCP-Server](https://www.modelscope.cn/mcp/servers/@antvis/mcp-server-chart)，可以在魔搭社区中配置 Host 资源后，在 src/agents/common/mcp.py 的 `MCP_SERVERS` 中添加 mcp-server，需要注意的是记得将 `type` 字段修改为 `transport`。 + +```python +# MCP Server configurations +MCP_SERVERS = { + "sequentialthinking": { + "url": "https://remote.mcpservers.org/sequentialthinking/mcp", + "transport": "streamable_http", + }, + + "mcp-server-chart": { + "url": "https://mcp.api-inference.modelscope.net/9993ae42524c4c/mcp", + "transport": "streamable_http", + } +} +``` + ### 服务端口说明 | 端口 | 服务 | 说明 | diff --git a/pyproject.toml b/pyproject.toml index a2adac3d..66ef321d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ dependencies = [ "tabulate>=0.9.0", "minio>=7.2.7", "Pillow>=10.5.0", + "pymysql>=1.1.0", ] [tool.ruff] line-length = 120 # 代码最大行宽 diff --git a/src/.env.template b/src/.env.template index 93424cb0..fe253aae 100644 --- a/src/.env.template +++ b/src/.env.template @@ -14,3 +14,12 @@ TOGETHER_API_KEY=tgp_v1_fPjW******irD6zesAn4 # 功能服务 TAVILY_API_KEY=tvly-3gR4ind9******JMOxw3E2LG # <<< 配置网络搜索 + + +# 基础配置示例 +MYSQL_HOST=192.168.1.100 +MYSQL_USER=username +MYSQL_PASSWORD=your_secure_password +MYSQL_DATABASE=database_name +MYSQL_PORT=3306 +MYSQL_CHARSET=utf8mb4 diff --git a/src/agents/chatbot/tools.py b/src/agents/chatbot/tools.py index 756f0564..7fdf4572 100644 --- a/src/agents/chatbot/tools.py +++ b/src/agents/chatbot/tools.py @@ -8,6 +8,7 @@ from langchain_core.tools import tool from src.agents.common.tools import get_buildin_tools from src.utils import logger from src.utils.minio_utils import upload_image_to_minio +from src.agents.common.toolkits.mysql import get_mysql_tools @tool @@ -30,6 +31,7 @@ def calculator(a: float, b: float, operation: str) -> float: logger.error(f"Calculator error: {e}") raise + @tool async def text_to_img_kolors(text: str) -> str: """（用来测试文件存储）使用Kolors模型生成图片， 会返回图片的URL""" @@ -42,12 +44,9 @@ async def text_to_img_kolors(text: str) -> str: "image_size": "512x512", "batch_size": 1, "num_inference_steps": 20, - "guidance_scale": 7.5 - } - headers = { - "Authorization": f"Bearer {os.getenv('SILICONFLOW_API_KEY')}", - "Content-Type": "application/json" + "guidance_scale": 7.5, } + headers = {"Authorization": f"Bearer {os.getenv('SILICONFLOW_API_KEY')}", "Content-Type": "application/json"} try: response = requests.post(url, json=payload, headers=headers) @@ -76,4 +75,5 @@ def get_tools() -> list[Any]: tools = get_buildin_tools() tools.append(calculator) tools.append(text_to_img_kolors) + tools.extend(get_mysql_tools()) return tools diff --git a/src/agents/common/base.py b/src/agents/common/base.py index 3d96cfb3..c784071f 100644 --- a/src/agents/common/base.py +++ b/src/agents/common/base.py @@ -53,8 +53,9 @@ class BaseAgent: context = self.context_schema.from_file(module_name=self.module_name, input_context=input_context) logger.debug(f"stream_messages: {context}") # TODO 的 Checkpointer 似乎还没有适配最新的 Context API + input_config = {"configurable": input_context, "recursion_limit": 100} async for msg, metadata in graph.astream( - {"messages": messages}, stream_mode="messages", context=context, config={"configurable": input_context} + {"messages": messages}, stream_mode="messages", context=context, config=input_config ): yield msg, metadata diff --git a/src/agents/common/mcp.py b/src/agents/common/mcp.py index c87d3858..5f132e14 100644 --- a/src/agents/common/mcp.py +++ b/src/agents/common/mcp.py @@ -24,6 +24,10 @@ MCP_SERVERS = { # "args": ["mcp-server-time"], # "transport": "stdio", # }, + # "mcp-server-chart": { + # "url": "https://mcp.api-inference.modelscope.net/9993ae42524c4c/mcp", + # "transport": "streamable_http", + # } } diff --git a/src/agents/common/toolkits/mysql/__init__.py b/src/agents/common/toolkits/mysql/__init__.py new file mode 100644 index 00000000..0480f381 --- /dev/null +++ b/src/agents/common/toolkits/mysql/__init__.py @@ -0,0 +1,3 @@ +from .tools import get_mysql_tools + +__all__ = ["get_mysql_tools"] diff --git a/src/agents/common/toolkits/mysql/connection.py b/src/agents/common/toolkits/mysql/connection.py new file mode 100644 index 00000000..cbb748c4 --- /dev/null +++ b/src/agents/common/toolkits/mysql/connection.py @@ -0,0 +1,200 @@ +import signal +import time +from typing import Any +from contextlib import contextmanager +import threading +import pymysql +from pymysql.cursors import DictCursor +from pymysql import MySQLError +from src.utils import logger + + +class MySQLConnectionManager: + """MySQL 数据库连接管理器""" + + def __init__(self, config: dict[str, Any]): + self.config = config + self.connection = None + self._lock = threading.Lock() + self.last_connection_time = 0 + self.max_connection_age = 3600 # 1小时后重新连接 + + def _get_connection(self) -> pymysql.Connection: + """获取数据库连接""" + current_time = time.time() + + # 检查连接是否过期或断开 + if ( + self.connection is None + or not self.connection.open + or current_time - self.last_connection_time > self.max_connection_age + ): + with self._lock: + # 双重检查 + if ( + self.connection is None + or not self.connection.open + or current_time - self.last_connection_time > self.max_connection_age + ): + # 关闭旧连接 + if self.connection and self.connection.open: + try: + self.connection.close() + except Exception as _: + pass + + # 创建新连接 + self.connection = self._create_connection() + self.last_connection_time = current_time + + return self.connection + + def _create_connection(self) -> pymysql.Connection: + """创建新的数据库连接""" + max_retries = 3 + for attempt in range(max_retries): + try: + connection = pymysql.connect( + host=self.config["host"], + user=self.config["user"], + password=self.config["password"], + database=self.config["database"], + port=self.config["port"], + charset=self.config.get("charset", "utf8mb4"), + cursorclass=DictCursor, + connect_timeout=10, + read_timeout=60, # 增加读取超时 + write_timeout=30, + autocommit=True, # 自动提交 + ) + logger.info(f"MySQL connection established successfully (attempt {attempt + 1})") + return connection + + except MySQLError as e: + logger.warning(f"Connection attempt {attempt + 1} failed: {e}") + if attempt < max_retries - 1: + time.sleep(2**attempt) # 指数退避 + else: + logger.error(f"Failed to connect to MySQL after {max_retries} attempts: {e}") + raise ConnectionError(f"MySQL connection failed: {e}") + + def test_connection(self) -> bool: + """测试连接是否有效""" + try: + if self.connection and self.connection.open: + # 执行简单查询测试连接 + with self.connection.cursor() as cursor: + cursor.execute("SELECT 1") + cursor.fetchone() + return True + except Exception as _: + pass + return False + + @contextmanager + def get_cursor(self): + """获取数据库游标的上下文管理器""" + max_retries = 2 + for attempt in range(max_retries): + try: + connection = self._get_connection() + cursor = connection.cursor() + + try: + yield cursor + connection.commit() + break # 成功，退出重试循环 + + except Exception as e: + connection.rollback() + + # 如果是连接错误，尝试重新连接 + if "MySQL" in str(e) or "connection" in str(e).lower(): + if attempt < max_retries - 1: + logger.warning(f"Connection error, retrying (attempt {attempt + 1}): {e}") + # 强制重新连接 + if self.connection: + try: + self.connection.close() + except Exception as _: + pass + self.connection = None + time.sleep(1) + continue + + raise e # 其他错误直接抛出 + + finally: + if cursor: + cursor.close() + + except Exception as e: + if attempt == max_retries - 1: + raise e # 最后一次尝试失败，抛出异常 + time.sleep(1) + + def close(self): + """关闭数据库连接""" + if self.connection: + self.connection.close() + self.connection = None + logger.info("MySQL connection closed") + + +class QueryTimeoutError(Exception): + """查询超时异常""" + + pass + + +class QueryResultTooLargeError(Exception): + """查询结果过大异常""" + + pass + + +def execute_query_with_timeout(connection: pymysql.Connection, sql: str, params: tuple = None, timeout: int = 10): + """带超时的查询执行""" + + def timeout_handler(_signum, _frame): + raise QueryTimeoutError(f"Query timeout after {timeout} seconds") + + # 设置信号处理 + old_handler = signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(timeout) + + try: + cursor = connection.cursor(DictCursor) + cursor.execute(sql, params or ()) + result = cursor.fetchall() + cursor.close() + return result + finally: + # 恢复原始信号处理 + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + +def limit_result_size(result: list, max_chars: int = 10000) -> list: + """限制结果大小""" + if not result: + return result + + # 计算结果的字符大小 + result_str = str(result) + if len(result_str) > max_chars: + # 返回部分结果并提示 + limited_result = [] + current_chars = 0 + for row in result: + row_str = str(row) + if current_chars + len(row_str) > max_chars: + break + limited_result.append(row) + current_chars += len(row_str) + + # 记录警告 + logger.warning(f"Query result truncated from {len(result)} to {len(limited_result)} rows due to size limit") + return limited_result + + return result diff --git a/src/agents/common/toolkits/mysql/exceptions.py b/src/agents/common/toolkits/mysql/exceptions.py new file mode 100644 index 00000000..96053b0d --- /dev/null +++ b/src/agents/common/toolkits/mysql/exceptions.py @@ -0,0 +1,34 @@ +class MySQLToolError(Exception): + """MySQL 工具基础异常""" + + pass + + +class MySQLConnectionError(MySQLToolError): + """MySQL 连接异常""" + + pass + + +class MySQLQueryError(MySQLToolError): + """MySQL 查询异常""" + + pass + + +class MySQLSecurityError(MySQLToolError): + """MySQL 安全异常""" + + pass + + +class MySQLTimeoutError(MySQLToolError): + """MySQL 超时异常""" + + pass + + +class MySQLResultTooLargeError(MySQLToolError): + """MySQL 结果过大异常""" + + pass diff --git a/src/agents/common/toolkits/mysql/security.py b/src/agents/common/toolkits/mysql/security.py new file mode 100644 index 00000000..dcb91556 --- /dev/null +++ b/src/agents/common/toolkits/mysql/security.py @@ -0,0 +1,87 @@ +import re + + +class MySQLSecurityChecker: + """MySQL 安全检查器""" + + # 允许的SQL操作 + ALLOWED_OPERATIONS = {"SELECT", "SHOW", "DESCRIBE", "EXPLAIN"} + + # 危险的关键词 + DANGEROUS_KEYWORDS = { + "DROP", + "DELETE", + "UPDATE", + "INSERT", + "CREATE", + "ALTER", + "TRUNCATE", + "REPLACE", + "LOAD", + "GRANT", + "REVOKE", + "SET", + "COMMIT", + "ROLLBACK", + "UNLOCK", + "KILL", + "SHUTDOWN", + } + + @classmethod + def validate_sql(cls, sql: str) -> bool: + """验证SQL语句的安全性""" + if not sql: + return False + + # 标准化SQL + sql_upper = sql.strip().upper() + + # 检查是否是允许的操作 + if not any(sql_upper.startswith(op) for op in cls.ALLOWED_OPERATIONS): + return False + + # 检查危险关键词 + for keyword in cls.DANGEROUS_KEYWORDS: + if keyword in sql_upper: + return False + + # 检查SQL注入模式 + sql_injection_patterns = [ + r"\bor\s+1\s*=\s*1\b", + r"\bunion\s+select\b", + r"\bexec\s*\(", + r"\bxp_cmdshell\b", + r"\bsleep\s*\(", + r"\bbenchmark\s*\(", + r"\bwaitfor\s+delay\b", + r"\b;\s*drop\b", + r"\b;\s*delete\b", + r"\b;\s*update\b", + r"\b;\s*insert\b", + ] + + for pattern in sql_injection_patterns: + if re.search(pattern, sql_upper, re.IGNORECASE): + return False + + return True + + @classmethod + def validate_table_name(cls, table_name: str) -> bool: + """验证表名的安全性""" + if not table_name: + return False + + # 检查表名只包含字母、数字、下划线 + return bool(re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table_name)) + + @classmethod + def validate_limit(cls, limit: int) -> bool: + """验证limit参数""" + return isinstance(limit, int) and 0 < limit <= 1000 + + @classmethod + def validate_timeout(cls, timeout: int) -> bool: + """验证timeout参数""" + return isinstance(timeout, int) and 1 <= timeout <= 60 diff --git a/src/agents/common/toolkits/mysql/tools.py b/src/agents/common/toolkits/mysql/tools.py new file mode 100644 index 00000000..cc13cdf5 --- /dev/null +++ b/src/agents/common/toolkits/mysql/tools.py @@ -0,0 +1,280 @@ +from typing import Annotated, Any +from langchain_core.tools import tool +from pydantic import BaseModel, Field + +from .connection import MySQLConnectionManager, limit_result_size +from .security import MySQLSecurityChecker +from .exceptions import MySQLConnectionError +from src.utils import logger + + +# 全局连接管理器实例 +_connection_manager: MySQLConnectionManager | None = None + + +def get_connection_manager() -> MySQLConnectionManager: + """获取全局连接管理器""" + global _connection_manager + if _connection_manager is None: + import os + + # 从环境变量中读取 MySQL 配置 + mysql_config = { + "host": os.getenv("MYSQL_HOST"), + "user": os.getenv("MYSQL_USER"), + "password": os.getenv("MYSQL_PASSWORD"), + "database": os.getenv("MYSQL_DATABASE"), + "port": int(os.getenv("MYSQL_PORT", "3306")), + "charset": "utf8mb4", + } + # 验证配置完整性 + required_keys = ["host", "user", "password", "database"] + for key in required_keys: + if not mysql_config[key]: + raise MySQLConnectionError(f"MySQL configuration missing required key: {key}") + + _connection_manager = MySQLConnectionManager(mysql_config) + return _connection_manager + + +class TableListModel(BaseModel): + """获取表名列表的参数模型""" + + pass + + +@tool(args_schema=TableListModel) +def mysql_list_tables() -> str: + """获取数据库中的所有表名 + + 这个工具用来列出当前数据库中所有的表名，帮助你了解数据库的结构。 + """ + try: + conn_manager = get_connection_manager() + + with conn_manager.get_cursor() as cursor: + # 获取表名 + cursor.execute("SHOW TABLES") + tables = cursor.fetchall() + + if not tables: + return "数据库中没有找到任何表" + + # 提取表名 + table_names = [] + for table in tables: + table_name = list(table.values())[0] + table_names.append(table_name) + + # 获取每个表的行数信息 + table_info = [] + for table_name in table_names: + try: + cursor.execute(f"SELECT COUNT(*) as count FROM `{table_name}`") + count_result = cursor.fetchone() + row_count = count_result["count"] + table_info.append(f"- {table_name} (约 {row_count} 行)") + except Exception: + table_info.append(f"- {table_name} (无法获取行数)") + + result = "数据库中的表:\n" + "\n".join(table_info) + logger.info(f"Retrieved {len(table_names)} tables from database") + return result + + except Exception as e: + error_msg = f"获取表名失败: {str(e)}" + logger.error(error_msg) + return error_msg + + +class TableDescribeModel(BaseModel): + """获取表结构的参数模型""" + + table_name: str = Field(description="要查询的表名", example="users") + + +@tool(args_schema=TableDescribeModel) +def mysql_describe_table(table_name: Annotated[str, "要查询结构的表名"]) -> str: + """获取指定表的详细结构信息 + + 这个工具用来查看表的字段信息、数据类型、是否允许NULL、默认值、键类型等。 + 帮助你了解表的结构，以便编写正确的SQL查询。 + """ + try: + # 验证表名安全性 + if not MySQLSecurityChecker.validate_table_name(table_name): + return "表名包含非法字符，请检查表名" + + conn_manager = get_connection_manager() + + with conn_manager.get_cursor() as cursor: + # 获取表结构 + cursor.execute(f"DESCRIBE `{table_name}`") + columns = cursor.fetchall() + + if not columns: + return f"表 {table_name} 不存在或没有字段" + + # 格式化输出 + result = f"表 `{table_name}` 的结构:\n\n" + result += "字段名\t\t类型\t\tNULL\t键\t默认值\t\t额外\n" + result += "-" * 80 + "\n" + + for col in columns: + field = col["Field"] or "" + type_str = col["Type"] or "" + null_str = col["Null"] or "" + key_str = col["Key"] or "" + default_str = col.get("Default") or "" + extra_str = col.get("Extra") or "" + + # 格式化输出 + result += f"{field:<16}\t{type_str:<16}\t{null_str:<8}\t{key_str:<4}\t{default_str:<16}\t{extra_str}\n" + + # 获取索引信息 + try: + cursor.execute(f"SHOW INDEX FROM `{table_name}`") + indexes = cursor.fetchall() + + if indexes: + result += "\n索引信息:\n" + index_dict = {} + for idx in indexes: + key_name = idx["Key_name"] + if key_name not in index_dict: + index_dict[key_name] = [] + index_dict[key_name].append(idx["Column_name"]) + + for key_name, columns in index_dict.items(): + result += f"- {key_name}: {', '.join(columns)}\n" + except Exception as e: + logger.warning(f"Failed to get index info for table {table_name}: {e}") + + logger.info(f"Retrieved structure for table {table_name}") + return result + + except Exception as e: + error_msg = f"获取表 {table_name} 结构失败: {str(e)}" + logger.error(error_msg) + return error_msg + + +class QueryModel(BaseModel): + """执行SQL查询的参数模型""" + + sql: str = Field(description="要执行的SQL查询语句（只能是SELECT语句）", example="SELECT * FROM users WHERE id = 1") + limit: int | None = Field(default=100, description="限制返回的最大行数，默认100，最大1000", ge=1, le=1000) + timeout: int | None = Field(default=10, description="查询超时时间（秒），默认10秒，最大60秒", ge=1, le=60) + + +@tool(args_schema=QueryModel) +def mysql_query( + sql: Annotated[str, "要执行的SQL查询语句（只能是SELECT语句）"], + limit: Annotated[int | None, "限制返回的最大行数，默认100，最大1000"] = 100, + timeout: Annotated[int | None, "查询超时时间（秒），默认10秒，最大60秒"] = 10, +) -> str: + """执行只读的SQL查询语句 + + 这个工具用来执行SQL查询并返回结果。支持复杂的SELECT查询，包括JOIN、GROUP BY等。 + 注意：只能执行查询操作，不能修改数据。 + + 参数: + - sql: SQL查询语句 + - limit: 返回结果的最大行数（防止结果过大） + - timeout: 查询超时时间（防止长时间运行的查询） + """ + try: + # 验证SQL安全性 + if not MySQLSecurityChecker.validate_sql(sql): + return "SQL语句包含不安全的操作或可能的注入攻击，请检查SQL语句" + + # 验证参数 + if not MySQLSecurityChecker.validate_limit(limit): + return "limit参数必须在1-1000之间" + + if not MySQLSecurityChecker.validate_timeout(timeout): + return "timeout参数必须在1-60之间" + + # 如果SQL中没有LIMIT，添加LIMIT子句 + if "LIMIT" not in sql.upper() and "limit" not in sql: + sql_upper = sql.strip().upper() + # 确保是SELECT语句才添加LIMIT + if sql_upper.startswith("SELECT"): + sql = f"{sql} LIMIT {limit}" + + conn_manager = get_connection_manager() + + with conn_manager.get_cursor() as cursor: + # 执行查询 + cursor.execute(sql) + result = cursor.fetchall() + + if not result: + return "查询执行成功，但没有返回任何结果" + + # 限制结果大小 + limited_result = limit_result_size(result, max_chars=10000) + + # 检查结果是否被截断 + if len(limited_result) < len(result): + warning = f"\n\n⚠️ 警告: 查询结果过大，只显示了前 {len(limited_result)} 行（共 {len(result)} 行）。\n" + warning += "建议使用更精确的查询条件或使用LIMIT子句来减少返回的数据量。" + else: + warning = "" + + # 格式化输出 + if limited_result: + # 获取列名 + columns = list(limited_result[0].keys()) + + # 计算每列的最大宽度 + col_widths = {} + for col in columns: + col_widths[col] = max(len(str(col)), max(len(str(row.get(col, ""))) for row in limited_result)) + col_widths[col] = min(col_widths[col], 50) # 限制最大宽度 + + # 构建表头 + header = "| " + " | ".join(f"{col:<{col_widths[col]}}" for col in columns) + " |" + separator = "|" + "|".join("-" * (col_widths[col] + 2) for col in columns) + "|" + + # 构建数据行 + rows = [] + for row in limited_result: + row_str = "| " + " | ".join(f"{str(row.get(col, '')):<{col_widths[col]}}" for col in columns) + " |" + rows.append(row_str) + + result_str = f"查询结果（共 {len(limited_result)} 行）:\n\n" + result_str += header + "\n" + separator + "\n" + result_str += "\n".join(rows[:50]) # 最多显示50行 + + if len(rows) > 50: + result_str += f"\n\n... 还有 {len(rows) - 50} 行未显示 ..." + + result_str += warning + + logger.info(f"Query executed successfully, returned {len(limited_result)} rows") + return result_str + else: + return "查询执行成功，但没有返回任何结果" + + except Exception as e: + error_msg = f"SQL查询执行失败: {str(e)}" + + # 提供更有用的错误信息 + if "timeout" in str(e).lower(): + error_msg += "\n\n💡 建议：查询超时了，请尝试以下方法：\n" + error_msg += "1. 减少查询的数据量（使用WHERE条件过滤）\n" + error_msg += "2. 使用LIMIT子句限制返回行数\n" + error_msg += "3. 增加timeout参数值（最大60秒）" + elif "table" in str(e).lower() and "doesn't exist" in str(e).lower(): + error_msg += "\n\n💡 建议：表不存在，请使用 mysql_list_tables 查看可用的表名" + elif "column" in str(e).lower() and "doesn't exist" in str(e).lower(): + error_msg += "\n\n💡 建议：列不存在，请使用 mysql_describe_table 查看表结构" + + logger.error(error_msg) + return error_msg + + +def get_mysql_tools() -> list[Any]: + """获取MySQL工具列表""" + return [mysql_list_tables, mysql_describe_table, mysql_query] diff --git a/src/agents/common/tools.py b/src/agents/common/tools.py index 87429583..766ada84 100644 --- a/src/agents/common/tools.py +++ b/src/agents/common/tools.py @@ -115,6 +115,10 @@ def get_buildin_tools() -> list: tools.extend(get_kb_based_tools()) tools.extend(get_static_tools()) + from src.agents.common.toolkits.mysql.tools import get_mysql_tools + + tools.extend(get_mysql_tools()) + except Exception as e: logger.error(f"Failed to get knowledge base retrievers: {e}") diff --git a/test/test_mysql_connection.py b/test/test_mysql_connection.py new file mode 100644 index 00000000..7a0285e6 --- /dev/null +++ b/test/test_mysql_connection.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +""" +MySQL 数据库连接验证脚本 +""" + +import os + +# 添加项目根目录到 Python 路径 +# sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent.parent.parent)) + + +def load_env_file(env_file="src/.env"): + """加载 .env 文件""" + env_vars = {} + if os.path.exists(env_file): + with open(env_file, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line and not line.startswith("#") and "=" in line: + key, value = line.split("=", 1) + env_vars[key] = value.strip("\"'") + return env_vars + + +def test_mysql_connection(): + """测试 MySQL 连接""" + print("=== MySQL 数据库连接验证 ===\n") + + # 加载环境变量 + env_vars = load_env_file() + + # 检查必需的环境变量 + required_vars = ["MYSQL_HOST", "MYSQL_USER", "MYSQL_PASSWORD", "MYSQL_DATABASE"] + missing_vars = [] + + for var in required_vars: + if var not in env_vars: + missing_vars.append(var) + + if missing_vars: + print(f"❌ 缺少必需的环境变量: {', '.join(missing_vars)}") + print("\n请在 .env 文件中设置以下变量:") + for var in missing_vars: + print(f" {var}=your_value") + return False + + # 显示配置信息（隐藏密码） + print("📋 数据库配置:") + print(f" Host: {env_vars.get('MYSQL_HOST', 'Not set')}") + print(f" User: {env_vars.get('MYSQL_USER', 'Not set')}") + print(f" Database: {env_vars.get('MYSQL_DATABASE', 'Not set')}") + print(f" Port: {env_vars.get('MYSQL_PORT', '3306')}") + print(f" Charset: {env_vars.get('MYSQL_CHARSET', 'utf8mb4')}") + print() + + try: + # 导入 MySQL 连接管理器 + from src.agents.common.toolkits.mysql.connection import MySQLConnectionManager + + # 创建连接配置 + mysql_config = { + "host": env_vars["MYSQL_HOST"], + "user": env_vars["MYSQL_USER"], + "password": env_vars["MYSQL_PASSWORD"], + "database": env_vars["MYSQL_DATABASE"], + "port": int(env_vars.get("MYSQL_PORT", "3306")), + "charset": env_vars.get("MYSQL_CHARSET", "utf8mb4"), + } + + # 创建连接管理器 + print("🔄 正在连接数据库...") + conn_manager = MySQLConnectionManager(mysql_config) + + # 测试连接 + with conn_manager.get_cursor() as cursor: + # 测试基本连接 + cursor.execute("SELECT 1 as test") + result = cursor.fetchone() + if result and result["test"] == 1: + print("✅ 数据库连接成功！") + + # 获取数据库版本 + cursor.execute("SELECT VERSION() as version") + version_info = cursor.fetchone() + print(f"📊 MySQL 版本: {version_info['version']}") + + # 获取当前数据库 + cursor.execute("SELECT DATABASE() as db_name") + db_info = cursor.fetchone() + print(f"🗄️ 当前数据库: {db_info['db_name']}") + + # 获取表数量 + cursor.execute("SHOW TABLES") + tables = cursor.fetchall() + print(f"📋 表数量: {len(tables)}") + + if tables: + print("\n📝 数据库表列表:") + for i, table in enumerate(tables[:10]): # 只显示前10个表 + table_name = list(table.values())[0] + print(f" {i + 1}. {table_name}") + + if len(tables) > 10: + print(f" ... 还有 {len(tables) - 10} 个表未显示") + + return True + else: + print("❌ 数据库连接测试失败") + return False + + except ImportError as e: + print(f"❌ 导入错误: {e}") + print("请确保已安装 pymysql 依赖") + return False + except Exception as e: + print(f"❌ 连接失败: {e}") + + # 提供故障排除建议 + print("\n💡 故障排除建议:") + print("1. 检查数据库服务是否正在运行") + print("2. 验证主机地址和端口是否正确") + print("3. 确认用户名和密码是否正确") + print("4. 检查数据库是否存在") + print("5. 确认网络连接和防火墙设置") + + return False + + +def test_tools(): + """测试 MySQL 工具是否正常工作""" + print("\n=== MySQL 工具测试 ===\n") + + try: + from src.agents.common.toolkits.mysql.tools import mysql_list_tables, mysql_describe_table, mysql_query + + print("✅ MySQL 工具导入成功") + + # 测试获取表名 + print("\n🔄 测试获取表名...") + result = mysql_list_tables.invoke({}) + if "失败" not in result and "错误" not in result: + print("✅ 获取表名工具正常") + else: + print(f"❌ 获取表名工具异常: {result}") + return False + + # 如果有表，测试获取表结构 + if "数据库中的表:" in result: + print("\n🔄 测试获取表结构...") + # 提取第一个表名 + lines = result.split("\n") + for line in lines: + if "- " in line and "(" in line: + table_name = line.split("- ")[1].split(" ")[0] + break + else: + table_name = None + + if table_name: + structure_result = mysql_describe_table.invoke({"table_name": table_name}) + if "失败" not in structure_result and "错误" not in structure_result: + print("✅ 获取表结构工具正常") + else: + print(f"❌ 获取表结构工具异常: {structure_result}") + return False + + # 测试简单查询 + print("\n🔄 测试SQL查询...") + query_result = mysql_query.invoke({"sql": f"SELECT COUNT(*) as total FROM `{table_name}`"}) + if "失败" not in query_result and "错误" not in query_result: + print("✅ SQL查询工具正常") + else: + print(f"❌ SQL查询工具异常: {query_result}") + return False + + return True + + except Exception as e: + print(f"❌ 工具测试失败: {e}") + return False + + +def main(): + """主函数""" + print("MySQL 数据库连接和工具验证脚本") + print("=" * 50) + + # 测试连接 + connection_ok = test_mysql_connection() + + if connection_ok: + # 测试工具 + tools_ok = test_tools() + + print("\n" + "=" * 50) + if connection_ok and tools_ok: + print("🎉 所有测试通过！MySQL 工具包可以正常使用") + else: + print("❌ 部分测试失败，请检查配置") + else: + print("\n" + "=" * 50) + print("❌ 数据库连接失败，请检查配置") + + +if __name__ == "__main__": + main() diff --git a/test/test_mysql_import.py b/test/test_mysql_import.py new file mode 100644 index 00000000..647f82f2 --- /dev/null +++ b/test/test_mysql_import.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 +""" +简单测试 MySQL 工具导入 +""" + +import sys + +sys.path.insert(0, ".") + +try: + from src.agents.common.toolkits.mysql.tools import mysql_list_tables, mysql_describe_table, mysql_query + + print("✅ MySQL 工具导入成功") + print(f"工具列表: {[tool.name for tool in [mysql_list_tables, mysql_describe_table, mysql_query]]}") + + # 检查工具描述 + for tool in [mysql_list_tables, mysql_describe_table, mysql_query]: + print(f" - {tool.name}: {tool.description[:60]}...") + +except Exception as e: + print(f"❌ 导入失败: {e}") + import traceback + + traceback.print_exc() diff --git a/web/src/components/ToolCallingResult/ToolResultRenderer.vue b/web/src/components/ToolCallingResult/ToolResultRenderer.vue index a478547a..17db38a9 100644 --- a/web/src/components/ToolCallingResult/ToolResultRenderer.vue +++ b/web/src/components/ToolCallingResult/ToolResultRenderer.vue @@ -25,6 +25,11 @@ :data="parsedData" /> + +
+ +
+