summary_generation.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. # agent/tools/summary_generation.py
  2. from langchain.tools import tool
  3. from typing import Dict, Any
  4. import pandas as pd
  5. from common.vanna_instance import get_vanna_instance
  6. from core.logging import get_agent_logger
  7. # Initialize logger
  8. logger = get_agent_logger("SummaryGeneration")
  9. @tool
  10. def generate_summary(question: str, query_result: Dict[str, Any], sql: str) -> Dict[str, Any]:
  11. """
  12. 为查询结果生成自然语言摘要。
  13. Args:
  14. question: 原始问题
  15. query_result: 查询结果数据
  16. sql: 执行的SQL语句
  17. Returns:
  18. 包含摘要结果的字典,格式:
  19. {
  20. "success": bool,
  21. "summary": str,
  22. "error": str或None
  23. }
  24. """
  25. try:
  26. logger.info(f"开始生成摘要,问题: {question}")
  27. if not query_result or not query_result.get("rows"):
  28. return {
  29. "success": True,
  30. "summary": "查询执行完成,但没有找到符合条件的数据。",
  31. "message": "无数据摘要"
  32. }
  33. # 重构DataFrame用于摘要生成
  34. df = _reconstruct_dataframe(query_result)
  35. if df is None or df.empty:
  36. return {
  37. "success": True,
  38. "summary": "查询执行完成,但数据为空。",
  39. "message": "空数据摘要"
  40. }
  41. # 调用Vanna生成摘要(thinking内容已在base_llm_chat.py中统一处理)
  42. vn = get_vanna_instance()
  43. summary = vn.generate_summary(question=question, df=df)
  44. if summary is None:
  45. # 生成默认摘要
  46. summary = _generate_default_summary(question, query_result, sql)
  47. logger.info(f"摘要生成成功: {summary[:100]}...")
  48. return {
  49. "success": True,
  50. "summary": summary,
  51. "message": "摘要生成成功"
  52. }
  53. except Exception as e:
  54. logger.error(f"摘要生成异常: {str(e)}")
  55. # 生成备用摘要
  56. fallback_summary = _generate_fallback_summary(question, query_result, sql)
  57. return {
  58. "success": True, # 即使异常也返回成功,因为有备用摘要
  59. "summary": fallback_summary,
  60. "message": f"使用备用摘要生成: {str(e)}"
  61. }
  62. def _reconstruct_dataframe(query_result: Dict[str, Any]) -> pd.DataFrame:
  63. """从查询结果重构DataFrame"""
  64. try:
  65. rows = query_result.get("rows", [])
  66. columns = query_result.get("columns", [])
  67. if not rows or not columns:
  68. return pd.DataFrame()
  69. return pd.DataFrame(rows, columns=columns)
  70. except Exception as e:
  71. logger.warning(f"DataFrame重构失败: {str(e)}")
  72. return pd.DataFrame()
  73. def _generate_default_summary(question: str, query_result: Dict[str, Any], sql: str) -> str:
  74. """生成默认摘要"""
  75. try:
  76. row_count = query_result.get("row_count", 0)
  77. columns = query_result.get("columns", [])
  78. if row_count == 0:
  79. return "查询执行完成,但没有找到符合条件的数据。"
  80. summary_parts = [f"根据您的问题「{question}」,查询返回了 {row_count} 条记录。"]
  81. if columns:
  82. summary_parts.append(f"数据包含以下字段:{', '.join(columns)}。")
  83. return ' '.join(summary_parts)
  84. except Exception:
  85. return f"查询执行完成,共返回 {query_result.get('row_count', 0)} 条记录。"
  86. def _generate_fallback_summary(question: str, query_result: Dict[str, Any], sql: str) -> str:
  87. """生成备用摘要"""
  88. row_count = query_result.get("row_count", 0)
  89. if row_count == 0:
  90. return "查询执行完成,但没有找到符合条件的数据。请检查查询条件是否正确。"
  91. return f"查询执行成功,共返回 {row_count} 条记录。数据已准备完毕,您可以查看详细结果。"