""" Layer 2: 自适应探索器 """ import json from typing import Any from dataclasses import dataclass, field from core.config import LLM_CONFIG from core.utils import get_llm_client, extract_json_object from core.sandbox 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 - 每条 SQL 必须包含聚合函数（COUNT/SUM/AVG/MIN/MAX）或 GROUP BY 或 LIMIT - 禁止 SELECT * - 用 ROUND 控制精度 - 合理使用 LIMIT（分组结果 15 行以内，时间序列 60 行以内） - 如果需要查看明细数据，必须加 LIMIT ## 探索策略 1. 第一轮：验证核心假设 2. 后续轮：基于已有结果追问 3. 不要重复查已经确认的事 4. 每轮要有新发现，否则就该结束""" @dataclass class ExplorationStep: """单步探索结果""" round_num: int = 0 reasoning: str = "" purpose: str = "" sql: str = "" action: str = "query" success: bool = False error: str | None = None columns: list[str] = field(default_factory=list) rows: list[dict] = field(default_factory=list) row_count: int = 0 @classmethod def from_decision(cls, round_num: int, decision: dict, result: dict) -> "ExplorationStep": return cls( round_num=round_num, reasoning=decision.get("reasoning", ""), purpose=decision.get("purpose", ""), sql=decision.get("sql", ""), action=decision.get("action", "query"), success=result.get("success", False), error=result.get("error"), columns=result.get("columns", []), rows=result.get("rows", []), 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: """自适应探索器""" def __init__(self, executor: SandboxExecutor): self.executor = executor self.client, self.model = get_llm_client(LLM_CONFIG) def explore( self, plan: dict, schema_text: str, max_rounds: int = 6, playbook_result: dict | None = None, ) -> list[ExplorationStep]: steps: list[ExplorationStep] = [] # 阶段 A: 预设查询 preset_context = "" if playbook_result and playbook_result.get("preset_queries"): preset_steps = self._run_preset_queries(playbook_result["preset_queries"]) steps.extend(preset_steps) preset_context = self._build_preset_context(preset_steps, playbook_result) # 阶段 B: 自适应探索 preset_used = len([s for s in steps if s.success]) remaining = max(1, max_rounds - preset_used) initial = ( f"## 分析计划\n```json\n{json.dumps(plan, ensure_ascii=False, indent=2)}\n```\n\n" f"## 数据库 Schema\n{schema_text}\n\n" ) # 注入历史上下文 prev_context = plan.pop("_prev_context", None) if prev_context: initial += f"## 历史分析参考\n{prev_context}\n\n" if preset_context: initial += ( f"## 预设分析结果（已执行）\n{preset_context}\n\n" f"请基于这些已有数据，决定是否需要进一步探索。\n" f"重点关注：预设结果中的异常、值得深挖的点。\n" f"如果预设结果已经足够，直接输出 done。" ) if playbook_result.get("exploration_hints"): initial += f"\n\n## 探索提示\n{playbook_result['exploration_hints']}" else: initial += "请开始第一轮探索。根据计划，先执行最关键的查询。" messages = [ {"role": "system", "content": EXPLORER_SYSTEM}, {"role": "user", "content": initial}, ] offset = len(steps) for round_num in range(offset + 1, offset + remaining + 1): print(f"\n 🔄 探索第 {round_num}/{max_rounds} 轮") decision = self._llm_decide(messages) reasoning = decision.get("reasoning", "") print(f" 💭 {reasoning[:80]}{'...' if len(reasoning) > 80 else ''}") if decision.get("action") == "done": print(f" ✅ 探索完成") steps.append(ExplorationStep.from_decision(round_num, decision, {"success": True})) break sql = decision.get("sql", "") if not sql: continue print(f" 📝 {decision.get('purpose', '')}") try: result = self.executor.execute(sql) except Exception as e: result = {"success": False, "error": str(e), "sql": sql} print(f" {'✅' if result['success'] else '❌'} {result.get('row_count', result.get('error', ''))}") steps.append(ExplorationStep.from_decision(round_num, decision, result)) messages.append({"role": "assistant", "content": json.dumps(decision, ensure_ascii=False)}) messages.append({"role": "user", "content": self._format_result(result)}) return steps def _run_preset_queries(self, preset_queries: list[dict]) -> list[ExplorationStep]: steps = [] for i, pq in enumerate(preset_queries, 1): sql, purpose = pq["sql"], pq.get("purpose", f"预设查询 {i}") print(f"\n 📌 预设查询 {i}/{len(preset_queries)}: {purpose}") try: result = self.executor.execute(sql) except Exception as e: result = {"success": False, "error": str(e), "sql": sql} decision = {"action": "query", "reasoning": f"[预设] {purpose}", "sql": sql, "purpose": purpose} steps.append(ExplorationStep.from_decision(i, decision, result)) print(f" {'✅' if result['success'] else '❌'} {result.get('row_count', result.get('error', ''))}") return steps def _build_preset_context(self, steps: list[ExplorationStep], playbook_result: dict) -> str: parts = [f"Playbook: {playbook_result.get('playbook_name', '未知')}"] for step in steps: if step.success: parts.append( f"### {step.purpose}\nSQL: `{step.sql}`\n" f"结果 ({step.row_count} 行): {json.dumps(step.rows[:15], ensure_ascii=False)}" ) else: parts.append(f"### {step.purpose}\nSQL: `{step.sql}`\n执行失败: {step.error}") return "\n\n".join(parts) def _llm_decide(self, messages: list[dict]) -> dict: response = self.client.chat.completions.create( model=self.model, messages=messages, temperature=0.2, max_tokens=1024, ) content = response.choices[0].message.content.strip() result = extract_json_object(content) return result if result else {"action": "done", "reasoning": f"无法解析: {content[:100]}"} def _format_result(self, result: dict) -> str: if not result.get("success"): return f"❌ 执行失败: {result.get('error', '未知错误')}" rows = result["rows"][:20] return ( f"查询结果：\n\n✅ 返回 {result['row_count']} 行\n" f"列: {result['columns']}\n数据:\n{json.dumps(rows, ensure_ascii=False, indent=2)}\n\n" f"请基于这个结果决定下一步。如果发现异常或值得深挖的点，继续查询。如果分析足够，输出 done。" )