File size: 1,475 Bytes
48d8f38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | 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)
# check response is str
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}
|