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 == "direct_reply": if decision.sql: raise RuntimeError("Agent 的自然语言回复不应包含 SQL") if not decision.reply.strip(): raise RuntimeError("Agent 的自然语言回复为空") 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 current.execution_mode != "sql" 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=", text) return text[:500] or type(exc).__name__