from typing import Dict, Any, List, Optional import logging from datetime import datetime import asyncio from concurrent.futures import ThreadPoolExecutor from langchain_openai import ChatOpenAI from th_agenter.core.context import UserContext from .smart_query import DatabaseQueryService from .postgresql_tool_manager import get_postgresql_tool from .mysql_tool_manager import get_mysql_tool from .table_metadata_service import TableMetadataService from ..core.config import get_settings # 配置日志 logger = logging.getLogger(__name__) class SmartWorkflowError(Exception): """智能工作流自定义异常""" pass class DatabaseConnectionError(SmartWorkflowError): """数据库连接异常""" pass class TableSchemaError(SmartWorkflowError): """表结构获取异常""" pass class SQLGenerationError(SmartWorkflowError): """SQL生成异常""" pass class QueryExecutionError(SmartWorkflowError): """查询执行异常""" pass class SmartDatabaseWorkflowManager: """ 智能数据库工作流管理器 负责协调数据库连接、表元数据获取、SQL生成、查询执行和AI总结的完整流程 """ def __init__(self, db=None): self.executor = ThreadPoolExecutor(max_workers=4) self.database_service = DatabaseQueryService() self.postgresql_tool = get_postgresql_tool() self.mysql_tool = get_mysql_tool() self.db = db self.table_metadata_service = TableMetadataService(db) if db else None async def initialize(self): from ..core.new_agent import new_agent self.llm = await new_agent() def _get_database_tool(self, db_type: str): """根据数据库类型获取对应的数据库工具""" if db_type.lower() == 'postgresql': return self.postgresql_tool elif db_type.lower() == 'mysql': return self.mysql_tool else: raise ValueError(f"不支持的数据库类型: {db_type}") async def _run_in_executor(self, func, *args): """在线程池中运行阻塞函数""" loop = asyncio.get_event_loop() return await loop.run_in_executor(self.executor, func, *args) def _convert_query_result_to_table_data(self, query_result: Dict[str, Any]) -> Dict[str, Any]: """ 将数据库查询结果转换为前端表格数据格式 参考Excel处理方式，以表格形式返回结果 """ try: data = query_result.get('data', []) columns = query_result.get('columns', []) row_count = query_result.get('row_count', 0) if not data or not columns: return { 'result_type': 'table', 'columns': [], 'data': [], 'total': 0, 'message': '查询未返回数据' } # 构建列定义 table_columns = [] for i, col_name in enumerate(columns): table_columns.append({ 'prop': f'col_{i}', 'label': str(col_name), 'width': 'auto' }) # 转换数据行 table_data = [] for row_index, row in enumerate(data): row_data = {'_index': str(row_index)} # 处理字典格式的行数据 if isinstance(row, dict): for i, col_name in enumerate(columns): col_prop = f'col_{i}' value = row.get(col_name) # 处理None值和特殊值 if value is None: row_data[col_prop] = '' elif isinstance(value, (int, float, str, bool)): row_data[col_prop] = str(value) else: row_data[col_prop] = str(value) else: # 处理列表格式的行数据（兼容性处理） for i, value in enumerate(row): col_prop = f'col_{i}' # 处理None值和特殊值 if value is None: row_data[col_prop] = '' elif isinstance(value, (int, float, str, bool)): row_data[col_prop] = str(value) else: row_data[col_prop] = str(value) table_data.append(row_data) return { 'result_type': 'table_data', 'columns': table_columns, 'data': table_data, 'total': row_count, 'message': f'查询成功，共返回 {row_count} 条记录' } except Exception as e: logger.error(f"转换查询结果异常: {str(e)}") return { 'result_type': 'error', 'columns': [], 'data': [], 'total': 0, 'message': f'结果转换失败: {str(e)}' } async def process_database_query_stream( self, user_query: str, user_id: int, database_config_id: int ): """ 流式处理数据库智能问数查询的主要工作流（基于保存的表元数据） 实时推送每个工作流步骤 新流程： 1. 根据database_config_id获取数据库配置并创建连接 2. 从系统数据库读取表元数据（只包含启用问答的表） 3. 根据表元数据生成SQL 4. 执行SQL查询 5. 查询数据后处理成表格形式 6. 生成数据总结 7. 返回结果 Args: user_query: 用户问题 user_id: 用户ID database_config_id: 数据库配置ID Yields: 包含工作流步骤或最终结果的字典 """ workflow_steps = [] try: logger.info(f"开始执行流式数据库查询工作流 - 用户ID: {user_id}, 数据库配置ID: {database_config_id}, 查询: {user_query[:50]}...") # 步骤1: 根据database_config_id获取数据库配置并创建连接 try: step_data = { 'type': 'workflow_step', 'step': 'database_connection', 'status': 'running', 'message': '正在建立数据库连接...', 'timestamp': datetime.now().isoformat() } yield step_data # 获取数据库配置并建立连接 connection_result = await self._connect_database(user_id, database_config_id) if not connection_result['success']: raise DatabaseConnectionError(connection_result['message']) step_data.update({ 'status': 'completed', 'message': '数据库连接成功', 'details': {'database': connection_result.get('database_name', 'Unknown')} }) yield step_data workflow_steps.append({ 'step': 'database_connection', 'status': 'completed', 'message': '数据库连接成功' }) except Exception as e: error_msg = f'数据库连接失败: {str(e)}' step_data = { 'type': 'workflow_step', 'step': 'database_connection', 'status': 'failed', 'message': error_msg, 'timestamp': datetime.now().isoformat() } yield step_data yield { 'type': 'error', 'message': error_msg, 'workflow_steps': workflow_steps } return # 步骤2: 从系统数据库读取表元数据（只包含启用问答的表） try: step_data = { 'type': 'workflow_step', 'step': 'table_metadata', 'status': 'running', 'message': '正在从系统数据库读取表元数据...', 'timestamp': datetime.now().isoformat() } yield step_data # 从系统数据库读取已保存的表元数据（只包含启用问答的表） tables_info = await self._get_saved_tables_metadata(user_id, database_config_id) step_data.update({ 'status': 'completed', 'message': f'成功读取 {len(tables_info)} 个启用问答的表元数据', 'details': {'table_count': len(tables_info), 'tables': list(tables_info.keys())} }) yield step_data workflow_steps.append({ 'step': 'table_metadata', 'status': 'completed', 'message': f'成功读取表元数据' }) except Exception as e: error_msg = f'获取表元数据失败: {str(e)}' step_data = { 'type': 'workflow_step', 'step': 'table_metadata', 'status': 'failed', 'message': error_msg, 'timestamp': datetime.now().isoformat() } yield step_data yield { 'type': 'error', 'message': error_msg, 'workflow_steps': workflow_steps } return # 步骤3: 根据表元数据生成SQL try: step_data = { 'type': 'workflow_step', 'step': 'sql_generation', 'status': 'running', 'message': '正在根据表元数据生成SQL查询...', 'timestamp': datetime.now().isoformat() } yield step_data # 根据表元数据选择相关表并生成SQL target_tables, target_schemas = await self._select_target_table(user_query, tables_info) step_data = { 'type': 'workflow_step', 'step': 'table_selected', 'status': 'completed', 'message': f'已经智能选择了相关表: {", ".join(target_tables)}', 'timestamp': datetime.now().isoformat() } yield step_data workflow_steps.append({ 'step': 'table_metadata', 'status': 'completed', 'message': f'已经智能选择了相关表: {", ".join(target_tables)}', }) sql_query = await self._generate_sql_query(user_query, target_tables, target_schemas) step_data.update({ 'status': 'completed', 'message': 'SQL查询生成成功', 'details': { 'target_tables': target_tables, 'generated_sql': sql_query[:100] + '...' if len(sql_query) > 100 else sql_query } }) yield step_data workflow_steps.append({ 'step': 'sql_generation', 'status': 'completed', 'message': 'SQL语句生成成功' }) except Exception as e: error_msg = f'SQL生成失败: {str(e)}' step_data = { 'type': 'workflow_step', 'step': 'sql_generation', 'status': 'failed', 'message': error_msg, 'timestamp': datetime.now().isoformat() } yield step_data yield { 'type': 'error', 'message': error_msg, 'workflow_steps': workflow_steps } return # 步骤4: 执行SQL查询 try: step_data = { 'type': 'workflow_step', 'step': 'query_execution', 'status': 'running', 'message': '正在执行SQL查询...', 'timestamp': datetime.now().isoformat() } yield step_data query_result = await self._execute_database_query(user_id, sql_query, database_config_id) step_data.update({ 'status': 'completed', 'message': f'查询执行成功，返回 {query_result.get("row_count", 0)} 条记录', 'details': {'row_count': query_result.get('row_count', 0)} }) yield step_data workflow_steps.append({ 'step': 'query_execution', 'status': 'completed', 'message': '查询执行成功' }) except Exception as e: error_msg = f'查询执行失败: {str(e)}' step_data = { 'type': 'workflow_step', 'step': 'query_execution', 'status': 'failed', 'message': error_msg, 'timestamp': datetime.now().isoformat() } yield step_data yield { 'type': 'error', 'message': error_msg, 'workflow_steps': workflow_steps } return # 步骤5: 查询数据后处理成表格形式（在步骤6中完成） # 步骤6: 生成数据总结 try: step_data = { 'type': 'workflow_step', 'step': 'ai_summary', 'status': 'running', 'message': '正在生成查询结果总结...', 'timestamp': datetime.now().isoformat() } yield step_data summary = await self._generate_database_summary(user_query, query_result, ', '.join(target_tables)) step_data.update({ 'status': 'completed', 'message': '总结生成完成', 'details': { 'tables_analyzed': target_tables, 'summary_length': len(summary) } }) yield step_data workflow_steps.append({ 'step': 'ai_summary', 'status': 'completed', 'message': '总结生成完成' }) except Exception as e: logger.warning(f'生成总结失败: {str(e)}') summary = '查询执行完成，但生成总结时出现问题。' workflow_steps.append({ 'step': 'ai_summary', 'status': 'warning', 'message': '总结生成失败，但查询成功' }) # 步骤7: 返回最终结果，且结果参考excel的处理方式，尽量以表格形式返回 try: step_data = { 'type': 'workflow_step', 'step': 'result_formatting', 'status': 'running', 'message': '正在格式化查询结果...', 'timestamp': datetime.now().isoformat() } yield step_data # 转换为表格格式 table_data = self._convert_query_result_to_table_data(query_result) step_data.update({ 'status': 'completed', 'message': '结果格式化完成' }) yield step_data workflow_steps.append({ 'step': 'result_formatting', 'status': 'completed', 'message': '结果格式化完成' }) # 返回最终结果 final_result = { 'type': 'final_result', 'success': True, 'data': { **table_data, 'generated_sql': sql_query, 'summary': summary, 'table_name': target_tables, 'query_result': query_result, 'metadata_source': 'saved_database' # 标记元数据来源 }, 'workflow_steps': workflow_steps, 'timestamp': datetime.now().isoformat() } yield final_result logger.info(f"数据库查询工作流完成 - 用户ID: {user_id}") except Exception as e: error_msg = f'结果格式化失败: {str(e)}' yield { 'type': 'error', 'message': error_msg, 'workflow_steps': workflow_steps } return except Exception as e: logger.error(f"数据库查询工作流异常: {str(e)}", exc_info=True) yield { 'type': 'error', 'message': f'系统异常: {str(e)}', 'workflow_steps': workflow_steps } async def _connect_database(self, user_id: int, database_config_id: int) -> Dict[str, Any]: """连接数据库（判断用户现有连接）""" try: # 获取数据库配置 from ..services.database_config_service import DatabaseConfigService config_service = DatabaseConfigService(self.db) config = config_service.get_config_by_id(database_config_id, user_id) if not config: return {'success': False, 'message': '数据库配置不存在'} # 根据数据库类型选择对应的工具 try: db_tool = self._get_database_tool(config.db_type) except ValueError as e: return {'success': False, 'message': str(e)} # 测试连接（如果已经有连接则直接复用） connection_config = { 'host': config.host, 'port': config.port, 'database': config.database, 'username': config.username, 'password': config_service._decrypt_password(config.password) } try: connection = db_tool._test_connection(connection_config) if connection['success'] == True: return { 'success': True, 'database_name': config.database, 'db_type': config.db_type, 'message': '连接成功' } else: return { 'success': False, 'database_name': config.database, 'db_type': config.db_type, 'message': '连接失败' } except Exception as e: return { 'success': False, 'message': f'连接失败: {str(e)}' } except Exception as e: logger.error(f"数据库连接异常: {str(e)}") return {'success': False, 'message': f'连接异常: {str(e)}'} async def _get_saved_tables_metadata(self, user_id: int, database_config_id: int) -> Dict[str, Dict[str, Any]]: """从系统数据库中读取已保存的表元数据""" try: if not self.table_metadata_service: raise TableSchemaError("表元数据服务未初始化") # 从数据库中获取表元数据 saved_metadata = self.table_metadata_service.get_user_table_metadata( user_id, database_config_id ) if not saved_metadata: raise TableSchemaError(f"未找到数据库配置ID {database_config_id} 的表元数据，请先在数据库管理页面收集表元数据") # 转换为所需格式 tables_metadata = {} for meta in saved_metadata: # 只处理启用问答的表 if meta.is_enabled_for_qa: tables_metadata[meta.table_name] = { 'table_name': meta.table_name, 'columns': meta.columns_info or [], 'primary_keys': meta.primary_keys or [], 'row_count': meta.row_count or 0, 'table_comment': meta.table_comment or '', 'qa_description': meta.qa_description or '', 'business_context': meta.business_context or '', 'from_saved_metadata': True # 标记来源 } if not tables_metadata: raise TableSchemaError("没有启用问答的表，请在数据库管理页面启用相关表的问答功能") logger.info(f"从系统数据库读取表元数据成功，共 {len(tables_metadata)} 个启用问答的表") return tables_metadata except Exception as e: logger.error(f"读取保存的表元数据异常: {str(e)}") raise TableSchemaError(f'读取表元数据失败: {str(e)}') async def _get_table_schema(self, user_id: int, table_name: str, database_config_id: int) -> Dict[str, Any]: """获取指定表结构""" try: # 获取数据库配置 from ..services.database_config_service import DatabaseConfigService config_service = DatabaseConfigService(self.db) config = config_service.get_config_by_id(database_config_id, user_id) if not config: raise TableSchemaError('数据库配置不存在') # 根据数据库类型选择对应的工具 try: db_tool = self._get_database_tool(config.db_type) except ValueError as e: raise TableSchemaError(str(e)) # 使用对应的数据库工具获取表结构 schema_result = await db_tool.describe_table(table_name) if schema_result.get('success'): return schema_result.get('schema', {}) else: raise TableSchemaError(schema_result.get('error', '获取表结构失败')) except Exception as e: logger.error(f"获取表结构异常: {str(e)}") raise TableSchemaError(f'获取表结构失败: {str(e)}') async def _select_target_table(self, user_query: str, tables_info: Dict[str, Dict]) -> tuple[List[str], List[Dict]]: """根据用户查询选择相关的表，支持返回多个表""" try: if len(tables_info) == 1: # 只有一个表，直接返回 table_name = list(tables_info.keys())[0] return [table_name], [tables_info[table_name]] # 多个表时，使用LLM选择相关的表 tables_summary = [] for table_name, schema in tables_info.items(): columns = schema.get('columns', []) column_names = [col.get('column_name', col.get('name', '')) for col in columns] qa_desc = schema.get('qa_description', '') business_ctx = schema.get('business_context', '') tables_summary.append(f"表名: {table_name}\n字段: {', '.join(column_names[:10])}\n表描述: {qa_desc}\n业务上下文: {business_ctx}") prompt = f""" 用户查询: {user_query} 可用的表: {chr(10).join(tables_summary)} 请根据用户查询选择相关的表，可以选择多个表。分析表之间可能的关联关系，返回所有相关的表名，用逗号分隔。 可以通过qa_description（表描述），business_context(表的业务上下文），以及column_names几个字段判断要使用哪些表。 注意：只返回表名列表，后面不要跟其他的内容。 例如直接输出: table1,table2,table3 """ response = await self.llm.ainvoke(prompt) selected_tables = [t.strip() for t in response.content.strip().split(',')] # 验证选择的表是否存在 valid_tables = [] valid_schemas = [] for table in selected_tables: if table in tables_info: valid_tables.append(table) valid_schemas.append(tables_info[table]) else: logger.warning(f"LLM选择的表 {table} 不存在") if valid_tables: return valid_tables, valid_schemas else: # 如果没有有效的表，选择第一个表 table_name = list(tables_info.keys())[0] logger.warning(f"没有找到有效的表，使用默认表 {table_name}") return [table_name], [tables_info[table_name]] except Exception as e: logger.error(f"选择目标表异常: {str(e)}") # 出现异常时选择第一个表 table_name = list(tables_info.keys())[0] return [table_name], [tables_info[table_name]] async def _generate_sql_query(self, user_query: str, table_names: List[str], table_schemas: List[Dict]) -> str: """生成SQL语句，支持多表关联查询""" try: # 构建所有表的结构信息 tables_info = [] for table_name, schema in zip(table_names, table_schemas): columns_info = [] for col in schema.get('columns', []): col_info = f"{col['column_name']} ({col['data_type']})" columns_info.append(col_info) table_info = f"表名: {table_name}\n" table_info += f"表描述: {schema.get('qa_description', '')}\n" table_info += f"业务上下文: {schema.get('business_context', '')}\n" table_info += "字段信息:\n" + "\n".join(columns_info) tables_info.append(table_info) schema_text = "\n\n".join(tables_info) prompt = f""" 基于以下表结构，将自然语言查询转换为SQL语句。如果需要关联多个表，请分析表之间的关系，使用合适的JOIN语法： {schema_text} 用户查询: {user_query} 请生成对应的SQL查询语句，要求： 1. 只返回SQL语句，不要包含其他解释 2. 如果查询涉及多个表，需要正确处理表之间的关联关系 3. 使用合适的JOIN类型（INNER JOIN、LEFT JOIN等） 4. 确保SELECT的字段来源明确，必要时使用表名前缀 """ # 使用LLM生成SQL response = await self.llm.ainvoke(prompt) sql_query = response.content.strip() # 清理SQL语句 if sql_query.startswith('```sql'): sql_query = sql_query[6:] if sql_query.endswith('```'): sql_query = sql_query[:-3] sql_query = sql_query.strip() logger.info(f"生成的SQL查询: {sql_query}") return sql_query except Exception as e: logger.error(f"SQL生成异常: {str(e)}") raise SQLGenerationError(f'SQL生成失败: {str(e)}') async def _execute_database_query(self, user_id: int, sql_query: str, database_config_id: int) -> Dict[str, Any]: """执行SQL语句""" try: # 获取数据库配置 from ..services.database_config_service import DatabaseConfigService config_service = DatabaseConfigService(self.db) config = config_service.get_config_by_id(database_config_id, user_id) if not config: raise QueryExecutionError('数据库配置不存在') # 根据数据库类型选择对应的工具 try: db_tool = self._get_database_tool(config.db_type) except ValueError as e: raise QueryExecutionError(str(e)) # 使用对应的数据库工具执行查询 if str(user_id) in db_tool.connections: query_result = db_tool._execute_query(db_tool.connections[str(user_id)]['connection'], sql_query) else: raise QueryExecutionError('请重新进行数据库连接') if query_result.get('success'): data = query_result.get('data', []) return { 'success': True, 'data': data, 'row_count': len(data), 'columns': query_result.get('columns', []), 'sql_query': sql_query } else: raise QueryExecutionError(query_result.get('error', '查询执行失败')) except Exception as e: logger.error(f"查询执行异常: {str(e)}") raise QueryExecutionError(f'查询执行失败: {str(e)}') async def _generate_database_summary(self, user_query: str, query_result: Dict, tables_str: str) -> str: """生成AI总结，支持多表查询结果""" try: data = query_result.get('data', []) row_count = query_result.get('row_count', 0) columns = query_result.get('columns', []) sql_query = query_result.get('sql_query', '') # 构建总结提示词 prompt = f""" 用户查询: {user_query} 涉及的表: {tables_str} 查询结果: 共 {row_count} 条记录 查询的字段: {', '.join(columns)} 执行的SQL: {sql_query} 前几条数据示例: {str(data[:3]) if data else '无数据'} 请基于以上信息，用中文生成一个简洁的查询结果总结，包括： 1. 查询涉及的表及其关系 2. 查询的主要发现和数据特征 3. 如果有关联查询，说明关联的结果特点 4. 最后对用户的问题进行回答 总结要求： 1. 语言简洁明了 2. 重点突出查询结果 3. 如果是多表查询，需要说明表之间的关系 4. 总结不超过300字 """ # 使用LLM生成总结 response = await self.llm.ainvoke(prompt) summary = response.content.strip() logger.info(f"生成的总结: {summary[:100]}...") return summary except Exception as e: logger.error(f"总结生成异常: {str(e)}") return f"查询完成，共返回 {query_result.get('row_count', 0)} 条记录。涉及的表: {tables_str}" async def process_database_query( self, user_query: str, user_id: int, database_config_id: int, table_name: Optional[str] = None, conversation_id: Optional[int] = None, is_new_conversation: bool = False ) -> Dict[str, Any]: """ 处理数据库智能问数查询的主要工作流（基于保存的表元数据） 新流程： 1. 根据database_config_id获取数据库配置 2. 创建数据库连接 3. 从系统数据库读取表元数据（只包含启用问答的表） 4. 根据表元数据生成SQL 5. 执行SQL查询 6. 查询数据后处理成表格形式 7. 生成数据总结 8. 返回结果 Args: user_query: 用户问题 user_id: 用户ID database_config_id: 数据库配置ID table_name: 表名（可选） conversation_id: 对话ID is_new_conversation: 是否为新对话 Returns: 包含查询结果的字典 """ try: logger.info(f"开始执行数据库查询工作流 - 用户ID: {user_id}, 数据库配置ID: {database_config_id}, 查询: {user_query[:50]}...") # 步骤1: 根据database_config_id获取数据库配置并创建连接 connection_result = await self._connect_database(user_id, database_config_id) if not connection_result['success']: raise DatabaseConnectionError(connection_result['message']) logger.info("数据库连接成功") # 步骤2: 从系统数据库读取表元数据（只包含启用问答的表） tables_info = await self._get_saved_tables_metadata(user_id, database_config_id) logger.info(f"表元数据读取完成 - 共{len(tables_info)}个启用问答的表") # 步骤3: 根据表元数据选择相关表并生成SQL target_tables, target_schemas = await self._select_target_table(user_query, tables_info) sql_query = await self._generate_sql_query(user_query, target_tables, target_schemas) logger.info(f"SQL生成完成 - 目标表: {', '.join(target_tables)}") # 步骤4: 执行SQL查询 query_result = await self._execute_database_query(user_id, sql_query, database_config_id) logger.info("查询执行完成") # 步骤5: 查询数据后处理成表格形式 table_data = self._convert_query_result_to_table_data(query_result) # 步骤6: 生成数据总结 summary = await self._generate_database_summary(user_query, query_result, ', '.join(target_tables)) # 步骤7: 返回结果 return { 'success': True, 'data': { **table_data, 'generated_sql': sql_query, 'summary': summary, 'table_names': target_tables, 'query_result': query_result, 'metadata_source': 'saved_database' # 标记元数据来源 } } except SmartWorkflowError as e: logger.error(f"数据库工作流异常: {str(e)}") return { 'success': False, 'error': str(e), 'error_type': type(e).__name__ } except Exception as e: logger.error(f"数据库工作流未知异常: {str(e)}", exc_info=True) return { 'success': False, 'error': f'系统异常: {str(e)}', 'error_type': 'SystemError' }