ZJU_czx
118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
import subprocess
|
|
import json
|
|
import base64
|
|
import pandas as pd
|
|
import aiofiles
|
|
import os
|
|
from typing import Any, Hashable
|
|
from pydantic import Field, model_validator
|
|
|
|
from app.llm import LLM
|
|
from app.tool.base import BaseTool
|
|
from app.logger import logger
|
|
|
|
|
|
class ChartVisualization(BaseTool):
|
|
name: str = "generate_data_visualization"
|
|
description: str = """Visualize a statistical chart using csv data and chart description. The tool accepts local csv data file path and description of the chart, and output a chart in png or html.
|
|
Note: Each tool call generates only one single chart.
|
|
"""
|
|
parameters: dict = {
|
|
"type": "object",
|
|
"properties": {
|
|
"csv_path": {
|
|
"type": "string",
|
|
"description": """file path of csv data with ".csv" in the end""",
|
|
},
|
|
"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: LLM = Field(default_factory=LLM, description="Language model instance")
|
|
|
|
@model_validator(mode="after")
|
|
def initialize_llm(self):
|
|
"""Initialize llm with default settings if not provided."""
|
|
if self.llm is None or not isinstance(self.llm, LLM):
|
|
self.llm = LLM(config_name=self.name.lower())
|
|
return self
|
|
|
|
async def execute(
|
|
self, csv_path: str, chart_description: str, output_type: str
|
|
) -> str:
|
|
logger.info(
|
|
f"📈 Chart Generation with data and description: {chart_description} with {csv_path} "
|
|
)
|
|
try:
|
|
df = pd.read_csv(csv_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 = csv_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 invoke_vmind(
|
|
self,
|
|
dict_data: list[dict[Hashable, Any]],
|
|
chart_description: str,
|
|
output_type: str,
|
|
):
|
|
llm_config = {
|
|
"base_url": self.llm.base_url,
|
|
"model": self.llm.model,
|
|
"api_key": self.llm.api_key,
|
|
}
|
|
vmind_params = {
|
|
"llm_config": llm_config,
|
|
"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}"}
|