{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# AI Data Chatbot\n", "Ask natural language questions about your dataset — get back tables and interactive charts.\n", "\n", "**Run all cells top to bottom, then use the chat widget at the bottom.**" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 1: Install dependencies (run once) ──────────────────────────────────\n", "# Uncomment and run this cell if you haven't installed the packages yet\n", "# !pip install anthropic pandas numpy plotly openpyxl python-dotenv ipywidgets" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 2: Imports ───────────────────────────────────────────────────────────\n", "import os, json, re, textwrap\n", "import pandas as pd\n", "import numpy as np\n", "import plotly.express as px\n", "import plotly.graph_objects as go\n", "import ipywidgets as widgets\n", "from IPython.display import display, clear_output, HTML\n", "import anthropic\n", "from dotenv import load_dotenv\n", "\n", "load_dotenv() # reads ANTHROPIC_API_KEY from .env file\n", "print('✓ Imports OK')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 3: API Key ───────────────────────────────────────────────────────────\n", "# Option A: reads from .env file automatically (recommended)\n", "# Option B: paste your key directly below\n", "# os.environ['ANTHROPIC_API_KEY'] = 'sk-ant-...'\n", "\n", "assert os.getenv('ANTHROPIC_API_KEY'), \"Set ANTHROPIC_API_KEY in .env or uncomment Option B above\"\n", "print('✓ API key found')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 4: Load Data ─────────────────────────────────────────────────────────\n", "# Default: load the sample employee dataset\n", "# To use your own file, change this path to any CSV or Excel file:\n", "DATA_PATH = 'sample_data/employees.csv'\n", "\n", "if DATA_PATH.endswith(('.xlsx', '.xls')):\n", " df = pd.read_excel(DATA_PATH)\n", "else:\n", " df = pd.read_csv(DATA_PATH)\n", "\n", "print(f'✓ Loaded {len(df):,} rows × {len(df.columns)} columns')\n", "print(f' Columns: {\", \".join(df.columns.tolist())}')\n", "df.head()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 5: Core helpers ──────────────────────────────────────────────────────\n", "\n", "def get_schema_info(df):\n", " lines = [f'Shape: {df.shape[0]} rows × {df.shape[1]} columns', 'Columns:']\n", " for col in df.columns:\n", " dtype = str(df[col].dtype)\n", " if pd.api.types.is_numeric_dtype(df[col]):\n", " lines.append(f' {col} ({dtype}): min={df[col].min()}, max={df[col].max()}, mean={df[col].mean():.1f}')\n", " else:\n", " uniq = df[col].nunique()\n", " sample = df[col].dropna().unique()[:5].tolist()\n", " lines.append(f' {col} ({dtype}): {uniq} unique values, e.g. {sample}')\n", " return '\\n'.join(lines)\n", "\n", "def strip_imports(code):\n", " return '\\n'.join(\n", " line for line in code.splitlines()\n", " if not line.strip().startswith(('import ', 'from '))\n", " )\n", "\n", "def safe_exec_pandas(code, df):\n", " code = strip_imports(code)\n", " ns = {'df': df.copy(), 'pd': pd, 'np': np,\n", " '__builtins__': {'len':len,'range':range,'print':print,'str':str,\n", " 'int':int,'float':float,'list':list,'dict':dict,\n", " 'zip':zip,'enumerate':enumerate,'sorted':sorted,\n", " 'min':min,'max':max,'sum':sum,'abs':abs,'round':round}}\n", " exec(code, ns)\n", " return ns.get('result_df')\n", "\n", "def safe_exec_viz(code, result_df):\n", " code = strip_imports(code)\n", " ns = {'result_df': result_df.copy(), 'px': px, 'go': go, 'pd': pd, 'np': np,\n", " '__builtins__': {'len':len,'range':range,'print':print,'str':str,\n", " 'int':int,'float':float,'list':list,'dict':dict,\n", " 'zip':zip,'enumerate':enumerate,'sorted':sorted,\n", " 'min':min,'max':max,'sum':sum,'abs':abs,'round':round}}\n", " exec(code, ns)\n", " return ns.get('fig')\n", "\n", "def fallback_chart(result_df):\n", " num_cols = result_df.select_dtypes(include='number').columns.tolist()\n", " cat_cols = result_df.select_dtypes(exclude='number').columns.tolist()\n", " if num_cols and cat_cols:\n", " return px.bar(result_df, x=cat_cols[0], y=num_cols[0], title='Results')\n", " elif num_cols:\n", " return px.bar(result_df, y=num_cols[0], title='Results')\n", " return None\n", "\n", "print('✓ Helpers ready')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Cell 6: LLM handler ───────────────────────────────────────────────────────\n", "\n", "SYSTEM_PROMPT = \"\"\"\n", "You are a data analyst assistant. The user will ask questions about a dataset.\n", "You must respond with ONLY a valid JSON object — no markdown, no explanation outside the JSON.\n", "\n", "JSON format:\n", "{\n", " \"explanation\": \"Plain-English summary of what you found\",\n", " \"pandas_code\": \"Python code using pandas. Input: `df`. Output: must assign a DataFrame to `result_df`.\",\n", " \"viz_code\": \"Python code using plotly. Input: `result_df`. Output: must assign a Plotly figure to `fig`. DO NOT include any import statements.\",\n", " \"viz_type\": \"bar | line | pie | scatter | heatmap | table_only | null\"\n", "}\n", "\n", "Rules:\n", "- pandas_code must always produce a `result_df` DataFrame\n", "- viz_code must always produce a `fig` Plotly figure (never use import statements)\n", "- If no chart makes sense, set viz_type to table_only and viz_code to empty string\n", "\"\"\"\n", "\n", "client = anthropic.Anthropic()\n", "conversation_history = []\n", "schema_info = get_schema_info(df)\n", "sample_rows = df.head(5).to_string()\n", "\n", "def ask(question):\n", " user_msg = f\"\"\"Dataset schema:\\n{schema_info}\\n\\nSample rows:\\n{sample_rows}\\n\\nQuestion: {question}\"\"\"\n", " conversation_history.append({'role': 'user', 'content': user_msg})\n", "\n", " response = client.messages.create(\n", " model='claude-opus-4-6',\n", " max_tokens=4096,\n", " thinking={'type': 'adaptive'},\n", " system=SYSTEM_PROMPT,\n", " messages=conversation_history,\n", " )\n", "\n", " raw = next((b.text for b in response.content if hasattr(b, 'text')), '')\n", " conversation_history.append({'role': 'assistant', 'content': raw})\n", "\n", " # Parse JSON\n", " text = re.sub(r'^```[\\w]*\\n?', '', raw.strip())\n", " text = re.sub(r'\\n?```$', '', text.strip())\n", " m = re.search(r'\\{[\\s\\S]*\\}', text)\n", " return json.loads(m.group(0) if m else text)\n", "\n", "print('✓ LLM handler ready')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ── Cell 7: Interactive Chat Widget ───────────────────────────────────────────\n\nchat_log = widgets.Output()\ntext_input = widgets.Text(\n placeholder='Ask a question about your data...',\n layout=widgets.Layout(width='75%')\n)\nsend_btn = widgets.Button(description='Ask', button_style='primary',\n layout=widgets.Layout(width='10%'))\nclear_btn = widgets.Button(description='Clear chat', button_style='warning',\n layout=widgets.Layout(width='12%'))\n\nEXAMPLES = [\n 'Compare 2022 vs 2023 highest paid employees by job title',\n 'Top 10 highest paid employees in 2023',\n 'Average salary by department for each year',\n 'Which department has highest salary growth 2021 to 2024?',\n 'Show salary distribution by location in 2023',\n]\nexample_btns = [widgets.Button(description=e[:55], layout=widgets.Layout(width='100%'))\n for e in EXAMPLES]\n\ndef run_query(question):\n with chat_log:\n print('\\n' + '-' * 60)\n display(HTML('You: ' + question))\n display(HTML('Thinking...'))\n\n try:\n result = ask(question)\n except Exception as e:\n with chat_log:\n clear_output(wait=True)\n display(HTML('Error calling LLM: ' + str(e)))\n return\n\n with chat_log:\n clear_output(wait=True)\n display(HTML('You: ' + question))\n display(HTML('Assistant: ' + result.get('explanation', '') + '
'))\n\n # Run pandas code\n pandas_code = result.get('pandas_code', '')\n try:\n result_df = safe_exec_pandas(pandas_code, df)\n if result_df is not None and not result_df.empty:\n display(HTML('Data Table:'))\n display(result_df)\n else:\n display(HTML('No data returned.'))\n return\n except Exception as e:\n display(HTML('Pandas error: ' + str(e)))\n display(HTML('
' + pandas_code + '
'))\n return\n\n # Run viz code\n viz_type = result.get('viz_type', '')\n viz_code = result.get('viz_code', '')\n if viz_type != 'table_only' and viz_code:\n try:\n fig = safe_exec_viz(viz_code, result_df)\n if fig:\n fig.show()\n else:\n raise ValueError('fig is None')\n except Exception:\n fig = fallback_chart(result_df)\n if fig:\n fig.show()\n\ndef on_send(b):\n q = text_input.value.strip()\n if q:\n text_input.value = ''\n run_query(q)\n\ndef on_example(b):\n run_query(b.description)\n\ndef on_clear(b):\n global conversation_history\n conversation_history = []\n with chat_log:\n clear_output()\n\nsend_btn.on_click(on_send)\nclear_btn.on_click(on_clear)\ntext_input.on_submit(on_send)\nfor btn in example_btns:\n btn.on_click(on_example)\n\ndisplay(HTML('

AI Data Chatbot

Type a question or click an example below:

'))\ndisplay(widgets.VBox(example_btns))\ndisplay(HTML('
'))\ndisplay(widgets.HBox([text_input, send_btn, clear_btn]))\ndisplay(chat_log)" } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }