添加公共工具:联网搜索(DuckDuckGo)和可视化图表(Mermaid)
All checks were successful
构建并部署 AI Agent 服务 / deploy (push) Successful in 10m38s
All checks were successful
构建并部署 AI Agent 服务 / deploy (push) Successful in 10m38s
This commit is contained in:
213
backend/app/core/visualization.py
Normal file
213
backend/app/core/visualization.py
Normal file
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
可视化图表公共工具 - 支持 Mermaid 和简单图表生成
|
||||
Visualization Public Utility - Support Mermaid and simple chart generation
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChartData:
|
||||
"""图表数据类"""
|
||||
title: str
|
||||
x_labels: List[str]
|
||||
y_values: List[float]
|
||||
chart_type: str = "bar" # bar, line, pie
|
||||
|
||||
|
||||
class VisualizationTool:
|
||||
"""可视化图表公共工具类"""
|
||||
|
||||
def generate_mermaid_chart(self, data: ChartData) -> str:
|
||||
"""
|
||||
生成 Mermaid 格式图表
|
||||
|
||||
Args:
|
||||
data: 图表数据
|
||||
|
||||
Returns:
|
||||
Mermaid 格式字符串
|
||||
"""
|
||||
if data.chart_type == "pie":
|
||||
return self._generate_mermaid_pie(data)
|
||||
elif data.chart_type == "bar":
|
||||
return self._generate_mermaid_bar(data)
|
||||
else:
|
||||
return self._generate_mermaid_bar(data)
|
||||
|
||||
def _generate_mermaid_pie(self, data: ChartData) -> str:
|
||||
"""生成 Mermaid 饼图"""
|
||||
lines = []
|
||||
lines.append("```mermaid")
|
||||
lines.append("pie showData")
|
||||
lines.append(f" title {data.title}")
|
||||
|
||||
for label, value in zip(data.x_labels, data.y_values):
|
||||
lines.append(f" \"{label}\" : {value}")
|
||||
|
||||
lines.append("```")
|
||||
return "\n".join(lines)
|
||||
|
||||
def _generate_mermaid_bar(self, data: ChartData) -> str:
|
||||
"""生成 Mermaid 柱状图(使用 xychart)"""
|
||||
lines = []
|
||||
lines.append("```mermaid")
|
||||
lines.append("xychart-beta")
|
||||
lines.append(f" title \"{data.title}\"")
|
||||
|
||||
# 构建 x-axis
|
||||
x_axis = " ".join([f"\"{label}\"" for label in data.x_labels])
|
||||
lines.append(f" x-axis [ {x_axis} ]")
|
||||
|
||||
# 构建 y-axis
|
||||
y_axis = " ".join([str(v) for v in data.y_values])
|
||||
lines.append(f" y-axis \"数值\" {y_axis}")
|
||||
|
||||
# 柱状图
|
||||
lines.append(f" bar [ {y_axis} ]")
|
||||
|
||||
lines.append("```")
|
||||
return "\n".join(lines)
|
||||
|
||||
def generate_matplotlib_chart(self, data: ChartData, output_path: Optional[str] = None) -> str:
|
||||
"""
|
||||
使用 matplotlib 生成图表
|
||||
|
||||
Args:
|
||||
data: 图表数据
|
||||
output_path: 输出文件路径(可选)
|
||||
|
||||
Returns:
|
||||
图表描述或生成的文件路径
|
||||
"""
|
||||
try:
|
||||
import matplotlib.pyplot as plt
|
||||
import io
|
||||
import base64
|
||||
|
||||
# 创建图表
|
||||
fig, ax = plt.subplots(figsize=(10, 6))
|
||||
|
||||
if data.chart_type == "bar":
|
||||
ax.bar(data.x_labels, data.y_values, color='#3498db')
|
||||
elif data.chart_type == "line":
|
||||
ax.plot(data.x_labels, data.y_values, marker='o', color='#2ecc71', linewidth=2)
|
||||
elif data.chart_type == "pie":
|
||||
ax.pie(data.y_values, labels=data.x_labels, autopct='%1.1f%%', startangle=90)
|
||||
ax.axis('equal')
|
||||
|
||||
ax.set_title(data.title, fontsize=14, pad=20)
|
||||
ax.grid(axis='y', alpha=0.3)
|
||||
|
||||
# 如果没有指定输出路径,返回 base64 编码
|
||||
if output_path is None:
|
||||
buf = io.BytesIO()
|
||||
plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
|
||||
buf.seek(0)
|
||||
img_str = base64.b64encode(buf.read()).decode('utf-8')
|
||||
plt.close()
|
||||
return f""
|
||||
else:
|
||||
plt.savefig(output_path, dpi=100, bbox_inches='tight')
|
||||
plt.close()
|
||||
return f"图表已保存至:{output_path}"
|
||||
|
||||
except ImportError:
|
||||
# 如果 matplotlib 未安装,返回 Mermaid 格式
|
||||
return self.generate_mermaid_chart(data)
|
||||
except Exception as e:
|
||||
print(f"生成图表出错:{e}")
|
||||
return self.generate_mermaid_chart(data)
|
||||
|
||||
def format_chart_result(self, mermaid_code: str, description: str = "") -> str:
|
||||
"""
|
||||
格式化图表输出
|
||||
|
||||
Args:
|
||||
mermaid_code: Mermaid 图表代码
|
||||
description: 图表描述
|
||||
|
||||
Returns:
|
||||
格式化后的 Markdown 文本
|
||||
"""
|
||||
lines = []
|
||||
lines.append("## 📊 可视化图表\n")
|
||||
|
||||
if description:
|
||||
lines.append(description)
|
||||
lines.append("")
|
||||
|
||||
lines.append(mermaid_code)
|
||||
lines.append("")
|
||||
lines.append("---")
|
||||
lines.append("💡 **图表说明**:")
|
||||
lines.append("- 上方为 Mermaid 格式图表,支持在 Markdown 中渲染")
|
||||
lines.append("- 如需更复杂图表,可进一步使用 matplotlib 生成")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def quick_chart_from_text(self, text: str, chart_type: str = "bar") -> str:
|
||||
"""
|
||||
从文本快速生成图表(大模型友好接口)
|
||||
|
||||
Args:
|
||||
text: 包含数据的文本(格式:标题,标签1:值1,标签2:值2,...)
|
||||
chart_type: 图表类型(bar, line, pie)
|
||||
|
||||
Returns:
|
||||
格式化后的图表输出
|
||||
"""
|
||||
try:
|
||||
# 解析文本格式
|
||||
parts = [p.strip() for p in text.split(',')]
|
||||
title = parts[0]
|
||||
|
||||
x_labels = []
|
||||
y_values = []
|
||||
|
||||
for part in parts[1:]:
|
||||
if ':' in part:
|
||||
label, val_str = part.split(':', 1)
|
||||
x_labels.append(label.strip())
|
||||
y_values.append(float(val_str.strip()))
|
||||
|
||||
data = ChartData(
|
||||
title=title,
|
||||
x_labels=x_labels,
|
||||
y_values=y_values,
|
||||
chart_type=chart_type
|
||||
)
|
||||
|
||||
mermaid = self.generate_mermaid_chart(data)
|
||||
return self.format_chart_result(mermaid)
|
||||
|
||||
except Exception as e:
|
||||
return f"解析图表数据出错:{e}\n\n请使用格式:标题,标签1:值1,标签2:值2,..."
|
||||
|
||||
|
||||
# 单例实例
|
||||
_visualization_tool = None
|
||||
|
||||
|
||||
def get_visualization_tool() -> VisualizationTool:
|
||||
"""获取可视化工具单例"""
|
||||
global _visualization_tool
|
||||
if _visualization_tool is None:
|
||||
_visualization_tool = VisualizationTool()
|
||||
return _visualization_tool
|
||||
|
||||
|
||||
def generate_chart(text: str, chart_type: str = "bar") -> str:
|
||||
"""
|
||||
便捷函数:快速生成图表
|
||||
|
||||
Args:
|
||||
text: 包含数据的文本(格式:标题,标签1:值1,标签2:值2,...)
|
||||
chart_type: 图表类型(bar, line, pie)
|
||||
|
||||
Returns:
|
||||
格式化后的图表输出
|
||||
"""
|
||||
tool = get_visualization_tool()
|
||||
return tool.quick_chart_from_text(text, chart_type)
|
||||
Reference in New Issue
Block a user