OpenManusHiDpiFix/app/tool/chart_visualization/chart_visualization.py

import subprocess
import json
import threading
import base64
import pandas as pd
import aiofiles
import os
from typing import Any, Dict, Hashable
from app.tool.base import BaseTool
from app.config import config


def extract_executable_code(code_str: str) -> str:
    """
    Extract executable code from function call's parameters

    Args:
        code_str (string): The python code generated by llm.

    Returns:
        String: Python code can execute directly.
    """
    lines = code_str.strip().splitlines()
    start_idx = -1
    end_idx = -1

    # Find first occurrence of ```
    for i, line in enumerate(lines):
        if "```" in line.strip() or '"""' in line.strip():
            start_idx = i
            break

    # Find last occurrence of ```
    for i in reversed(range(len(lines))):
        if "```" in line.strip() or '"""' in line.strip():
            end_idx = i
            break

    if start_idx != -1 and end_idx != -1 and start_idx < end_idx:
        lines = lines[start_idx + 1 : end_idx]
    elif start_idx != -1:
        lines = lines[start_idx + 1 :]
    elif end_idx != -1:
        lines = lines[:end_idx]

    return "\n".join(lines)


class ChartVisualization(BaseTool):
    name: str = "generate_data_visualization"
    description: str = """Visualize a statistical chart using csv data and chart description. The tool accepts code to generate csv data and description of the chart, and output a chart in png or html.
Note: Each tool call generates a single chart.
"""
    parameters: dict = {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": """Python code template EXCLUSIVELY for CSV generation. MUST CONTAIN:
1. Data loading logic (handle dataframe/dict/file/url/json)
2. Data processing (cleaning/transformation)
3. CSV saving with path print (Only csv path)
""",
                # example
                "examples": [
                    """import pandas as pd
# Create safe output directory if not exists
output_dir = './data/chart_generation_temp'
os.makedirs(output_dir, exist_ok=True)

# Data loading
# Add your data loding logic here

# Data processing placeholder
# Add your cleaning/transformation logic here

# Final: Generic output handling
output_path = os.path.join(output_dir, 'csv_file_name.csv')
df.to_csv(output_path, index=False)

print(output_path)"""
                ],
            },
            "chart_description": {
                "type": "string",
                "description": "The chart title or description should be concise and clear",
                "examples": ["Product sales distribution", "Monthly revenue trend"],
            },
            "output_type": {
                "description": "Rendering format (html=interactive)",
                "type": "string",
                "default": "html",
                "enum": ["png", "html"],
            },
        },
        "required": ["code", "chart_description"],
    }
    llm: dict = config.llm["default"]

    async def execute(self, code: str, chart_description: str, output_type: str) -> str:
        code_result = await self.execute_code(code=code)
        if "success" in code_result and code_result["success"] is False:
            return code_result
        if code_result["observation"].startswith("Error"):
            return {"observation": code_result["observation"], "success": False}

        try:
            data_path = (
                code_result["observation"].replace("\n", "").replace("\r", "").strip()
            )
            if not data_path.endswith(".csv"):
                return {
                    "observation": "Error: Code should ONLY output CSV data path",
                    "success": False,
                }
            df = pd.read_csv(data_path)
            df = df.astype(object)
            df = df.where(pd.notnull(df), None)
            data_dict_list = df.to_json(orient="records", force_ascii=False)
            result = await self.invoke_vmind(
                data_dict_list, chart_description, output_type
            )
            if "error" in result:
                return {
                    "observation": f"Error: {result["error"]}",
                    "success": False,
                }
            chart_file_path = data_path.replace(".csv", f".{output_type}")
            while os.path.exists(chart_file_path):
                chart_file_path = chart_file_path.replace(
                    f".{output_type}", f"_new.{output_type}"
                )
            if output_type == "png":
                byte_data = base64.b64decode(result["res"])
                async with aiofiles.open(chart_file_path, "wb") as file:
                    await file.write(byte_data)
            else:
                async with aiofiles.open(
                    chart_file_path, "w", encoding="utf-8"
                ) as file:
                    await file.write(result["res"])
            return {"observation": f"chart successfully saved to {chart_file_path}"}
        except Exception as e:
            return {
                "observation": f"Error: {e}",
                "success": False,
            }

    async def execute_code(
        self,
        code: str,
        timeout: int = 5,
    ) -> Dict:
        """
        Executes the provided Python code with a timeout.

        Args:
            code (str): The Python code to execute.
            timeout (int): Execution timeout in seconds.

        Returns:
            Dict: Contains 'output' with execution output or error message and 'success' status.
        """
        result = {"observation": ""}
        be_extracted_code = extract_executable_code(code)

        def run_code():
            try:
                safe_globals = {"__builtins__": dict(__builtins__)}

                import sys
                from io import StringIO

                output_buffer = StringIO()
                sys.stdout = output_buffer

                exec(be_extracted_code, safe_globals, {})

                sys.stdout = sys.__stdout__

                result["observation"] = output_buffer.getvalue()

            except Exception as e:
                result["observation"] = str(e)
                result["success"] = False

        thread = threading.Thread(target=run_code)
        thread.start()
        thread.join(timeout)

        if thread.is_alive():
            return {
                "observation": f"Execution timeout after {timeout} seconds",
                "success": False,
            }

        return result

    async def invoke_vmind(
        self,
        dict_data: list[dict[Hashable, Any]],
        chart_description: str,
        output_type: str,
    ):
        vmind_options = {
            "url": self.llm.base_url + "/chat/completions",
            "model": self.llm.model,
            "headers": {
                "api-key": self.llm.api_key,
                "Authorization": f"Bearer {self.llm.api_key}",
            },
        }
        vmind_params = {
            "options": vmind_options,
            "user_prompt": chart_description,
            "dataset": dict_data,
            "output_type": output_type,
        }
        process = subprocess.run(
            ["npx", "ts-node", "src/chartVisualize.ts"],
            input=json.dumps(vmind_params),
            capture_output=True,
            text=True,
            encoding="utf-8",
            cwd=os.path.dirname(__file__),
        )
        if process.returncode == 0:
            return json.loads(process.stdout)
        else:
            return {"error": f"Node.js Error: {process.stderr}"}