Build an AI Data Analyst That Thinks Like a Senior Analyst

import pandas as pd
from IPython.display import display
orders = pd.read_csv("online_orders.csv")
print(f"Loaded {len(orders):,} rows and {len(orders.columns)} columns.")
display(orders.head())

Loaded 29 rows and 6 columns.

schema_preview = pd.DataFrame({
    "column": orders.columns,
    "dtype": orders.dtypes.astype(str).values,
    "missing_values": orders.isna().sum().values,
})
display(schema_preview)

import duckdb
con = duckdb.connect()
con.register("online_orders", orders)
preview = con.execute("""
    SELECT
        promotion_id,
        COUNT(*) AS n_orders,
        SUM(units_sold) AS total_units,
        SUM(cost_in_dollars * units_sold) AS total_revenue,
        ROUND(AVG(units_sold), 2) AS avg_units_per_order
    FROM online_orders
    GROUP BY promotion_id
    ORDER BY avg_units_per_order DESC
""").df()
display(preview)

class LLMClient:
    def __init__(self, client, model, provider):
        self.client = client
        self.model = model
        self.provider = provider

    def complete(self, prompt):
        if self.provider == "anthropic":
            response = self.client.messages.create(
                model=self.model,
                max_tokens=1024,
                messages=[{"role": "user", "content": prompt}],
            )

            for block in response.content:
                if block.type == "text":
                    return block.text

            raise ValueError("No text block found in Claude's response.")

        if self.provider == "openai":
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
            )
            return response.choices[0].message.content

        raise ValueError(f"Unsupported provider: {self.provider}")

import json
import re
def parse_json(text):
    text = text.strip()

    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
        text = re.sub(r"\s*```$", "", text)

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    candidates = []
    object_match = re.search(r"\{.*\}", text, re.DOTALL)
    array_match = re.search(r"\[.*\]", text, re.DOTALL)
    if object_match:
        candidates.append(object_match)
    if array_match:
        candidates.append(array_match)
    candidates.sort(key=lambda match: match.start())
    for match in candidates:
        try:
            return json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
    raise ValueError(f"No valid JSON found in model output:\n{text}")

class SeniorAnalyst:
    MIN_SUPPORT = 3  # 仅在组别订单数达到此阈值后才信任结果

    def __init__(self, llm, table_name, dataframe):
        self.llm = llm
        self.table_name = table_name
        self.con = duckdb.connect()
        self.con.register(table_name, dataframe)
        self.schema = self.con.execute(f"DESCRIBE {table_name}").df()

    def understand_business_context(self, question):
        row_count = self.con.execute(
            f"SELECT COUNT(*) FROM {self.table_name}"
        ).fetchone()[0]

        columns = self.schema[
            ["column_name", "column_type"]
        ].to_dict("records")
        prompt = f"""你是一位资深数据分析师。利益相关者提出问题："{question}"
表：{self.table_name}
列：{columns}
行数：{row_count}

用这张表实际能回答的表述方式重述利益相关者的问题。

同时说明表的粒度（一行代表什么），并列出你已经看到的任何限制条件：样本量、时间范围、缺失维度、缺失上下文等。

仅返回 JSON： {{"restated_question": "...", "grain": "...",
"limitations": ["...", "..."]}}"""
        context = parse_json(self.llm.complete(prompt))
        self.context = context
        return context

def generate_hypotheses(self, n=2):
    columns = list(self.schema["column_name"])
    prompt = f"""Business context: {self.context}
Propose {n} specific, testable hypotheses that would help answer the
restated question, using only columns in: {columns}.
Each hypothesis should be something we can test using SQL.
Return JSON only: [{{"hypothesis": "...", "why": "..."}}, ...]""
    hypotheses = parse_json(self.llm.complete(prompt))
    self.hypotheses = hypotheses
    return hypotheses

def plan_sql(self, hypothesis):
    columns = list(self.schema["column_name"])
    prompt = f"""Table: {self.table_name}
