250 lines
8.1 KiB
Python
250 lines
8.1 KiB
Python
"""
|
||
Text2SQL Few-shot 集成示例
|
||
|
||
展示如何在现有系统中集成经验数据集few-shot功能
|
||
"""
|
||
|
||
from typing import Optional, List
|
||
|
||
|
||
def enhance_orchestrator_with_fewshot(orchestrator_class):
|
||
"""
|
||
装饰器：为Text2SQLOrchestrator添加few-shot功能
|
||
|
||
用法：
|
||
from agents.orchestrator import Text2SQLOrchestrator
|
||
EnhancedOrchestrator = enhance_orchestrator_with_fewshot(Text2SQLOrchestrator)
|
||
orch = EnhancedOrchestrator(..., fewshot_enabled=True, fewshot_top_k=3)
|
||
"""
|
||
class FewShotEnhancedOrchestrator(orchestrator_class):
|
||
"""增强版Orchestrator，支持few-shot"""
|
||
|
||
def __init__(
|
||
self,
|
||
*args,
|
||
fewshot_enabled: bool = True,
|
||
fewshot_samples_path: Optional[str] = None,
|
||
fewshot_top_k: int = 3,
|
||
fewshot_min_rating: int = 7,
|
||
**kwargs
|
||
):
|
||
super().__init__(*args, **kwargs)
|
||
|
||
self.fewshot_enabled = fewshot_enabled
|
||
self.fewshot_top_k = fewshot_top_k
|
||
self.fewshot_min_rating = fewshot_min_rating
|
||
self.fewshot_selector = None
|
||
|
||
if fewshot_enabled:
|
||
from utils.fewshot_selector import FewShotSelector
|
||
path = fewshot_samples_path or "data/experiences/all_samples.jsonl"
|
||
self.fewshot_selector = FewShotSelector(path)
|
||
print(f"✅ Few-shot已启用: top_k={fewshot_top_k}, min_rating={fewshot_min_rating}")
|
||
|
||
def _generate_sql_with_fewshot(
|
||
self,
|
||
question: str,
|
||
schema_str: str,
|
||
dialect: str = "tsql"
|
||
) -> str:
|
||
"""带few-shot的SQL生成"""
|
||
from llm.deepseek_client import DeepSeekConfig
|
||
|
||
# 获取few-shot示例
|
||
examples_prompt = ""
|
||
if self.fewshot_selector:
|
||
examples = self.fewshot_selector.select(
|
||
question=question,
|
||
top_k=self.fewshot_top_k,
|
||
min_rating=self.fewshot_min_rating,
|
||
exclude_qids=[] # 可排除当前问题（如果已存在）
|
||
)
|
||
if examples:
|
||
examples_prompt = "\n\n".join([
|
||
f"示例 {i+1}：\n问题：{ex.question_zh}\nSQL：\n{ex.sql}"
|
||
for i, ex in enumerate(examples)
|
||
])
|
||
examples_prompt = f"参考以下相似示例：\n\n{examples_prompt}\n\n"
|
||
|
||
# 构造增强的Prompt
|
||
from config.prompts import SQL_GENERATOR_SYSTEM, SQL_GENERATOR_USER
|
||
|
||
user_content = f"""{examples_prompt}
|
||
当前问题：{question}
|
||
|
||
Schema信息：
|
||
{schema_str}
|
||
|
||
数据库方言：{dialect}
|
||
|
||
请参考上述示例的SQL编写风格，生成**有用 SQL**：
|
||
- 与示例保持相同的版式（大写关键字、4空格缩进、PascalCase别名）
|
||
- 根据当前Schema选择正确的表和字段
|
||
- 输出纯SQL，不要其他内容"""
|
||
|
||
messages = [
|
||
{"role": "system", "content": SQL_GENERATOR_SYSTEM},
|
||
{"role": "user", "content": user_content}
|
||
]
|
||
|
||
# 调用LLM
|
||
config = DeepSeekConfig(
|
||
api_key=self.deepseek.config.api_key,
|
||
base_url=self.deepseek.config.base_url,
|
||
model_name=self.deepseek.config.model_name,
|
||
temperature=0.3,
|
||
max_tokens=4096
|
||
)
|
||
|
||
from llm.deepseek_client import DeepSeekClient
|
||
client = DeepSeekClient(config)
|
||
response = client.chat(messages)
|
||
sql = response.content.strip()
|
||
|
||
# 清理代码块
|
||
if "```sql" in sql:
|
||
sql = sql[sql.find("```sql")+6:sql.find("```", sql.find("```sql")+6)].strip()
|
||
elif "```" in sql:
|
||
sql = sql[sql.find("```")+3:sql.find("```", sql.find("```")+3)].strip()
|
||
|
||
from utils.sql_parser import normalize_sql_for_dialect
|
||
return normalize_sql_for_dialect(sql, dialect)
|
||
|
||
# 重写 _generate_sql 方法
|
||
def _generate_sql(
|
||
self,
|
||
question: str,
|
||
schema_str: str,
|
||
dialect: str = "tsql"
|
||
) -> str:
|
||
if self.fewshot_enabled and self.fewshot_selector:
|
||
return self._generate_sql_with_fewshot(question, schema_str, dialect)
|
||
else:
|
||
return super()._generate_sql(question, schema_str, dialect)
|
||
|
||
return FewShotEnhancedOrchestrator
|
||
|
||
|
||
# ========== 集成步骤 ==========
|
||
|
||
def integrate_fewshot():
|
||
"""
|
||
集成few-shot功能的完整步骤
|
||
"""
|
||
print("=" * 60)
|
||
print("Few-shot 集成指南")
|
||
print("=" * 60)
|
||
|
||
steps = """
|
||
**步骤1：确保经验数据集已生成**
|
||
```bash
|
||
cd text2sql_agent_camel
|
||
python scripts/parse_examples.py
|
||
```
|
||
生成: data/experiences/all_samples.jsonl
|
||
|
||
**步骤2：修改 agents/orchestrator.py**
|
||
|
||
在文件开头添加：
|
||
```python
|
||
from utils.fewshot_selector import FewShotSelector
|
||
```
|
||
|
||
修改 Text2SQLOrchestrator.__init__：
|
||
```python
|
||
def __init__(self, ..., fewshot_enabled=True, fewshot_top_k=3, ...):
|
||
# ... 原有代码 ...
|
||
self.fewshot_enabled = fewshot_enabled
|
||
self.fewshot_top_k = fewshot_top_k
|
||
self.fewshot_selector = None
|
||
|
||
if fewshot_enabled:
|
||
self.fewshot_selector = FewShotSelector(
|
||
os.getenv("FEWSHOT_DATA_PATH", "data/experiences/all_samples.jsonl")
|
||
)
|
||
logger.info(f"Few-shot已启用: top_k={fewshot_top_k}")
|
||
```
|
||
|
||
修改 _generate_sql 方法（约第251行）：
|
||
```python
|
||
def _generate_sql(self, question: str, schema_str: str, dialect: str = "tsql") -> str:
|
||
# 添加few-shot增强
|
||
if self.fewshot_enabled and self.fewshot_selector:
|
||
examples = self.fewshot_selector.select(
|
||
question=question,
|
||
top_k=self.fewshot_top_k,
|
||
min_rating=7,
|
||
exclude_qids=[] # 可选：排除当前问题对应的QID
|
||
)
|
||
if examples:
|
||
examples_prompt = "\\n\\n".join([
|
||
f"示例（{ex.qid}）：\\n问题：{ex.question_zh}\\nSQL：\\n{ex.sql}"
|
||
for ex in examples
|
||
])
|
||
# 将examples_prompt注入到用户消息中
|
||
# （需要调整 deepseek_client 调用逻辑）
|
||
# ... 原有生成逻辑 ...
|
||
```
|
||
|
||
**步骤3：配置环境变量（可选）**
|
||
```bash
|
||
# .env 文件添加
|
||
FEWSHOT_ENABLED=true
|
||
FEWSHOT_TOP_K=3
|
||
FEWSHOT_MIN_RATING=7
|
||
FEWSHOT_DATA_PATH=./data/experiences/all_samples.jsonl
|
||
```
|
||
|
||
**步骤4：测试效果**
|
||
```bash
|
||
# 对比测试
|
||
python main.py "查询2024年1月的销售额" --verbose
|
||
# 观察日志中的 "Few-shot已选择: Q31, Q6, Q44"
|
||
```
|
||
|
||
**步骤5：A/B测试评估**
|
||
```python
|
||
# 创建对比脚本
|
||
from agents.orchestrator import Text2SQLOrchestrator
|
||
|
||
# 无few-shot
|
||
orch_baseline = Text2SQLOrchestrator(..., fewshot_enabled=False)
|
||
result1 = orch_baseline.generate(question)
|
||
|
||
# 有few-shot
|
||
orch_fewshot = Text2SQLOrchestrator(..., fewshot_enabled=True, fewshot_top_k=3)
|
||
result2 = orch_fewshot.generate(question)
|
||
|
||
print(f"Baseline: {result1.sql[:100]}")
|
||
print(f"Few-shot: {result2.sql[:100]}")
|
||
```
|
||
"""
|
||
print(steps)
|
||
|
||
|
||
def quick_test_fewshot():
|
||
"""快速测试few-shot选择器"""
|
||
from utils.fewshot_selector import FewShotSelector
|
||
|
||
selector = FewShotSelector("data/experiences/all_samples.jsonl")
|
||
|
||
test_questions = [
|
||
"查询2024年1月的销售额",
|
||
"统计所有活跃账户数量",
|
||
"按对手方列出未结算交易",
|
||
"查询今天的现金余额",
|
||
]
|
||
|
||
print("\nFew-shot选择测试:\n")
|
||
for q in test_questions:
|
||
print(f"问题: {q}")
|
||
examples = selector.select(q, top_k=2, min_rating=8)
|
||
for ex in examples:
|
||
print(f" [{ex.qid}] {ex.question_zh[:40]}... (rating={ex.rating})")
|
||
print()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
integrate_fewshot()
|
||
# quick_test_fewshot()
|