238 lines
8.9 KiB
Python
238 lines
8.9 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
自动集成few-shot功能到 Text2SQL 系统
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
python scripts/patch_fewshot.py --enable --top-k 3 --min-rating 7
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import re
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
def patch_orchestrator(
|
|||
|
|
enable: bool = True,
|
|||
|
|
top_k: int = 3,
|
|||
|
|
min_rating: int = 7,
|
|||
|
|
data_path: str = "data/experiences/all_samples.jsonl"
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
修改 agents/orchestrator.py 添加few-shot支持
|
|||
|
|
"""
|
|||
|
|
orch_path = Path(__file__).parent.parent / "agents" / "orchestrator.py"
|
|||
|
|
|
|||
|
|
if not orch_path.exists():
|
|||
|
|
print(f"❌ 文件不存在: {orch_path}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
content = orch_path.read_text(encoding='utf-8')
|
|||
|
|
|
|||
|
|
# 1. 在文件头部添加import
|
|||
|
|
import_line = "from utils.fewshot_selector import FewShotSelector"
|
|||
|
|
if import_line not in content:
|
|||
|
|
# 找到合适的插入点（在其他import之后）
|
|||
|
|
last_import = content.rfind("import ")
|
|||
|
|
if last_import != -1:
|
|||
|
|
insert_pos = content.find("\n", last_import) + 1
|
|||
|
|
content = content[:insert_pos] + import_line + "\n" + content[insert_pos:]
|
|||
|
|
print("✅ 添加 import: from utils.fewshot_selector import FewShotSelector")
|
|||
|
|
|
|||
|
|
# 2. 修改 __init__ 方法
|
|||
|
|
init_pattern = r'def __init__\([\s\S]*?\):'
|
|||
|
|
init_match = re.search(init_pattern, content)
|
|||
|
|
|
|||
|
|
if init_match:
|
|||
|
|
init_text = content[init_match.start():init_match.end()]
|
|||
|
|
|
|||
|
|
# 检查是否已添加few-shot参数
|
|||
|
|
if "fewshot_enabled" not in init_text:
|
|||
|
|
# 在参数列表末尾添加（在max_retry之后）
|
|||
|
|
new_params = """,
|
|||
|
|
# Few-shot 配置
|
|||
|
|
fewshot_enabled: bool = True,
|
|||
|
|
fewshot_samples_path: Optional[str] = None,
|
|||
|
|
fewshot_top_k: int = 3,
|
|||
|
|
fewshot_min_rating: int = 7,"""
|
|||
|
|
|
|||
|
|
# 找到max_retry参数后插入
|
|||
|
|
max_retry_pos = init_text.find("max_retry:")
|
|||
|
|
if max_retry_pos != -1:
|
|||
|
|
comma_pos = init_text.find(",", max_retry_pos) + 1
|
|||
|
|
new_init = (
|
|||
|
|
init_text[:comma_pos] +
|
|||
|
|
new_params +
|
|||
|
|
init_text[comma_pos:]
|
|||
|
|
)
|
|||
|
|
content = content[:init_match.start()] + new_init + content[init_match.end():]
|
|||
|
|
print("✅ 修改 __init__: 添加few-shot参数")
|
|||
|
|
|
|||
|
|
# 在__init__方法体内添加few-shot初始化代码
|
|||
|
|
# 找到"self.max_retry = max_retry"这一行
|
|||
|
|
max_retry_assign = "self.max_retry = max_retry"
|
|||
|
|
if max_retry_assign in content[init_match.start():init_match.end()]:
|
|||
|
|
init_body_start = init_match.end()
|
|||
|
|
init_body = content[init_body_start:init_body_start + 2000] # 读取前2000字符
|
|||
|
|
|
|||
|
|
if "self.fewshot_selector" not in init_body:
|
|||
|
|
fewshot_init = '''
|
|||
|
|
|
|||
|
|
# Few-shot 初始化
|
|||
|
|
self.fewshot_enabled = fewshot_enabled
|
|||
|
|
self.fewshot_top_k = fewshot_top_k
|
|||
|
|
self.fewshot_min_rating = fewshot_min_rating
|
|||
|
|
self.fewshot_selector = None
|
|||
|
|
|
|||
|
|
if self.fewshot_enabled:
|
|||
|
|
try:
|
|||
|
|
from utils.fewshot_selector import FewShotSelector
|
|||
|
|
path = fewshot_samples_path or os.getenv(
|
|||
|
|
"FEWSHOT_DATA_PATH",
|
|||
|
|
"./data/experiences/all_samples.jsonl"
|
|||
|
|
)
|
|||
|
|
self.fewshot_selector = FewShotSelector(path)
|
|||
|
|
logger.info(
|
|||
|
|
f"Few-shot已启用: top_k={fewshot_top_k}, "
|
|||
|
|
f"min_rating={fewshot_min_rating}"
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Few-shot加载失败: {e}，将使用标准生成")
|
|||
|
|
self.fewshot_enabled = False
|
|||
|
|
'''
|
|||
|
|
# 插入到self.max_retry赋值之后
|
|||
|
|
insert_pos = content[init_match.start():].find(max_retry_assign)
|
|||
|
|
if insert_pos != -1:
|
|||
|
|
absolute_pos = init_match.start() + insert_pos + len(max_retry_assign) + 1
|
|||
|
|
content = content[:absolute_pos] + fewshot_init + content[absolute_pos:]
|
|||
|
|
print("✅ 修改 __init__: 添加few-shot初始化")
|
|||
|
|
|
|||
|
|
# 3. 修改 _generate_sql 方法
|
|||
|
|
generate_pattern = r'def _generate_sql\([\s\S]*?\) -> str:'
|
|||
|
|
generate_match = re.search(generate_pattern, content)
|
|||
|
|
|
|||
|
|
if generate_match:
|
|||
|
|
generate_text = content[generate_match.start():generate_match.end()]
|
|||
|
|
|
|||
|
|
if "if self.fewshot_enabled" not in generate_text:
|
|||
|
|
# 在方法开头注入few-shot逻辑
|
|||
|
|
fewshot_logic = ''' # Few-shot 增强
|
|||
|
|
if self.fewshot_enabled and self.fewshot_selector:
|
|||
|
|
try:
|
|||
|
|
examples = self.fewshot_selector.select(
|
|||
|
|
question=question,
|
|||
|
|
top_k=self.fewshot_top_k,
|
|||
|
|
min_rating=self.fewshot_min_rating
|
|||
|
|
)
|
|||
|
|
if examples:
|
|||
|
|
examples_prompt = "\\n".join([
|
|||
|
|
f"示例 {i+1}：\\n问题：{ex.question_zh}\\nSQL：\\n{ex.sql}"
|
|||
|
|
for i, ex in enumerate(examples)
|
|||
|
|
])
|
|||
|
|
examples_prompt = (
|
|||
|
|
f"参考以下相似问题的SQL示例：\\n\\n{examples_prompt}\\n\\n"
|
|||
|
|
f"请学习上述示例的SQL编写风格和模式，为当前问题生成SQL。"
|
|||
|
|
)
|
|||
|
|
# 注入到schema_str前或user_content中
|
|||
|
|
# 这里选择在schema_str前添加
|
|||
|
|
schema_str = f"{examples_prompt}Schema信息：\\n{schema_str}"
|
|||
|
|
logger.debug(f"已注入 {len(examples)} 个few-shot示例")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Few-shot检索失败: {e}")
|
|||
|
|
|
|||
|
|
'''
|
|||
|
|
# 插入到方法开头（"dialect_label = dialect" 之前）
|
|||
|
|
dialect_pos = content[generate_match.start():].find("dialect_label = dialect")
|
|||
|
|
if dialect_pos != -1:
|
|||
|
|
absolute_pos = generate_match.start() + dialect_pos
|
|||
|
|
content = content[:absolute_pos] + fewshot_logic + content[absolute_pos:]
|
|||
|
|
print("✅ 修改 _generate_sql: 注入few-shot逻辑")
|
|||
|
|
|
|||
|
|
# 保存修改
|
|||
|
|
backup_path = orch_path.with_suffix('.py.bak')
|
|||
|
|
orch_path.write_text(content, encoding='utf-8')
|
|||
|
|
print(f"✅ 已修改 {orch_path}")
|
|||
|
|
print(f" 备份文件: {backup_path}")
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def create_env_config(top_k: int, min_rating: int):
|
|||
|
|
"""更新.env文件添加few-shot配置"""
|
|||
|
|
env_path = Path(__file__).parent.parent / ".env"
|
|||
|
|
|
|||
|
|
if not env_path.exists():
|
|||
|
|
print(f"⚠️ .env文件不存在: {env_path}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
content = env_path.read_text(encoding='utf-8')
|
|||
|
|
|
|||
|
|
additions = f"""
|
|||
|
|
# Few-shot 配置
|
|||
|
|
FEWSHOT_ENABLED=true
|
|||
|
|
FEWSHOT_TOP_K={top_k}
|
|||
|
|
FEWSHOT_MIN_RATING={min_rating}
|
|||
|
|
FEWSHOT_DATA_PATH=./data/experiences/all_samples.jsonl
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 检查是否已存在
|
|||
|
|
if "FEWSHOT_ENABLED" not in content:
|
|||
|
|
with env_path.open('a', encoding='utf-8') as f:
|
|||
|
|
f.write(additions)
|
|||
|
|
print(f"✅ 已添加Few-shot配置到 .env")
|
|||
|
|
else:
|
|||
|
|
print(f"ℹ️ .env 中已存在FEWSHOT配置，跳过")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description="集成Few-shot功能到Text2SQL系统")
|
|||
|
|
parser.add_argument("--enable", action="store_true", help="启用few-shot")
|
|||
|
|
parser.add_argument("--top-k", type=int, default=3, help="每次使用的示例数量")
|
|||
|
|
parser.add_argument("--min-rating", type=int, default=7, help="示例最低评分")
|
|||
|
|
parser.add_argument("--env-only", action="store_true", help="仅修改.env")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
if not args.enable and not args.env_only:
|
|||
|
|
print("请使用 --enable 启用few-shot，或 --env-only 仅修改环境变量")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print("=" * 60)
|
|||
|
|
print("Text2SQL Few-shot 集成工具")
|
|||
|
|
print("=" * 60)
|
|||
|
|
|
|||
|
|
if args.env_only:
|
|||
|
|
create_env_config(args.top_k, args.min_rating)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
if args.enable:
|
|||
|
|
print(f"\n配置: top_k={args.top_k}, min_rating={args.min_rating}")
|
|||
|
|
|
|||
|
|
# 1. 修改orchestrator.py
|
|||
|
|
success = patch_orchestrator(
|
|||
|
|
enable=True,
|
|||
|
|
top_k=args.top_k,
|
|||
|
|
min_rating=args.min_rating
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if success:
|
|||
|
|
# 2. 更新.env
|
|||
|
|
create_env_config(args.top_k, args.min_rating)
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 60)
|
|||
|
|
print("✅ Few-shot集成完成！")
|
|||
|
|
print("=" * 60)
|
|||
|
|
print(f"\n下一步:")
|
|||
|
|
print(f"1. 确保经验数据集已生成:")
|
|||
|
|
print(f" python scripts/parse_examples.py")
|
|||
|
|
print(f"\n2. 测试系统:")
|
|||
|
|
print(f" python main.py \"查询2024年1月的销售额\" --verbose")
|
|||
|
|
print(f"\n3. 查看日志中的few-shot选择:")
|
|||
|
|
print(f" Few-shot已选择: Q31, Q6, Q44")
|
|||
|
|
print(f"\n4. 调整参数:")
|
|||
|
|
print(f" 编辑 .env 修改 FEWSHOT_TOP_K 和 FEWSHOT_MIN_RATING")
|
|||
|
|
else:
|
|||
|
|
print("❌ 集成失败")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|