717 lines
25 KiB
Python
717 lines
25 KiB
Python
import pandas as pd
|
||
import pymysql
|
||
import psycopg2
|
||
import tempfile
|
||
import os
|
||
from typing import Dict, Any, List
|
||
from datetime import datetime
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
from langchain_community.chat_models import ChatZhipuAI
|
||
from langchain_core.messages import HumanMessage
|
||
from loguru import logger
|
||
|
||
# 在 SmartQueryService 类中添加方法
|
||
|
||
from .table_metadata_service import TableMetadataService
|
||
|
||
class SmartQueryService:
|
||
"""
|
||
智能问数服务基类
|
||
"""
|
||
def __init__(self):
|
||
self.executor = ThreadPoolExecutor(max_workers=4)
|
||
self.table_metadata_service = None
|
||
|
||
def set_db_session(self, db_session):
|
||
"""设置数据库会话"""
|
||
self.table_metadata_service = TableMetadataService(db_session)
|
||
|
||
async def _run_in_executor(self, func, *args):
|
||
"""在线程池中运行阻塞函数"""
|
||
loop = asyncio.get_event_loop()
|
||
return await loop.run_in_executor(self.executor, func, *args)
|
||
|
||
class ExcelAnalysisService(SmartQueryService):
|
||
"""
|
||
Excel数据分析服务
|
||
"""
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.user_dataframes = {} # 存储用户的DataFrame
|
||
|
||
def analyze_dataframe(self, df: pd.DataFrame, filename: str) -> Dict[str, Any]:
|
||
"""
|
||
分析DataFrame并返回基本信息
|
||
"""
|
||
try:
|
||
# 基本统计信息
|
||
rows, columns = df.shape
|
||
|
||
# 列信息
|
||
column_info = []
|
||
for col in df.columns:
|
||
col_info = {
|
||
'name': col,
|
||
'dtype': str(df[col].dtype),
|
||
'null_count': int(df[col].isnull().sum()),
|
||
'unique_count': int(df[col].nunique())
|
||
}
|
||
|
||
# 如果是数值列，添加统计信息
|
||
if pd.api.types.is_numeric_dtype(df[col]):
|
||
df.fillna({col:0}) #数值列，将空值补0
|
||
col_info.update({
|
||
'mean': float(df[col].mean()) if not df[col].isnull().all() else None,
|
||
'std': float(df[col].std()) if not df[col].isnull().all() else None,
|
||
'min': float(df[col].min()) if not df[col].isnull().all() else None,
|
||
'max': float(df[col].max()) if not df[col].isnull().all() else None
|
||
})
|
||
|
||
column_info.append(col_info)
|
||
|
||
# 数据预览（前5行）
|
||
preview_data = df.head().fillna('').to_dict('records')
|
||
|
||
# 数据质量检查
|
||
quality_issues = []
|
||
|
||
# 检查缺失值
|
||
missing_cols = df.columns[df.isnull().any()].tolist()
|
||
if missing_cols:
|
||
quality_issues.append({
|
||
'type': 'missing_values',
|
||
'description': f'以下列存在缺失值: {", ".join(map(str, missing_cols))}',
|
||
'columns': missing_cols
|
||
})
|
||
|
||
# 检查重复行
|
||
duplicate_count = df.duplicated().sum()
|
||
if duplicate_count > 0:
|
||
quality_issues.append({
|
||
'type': 'duplicate_rows',
|
||
'description': f'发现 {duplicate_count} 行重复数据',
|
||
'count': int(duplicate_count)
|
||
})
|
||
|
||
return {
|
||
'filename': filename,
|
||
'rows': rows,
|
||
'columns': columns,
|
||
'column_names': [str(col) for col in df.columns.tolist()],
|
||
'column_info': column_info,
|
||
'preview': preview_data,
|
||
'quality_issues': quality_issues,
|
||
'memory_usage': f"{df.memory_usage(deep=True).sum() / 1024 / 1024:.2f} MB"
|
||
}
|
||
|
||
except Exception as e:
|
||
print(e)
|
||
raise Exception(f"DataFrame分析失败: {str(e)}")
|
||
|
||
def _create_pandas_agent(self, df: pd.DataFrame):
|
||
"""
|
||
创建pandas代理
|
||
"""
|
||
try:
|
||
# 使用智谱AI作为LLM
|
||
llm = ChatZhipuAI(
|
||
model="glm-4",
|
||
api_key=os.getenv("ZHIPUAI_API_KEY"),
|
||
temperature=0.1
|
||
)
|
||
agent = None
|
||
logger.error('创建pandas代理失败 - 暂屏蔽处理')
|
||
|
||
# # 创建pandas代理
|
||
# agent = create_pandas_dataframe_agent(
|
||
# llm=llm,
|
||
# df=df,
|
||
# verbose=True,
|
||
# return_intermediate_steps=True,
|
||
# handle_parsing_errors=True,
|
||
# max_iterations=3,
|
||
# early_stopping_method="force",
|
||
# allow_dangerous_code=True # 允许执行代码以支持数据分析
|
||
# )
|
||
|
||
return agent
|
||
|
||
except Exception as e:
|
||
raise Exception(f"创建pandas代理失败: {str(e)}")
|
||
|
||
def _execute_pandas_query(self, agent, query: str) -> Dict[str, Any]:
|
||
"""
|
||
执行pandas查询
|
||
"""
|
||
try:
|
||
# 执行查询
|
||
# 使用invoke方法来处理有多个输出键的情况
|
||
agent_result = agent.invoke({"input": query})
|
||
# 提取主要结果
|
||
result = agent_result.get('output', agent_result)
|
||
|
||
# 解析结果
|
||
if isinstance(result, pd.DataFrame):
|
||
# 如果结果是DataFrame
|
||
data = result.fillna('').to_dict('records')
|
||
columns = result.columns.tolist()
|
||
total = len(result)
|
||
|
||
return {
|
||
'data': data,
|
||
'columns': columns,
|
||
'total': total,
|
||
'result_type': 'dataframe'
|
||
}
|
||
else:
|
||
# 如果结果是其他类型（字符串、数字等）
|
||
return {
|
||
'data': [{'result': str(result)}],
|
||
'columns': ['result'],
|
||
'total': 1,
|
||
'result_type': 'scalar'
|
||
}
|
||
|
||
except Exception as e:
|
||
raise Exception(f"pandas查询执行失败: {str(e)}")
|
||
|
||
async def execute_natural_language_query(
|
||
self,
|
||
query: str,
|
||
user_id: int,
|
||
page: int = 1,
|
||
page_size: int = 20
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
执行自然语言查询
|
||
"""
|
||
try:
|
||
# 查找用户的临时文件
|
||
temp_dir = tempfile.gettempdir()
|
||
user_files = [f for f in os.listdir(temp_dir)
|
||
if f.startswith(f"excel_{user_id}_") and f.endswith('.pkl')]
|
||
|
||
if not user_files:
|
||
return {
|
||
'success': False,
|
||
'message': '未找到上传的Excel文件，请先上传文件'
|
||
}
|
||
|
||
# 使用最新的文件
|
||
latest_file = sorted(user_files)[-1]
|
||
file_path = os.path.join(temp_dir, latest_file)
|
||
|
||
# 加载DataFrame
|
||
df = pd.read_pickle(file_path)
|
||
|
||
# 创建pandas代理
|
||
agent = self._create_pandas_agent(df)
|
||
|
||
# 执行查询
|
||
query_result = await self._run_in_executor(
|
||
self._execute_pandas_query, agent, query
|
||
)
|
||
|
||
# 分页处理
|
||
total = query_result['total']
|
||
start_idx = (page - 1) * page_size
|
||
end_idx = start_idx + page_size
|
||
|
||
paginated_data = query_result['data'][start_idx:end_idx]
|
||
|
||
# 生成AI总结
|
||
summary = await self._generate_summary(query, query_result, df)
|
||
|
||
return {
|
||
'success': True,
|
||
'data': {
|
||
'data': paginated_data,
|
||
'columns': query_result['columns'],
|
||
'total': total,
|
||
'page': page,
|
||
'page_size': page_size,
|
||
'generated_code': f"# 基于自然语言查询: {query}\n# 使用LangChain Pandas代理执行",
|
||
'summary': summary,
|
||
'result_type': query_result['result_type']
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'message': f"查询执行失败: {str(e)}"
|
||
}
|
||
|
||
async def _generate_summary(self, query: str, result: Dict[str, Any], df: pd.DataFrame) -> str:
|
||
"""
|
||
生成AI总结
|
||
"""
|
||
try:
|
||
llm = ChatZhipuAI(
|
||
model="glm-4",
|
||
api_key=os.getenv("ZHIPUAI_API_KEY"),
|
||
temperature=0.3
|
||
)
|
||
|
||
# 构建总结提示
|
||
prompt = f"""
|
||
用户查询: {query}
|
||
|
||
数据集信息:
|
||
- 总行数: {len(df)}
|
||
- 总列数: {len(df.columns)}
|
||
- 列名: {', '.join(str(col) for col in df.columns.tolist())}
|
||
|
||
查询结果:
|
||
- 结果类型: {result['result_type']}
|
||
- 结果行数: {result['total']}
|
||
- 结果列数: {len(result['columns'])}
|
||
|
||
请基于以上信息，用中文生成一个简洁的分析总结，包括:
|
||
1. 查询的主要目的
|
||
2. 关键发现
|
||
3. 数据洞察
|
||
4. 建议的后续分析方向
|
||
|
||
总结应该专业、准确、易懂，控制在200字以内。
|
||
"""
|
||
|
||
response = await self._run_in_executor(
|
||
lambda: llm.invoke([HumanMessage(content=prompt)])
|
||
)
|
||
|
||
return response.content
|
||
|
||
except Exception as e:
|
||
return f"查询已完成，但生成总结时出现错误: {str(e)}"
|
||
|
||
class DatabaseQueryService(SmartQueryService):
|
||
"""
|
||
数据库查询服务
|
||
"""
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.user_connections = {} # 存储用户的数据库连接信息
|
||
|
||
def _create_connection(self, config: Dict[str, str]):
|
||
"""
|
||
创建数据库连接
|
||
"""
|
||
db_type = config['type'].lower()
|
||
|
||
try:
|
||
if db_type == 'mysql':
|
||
connection = pymysql.connect(
|
||
host=config['host'],
|
||
port=int(config['port']),
|
||
user=config['username'],
|
||
password=config['password'],
|
||
database=config['database'],
|
||
charset='utf8mb4'
|
||
)
|
||
elif db_type == 'postgresql':
|
||
connection = psycopg2.connect(
|
||
host=config['host'],
|
||
port=int(config['port']),
|
||
user=config['username'],
|
||
password=config['password'],
|
||
database=config['database']
|
||
)
|
||
|
||
return connection
|
||
|
||
except Exception as e:
|
||
raise Exception(f"数据库连接失败: {str(e)}")
|
||
|
||
async def test_connection(self, config: Dict[str, str]) -> bool:
|
||
"""
|
||
测试数据库连接
|
||
"""
|
||
try:
|
||
connection = await self._run_in_executor(self._create_connection, config)
|
||
connection.close()
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
async def connect_database(self, config: Dict[str, str], user_id: int) -> Dict[str, Any]:
|
||
"""
|
||
连接数据库并获取表列表
|
||
"""
|
||
try:
|
||
connection = await self._run_in_executor(self._create_connection, config)
|
||
|
||
# 获取表列表
|
||
tables = await self._run_in_executor(self._get_tables, connection, config['type'])
|
||
|
||
# 存储连接信息
|
||
self.user_connections[user_id] = {
|
||
'config': config,
|
||
'connection': connection,
|
||
'connected_at': datetime.now()
|
||
}
|
||
|
||
return {
|
||
'success': True,
|
||
'data': {
|
||
'tables': tables,
|
||
'database_type': config['type'],
|
||
'database_name': config['database']
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'message': f"数据库连接失败: {str(e)}"
|
||
}
|
||
|
||
def _get_tables(self, connection, db_type: str) -> List[str]:
|
||
"""
|
||
获取数据库表列表
|
||
"""
|
||
cursor = connection.cursor()
|
||
|
||
try:
|
||
if db_type.lower() == 'mysql':
|
||
cursor.execute("SHOW TABLES")
|
||
tables = [row[0] for row in cursor.fetchall()]
|
||
elif db_type.lower() == 'postgresql':
|
||
cursor.execute("""
|
||
SELECT table_name
|
||
FROM information_schema.tables
|
||
WHERE table_schema = 'public'
|
||
""")
|
||
tables = [row[0] for row in cursor.fetchall()]
|
||
|
||
elif db_type.lower() == 'sqlserver':
|
||
cursor.execute("""
|
||
SELECT table_name
|
||
FROM information_schema.tables
|
||
WHERE table_type = 'BASE TABLE'
|
||
""")
|
||
tables = [row[0] for row in cursor.fetchall()]
|
||
else:
|
||
tables = []
|
||
|
||
return tables
|
||
|
||
finally:
|
||
cursor.close()
|
||
|
||
async def get_table_schema(self, table_name: str, user_id: int) -> Dict[str, Any]:
|
||
"""
|
||
获取表结构
|
||
"""
|
||
try:
|
||
if user_id not in self.user_connections:
|
||
return {
|
||
'success': False,
|
||
'message': '数据库连接已断开，请重新连接'
|
||
}
|
||
|
||
connection = self.user_connections[user_id]['connection']
|
||
db_type = self.user_connections[user_id]['config']['type']
|
||
|
||
schema = await self._run_in_executor(
|
||
self._get_table_schema, connection, table_name, db_type
|
||
)
|
||
|
||
return {
|
||
'success': True,
|
||
'data': {
|
||
'schema': schema,
|
||
'table_name': table_name
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'message': f"获取表结构失败: {str(e)}"
|
||
}
|
||
|
||
def _get_table_schema(self, connection, table_name: str, db_type: str) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取表结构信息
|
||
"""
|
||
cursor = connection.cursor()
|
||
|
||
try:
|
||
if db_type.lower() == 'mysql':
|
||
cursor.execute(f"DESCRIBE {table_name}")
|
||
columns = cursor.fetchall()
|
||
schema = [{
|
||
'column_name': col[0],
|
||
'data_type': col[1],
|
||
'is_nullable': 'YES' if col[2] == 'YES' else 'NO',
|
||
'column_key': col[3],
|
||
'column_default': col[4]
|
||
} for col in columns]
|
||
elif db_type.lower() == 'postgresql':
|
||
cursor.execute("""
|
||
SELECT column_name, data_type, is_nullable, column_default
|
||
FROM information_schema.columns
|
||
WHERE table_name = %s
|
||
ORDER BY ordinal_position
|
||
""", (table_name,))
|
||
columns = cursor.fetchall()
|
||
schema = [{
|
||
'column_name': col[0],
|
||
'data_type': col[1],
|
||
'is_nullable': col[2],
|
||
'column_default': col[3]
|
||
} for col in columns]
|
||
|
||
else:
|
||
schema = []
|
||
|
||
return schema
|
||
|
||
finally:
|
||
cursor.close()
|
||
|
||
async def execute_natural_language_query(
|
||
self,
|
||
query: str,
|
||
table_name: str,
|
||
user_id: int,
|
||
page: int = 1,
|
||
page_size: int = 20
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
执行自然语言数据库查询
|
||
"""
|
||
try:
|
||
if user_id not in self.user_connections:
|
||
return {
|
||
'success': False,
|
||
'message': '数据库连接已断开，请重新连接'
|
||
}
|
||
|
||
connection = self.user_connections[user_id]['connection']
|
||
|
||
# 这里应该集成MCP服务来将自然语言转换为SQL
|
||
# 目前先使用简单的实现
|
||
sql_query = await self._convert_to_sql(query, table_name, connection)
|
||
|
||
# 执行SQL查询
|
||
result = await self._run_in_executor(
|
||
self._execute_sql_query, connection, sql_query, page, page_size
|
||
)
|
||
|
||
# 生成AI总结
|
||
summary = await self._generate_db_summary(query, result, table_name)
|
||
|
||
result['generated_code'] = sql_query
|
||
result['summary'] = summary
|
||
|
||
return {
|
||
'success': True,
|
||
'data': result
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
'success': False,
|
||
'message': f"数据库查询执行失败: {str(e)}"
|
||
}
|
||
|
||
async def _convert_to_sql(self, query: str, table_name: str, connection) -> str:
|
||
"""
|
||
将自然语言转换为SQL查询
|
||
TODO: 集成MCP服务
|
||
"""
|
||
# 这是一个简化的实现，实际应该使用MCP服务
|
||
# 根据常见的查询模式生成SQL
|
||
|
||
query_lower = query.lower()
|
||
|
||
if '所有' in query or '全部' in query or 'all' in query_lower:
|
||
return f"SELECT * FROM {table_name} LIMIT 100"
|
||
elif '统计' in query or '总数' in query or 'count' in query_lower:
|
||
return f"SELECT COUNT(*) as total_count FROM {table_name}"
|
||
elif '最近' in query or 'recent' in query_lower:
|
||
return f"SELECT * FROM {table_name} ORDER BY id DESC LIMIT 10"
|
||
elif '分组' in query or 'group' in query_lower:
|
||
# 简单的分组查询，需要根据实际表结构调整
|
||
return f"SELECT COUNT(*) as count FROM {table_name} GROUP BY id LIMIT 10"
|
||
else:
|
||
# 默认查询
|
||
return f"SELECT * FROM {table_name} LIMIT 20"
|
||
|
||
def _execute_sql_query(self, connection, sql_query: str, page: int, page_size: int) -> Dict[str, Any]:
|
||
"""
|
||
执行SQL查询
|
||
"""
|
||
cursor = connection.cursor()
|
||
|
||
try:
|
||
# 执行查询
|
||
cursor.execute(sql_query)
|
||
|
||
# 获取列名
|
||
columns = [desc[0] for desc in cursor.description] if cursor.description else []
|
||
|
||
# 获取所有结果
|
||
all_results = cursor.fetchall()
|
||
total = len(all_results)
|
||
|
||
# 分页
|
||
start_idx = (page - 1) * page_size
|
||
end_idx = start_idx + page_size
|
||
paginated_results = all_results[start_idx:end_idx]
|
||
|
||
# 转换为字典格式
|
||
data = []
|
||
for row in paginated_results:
|
||
row_dict = {}
|
||
for i, value in enumerate(row):
|
||
if i < len(columns):
|
||
row_dict[columns[i]] = value
|
||
data.append(row_dict)
|
||
|
||
return {
|
||
'data': data,
|
||
'columns': columns,
|
||
'total': total,
|
||
'page': page,
|
||
'page_size': page_size
|
||
}
|
||
|
||
finally:
|
||
cursor.close()
|
||
|
||
async def _generate_db_summary(self, query: str, result: Dict[str, Any], table_name: str) -> str:
|
||
"""
|
||
生成数据库查询总结
|
||
"""
|
||
try:
|
||
llm = ChatZhipuAI(
|
||
model="glm-4",
|
||
api_key=os.getenv("ZHIPUAI_API_KEY"),
|
||
temperature=0.3
|
||
)
|
||
|
||
prompt = f"""
|
||
用户查询: {query}
|
||
目标表: {table_name}
|
||
|
||
查询结果:
|
||
- 结果行数: {result['total']}
|
||
- 结果列数: {len(result['columns'])}
|
||
- 列名: {', '.join(result['columns'])}
|
||
|
||
请基于以上信息，用中文生成一个简洁的数据库查询分析总结，包括:
|
||
1. 查询的主要目的
|
||
2. 关键数据发现
|
||
3. 数据特征分析
|
||
4. 建议的后续查询方向
|
||
|
||
总结应该专业、准确、易懂，控制在200字以内。
|
||
"""
|
||
|
||
response = await self._run_in_executor(
|
||
lambda: llm.invoke([HumanMessage(content=prompt)])
|
||
)
|
||
|
||
return response.content
|
||
|
||
except Exception as e:
|
||
return f"查询已完成，但生成总结时出现错误: {str(e)}"
|
||
|
||
# 在 SmartQueryService 类中添加方法
|
||
|
||
from .table_metadata_service import TableMetadataService
|
||
|
||
class SmartQueryService:
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.table_metadata_service = None
|
||
|
||
def set_db_session(self, db_session):
|
||
"""设置数据库会话"""
|
||
self.table_metadata_service = TableMetadataService(db_session)
|
||
|
||
async def get_database_context(self, user_id: int, query: str) -> str:
|
||
"""获取数据库上下文信息用于问答"""
|
||
if not self.table_metadata_service:
|
||
return ""
|
||
|
||
try:
|
||
# 获取用户的表元数据
|
||
table_metadata_list = self.table_metadata_service.get_user_table_metadata(user_id)
|
||
|
||
if not table_metadata_list:
|
||
return ""
|
||
|
||
# 构建数据库上下文
|
||
context_parts = []
|
||
context_parts.append("=== 数据库表信息 ===")
|
||
|
||
for metadata in table_metadata_list:
|
||
table_info = []
|
||
table_info.append(f"表名: {metadata.table_name}")
|
||
|
||
if metadata.table_comment:
|
||
table_info.append(f"表描述: {metadata.table_comment}")
|
||
|
||
if metadata.qa_description:
|
||
table_info.append(f"业务说明: {metadata.qa_description}")
|
||
|
||
# 添加列信息
|
||
if metadata.columns_info:
|
||
columns = []
|
||
for col in metadata.columns_info:
|
||
col_desc = f"{col['column_name']} ({col['data_type']})"
|
||
if col.get('column_comment'):
|
||
col_desc += f" - {col['column_comment']}"
|
||
columns.append(col_desc)
|
||
table_info.append(f"字段: {', '.join(columns)}")
|
||
|
||
# 添加示例数据
|
||
if metadata.sample_data:
|
||
table_info.append(f"示例数据: {metadata.sample_data[:2]}")
|
||
|
||
table_info.append(f"总行数: {metadata.row_count}")
|
||
|
||
context_parts.append("\n".join(table_info))
|
||
context_parts.append("---")
|
||
|
||
return "\n".join(context_parts)
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取数据库上下文失败: {str(e)}")
|
||
return ""
|
||
|
||
async def execute_smart_query(self, query: str, user_id: int, **kwargs) -> Dict[str, Any]:
|
||
"""执行智能查询（集成表元数据）"""
|
||
try:
|
||
# 获取数据库上下文
|
||
db_context = await self.get_database_context(user_id, query)
|
||
|
||
# 构建增强的提示词
|
||
enhanced_prompt = f"""
|
||
{db_context}
|
||
|
||
用户问题: {query}
|
||
|
||
请基于上述数据库表信息，生成相应的SQL查询语句。
|
||
注意：
|
||
1. 使用准确的表名和字段名
|
||
2. 考虑数据类型和约束
|
||
3. 参考示例数据理解数据格式
|
||
4. 生成高效的查询语句
|
||
"""
|
||
|
||
# 调用原有的查询逻辑
|
||
return await super().execute_smart_query(enhanced_prompt, user_id, **kwargs)
|
||
|
||
except Exception as e:
|
||
logger.error(f"智能查询失败: {str(e)}")
|
||
return {
|
||
'success': False,
|
||
'message': f"查询失败: {str(e)}"
|
||
} |