Spaces:
Running
Running
Download src/data/generate_text2sql_dataset.py from daisysooyeon/SAP-ERP-AI-Agent: direct link, hf CLI and curl.
- Browser
- Download file 28.1 kB
-
https://huggingface.co/spaces/daisysooyeon/SAP-ERP-AI-Agent/resolve/main/src/data/generate_text2sql_dataset.py
- Command line
-
hf download hf://spaces/daisysooyeon/SAP-ERP-AI-Agent/src/data/generate_text2sql_dataset.py
-
curl -L -o generate_text2sql_dataset.py https://huggingface.co/spaces/daisysooyeon/SAP-ERP-AI-Agent/resolve/main/src/data/generate_text2sql_dataset.py
28.1 kB
| """ | |
| src/data/generate_text2sql_dataset.py | |
| ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| Text-to-SQL νκ°μ© ν μ€νΈ μΌμ΄μ€ μμ±κΈ° | |
| νλ¦: | |
| 1. data/sap_erp.db (SQLite)μμ λ€μν μλλ¦¬μ€ μΌμ΄μ€λ₯Ό μ§μ μνλ§ | |
| - WBSTA = A / B / C μΌμ΄μ€ | |
| - MARD JOIN κ°λ₯ν μΌμ΄μ€ (available_stock μμ) | |
| - NULL μΌμ΄μ€ (no-join row) | |
| - μ‘΄μ¬νμ§ μλ order_id (negative case) | |
| 2. μνλ§ν μ€μ DB κ°μΌλ‘ expected_values μλ κ³μ° | |
| 3. κΈ°μ‘΄ text2sql_test_cases.json ν¬λ§·μΌλ‘ μ μ₯ | |
| μ€ν: | |
| python -m src.data.generate_text2sql_dataset | |
| python -m src.data.generate_text2sql_dataset --n-samples 30 --output data/eval/text2sql_test_cases_gen.json | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import logging | |
| import random | |
| import sqlite3 | |
| import sys | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| from typing import Any | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| sys.stdout.reconfigure(encoding="utf-8", errors="replace") | |
| _ROOT_FOR_LOG = Path(__file__).resolve().parent.parent.parent | |
| _LOG_DIR = _ROOT_FOR_LOG / "logs" | |
| _LOG_DIR.mkdir(parents=True, exist_ok=True) | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| handlers=[ | |
| logging.FileHandler(_LOG_DIR / "data_generation.log", encoding="utf-8"), | |
| logging.StreamHandler(sys.stdout), | |
| ], | |
| force=True, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| _ROOT = Path(__file__).resolve().parent.parent.parent | |
| DEFAULT_DB = str(_ROOT / "data" / "sap_erp.db") | |
| DEFAULT_OUTPUT = str(_ROOT / "data" / "eval" / "text2sql_test_cases_gen.json") | |
| REQUIRED_COLUMNS = ["material_name", "quantity", "delivery_status", "delivery_date", "available_stock"] | |
| # negative caseμ μ¬μ©ν κ°μ§ order_id | |
| NEGATIVE_ORDER_ID = "9999999999" | |
| # --------------------------------------------------------------------------- | |
| # λ°μ΄ν° ꡬ쑰 | |
| # --------------------------------------------------------------------------- | |
| class TestCase: | |
| id: str | |
| description: str | |
| order_id: str | |
| item_no: int | |
| expected_columns: list[str] | |
| expected_values: dict[str, Any] | None | |
| golden_sql: str # μ λ΅(λ νΌλ°μ€) SQL β νλ‘λμ λꡬμ _hardcoded_query κ²°κ³Ό | |
| notes: str | |
| # --------------------------------------------------------------------------- | |
| # DB μνλ§ μΏΌλ¦¬ λͺ¨μ | |
| # --------------------------------------------------------------------------- | |
| # κ° μλ리μ€λ₯Ό (description_template, notes, sql) ννλ‘ μ μ | |
| # SQLμ (VBELN, POSNR, quantity, delivery_status, delivery_date, available_stock) λ°ν | |
| _SAMPLING_QUERIES: list[tuple[str, str, str]] = [ | |
| # ββ μΌμ΄μ€ 1: WBSTA = A (λ―Έμ²λ¦¬) βββββββββββββββββββββββββββββββββββββββββ | |
| ( | |
| "Order with WBSTA=A (goods movement not yet processed)", | |
| "WBSTA=A means goods movement not yet started", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbup.WBSTA = 'A' | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 2: WBSTA = B (λΆλΆ μ²λ¦¬) βββββββββββββββββββββββββββββββββββββββ | |
| ( | |
| "Order with WBSTA=B (partially processed)", | |
| "WBSTA=B means goods movement partially processed", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbup.WBSTA = 'B' | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 3: WBSTA = C (μμ μ²λ¦¬) βββββββββββββββββββββββββββββββββββββββ | |
| ( | |
| "Order with WBSTA=C (fully processed)", | |
| "WBSTA=C means goods movement fully completed", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbup.WBSTA = 'C' | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 4: MARD JOIN κ°λ₯ (available_stock μμ) βββββββββββββββββββββ | |
| ( | |
| "Order with available stock via MARD JOIN", | |
| "MARD.LABST (unrestricted stock) should be present", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| mard.LABST AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| LEFT JOIN MARD mard ON vbap.MATNR = mard.MATNR | |
| WHERE mard.LABST IS NOT NULL AND mard.LABST > 0 | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 5: delivery_dateλ§ μκ³ stock μμ ββββββββββββββββββββββββββββ | |
| ( | |
| "Order with delivery date only (no stock join)", | |
| "MARD join yields NULL β delivery_date must be from VBEP", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| LEFT JOIN MARD mard ON vbap.MATNR = mard.MATNR | |
| WHERE vbep.EDATU IS NOT NULL AND mard.LABST IS NULL | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 6: λΉκΈ°λ³Έ POSNR (POSNR != 10) ββββββββββββββββββββββββββββββββ | |
| ( | |
| "Order item with non-default POSNR (item number != 10)", | |
| "Tests that POSNR is handled correctly for non-standard line items", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| mard.LABST AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| LEFT JOIN MARD mard ON vbap.MATNR = mard.MATNR | |
| WHERE vbap.POSNR != 10 | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 7: μΌλ° κΈ°λ³Έ μΌμ΄μ€ (νΉμ 쑰건 μμ) ββββββββββββββββββββββββββ | |
| ( | |
| "Standard order β baseline with all joins present", | |
| "General case: material name, delivery status, and delivery date all populated", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE makt.MAKTX IS NOT NULL | |
| AND vbup.WBSTA IS NOT NULL | |
| AND vbep.EDATU IS NOT NULL | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 8: λ©κΈ° μ΄κ³Ό λ―Έμλ£ μ£Όλ¬Έ (overdue) ββββββββββββββββββββββββββββ | |
| ( | |
| "Overdue order β delivery date in the past, goods movement incomplete", | |
| "EDATU < today and WBSTA in (A, B) β tests time-sensitive filtering", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbep.EDATU < strftime('%Y%m%d', 'now') | |
| AND vbup.WBSTA IN ('A', 'B') | |
| AND vbep.EDATU IS NOT NULL | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 9: λ―Έλ λ©κΈ° μμ μ£Όλ¬Έ ββββββββββββββββββββββββββββββββββββββββ | |
| ( | |
| "Upcoming delivery β delivery date in the future, not yet processed", | |
| "EDATU >= today and WBSTA = A β tests future-date query handling", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbep.EDATU >= strftime('%Y%m%d', 'now') | |
| AND vbup.WBSTA = 'A' | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 10: μ¬κ³ κ° μ£Όλ¬Έλ μ΄μμΈ κ²½μ° (stock surplus) ββββββββββββββββββ | |
| ( | |
| "Order where available stock exceeds ordered quantity", | |
| "MARD.LABST > VBAP.KWMENG β tests stock-vs-quantity comparison logic", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| mard.LABST AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| LEFT JOIN MARD mard ON vbap.MATNR = mard.MATNR | |
| WHERE mard.LABST > vbap.KWMENG | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 11: λ€μ€ λΌμΈ μμ΄ν μ£Όλ¬Έ βββββββββββββββββββββββββββββββββββββ | |
| ( | |
| "Multi-line order β sales order containing more than one item", | |
| "Tests item disambiguation when VBELN has multiple POSNRs", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| NULL AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| WHERE vbap.VBELN IN ( | |
| SELECT VBELN FROM VBAP GROUP BY VBELN HAVING COUNT(DISTINCT POSNR) > 1 | |
| ) | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| # ββ μΌμ΄μ€ 12: μ¬κ³ μμ (LABST = 0 λλ JOIN λΆκ°) ββββββββββββββββββββββ | |
| ( | |
| "Order with zero or no available stock", | |
| "MARD.LABST = 0 or NULL β tests zero-stock edge case handling", | |
| """ | |
| SELECT | |
| vbap.VBELN AS order_id, | |
| vbap.POSNR AS item_no, | |
| makt.MAKTX AS material_name, | |
| vbap.KWMENG AS quantity, | |
| vbup.WBSTA AS delivery_status, | |
| vbep.EDATU AS delivery_date, | |
| COALESCE(mard.LABST, 0.0) AS available_stock | |
| FROM VBAP vbap | |
| LEFT JOIN MAKT makt ON vbap.MATNR = makt.MATNR AND makt.SPRAS = 'E' | |
| LEFT JOIN VBUP vbup ON CAST(vbap.VBELN AS INTEGER) = CAST(vbup.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbup.POSNR AS INTEGER) | |
| LEFT JOIN VBEP vbep ON CAST(vbap.VBELN AS INTEGER) = CAST(vbep.VBELN AS INTEGER) AND CAST(vbap.POSNR AS INTEGER) = CAST(vbep.POSNR AS INTEGER) | |
| LEFT JOIN MARD mard ON vbap.MATNR = mard.MATNR | |
| WHERE (mard.LABST IS NULL OR mard.LABST = 0) | |
| AND vbap.KWMENG > 0 | |
| AND makt.MAKTX IS NOT NULL | |
| ORDER BY RANDOM() | |
| LIMIT {limit} | |
| """, | |
| ), | |
| ] | |
| # --------------------------------------------------------------------------- | |
| # DB μ°κ²° ν¬νΌ | |
| # --------------------------------------------------------------------------- | |
| def _query(db_path: str, sql: str, conn: sqlite3.Connection | None = None) -> list[dict]: | |
| """SQL μ€ν β dict 리μ€νΈ. connμ΄ μ£Όμ΄μ§λ©΄ κ·Έ μ°κ²°μ μ¬μ¬μ©(λ«μ§ μμ), | |
| μμΌλ©΄ μΌνμ± μ°κ²°μ μ΄κ³ λ«λλ€(νμνΈν). λ°°μΉ μνλ§μμ connμ μ¬μ¬μ©ν΄ | |
| μ°κ²° μ€ν λΉμ© + SQLite νμ΄μ§ μΊμ μ½λμ€ννΈλ₯Ό νΌνλ€.""" | |
| own = conn is None | |
| if own: | |
| conn = sqlite3.connect(db_path) | |
| conn.row_factory = sqlite3.Row | |
| try: | |
| rows = conn.execute(sql).fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| if own: | |
| conn.close() | |
| def _edatu_to_iso(edatu: str | None) -> str | None: | |
| """YYYYMMDD β YYYY-MM-DD λ³ν (μ΄λ―Έ ISOλ©΄ κ·Έλλ‘)""" | |
| if edatu and len(edatu) == 8: | |
| return f"{edatu[:4]}-{edatu[4:6]}-{edatu[6:]}" | |
| return edatu | |
| # μνλ§Β·μΊλ Έλ컬 μΏΌλ¦¬κ° μ°λ CAST(...) μ‘°μΈ ν€μ λν ννμ μΈλ±μ€. | |
| # μ΄κ² μμΌλ©΄ CAST μλ³ μ‘°μΈμ΄ μΈλ±μ€λ₯Ό λͺ» ν VBAPΓVBUP μ€μ²©λ£¨ν(μμμ΅ λΉκ΅)λ‘ νλ°νλ€ | |
| # (μλλ¦¬μ€ μΏΌλ¦¬ 1ν ~254μ΄ κ΄μΈ‘). μΈλ±μ€κ° μμΌλ©΄ ~0.2μ΄λ‘ λ¨μ΄μ§λ€. λ©±λ± μμ±. | |
| _PERF_INDEXES = [ | |
| "CREATE INDEX IF NOT EXISTS ix_vbup_norm ON VBUP(CAST(VBELN AS INTEGER), CAST(POSNR AS INTEGER))", | |
| "CREATE INDEX IF NOT EXISTS ix_vbep_norm ON VBEP(CAST(VBELN AS INTEGER), CAST(POSNR AS INTEGER))", | |
| "CREATE INDEX IF NOT EXISTS ix_mard_mw ON MARD(MATNR, WERKS)", | |
| "CREATE INDEX IF NOT EXISTS ix_makt_m ON MAKT(MATNR)", | |
| ] | |
| def _ensure_perf_indexes(conn: sqlite3.Connection) -> None: | |
| """μνλ§ μ±λ₯μ© μΈλ±μ€λ₯Ό λ©±λ± μμ±. μ€ν¨ν΄λ(μ: μ½κΈ°μ μ©/λ½) νλ¦μ λμ§ μλλ€.""" | |
| try: | |
| for ddl in _PERF_INDEXES: | |
| conn.execute(ddl) | |
| conn.commit() | |
| except sqlite3.Error as e: | |
| logger.warning("Perf index creation skipped: %s", e) | |
| def _canonical_query_and_row( | |
| db_path: str, order_id: str, item_no: str | int, | |
| conn: sqlite3.Connection | None = None, | |
| ) -> tuple[str, dict | None]: | |
| """expected_values + golden_sqlμ λ¨μΌ μΆμ²(source of truth). | |
| νλ‘λμ ν둬ννΈκ° LLMμκ² μ§μνλ μλ―Έμ 1:1λ‘ μΌμΉνλ λ νΌλ°μ€ 쿼리 | |
| (src.tools.text2sql._hardcoded_query)λ₯Ό κ·Έλλ‘ μ€νν΄ μ λ΅ νμ λ§λ λ€. | |
| μλλ¦¬μ€ SQLμ (order_id, item_no) λ€μμ± μ νμλ§ μ°κ³ , κ°Β·μ λ΅SQLμ μ¬κΈ°μ μ°μΆνλ€. | |
| μ΄λ κ² ν΄μΌ 골λ (κ°+SQL)κ³Ό ν둬ννΈκ° κ°μ κ²°μ μ μ μλ₯Ό 곡μ νλ€. | |
| connμ΄ μ£Όμ΄μ§λ©΄ μ¬μ¬μ©νλ€(λ°°μΉ μνλ§ μ±λ₯). | |
| Returns: (golden_sql, μ λ΅ ν or None) | |
| """ | |
| from src.tools.text2sql import _hardcoded_query | |
| sql = _hardcoded_query(str(order_id), str(item_no)) | |
| rows = _query(db_path, sql, conn=conn) | |
| return sql, (rows[0] if rows else None) | |
| # --------------------------------------------------------------------------- | |
| # ν μ€νΈ μΌμ΄μ€ λΉλ | |
| # --------------------------------------------------------------------------- | |
| def _build_expected_values(row: dict) -> dict | None: | |
| """DB νμμ expected_values dictλ₯Ό ꡬμ±ν©λλ€.""" | |
| ev: dict = {} | |
| if row.get("material_name"): | |
| ev["material_name"] = row["material_name"] | |
| if row.get("quantity") is not None: | |
| ev["quantity"] = float(row["quantity"]) | |
| if row.get("delivery_status"): | |
| ev["delivery_status"] = row["delivery_status"] | |
| if row.get("delivery_date"): | |
| ev["delivery_date"] = _edatu_to_iso(row["delivery_date"]) | |
| if row.get("available_stock") is not None: | |
| ev["available_stock"] = float(row["available_stock"]) | |
| return ev if ev else None | |
| def _sample_cases( | |
| db_path: str, | |
| n_per_scenario: int, | |
| seed: int = 42, | |
| ) -> list[TestCase]: | |
| """κ° μνλ§ μλ리μ€μμ n_per_scenarioκ°μ© μΌμ΄μ€λ₯Ό μμ§ν©λλ€.""" | |
| random.seed(seed) | |
| cases: list[TestCase] = [] | |
| # μ°κ²° 1κ° μ¬μ¬μ©(264+ μ€ν λ°©μ§ + νμ΄μ§ μΊμ μ μ§) + CAST μ‘°μΈ μΈλ±μ€ λ©±λ± μμ±. | |
| # μΈλ±μ€κ° μμΌλ©΄ μλλ¦¬μ€ μΏΌλ¦¬μ CAST μλ³ μ‘°μΈμ΄ μ€μ²©λ£¨νλ‘ νλ°νλ€(1ν ~254μ΄). | |
| conn = sqlite3.connect(db_path) | |
| conn.row_factory = sqlite3.Row | |
| _ensure_perf_indexes(conn) | |
| try: | |
| for desc_tmpl, notes, sql_tmpl in _SAMPLING_QUERIES: | |
| # λλνκ² μΏΌλ¦¬ ν λλ€ μ ν | |
| sql = sql_tmpl.format(limit=n_per_scenario * 5) | |
| try: | |
| rows = _query(db_path, sql, conn=conn) | |
| except Exception as e: | |
| logger.warning("Sampling failed (%s): %s", desc_tmpl[:40], e) | |
| continue | |
| sampled = random.sample(rows, min(n_per_scenario, len(rows))) | |
| for row in sampled: | |
| order_id = str(row.get("order_id", "")).strip() | |
| item_no = int(row.get("item_no", 10)) | |
| # expected_valuesΒ·golden_sqlμ μλλ¦¬μ€ νμ΄ μλλΌ λ νΌλ°μ€(μΊλ Έλ컬) 쿼리μμ μ°μΆ | |
| golden_sql, canon = _canonical_query_and_row(db_path, order_id, item_no, conn=conn) | |
| cases.append(TestCase( | |
| id="", # λμ€μ λΆμ¬ | |
| description=desc_tmpl, | |
| order_id=order_id, | |
| item_no=item_no, | |
| expected_columns=REQUIRED_COLUMNS, | |
| expected_values=_build_expected_values(canon) if canon else None, | |
| golden_sql=golden_sql, | |
| notes=notes, | |
| )) | |
| finally: | |
| conn.close() | |
| return cases | |
| def _add_negative_case(cases: list[TestCase]) -> None: | |
| """μ‘΄μ¬νμ§ μλ μ£Όλ¬Έλ²νΈ μΌμ΄μ€(expected_values=None)λ₯Ό μΆκ°ν©λλ€. | |
| μ λ΅ SQLμ μ ν¨νμ§λ§ κ²°κ³Ό νμ΄ μλ μΌμ΄μ€λ₯Ό κ²μ¦νλ€.""" | |
| from src.tools.text2sql import _hardcoded_query | |
| cases.append(TestCase( | |
| id="", | |
| description="Non-existent order β should return no row (None result)", | |
| order_id=NEGATIVE_ORDER_ID, | |
| item_no=10, | |
| expected_columns=REQUIRED_COLUMNS, | |
| expected_values=None, | |
| golden_sql=_hardcoded_query(NEGATIVE_ORDER_ID, "10"), | |
| notes="SQL must execute without error even if no rows returned", | |
| )) | |
| # --------------------------------------------------------------------------- | |
| # λ©μΈ μμ± ν¨μ | |
| # --------------------------------------------------------------------------- | |
| def generate_text2sql_dataset( | |
| n_samples: int = 20, | |
| output_path: str = DEFAULT_OUTPUT, | |
| db_path: str = DEFAULT_DB, | |
| seed: int = 42, | |
| append: bool = False, | |
| ) -> list[dict]: | |
| """ | |
| SQLite DBλ₯Ό μνλ§νμ¬ text2sql ν μ€νΈ μΌμ΄μ€λ₯Ό μμ±ν©λλ€. | |
| Parameters | |
| ---------- | |
| n_samples : μλ리μ€λΉ μν μ (μ€μ μ΄ μΌμ΄μ€ μ β n_samples Γ μλλ¦¬μ€ μ) | |
| output_path : μΆλ ₯ JSON κ²½λ‘ | |
| db_path : sap_erp.db κ²½λ‘ | |
| seed : λλ€ μλ (μ¬νμ±) | |
| append : κΈ°μ‘΄ νμΌμ μΆκ° μ¬λΆ | |
| """ | |
| if not Path(db_path).exists(): | |
| raise FileNotFoundError( | |
| f"DB file not found: {db_path}\n" | |
| "Check the paths.erp_db configuration in configs.yaml." | |
| ) | |
| logger.info("Starting DB sampling: %s", db_path) | |
| cases = _sample_cases(db_path, n_samples, seed) | |
| _add_negative_case(cases) | |
| # μ€λ³΅ (order_id, item_no) μ κ±° | |
| seen: set[tuple] = set() | |
| deduped: list[TestCase] = [] | |
| for c in cases: | |
| key = (c.order_id, c.item_no) | |
| if key not in seen: | |
| seen.add(key) | |
| deduped.append(c) | |
| cases = deduped | |
| # id λΆμ¬ | |
| start_idx = 1 | |
| if append: | |
| existing_path = Path(output_path) | |
| if existing_path.exists(): | |
| try: | |
| existing = json.loads(existing_path.read_text(encoding="utf-8")) | |
| start_idx = len(existing) + 1 | |
| except Exception: | |
| pass | |
| for idx, c in enumerate(cases, start=start_idx): | |
| c.id = f"tc_{idx:03d}" | |
| # dict λ³ν | |
| records = [asdict(c) for c in cases] | |
| # append μ²λ¦¬ | |
| if append: | |
| existing_path = Path(output_path) | |
| if existing_path.exists(): | |
| try: | |
| existing = json.loads(existing_path.read_text(encoding="utf-8")) | |
| records = existing + records | |
| except Exception as e: | |
| logger.warning("Failed to read existing file: %s", e) | |
| out = Path(output_path) | |
| out.parent.mkdir(parents=True, exist_ok=True) | |
| out.write_text(json.dumps(records, indent=2, ensure_ascii=False), encoding="utf-8") | |
| # Summary | |
| print("\n" + "=" * 60) | |
| print(" Text-to-SQL Dataset Generation Results") | |
| print("=" * 60) | |
| print(f" DB Path : {db_path}") | |
| print(f" Scenarios : {len(_SAMPLING_QUERIES) + 1}") | |
| print(f" Generated Cases : {len(cases)}") | |
| print(f" Output File : {output_path}") | |
| print("=" * 60) | |
| return records | |
| # --------------------------------------------------------------------------- | |
| # CLI | |
| # --------------------------------------------------------------------------- | |
| def _parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Text-to-SQL νκ°μ© ν μ€νΈ μΌμ΄μ€ μμ±") | |
| parser.add_argument( | |
| "--n-samples", type=int, default=10, | |
| help="μλ리μ€λΉ μν μ (κΈ°λ³Έ: 10, μ΄ μΌμ΄μ€ β 10 Γ μλλ¦¬μ€ μ)", | |
| ) | |
| parser.add_argument( | |
| "--output", default=DEFAULT_OUTPUT, | |
| help=f"μΆλ ₯ JSON κ²½λ‘ (κΈ°λ³Έ: {DEFAULT_OUTPUT})", | |
| ) | |
| parser.add_argument( | |
| "--db", default=DEFAULT_DB, | |
| help=f"SQLite DB κ²½λ‘ (κΈ°λ³Έ: {DEFAULT_DB})", | |
| ) | |
| parser.add_argument( | |
| "--seed", type=int, default=42, | |
| help="λλ€ μλ (κΈ°λ³Έ: 42)", | |
| ) | |
| parser.add_argument( | |
| "--append", action="store_true", | |
| help="κΈ°μ‘΄ μΆλ ₯ νμΌμ κ²°κ³Όλ₯Ό μΆκ°", | |
| ) | |
| return parser.parse_args() | |
| if __name__ == "__main__": | |
| args = _parse_args() | |
| generate_text2sql_dataset( | |
| n_samples=args.n_samples, | |
| output_path=args.output, | |
| db_path=args.db, | |
| seed=args.seed, | |
| append=args.append, | |
| ) | |