SAP-ERP-AI-Agent / src /data /generate_text2sql_dataset.py
daisysooyeon's picture
deploy: SAP ERP AI Agent (HF Spaces docker)
50efdc6
Raw History Blame Contribute Delete
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"
# ---------------------------------------------------------------------------
# 데이터 ꡬ쑰
# ---------------------------------------------------------------------------
@dataclass
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,
)