data-query-agent/src/data_query_agent/service.py

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=<redacted>", text)
        return text[:500] or type(exc).__name__