|
|
@@ -0,0 +1,308 @@
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import asyncio
|
|
|
+import logging
|
|
|
+import re
|
|
|
+import uuid
|
|
|
+from datetime import datetime
|
|
|
+from pathlib import Path
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+import pandas as pd
|
|
|
+
|
|
|
+from .codex_runtime import CodexRuntime
|
|
|
+from .commands import HELP_TEXT, SKILLS_TEXT, parse_command
|
|
|
+from .config import Settings
|
|
|
+from .feishu import FeishuApi, FeishuLongConnection
|
|
|
+from .models import IncomingMessage, QueryAnalysis, QueryDecision
|
|
|
+from .odps_client import ODPSClient, validate_for_odps
|
|
|
+from .reports import build_profile, write_json
|
|
|
+from .skill_executor import SkillExecutor
|
|
|
+from .sql_guard import SQLGuard
|
|
|
+from .state import StateStore
|
|
|
+
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
+
|
|
|
+class DataQueryService:
|
|
|
+ def __init__(self, settings: Settings) -> None:
|
|
|
+ self.settings = settings
|
|
|
+ settings.runtime_dir.mkdir(parents=True, exist_ok=True)
|
|
|
+ (settings.runtime_dir / "runs").mkdir(parents=True, exist_ok=True)
|
|
|
+ self.state = StateStore(settings.db_path, settings.conversation_idle_hours)
|
|
|
+ self.codex = CodexRuntime(settings)
|
|
|
+ self.odps = ODPSClient(settings)
|
|
|
+ self.guard = SQLGuard(settings.odps_allowed_projects)
|
|
|
+ self.skills = SkillExecutor(settings)
|
|
|
+ self.feishu = FeishuApi(settings)
|
|
|
+ self.queue: asyncio.Queue[IncomingMessage] = asyncio.Queue(maxsize=1000)
|
|
|
+ self.query_slots = asyncio.Semaphore(settings.query_concurrency)
|
|
|
+ self._locks: dict[str, asyncio.Lock] = {}
|
|
|
+ self._listener: FeishuLongConnection | None = None
|
|
|
+ self._worker_tasks: list[asyncio.Task[Any]] = []
|
|
|
+ self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
+
|
|
|
+ def _authorized(self, message: IncomingMessage) -> bool:
|
|
|
+ return (
|
|
|
+ message.message_type in {"text", "post"}
|
|
|
+ and bool(message.text)
|
|
|
+ and bool(message.message_id)
|
|
|
+ and bool(message.chat_id)
|
|
|
+ and bool(message.sender_open_id)
|
|
|
+ and message.chat_id in self.settings.allowed_chat_ids
|
|
|
+ and message.chat_type == "group"
|
|
|
+ and message.mentioned_bot
|
|
|
+ )
|
|
|
+
|
|
|
+ def _enqueue_from_websocket(self, message: IncomingMessage) -> None:
|
|
|
+ if not self._authorized(message) or self._loop is None:
|
|
|
+ return
|
|
|
+
|
|
|
+ def enqueue() -> None:
|
|
|
+ try:
|
|
|
+ self.queue.put_nowait(message)
|
|
|
+ except asyncio.QueueFull:
|
|
|
+ logger.error("Inbound queue full; message_id=%s dropped", message.message_id)
|
|
|
+
|
|
|
+ self._loop.call_soon_threadsafe(enqueue)
|
|
|
+
|
|
|
+ async def start(self) -> None:
|
|
|
+ self._loop = asyncio.get_running_loop()
|
|
|
+ bot_open_id = await self.feishu.bot_open_id()
|
|
|
+ self._listener = FeishuLongConnection(self.settings, bot_open_id, self._enqueue_from_websocket)
|
|
|
+ self._listener.start()
|
|
|
+ worker_count = max(2, self.settings.query_concurrency * 2)
|
|
|
+ self._worker_tasks = [asyncio.create_task(self._worker(index)) for index in range(worker_count)]
|
|
|
+ logger.info("Data query agent started workers=%d", worker_count)
|
|
|
+
|
|
|
+ async def stop(self) -> None:
|
|
|
+ for task in self._worker_tasks:
|
|
|
+ task.cancel()
|
|
|
+ await asyncio.gather(*self._worker_tasks, return_exceptions=True)
|
|
|
+ await self.feishu.close()
|
|
|
+ self.state.close()
|
|
|
+
|
|
|
+ async def _worker(self, index: int) -> None:
|
|
|
+ while True:
|
|
|
+ message = await self.queue.get()
|
|
|
+ try:
|
|
|
+ await self.handle(message)
|
|
|
+ except asyncio.CancelledError:
|
|
|
+ raise
|
|
|
+ except Exception:
|
|
|
+ logger.exception("Unhandled message failure worker=%d message_id=%s", index, message.message_id)
|
|
|
+ finally:
|
|
|
+ self.queue.task_done()
|
|
|
+
|
|
|
+ async def handle(self, message: IncomingMessage) -> None:
|
|
|
+ if not self._authorized(message):
|
|
|
+ return
|
|
|
+ if not self.state.claim_message(message.message_id, message.conversation_key):
|
|
|
+ logger.info("Duplicate Feishu message ignored message_id=%s", message.message_id)
|
|
|
+ return
|
|
|
+ lock = self._locks.setdefault(message.conversation_key, asyncio.Lock())
|
|
|
+ async with lock:
|
|
|
+ try:
|
|
|
+ command = parse_command(message.text)
|
|
|
+ if command:
|
|
|
+ await self._handle_command(message, command)
|
|
|
+ self.state.finish_message(message.message_id, "command")
|
|
|
+ return
|
|
|
+ await self.feishu.reply_text(message.message_id, "收到,正在处理你的请求…")
|
|
|
+ async with self.query_slots:
|
|
|
+ await self._handle_query(message)
|
|
|
+ self.state.finish_message(message.message_id, "completed")
|
|
|
+ except Exception as exc:
|
|
|
+ safe_error = self._safe_error(exc)
|
|
|
+ logger.error("Message handling failed message_id=%s error=%s", message.message_id, safe_error)
|
|
|
+ self.state.finish_message(message.message_id, "failed", safe_error)
|
|
|
+ self.state.fail_active_runs(message.message_id, safe_error)
|
|
|
+ try:
|
|
|
+ await self.feishu.reply_text(message.message_id, f"查询失败:{safe_error}")
|
|
|
+ except Exception as reply_exc:
|
|
|
+ logger.error(
|
|
|
+ "Failed to send error reply message_id=%s error=%s",
|
|
|
+ message.message_id,
|
|
|
+ self._safe_error(reply_exc),
|
|
|
+ )
|
|
|
+
|
|
|
+ async def _handle_command(self, message: IncomingMessage, command: str) -> None:
|
|
|
+ if command == "help":
|
|
|
+ await self.feishu.reply_text(message.message_id, HELP_TEXT)
|
|
|
+ return
|
|
|
+ if command == "skills":
|
|
|
+ await self.feishu.reply_text(message.message_id, SKILLS_TEXT)
|
|
|
+ return
|
|
|
+ conversation = self.state.reset_conversation(message.conversation_key, new_session=command == "new")
|
|
|
+ if command == "new":
|
|
|
+ text = f"已开启新会话({conversation.session_id[:8]}),下一条问题将使用全新上下文。"
|
|
|
+ else:
|
|
|
+ text = "已清空当前上下文,下一条问题将重新开始。"
|
|
|
+ await self.feishu.reply_text(message.message_id, text)
|
|
|
+
|
|
|
+ async def _handle_query(self, message: IncomingMessage) -> None:
|
|
|
+ conversation = self.state.get_conversation(message.conversation_key)
|
|
|
+ thread_id, decision = await self.codex.plan(conversation.thread_id, message.text)
|
|
|
+ self.state.set_thread(message.conversation_key, thread_id)
|
|
|
+ if decision.status != "ready":
|
|
|
+ await self.feishu.reply_text(message.message_id, decision.reply)
|
|
|
+ return
|
|
|
+ if decision.execution_mode == "sql" and not decision.sql:
|
|
|
+ raise RuntimeError("Agent 返回 ready 但没有 SQL")
|
|
|
+ if decision.execution_mode == "skill_script" and decision.selected_skill != "query-user-behavior-path":
|
|
|
+ raise RuntimeError("Agent 请求了未授权的 Skill 脚本")
|
|
|
+
|
|
|
+ run_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
|
|
|
+ run_dir = self.settings.runtime_dir / "runs" / run_id
|
|
|
+ run_dir.mkdir(parents=True, exist_ok=False)
|
|
|
+ self.state.create_run(run_id, message.conversation_key, message.message_id, run_dir)
|
|
|
+ write_json(run_dir / "request.json", {
|
|
|
+ "run_id": run_id,
|
|
|
+ "message_id": message.message_id,
|
|
|
+ "conversation_key": message.conversation_key,
|
|
|
+ "session_id": conversation.session_id,
|
|
|
+ "question": message.text,
|
|
|
+ "selected_skill": decision.selected_skill,
|
|
|
+ "execution_mode": decision.execution_mode,
|
|
|
+ "parameters": decision.parameters.model_dump(),
|
|
|
+ })
|
|
|
+
|
|
|
+ await self.feishu.reply_text(
|
|
|
+ message.message_id,
|
|
|
+ f"已匹配 Skill:{decision.selected_skill},正在执行查询…",
|
|
|
+ )
|
|
|
+
|
|
|
+ if decision.execution_mode == "skill_script":
|
|
|
+ self.state.update_run(run_id, "running", metadata={"skill": decision.selected_skill})
|
|
|
+ artifact = await self.skills.run_user_timeline(decision.parameters, run_dir)
|
|
|
+ await self._publish_result(
|
|
|
+ message,
|
|
|
+ run_id,
|
|
|
+ thread_id,
|
|
|
+ decision.title,
|
|
|
+ decision.selected_skill,
|
|
|
+ artifact.dataframe,
|
|
|
+ artifact.xlsx_path,
|
|
|
+ artifact.instance_id,
|
|
|
+ artifact.truncated,
|
|
|
+ )
|
|
|
+ return
|
|
|
+
|
|
|
+ current = decision
|
|
|
+ result = None
|
|
|
+ repairs: list[dict[str, Any]] = []
|
|
|
+ for attempt in range(self.settings.query_max_repairs + 1):
|
|
|
+ sql = (current.sql or "").strip().rstrip(";")
|
|
|
+ (run_dir / "query.sql").write_text(sql + "\n", encoding="utf-8")
|
|
|
+ write_json(run_dir / "query_plan.json", {
|
|
|
+ "title": current.title,
|
|
|
+ "selected_skill": current.selected_skill,
|
|
|
+ "execution_mode": current.execution_mode,
|
|
|
+ "parameters": current.parameters.model_dump(),
|
|
|
+ "assumptions": current.assumptions,
|
|
|
+ "attempt": attempt,
|
|
|
+ "repairs": repairs,
|
|
|
+ })
|
|
|
+ try:
|
|
|
+ if current.selected_skill == "odps-product-efficiency-report":
|
|
|
+ self.guard.validate_product_efficiency_contract(sql, current.parameters.data_mode)
|
|
|
+ await validate_for_odps(sql, self.guard, self.odps)
|
|
|
+ self.state.update_run(run_id, "running")
|
|
|
+ result = await self.odps.execute(sql)
|
|
|
+ break
|
|
|
+ except Exception as exc:
|
|
|
+ if attempt >= self.settings.query_max_repairs:
|
|
|
+ self.state.update_run(run_id, "failed", metadata={"attempts": attempt + 1})
|
|
|
+ raise
|
|
|
+ repairs.append({"attempt": attempt + 1, "error": self._safe_error(exc)})
|
|
|
+ thread_id, current = await self.codex.repair(thread_id, sql, str(exc), attempt + 1)
|
|
|
+ self.state.set_thread(message.conversation_key, thread_id)
|
|
|
+ if current.status != "ready" or not current.sql:
|
|
|
+ raise RuntimeError(current.reply or "Agent 无法修复 SQL")
|
|
|
+
|
|
|
+ if result is None:
|
|
|
+ raise RuntimeError("查询未返回结果")
|
|
|
+ frame = result.dataframe
|
|
|
+ if not isinstance(frame, pd.DataFrame):
|
|
|
+ frame = pd.DataFrame(frame)
|
|
|
+ info = {
|
|
|
+ "run_id": run_id,
|
|
|
+ "原始问题": message.text,
|
|
|
+ "匹配 Skill": current.selected_skill,
|
|
|
+ "ODPS instance_id": result.instance_id,
|
|
|
+ "返回行数": len(frame.index),
|
|
|
+ "是否截断": result.truncated,
|
|
|
+ "生成时间": datetime.now().isoformat(timespec="seconds"),
|
|
|
+ }
|
|
|
+ artifact = await self.skills.format_report(
|
|
|
+ current.selected_skill,
|
|
|
+ current.parameters,
|
|
|
+ run_dir,
|
|
|
+ frame,
|
|
|
+ current.sql or "",
|
|
|
+ info,
|
|
|
+ )
|
|
|
+ await self._publish_result(
|
|
|
+ message,
|
|
|
+ run_id,
|
|
|
+ thread_id,
|
|
|
+ current.title,
|
|
|
+ current.selected_skill,
|
|
|
+ artifact.dataframe,
|
|
|
+ artifact.xlsx_path,
|
|
|
+ result.instance_id,
|
|
|
+ result.truncated,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def _publish_result(
|
|
|
+ self,
|
|
|
+ message: IncomingMessage,
|
|
|
+ run_id: str,
|
|
|
+ thread_id: str,
|
|
|
+ title: str,
|
|
|
+ selected_skill: str,
|
|
|
+ frame: pd.DataFrame,
|
|
|
+ xlsx_path: Path,
|
|
|
+ instance_id: str,
|
|
|
+ truncated: bool,
|
|
|
+ ) -> None:
|
|
|
+ profile = build_profile(frame, truncated)
|
|
|
+ try:
|
|
|
+ thread_id, analysis = await self.codex.analyze(thread_id, message.text, profile)
|
|
|
+ self.state.set_thread(message.conversation_key, thread_id)
|
|
|
+ except Exception as exc:
|
|
|
+ logger.error("Codex result analysis failed; using deterministic summary: %s", self._safe_error(exc))
|
|
|
+ analysis = QueryAnalysis(
|
|
|
+ summary=f"查询完成,共返回 {len(frame.index)} 行、{len(frame.columns)} 列。",
|
|
|
+ highlights=[],
|
|
|
+ caveats=["结果达到行数上限,表格仅包含前若干行。"] if truncated else [],
|
|
|
+ )
|
|
|
+ summary_payload = analysis.model_dump()
|
|
|
+ summary_payload["profile"] = profile
|
|
|
+ run_dir = self.settings.runtime_dir / "runs" / run_id
|
|
|
+ write_json(run_dir / "summary.json", summary_payload)
|
|
|
+ self.state.update_run(
|
|
|
+ run_id,
|
|
|
+ "publishing",
|
|
|
+ instance_id=instance_id,
|
|
|
+ metadata={"skill": selected_skill, "rows": len(frame.index), "truncated": truncated},
|
|
|
+ )
|
|
|
+ await self.feishu.reply_text(message.message_id, "查询完成,正在生成并发布飞书表格…")
|
|
|
+ url = await self.feishu.publish_sheet(xlsx_path, title)
|
|
|
+ self.state.update_run(
|
|
|
+ run_id,
|
|
|
+ "completed",
|
|
|
+ feishu_url=url,
|
|
|
+ metadata={"skill": selected_skill, "rows": len(frame.index), "truncated": truncated},
|
|
|
+ )
|
|
|
+ lines = [analysis.summary]
|
|
|
+ lines.extend(f"- {item}" for item in analysis.highlights)
|
|
|
+ lines.extend(f"- 注意:{item}" for item in analysis.caveats)
|
|
|
+ lines.append(f"\n返回 **{len(frame.index)}** 行" + ("(已截断)" if truncated else ""))
|
|
|
+ await self.feishu.reply_card(message.message_id, title, "\n".join(lines)[:5000], url)
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _safe_error(exc: Exception) -> str:
|
|
|
+ text = str(exc).replace("\n", " ")
|
|
|
+ text = re.sub(r"(?i)(access[_ -]?key|secret|token)\s*[:=]\s*\S+", r"\1=<redacted>", text)
|
|
|
+ return text[:500] or type(exc).__name__
|