modify_config.py

import yaml
import re

def modify_config_yaml(input_file, output_file=None):
    """
    修改配置文件：
    1. 在 init_codes 中添加 _fk 表
    2. 在 Phase 2 的 "Before ANY filtering" 部分后添加外键说明
    """
    if output_file is None:
        output_file = input_file

    # 读取 YAML
    with open(input_file, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)

    # ============= 修改 init_codes =============
    init_codes = config.get('init_codes', [])

    # 找到导入表的那一行
    table_import_line = None
    table_import_index = None

    for i, line in enumerate(init_codes):
        if 'from database.' in line and 'import' in line:
            table_import_line = line
            table_import_index = i
            break

    if table_import_line:
        # 提取表名
        match = re.search(r'import\s+(.+)$', table_import_line)
        if match:
            tables_str = match.group(1)
            tables = [t.strip() for t in tables_str.split(',')]

            # 为每个表添加 _fk 后缀
            fk_tables = [f"{table}_fk" for table in tables]

            # 合并原表名和外键表名
            all_tables = tables + fk_tables

            # 重新构建导入语句
            import_path = re.search(r'from\s+([\w.]+)\s+import', table_import_line).group(1)
            new_import_line = f"from {import_path} import {', '.join(all_tables)}"

            # 更新 init_codes
            init_codes[table_import_index] = new_import_line
            config['init_codes'] = init_codes

            print(f"✓ 已添加外键表: {', '.join(fk_tables)}")

    # ============= 修改 prompt =============
    prompt = config.get('prompt', '')

    # 构建要插入的外键说明（两段话，简洁明了）
    fk_instruction = """
**Before performing any JOIN operations:**
- First check the corresponding `_fk` DataFrame (e.g., `cards_fk` for `cards` table) to understand foreign key relationships
- The `_fk` DataFrames show: which columns (`from`) reference which tables (`table`) and their target columns (`to`)
"""

    # 找到插入位置：在 "Before ANY filtering:" 这一段的末尾
    # 定位策略：找到 "# Step 3:" 这一行，在它前面插入

    if "# Step 3: df[df['name'] == 'Superman']  # Now safe to filter exactly" in prompt:
        # 在这行之前插入外键说明
        insertion_point = "# Step 3: df[df['name'] == 'Superman']  # Now safe to filter exactly"
        prompt = prompt.replace(
            insertion_point,
            fk_instruction + "\n" + insertion_point
        )
        config['prompt'] = prompt
        print("✓ 已在 Phase 2 末尾添加外键说明")
    else:
        print("⚠ 未找到插入位置标记，请手动检查")

    # 保存修改后的 YAML
    with open(output_file, 'w', encoding='utf-8') as f:
        yaml.dump(config, f, allow_unicode=True, sort_keys=False, width=1000)

    print(f"\n✓ 配置文件已保存到: {output_file}")

    # 显示修改摘要
    print("\n" + "="*60)
    print("修改摘要:")
    print("="*60)
    print(f"1. init_codes 已更新:")
    for line in config['init_codes']:
        if 'import' in line:
            print(f"   {line}")
    print(f"\n2. prompt 中 Phase 2 已添加外键使用说明")

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("使用方法:")
        print("  python modify_yaml.py <config.yaml>                  # 原地修改")
        print("  python modify_yaml.py <config.yaml> <output.yaml>    # 保存到新文件")
        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None

    modify_config_yaml(input_file, output_file)