Spaces:
Running
Running
File size: 5,009 Bytes
44de337 00c47cd 44de337 ca7adb9 44de337 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import os
import textwrap
import nbformat
from nbformat import v4 as nbf
import logging
import json
from datetime import datetime, timezone
import pandas as pd
logger = logging.getLogger(__name__)
class Notebook:
def __init__(self, data):
self.nb = nbf.new_notebook()
self.nb.cells = []
self.data = data
self.output_path = "tmp/notebook.ipynb"
self.imports_added = False
def build_notebook(self):
success = 0
metrics = self.data
logger.info(f"Start building notebook with {len(metrics)} metrics")
for i, metric in enumerate(metrics):
logger.info(f"Processing metric {i+1} of {len(metrics)}")
metric_result = json.loads(metric.get("body"))
table_data = metric_result.get("table_data")
query_config = metric_result.get("sql_config")
chart_data = metric_result.get("chart_data", {})
chart_config = metric_result.get("chart_config", {})
chart_type = metric_result.get("chart_config", {}).get("type", {})
if not self.nb.cells:
self._add_title()
self._add_metric_header(config=chart_config, index=i)
self._add_query_details(query_config=query_config)
if not self.imports_added:
self._import_libs()
self.imports_added = True
self._add_table(raw_table=table_data)
if chart_data and chart_config:
self._chart_code(chart_type=chart_type, chart_data=chart_data, chart_config=chart_config)
def export_notebook(self):
self.build_notebook()
logger.info(f"Exporting notebook to {self.output_path}")
os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
with open(self.output_path, 'w') as f:
nbformat.write(self.nb, f)
logger.info(f"Notebook exported successfully to {self.output_path}")
return self.output_path
def _import_libs(self):
self.nb.cells.append(nbf.new_markdown_cell("### 📦 Imports"))
self.nb.cells.append(nbf.new_code_cell("import pandas as pd\nimport plotly.express as px"))
def _add_title(self):
self.nb.cells.append(nbf.new_markdown_cell("# Metrics Dashboard"))
def _add_metric_header(self, config, index):
self.nb.cells.append(nbf.new_markdown_cell(f"## 📈 Metric {index+1} {config.get('title', '')}\n**Timestamp:** {datetime.now(timezone.utc).isoformat()}"))
def _chart_code(self, chart_type, chart_data, chart_config):
df_code = f"df = pd.DataFrame({chart_data.get(chart_type).get('data')})"
chart_dispatch = {
"bar": f"""
fig = px.bar(df, x='x', y='y', title='{chart_config['title']}',
labels={{'y': '{chart_config['y_axis_label']}', 'x': '{chart_config['x_axis_label']}'}}, color='x', template='plotly_white')
""",
"line": f"""
fig = px.line(df, x='x', y='y', title='{chart_config['title']}',
labels={{'y': '{chart_config['y_axis_label']}', 'x': '{chart_config['x_axis_label']}'}}, template='plotly_white')
""",
"pie": f"""
fig = px.pie(df, values='y', names='x', title='{chart_config['title']}',
labels={{'y': '{chart_config['y_axis_label']}'}}, template='plotly_white', hole=0.3)
fig.update_traces(textposition='inside', textfont_size=12)
""",
"hist": f"""
df['bin_center'] = (df['bin_start'] + df['bin_end']) / 2
fig = px.bar(df, x='bin_center', y='frequency', title='{chart_config['title']}',
labels={{'frequency': '{chart_config['y_axis_label']}', 'bin_center': '{chart_config['x_axis_label']}'}},
template='plotly_white')
"""
}
chart_code = f"{df_code}\n{textwrap.dedent(chart_dispatch.get(chart_type))}\nfig.update_layout(showlegend=False)\nfig.show()"
self.nb.cells.append(nbf.new_markdown_cell(f"### 📊 Chart"))
self.nb.cells.append(nbf.new_code_cell(chart_code))
def _add_query_details(self, query_config):
sql_clean = query_config['sql_query'].replace("\\n", "\n").strip()
self.nb.cells.append(nbf.new_markdown_cell(f"**User Query:** {query_config['text_query']}\n"))
self.nb.cells.append(nbf.new_markdown_cell(f"### 🔍 Query Details\n**SQL:**\n```sql\n{sql_clean}\n```\n**Explanation:** {query_config['explanation']}"))
def _add_table(self, raw_table):
table_data = pd.DataFrame(raw_table).fillna("NAN").to_dict(orient="records")
raw_code = f'raw_table = pd.DataFrame({table_data})\nraw_table.head()'
self.nb.cells.append(nbf.new_markdown_cell("### Table"))
self.nb.cells.append(nbf.new_code_cell(raw_code)) |