| from pandasai import Agent |
| from pydantic import BaseModel |
| import json |
| from web.file import FileBean, get_dataframes |
| from base.no_parser import NoParser |
| from pathlib import Path |
| from base.init import llm |
| from web.file import FileBean, get_dataframes |
| import matplotlib |
| import matplotlib.pyplot as plt |
| plt.rcParams['font.family'] = ['Heiti TC', 'Arial Unicode MS'] |
| matplotlib.use('Agg') |
| config = {"llm": llm, "verbose": True, "open_charts": False, |
| "response_parser": NoParser, "save_charts": True} |
|
|
| class BaseReq(BaseModel): |
| prompt: str |
|
|
|
|
| class PlotReq(BaseReq): |
| fileList: list[FileBean] |
|
|
| def gen_plot(req: PlotReq): |
| dfs = [df for fb in req.fileList for df in get_dataframes(fb).values()] |
| lake = Agent( |
| dfs, config={**config, "save_charts_path": f".data/static/plot"}) |
| response = lake.chat(req.prompt) |
| |
| if isinstance(response, str): |
| return {"resultType": 'error', "resultValueStr": response} |
| result_type = response['type'] |
| if result_type == 'dataframe': |
| v_dict = response['value'].to_dict(orient='records') |
| result_value_str = json.dumps(v_dict) |
| elif result_type == 'plot': |
| file_path = Path(response["value"]) |
| result_value_str = file_path.name |
| elif result_type == 'number': |
| result_value_str = str(response["value"]) |
| else: |
| result_value_str = response["value"] |
| return {"resultType": result_type, "resultValueStr": result_value_str} |
|
|
