Spaces:
Sleeping
Sleeping
| from typing import List | |
| import plotly.io as pio | |
| import plotly.express as px | |
| import pandas as pd | |
| from utils import TEMP_DIR | |
| import os | |
| import ast | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| root_url = os.getenv("ROOT_URL", "") | |
| def _write_chart(fig, chart_path, chart_url): | |
| """Write a Plotly figure to disk and return a responsive iframe HTML string.""" | |
| pio.write_html(fig, chart_path, full_html=False, config={"responsive": True}) | |
| return ( | |
| 'Please display this iframe: ' | |
| '<div style="width:100%;overflow-x:auto;">' | |
| '<iframe style="width:100%;min-width:400px;" height="500" ' | |
| f'src="{chart_url}" frameborder="0" allowfullscreen>' | |
| '</iframe></div>' | |
| ) | |
| def llm_chart_data_scrub(data, layout): | |
| #Processing data to account for variation from LLM | |
| data_list = [] | |
| layout_dict = {} | |
| if isinstance(data, list): | |
| data_list = data | |
| else: | |
| data_list.append(data) | |
| false_replace = [':false', ': false'] | |
| false_value = ':False' | |
| true_replace = [':true', ': true'] | |
| true_value = ':True' | |
| data_dict = {} | |
| for data_obj in data_list: | |
| if isinstance(data_obj, str): | |
| data_obj = data_obj.replace("\n", "") | |
| for replace in false_replace: | |
| data_obj = data_obj.replace(replace, false_value) | |
| for replace in true_replace: | |
| data_obj = data_obj.replace(replace, true_value) | |
| print(data_obj) | |
| data_dict = ast.literal_eval(data_obj) | |
| else: | |
| data_dict = data_obj | |
| if layout and isinstance(layout, list): | |
| layout_obj = layout[0] | |
| else: | |
| layout_obj = layout | |
| if layout_obj and isinstance(layout_obj, str): | |
| for replace in false_replace: | |
| layout_obj = layout_obj.replace(replace, false_value) | |
| for replace in true_replace: | |
| layout_obj = layout_obj.replace(replace, true_value) | |
| print(layout_obj) | |
| layout_dict = ast.literal_eval(layout_obj) | |
| else: | |
| layout_dict = layout_obj | |
| return data_dict, layout_dict | |
| def scatter_chart_fig(df, x_column: List[str], y_column: str, category: str="", trendline: str="", | |
| trendline_options: List[dict]=[{}], marginal_x: str="", marginal_y: str="", | |
| size: str=""): | |
| function_args = {"data_frame":df, "x":x_column, "y":y_column} | |
| if category: | |
| function_args["color"] = category | |
| if trendline: | |
| function_args["trendline"] = trendline | |
| if marginal_x: | |
| function_args["marginal_x"] = marginal_x | |
| if marginal_y: | |
| function_args["marginal_y"] = marginal_y | |
| if size: | |
| df.loc[df[size] < 0, size] = 0 | |
| function_args["size"] = size | |
| if trendline_options: | |
| trendline_options_dict = {} | |
| if trendline_options and isinstance(trendline_options, list): | |
| trendline_options_obj = trendline_options[0] | |
| else: | |
| trendline_options_obj = trendline_options | |
| if trendline_options_obj and isinstance(trendline_options_obj, str): | |
| trendline_options_dict = ast.literal_eval(trendline_options_obj) | |
| else: | |
| trendline_options_dict = trendline_options_obj | |
| function_args["trendline_options"] = trendline_options_dict | |
| fig = px.scatter(**function_args) | |
| return fig | |
| def scatter_chart_generation_func(x_column: List[str], y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], | |
| category: str="", trendline: str="", trendline_options: List[dict]=[{}], marginal_x: str="", marginal_y: str="", | |
| size: str="", **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| initial_graph = scatter_chart_fig(df, x_column=x_column, y_column=y_column, | |
| category=category, trendline=trendline, trendline_options=trendline_options, | |
| marginal_x=marginal_x, marginal_y=marginal_y, size=size) | |
| fig = initial_graph.to_dict() | |
| print(data) | |
| print(layout) | |
| data_dict,layout_dict = llm_chart_data_scrub(data,layout) | |
| #Applying stylings and settings generated from LLM | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| data_ignore = ["x","y","type"] | |
| if size: | |
| data_ignore.append("marker") | |
| for key, value in data_dict.items(): | |
| if key not in data_ignore: | |
| for data_item in fig["data"]: | |
| data_item[key] = value | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("SCATTER PLOT ERROR") | |
| print(e) | |
| reply = f"""There was an error generating the Plotly Scatter Plot from {x_column}, {y_column}, {data}, and {layout} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| return {"reply": reply} | |
| def line_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], | |
| category: str="", **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| function_args = {"data_frame":df, "x":x_column, "y":y_column} | |
| if category: | |
| function_args["color"] = category | |
| initial_graph = px.line(**function_args) | |
| fig = initial_graph.to_dict() | |
| data_dict,layout_dict = llm_chart_data_scrub(data,layout) | |
| print(data_dict) | |
| print(layout_dict) | |
| #Applying stylings and settings generated from LLM | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| for key, value in data_dict.items(): | |
| if key not in ["x","y","type"]: | |
| for data_item in fig["data"]: | |
| data_item[key] = value | |
| print(fig) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("LINE CHART ERROR") | |
| print(e) | |
| reply = f"""There was an error generating the Plotly Line Chart from {x_column}, {y_column}, {data}, and {layout} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| return {"reply": reply} | |
| def bar_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], | |
| category: str="", facet_row: str="", facet_col: str="", **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| function_args = {"data_frame":df, "x":x_column, "y":y_column} | |
| if category: | |
| function_args["color"] = category | |
| if facet_row: | |
| function_args["facet_row"] = facet_row | |
| if facet_col: | |
| function_args["facet_col"] = facet_col | |
| initial_graph = px.bar(**function_args) | |
| fig = initial_graph.to_dict() | |
| data_dict,layout_dict = llm_chart_data_scrub(data,layout) | |
| print(data_dict) | |
| print(layout_dict) | |
| #Applying stylings and settings generated from LLM | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| for key, value in data_dict.items(): | |
| if key not in ["x","y","type"]: | |
| for data_item in fig["data"]: | |
| data_item[key] = value | |
| print(fig) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("BAR CHART ERROR") | |
| print(e) | |
| reply = f"""There was an error generating the Plotly Bar Chart from {x_column}, {y_column}, {data}, and {layout} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| return {"reply": reply} | |
| def pie_chart_generation_func(values: str, names: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| function_args = {"data_frame":df, "values":values, "names":names} | |
| initial_graph = px.pie(**function_args) | |
| fig = initial_graph.to_dict() | |
| data_dict,layout_dict = llm_chart_data_scrub(data,layout) | |
| print(data_dict) | |
| print(layout_dict) | |
| #Applying stylings and settings generated from LLM | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| for key, value in data_dict.items(): | |
| if key not in ["x","y","type"]: | |
| for data_item in fig["data"]: | |
| data_item[key] = value | |
| print(fig) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("PIE CHART ERROR") | |
| print(e) | |
| reply = f"""There was an error generating the Plotly Pie Chart from {values}, {names}, {data}, and {layout} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| return {"reply": reply} | |
| def histogram_generation_func(x_column: str, session_hash, session_folder, y_column: str="", data: List[dict]=[{}], layout: List[dict]=[{}], histnorm: str="", category: str="", | |
| histfunc: str="", **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| print(x_column) | |
| function_args = {"data_frame":df, "x":x_column} | |
| if y_column: | |
| function_args["y"] = y_column | |
| if histnorm: | |
| function_args["histnorm"] = histnorm | |
| if category: | |
| function_args["color"] = category | |
| if histfunc: | |
| function_args["histfunc"] = histfunc | |
| initial_graph = px.histogram(**function_args) | |
| fig = initial_graph.to_dict() | |
| data_dict,layout_dict = llm_chart_data_scrub(data,layout) | |
| print(data_dict) | |
| print(layout_dict) | |
| #Applying stylings and settings generated from LLM | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| for key, value in data_dict.items(): | |
| if key not in ["x","y","type"]: | |
| for data_item in fig["data"]: | |
| data_item[key] = value | |
| print(fig) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("HISTOGRAM ERROR") | |
| print(e) | |
| reply = f"""There was an error generating the Plotly Histogram from {x_column}. | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| return {"reply": reply} | |
| def box_chart_generation_func(y_column: str, session_hash, session_folder, | |
| x_column: str="", category: str="", | |
| layout: List[dict]=[{}], **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| function_args = {"data_frame": df, "y": y_column} | |
| if x_column: | |
| function_args["x"] = x_column | |
| if category: | |
| function_args["color"] = category | |
| initial_graph = px.box(**function_args) | |
| fig = initial_graph.to_dict() | |
| _, layout_dict = llm_chart_data_scrub({}, layout) | |
| if layout_dict: | |
| fig["layout"] = layout_dict | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("BOX CHART ERROR") | |
| print(e) | |
| return {"reply": f"There was an error generating the box plot. Error: {e}. You should probably try again."} | |
| def correlation_heatmap_func(session_hash, session_folder, columns: List[str]=[], **kwargs): | |
| try: | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| numeric_df = df[columns].select_dtypes(include='number') if columns else df.select_dtypes(include='number') | |
| if numeric_df.shape[1] < 2: | |
| return {"reply": "At least two numeric columns are needed for a correlation matrix. Please refine your query to include more numeric columns."} | |
| corr = numeric_df.corr().round(3) | |
| fig = px.imshow( | |
| corr, | |
| text_auto='.2f', | |
| color_continuous_scale='RdBu_r', | |
| zmin=-1, | |
| zmax=1, | |
| title='Correlation Matrix', | |
| aspect='auto', | |
| ) | |
| fig.update_layout(font=dict(family='Inter, system-ui, sans-serif')) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("CORRELATION HEATMAP ERROR") | |
| print(e) | |
| return {"reply": f"There was an error generating the correlation heatmap. Error: {e}. You should probably try again."} | |
| def rolling_stats_func(x_column: str, y_column: str, session_hash, session_folder, | |
| window: int = 7, stats: List[str] = ["mean"], | |
| layout: List[dict] = [{}], category: str = "", **kwargs): | |
| try: | |
| import plotly.graph_objects as go | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| chart_path = f'{dir_path}/chart.html' | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| try: | |
| df[x_column] = pd.to_datetime(df[x_column]) | |
| except Exception: | |
| pass | |
| df = df.sort_values(x_column) | |
| valid_stats = {"mean", "std", "min", "max"} | |
| selected_stats = [s for s in stats if s in valid_stats] or ["mean"] | |
| fig = go.Figure() | |
| groups = df[category].unique().tolist() if category and category in df.columns else [None] | |
| for group in groups: | |
| group_df = df[df[category] == group] if group is not None else df | |
| prefix = f"{group} — " if group is not None else "" | |
| fig.add_trace(go.Scatter( | |
| x=group_df[x_column].values, y=group_df[y_column].values, | |
| mode="lines", name=f"{prefix}{y_column} (raw)", | |
| opacity=0.35, line=dict(width=1) | |
| )) | |
| rolling_obj = group_df[y_column].rolling(window) | |
| for stat in selected_stats: | |
| rolled = getattr(rolling_obj, stat)() | |
| fig.add_trace(go.Scatter( | |
| x=group_df[x_column].values, y=rolled.values, | |
| mode="lines", name=f"{prefix}Rolling {stat.capitalize()} (w={window})", | |
| line=dict(width=2.5) | |
| )) | |
| fig.update_layout( | |
| title=f"Rolling Statistics (window={window}) — {y_column}", | |
| xaxis_title=x_column, | |
| yaxis_title=y_column, | |
| font=dict(family="Inter, system-ui, sans-serif"), | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), | |
| ) | |
| _, layout_dict = llm_chart_data_scrub({}, layout) | |
| if layout_dict: | |
| fig.update_layout(**layout_dict) | |
| chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html' | |
| return {"reply": _write_chart(fig, chart_path, chart_url)} | |
| except Exception as e: | |
| print("ROLLING STATS ERROR") | |
| print(e) | |
| return {"reply": f"There was an error generating the rolling statistics chart. Error: {e}. You should probably try again."} | |
| def table_generation_func(session_hash, session_folder, **kwargs): | |
| print("TABLE GENERATION") | |
| try: | |
| from html import escape | |
| dir_path = TEMP_DIR / str(session_hash) / str(session_folder) | |
| csv_query_path = f'{dir_path}/query.csv' | |
| df = pd.read_csv(csv_query_path) | |
| total_rows = len(df) | |
| max_rows = 200 | |
| if total_rows > max_rows: | |
| df = df.head(max_rows) | |
| note = (f'<p class="vda-table-note">Showing first {max_rows} of {total_rows} rows' | |
| ' — refine your query to see more specific results.</p>') | |
| else: | |
| note = '' | |
| header_cells = ''.join(f'<th>{escape(str(col))}</th>' for col in df.columns) | |
| row_html = [ | |
| '<tr>' + ''.join(f'<td>{escape(str(val))}</td>' for val in row) + '</tr>' | |
| for _, row in df.iterrows() | |
| ] | |
| style = ( | |
| '<style>' | |
| '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}' | |
| '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}' | |
| '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}' | |
| '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}' | |
| '.vda-table tbody tr:nth-child(even){background:#f8fafc;}' | |
| '.vda-table tbody tr:last-child td{border-bottom:none;}' | |
| '.vda-table-note{font-size:12px;color:#6b7280;margin:4px 0 0;text-align:right;}' | |
| '</style>' | |
| ) | |
| table = ( | |
| '<div class="vda-table-wrap"><table class="vda-table">' | |
| f'<thead><tr>{header_cells}</tr></thead>' | |
| '<tbody>' + '\n'.join(row_html) + '</tbody>' | |
| '</table></div>' | |
| ) | |
| return {"reply": style + table + note} | |
| except Exception as e: | |
| print("TABLE ERROR") | |
| print(e) | |
| return {"reply": f"There was an error generating the table. Error: {e}. You should probably try again."} | |