Columns: {columns}
Hypothesis to test: {hypothesis['hypothesis']}
Write one DuckDB SQL query that tests this hypothesis.
Use only the available columns, do not invent columns, and if the query
groups rows, include a COUNT(*) column named n_orders so the result can
be checked for sample size before anyone trusts it.
Return JSON only: {{"sql": "...", "purpose": "..."}}"""
    plan = parse_json(self.llm.complete(prompt))
    return plan

生成的SQL:
    WITH promo_sums AS (
        SELECT promotion_id, SUM(units_sold) AS total_units, COUNT(*) AS n_orders
        FROM online_orders
        GROUP BY promotion_id
    ),
    ranked AS (
        SELECT promotion_id, total_units, n_orders,
               RANK() OVER (ORDER BY total_units DESC) AS rnk
        FROM promo_sums
    )
    SELECT
        r1.promotion_id AS top_promotion_id,
        r1.total_units AS top_total_units,
        r1.n_orders AS top_n_orders,
        r2.promotion_id AS second_promotion_id,
        r2.total_units AS second_total_units,
        r2.n_orders AS second_n_orders,
        (r1.total_units - r2.total_units) * 1.0 / r2.total_units AS pct_difference
    FROM ranked r1
    JOIN ranked r2 ON r2.rnk = 2
    WHERE r1.rnk = 1

   'purpose': 'Identify the promotion_id with the highest total units
     sold and compare it to the second-highest to test whether it exceeds
     it by at least 20%, including order counts to assess statistical
     support.'

def validate(self, sql_plan):
    result = self.con.execute(sql_plan["sql"]).df()
    if "n_orders" in result.columns:
        result["low_confidence"] = result["n_orders"] < self.MIN_SUPPORT
    else:
        result["low_confidence"] = False
    return result

def summarize(self, hypothesis, validated_result):
        flagged = validated_result[validated_result["low_confidence"]]
        prompt = f"""假设：{hypothesis['hypothesis']}
    查询结果:
    {validated_result.to_string(index=False)}
    被标记为低置信度的行支持的订单数量少于{self.MIN_SUPPORT}个
    不应作为结论依据。
    低置信度行：{flagged.to_dict('records')}
    根据结果撰写简洁的3到4句高管摘要
    结论仅基于展示的数据，明确避免
    将低置信度行作为主要结论，并且不要
    虚构数据不支持的解释。"""
        return self.llm.complete(prompt)

def recommend(self, summary):
        prompt = f"""高管摘要：{summary}
    根据摘要支持的内容提出2到3个具体业务建议
    建议必须基于证据，不得依赖低置信度数据
    或虚构事实，若证据薄弱，应建议进一步分析
    而不是假装结论已明确。"""
        return self.llm.complete(prompt)

def run(self, question):
        context = self.understand_business_context(question)
        hypotheses = self.generate_hypotheses()
        top_hypothesis = hypotheses[0]
        plan = self.plan_sql(top_hypothesis)
        validated = self.validate(plan)
        summary = self.summarize(top_hypothesis, validated)
        recommendation = self.recommend(summary)
        return {
            "context": context,
            "hypotheses": hypotheses,
            "sql_plan": plan,
            "validated_result": validated,
            "summary": summary,
            "recommendation": recommendation,
        }

PROVIDER = "anthropic"
    API_KEY = "YOUR_API_KEY_HERE"
    ANTHROPIC_MODEL = "claude-sonnet-5"
    OPENAI_MODEL = "gpt-4o"
    if API_KEY == "YOUR_API_KEY_HERE":
        raise ValueError(
            "在运行LLM部分之前，请将真实的API密钥粘贴到API_KEY中。"
        )
    if PROVIDER.lower() == "anthropic":
        from anthropic import Anthropic
        client = Anthropic(api_key=API_KEY)
        llm = LLMClient(client=client, model=ANTHROPIC_MODEL, provider="anthropic")
    elif PROVIDER.lower() == "openai":
        from openai import OpenAI
        client = OpenAI(api_key=API_KEY)
        llm = LLMClient(client=client, model=OPENAI_MODEL, provider="openai")
    else:
        raise ValueError("PROVIDER必须是'openai'或'anthropic'之一。")
    analyst = SeniorAnalyst(llm, "online_orders", orders)
    result = analyst.run("我们应该增加哪个促销活动？")
    print(result["summary"])
    print(result["recommendation"])