ailine/backend/app/core/visualization.py

"""
可视化图表公共工具 - 支持 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"![{data.title}](data:image/png;base64,{img_str})"
            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)