Chain-of-Table: рассуждения через трансформацию таблиц

COT_TABLE_SYSTEM = """Ты аналитик данных. Отвечай на вопросы по таблицам,
используя метод Chain-of-Table.

Доступные операции:
1. f_select_column(columns: list) — оставить только указанные столбцы
2. f_select_row(condition: str) — отфильтровать строки по условию
3. f_group_by(column: str, agg: str) — сгруппировать (agg: count/sum/avg/min/max)
4. f_sort_by(column: str, order: str) — отсортировать (asc/desc)
5. f_add_col(name: str, expression: str) — добавить вычисляемый столбец

Правила:
- На каждом шаге выполняй РОВНО ОДНУ операцию
- После каждой операции покажи результат в виде markdown-таблицы
- Таблица каждого шага — вход для следующего шага
- В конце сформулируй ответ на основе финальной таблицы
- Не пропускай шаги, не объединяй операции"""

def table_to_markdown(headers: list[str], rows: list[list]) -> str:
    """Конвертирует данные в Markdown-таблицу."""
    header_line = "| " + " | ".join(headers) + " |"
    separator = "| " + " | ".join("---" for _ in headers) + " |"
    row_lines = [
        "| " + " | ".join(str(cell) for cell in row) + " |"
        for row in rows
    ]
    return "\n".join([header_line, separator] + row_lines)


def build_cot_table_prompt(table_md: str, question: str) -> str:
    """Собирает промпт для Chain-of-Table."""
    return f"""Таблица:
{table_md}

Вопрос: {question}

Используй метод Chain-of-Table. Выполняй операции по одной,
показывая промежуточную таблицу после каждого шага."""

import pandas as pd
from openai import OpenAI

client = OpenAI()

PLANNER_PROMPT = """Ты планируешь цепочку операций над таблицей.
Верни JSON-массив операций. Каждая операция — объект:

{"op": "select_row", "condition": "column == 'value'"}
{"op": "select_column", "columns": ["col1", "col2"]}
{"op": "sort_by", "column": "col1", "order": "desc"}
{"op": "group_by", "column": "col1", "agg": {"col2": "sum"}}
{"op": "add_col", "name": "new_col", "expression": "col1 - col2"}

Верни ТОЛЬКО JSON-массив, без пояснений."""


def plan_operations(table_md: str, question: str) -> list[dict]:
    """LLM планирует цепочку операций."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": PLANNER_PROMPT},
            {"role": "user", "content": f"Таблица:\n{table_md}\n\nВопрос: {question}"}
        ],
        temperature=0,
        response_format={"type": "json_object"}
    )
    import json
    result = json.loads(response.choices[0].message.content)
    return result if isinstance(result, list) else result.get("operations", [])


def execute_operations(df: pd.DataFrame, operations: list[dict]) -> pd.DataFrame:
    """Выполняет спланированные операции в pandas."""
    for op in operations:
        match op["op"]:
            case "select_row":
                df = df.query(op["condition"])
            case "select_column":
                df = df[op["columns"]]
            case "sort_by":
                ascending = op.get("order", "asc") == "asc"
                df = df.sort_values(op["column"], ascending=ascending)
            case "group_by":
                df = df.groupby(op["column"]).agg(op["agg"]).reset_index()
            case "add_col":
                df[op["name"]] = df.eval(op["expression"])
    return df


# Использование
df = pd.read_csv("sales.csv")
table_md = df.to_markdown(index=False)
question = "Какой регион принёс максимальную выручку в Q4?"

operations = plan_operations(table_md, question)
# [{"op": "select_column", "columns": ["Регион", "Q4"]},
#  {"op": "group_by", "column": "Регион", "agg": {"Q4": "sum"}},
#  {"op": "sort_by", "column": "Q4", "order": "desc"}]

result = execute_operations(df, operations)
print(result.head(1))

SQL_PLANNER_PROMPT = """Ты аналитик. На основе вопроса и таблицы
составь SQL-запрос. Таблица называется `data`.

Используй подход Chain-of-Table: разбей запрос на шаги
с помощью CTE (WITH ... AS).

Верни ТОЛЬКО SQL, без пояснений."""


def generate_sql(table_md: str, question: str) -> str:
    """LLM генерирует SQL через Chain-of-Table."""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SQL_PLANNER_PROMPT},
            {"role": "user", "content": f"Таблица:\n{table_md}\n\nВопрос: {question}"}
        ],
        temperature=0
    )
    sql = response.choices[0].message.content
    # Очистка от markdown-блока
    if "```sql" in sql:
        sql = sql.split("```sql")[1].split("```")[0]
    return sql.strip()


# Результат — пошаговый SQL через CTE:
# WITH step1 AS (
#     SELECT * FROM data WHERE region = 'Москва'
# ),
# step2 AS (
#     SELECT manager, q4 - q1 AS growth FROM step1
# )
# SELECT manager, growth FROM step2 ORDER BY growth DESC LIMIT 1;