""" 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, )