239 lines
7.2 KiB
Python
239 lines
7.2 KiB
Python
|
"""
|
|||
|
|
Layer 2: 自适应探索器
|
|||
|
|
基于分析计划 + 已有发现，动态决定下一步查什么。
|
|||
|
|
多轮迭代，直到 AI 判断"够了"或达到上限。
|
|||
|
|
"""
|
|||
|
|
import json
|
|||
|
|
import re
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import openai
|
|||
|
|
from config import LLM_CONFIG
|
|||
|
|
from sandbox_executor import SandboxExecutor
|
|||
|
|
|
|||
|
|
|
|||
|
|
EXPLORER_SYSTEM = """你是一个数据分析执行者。你的上级给了你一个分析计划，你需要通过迭代执行 SQL 来完成分析。
|
|||
|
|
|
|||
|
|
## 你的工作方式
|
|||
|
|
每一轮你看到：
|
|||
|
|
1. 分析计划（上级给的目标）
|
|||
|
|
2. 数据库 Schema（表结构、数据画像）
|
|||
|
|
3. 之前的探索历史（查过什么、得到什么结果）
|
|||
|
|
|
|||
|
|
你决定下一步：
|
|||
|
|
- 输出一条 SQL 继续探索
|
|||
|
|
- 或者输出 done 表示分析足够
|
|||
|
|
|
|||
|
|
## 输出格式（严格 JSON）
|
|||
|
|
```json
|
|||
|
|
{
|
|||
|
|
"action": "query",
|
|||
|
|
"reasoning": "为什么要做这个查询",
|
|||
|
|
"sql": "SELECT ...",
|
|||
|
|
"purpose": "这个查询的目的"
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
或：
|
|||
|
|
```json
|
|||
|
|
{
|
|||
|
|
"action": "done",
|
|||
|
|
"reasoning": "为什么分析已经足够"
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## SQL 规则
|
|||
|
|
- 只用 SELECT，必须有聚合函数或 GROUP BY
|
|||
|
|
- 禁止 SELECT *
|
|||
|
|
- 用 ROUND 控制精度
|
|||
|
|
- 合理使用 LIMIT（分组结果 15 行以内，时间序列 60 行以内）
|
|||
|
|
|
|||
|
|
## 探索策略
|
|||
|
|
1. 第一轮：验证核心假设（计划中最关键的查询）
|
|||
|
|
2. 后续轮：基于已有结果追问
|
|||
|
|
- 发现离群值 → 追问为什么
|
|||
|
|
- 发现异常比例 → 追问细分维度
|
|||
|
|
- 结果平淡 → 换个角度试试
|
|||
|
|
3. 不要重复查已经确认的事
|
|||
|
|
4. 每轮要有新发现，否则就该结束"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
EXPLORER_CONTINUE = """查询结果：
|
|||
|
|
|
|||
|
|
{result_text}
|
|||
|
|
|
|||
|
|
请基于这个结果决定下一步。如果发现异常或值得深挖的点，继续查询。如果分析足够，输出 done。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ExplorationStep:
|
|||
|
|
"""单步探索结果"""
|
|||
|
|
def __init__(self, round_num: int, decision: dict, result: dict):
|
|||
|
|
self.round_num = round_num
|
|||
|
|
self.reasoning = decision.get("reasoning", "")
|
|||
|
|
self.purpose = decision.get("purpose", "")
|
|||
|
|
self.sql = decision.get("sql", "")
|
|||
|
|
self.action = decision.get("action", "query")
|
|||
|
|
self.success = result.get("success", False)
|
|||
|
|
self.error = result.get("error")
|
|||
|
|
self.columns = result.get("columns", [])
|
|||
|
|
self.rows = result.get("rows", [])
|
|||
|
|
self.row_count = result.get("row_count", 0)
|
|||
|
|
|
|||
|
|
def to_dict(self) -> dict:
|
|||
|
|
d = {
|
|||
|
|
"round": self.round_num,
|
|||
|
|
"action": self.action,
|
|||
|
|
"reasoning": self.reasoning,
|
|||
|
|
"purpose": self.purpose,
|
|||
|
|
"sql": self.sql,
|
|||
|
|
"success": self.success,
|
|||
|
|
}
|
|||
|
|
if self.success:
|
|||
|
|
d["result"] = {
|
|||
|
|
"columns": self.columns,
|
|||
|
|
"rows": self.rows,
|
|||
|
|
"row_count": self.row_count,
|
|||
|
|
}
|
|||
|
|
else:
|
|||
|
|
d["result"] = {"error": self.error}
|
|||
|
|
return d
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Explorer:
|
|||
|
|
"""自适应探索器：多轮迭代执行 SQL"""
|
|||
|
|
|
|||
|
|
def __init__(self, executor: SandboxExecutor):
|
|||
|
|
self.executor = executor
|
|||
|
|
self.client = openai.OpenAI(
|
|||
|
|
api_key=LLM_CONFIG["api_key"],
|
|||
|
|
base_url=LLM_CONFIG["base_url"],
|
|||
|
|
)
|
|||
|
|
self.model = LLM_CONFIG["model"]
|
|||
|
|
|
|||
|
|
def explore(
|
|||
|
|
self,
|
|||
|
|
plan: dict,
|
|||
|
|
schema_text: str,
|
|||
|
|
max_rounds: int = 6,
|
|||
|
|
) -> list[ExplorationStep]:
|
|||
|
|
"""
|
|||
|
|
执行探索循环
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
plan: Planner 生成的分析计划
|
|||
|
|
schema_text: Schema 文本描述
|
|||
|
|
max_rounds: 最大探索轮数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
探步列表
|
|||
|
|
"""
|
|||
|
|
steps: list[ExplorationStep] = []
|
|||
|
|
|
|||
|
|
# 构建初始消息
|
|||
|
|
messages = [
|
|||
|
|
{"role": "system", "content": EXPLORER_SYSTEM},
|
|||
|
|
{
|
|||
|
|
"role": "user",
|
|||
|
|
"content": (
|
|||
|
|
f"## 分析计划\n```json\n{json.dumps(plan, ensure_ascii=False, indent=2)}\n```\n\n"
|
|||
|
|
f"## 数据库 Schema\n{schema_text}\n\n"
|
|||
|
|
f"请开始第一轮探索。根据计划，先执行最关键的查询。"
|
|||
|
|
),
|
|||
|
|
},
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
for round_num in range(1, max_rounds + 1):
|
|||
|
|
print(f"\n 🔄 探索第 {round_num}/{max_rounds} 轮")
|
|||
|
|
|
|||
|
|
# LLM 决策
|
|||
|
|
decision = self._llm_decide(messages)
|
|||
|
|
action = decision.get("action", "query")
|
|||
|
|
reasoning = decision.get("reasoning", "")
|
|||
|
|
|
|||
|
|
print(f" 💭 {reasoning[:80]}{'...' if len(reasoning) > 80 else ''}")
|
|||
|
|
|
|||
|
|
if action == "done":
|
|||
|
|
print(f" ✅ 探索完成")
|
|||
|
|
steps.append(ExplorationStep(round_num, decision, {"success": True}))
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# 执行 SQL
|
|||
|
|
sql = decision.get("sql", "")
|
|||
|
|
purpose = decision.get("purpose", "")
|
|||
|
|
|
|||
|
|
if not sql:
|
|||
|
|
print(f" ⚠️ 未生成 SQL，跳过")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
print(f" 📝 {purpose}")
|
|||
|
|
result = self.executor.execute(sql)
|
|||
|
|
|
|||
|
|
if result["success"]:
|
|||
|
|
print(f" ✅ {result['row_count']} 行")
|
|||
|
|
else:
|
|||
|
|
print(f" ❌ {result['error']}")
|
|||
|
|
|
|||
|
|
step = ExplorationStep(round_num, decision, result)
|
|||
|
|
steps.append(step)
|
|||
|
|
|
|||
|
|
# 更新对话历史
|
|||
|
|
messages.append({
|
|||
|
|
"role": "assistant",
|
|||
|
|
"content": json.dumps(decision, ensure_ascii=False),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
result_text = self._format_result(result)
|
|||
|
|
messages.append({
|
|||
|
|
"role": "user",
|
|||
|
|
"content": EXPLORER_CONTINUE.format(result_text=result_text),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return steps
|
|||
|
|
|
|||
|
|
def _llm_decide(self, messages: list[dict]) -> dict:
|
|||
|
|
"""LLM 决策"""
|
|||
|
|
response = self.client.chat.completions.create(
|
|||
|
|
model=self.model,
|
|||
|
|
messages=messages,
|
|||
|
|
temperature=0.2,
|
|||
|
|
max_tokens=1024,
|
|||
|
|
)
|
|||
|
|
content = response.choices[0].message.content.strip()
|
|||
|
|
return self._extract_json(content)
|
|||
|
|
|
|||
|
|
def _format_result(self, result: dict) -> str:
|
|||
|
|
"""格式化查询结果"""
|
|||
|
|
if not result.get("success"):
|
|||
|
|
return f"❌ 执行失败: {result.get('error', '未知错误')}"
|
|||
|
|
|
|||
|
|
rows = result["rows"][:20]
|
|||
|
|
return (
|
|||
|
|
f"✅ 查询成功，返回 {result['row_count']} 行\n"
|
|||
|
|
f"列: {result['columns']}\n"
|
|||
|
|
f"数据:\n{json.dumps(rows, ensure_ascii=False, indent=2)}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _extract_json(self, text: str) -> dict:
|
|||
|
|
"""从 LLM 输出提取 JSON"""
|
|||
|
|
try:
|
|||
|
|
return json.loads(text)
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
for pattern in [r'```json\s*\n(.*?)\n```', r'```\s*\n(.*?)\n```']:
|
|||
|
|
match = re.search(pattern, text, re.DOTALL)
|
|||
|
|
if match:
|
|||
|
|
try:
|
|||
|
|
return json.loads(match.group(1))
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', text, re.DOTALL)
|
|||
|
|
if match:
|
|||
|
|
try:
|
|||
|
|
return json.loads(match.group())
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return {"action": "done", "reasoning": f"无法解析输出: {text[:100]}"}
|