Agent/examples/auto_put_ad_mini/tools/feishu_doc.py

"""
飞书文档工具 — auto_put_ad_mini

职责：
  - 将本地 xlsx 文件上传并导入为飞书在线表格
  - 设置文档权限（任何人获得链接可查看 - 支持不同主体访问）
  - 通过 IM 发送在线表格链接

飞书 Drive API（通过 httpx 直连）：
  - 上传素材：POST /drive/v1/medias/upload_all
  - 创建导入任务：POST /drive/v1/import_tasks
  - 查询导入结果：GET /drive/v1/import_tasks/{ticket}
  - 设置权限：PATCH /drive/v1/permissions/{token}/public

飞书 IM（通过框架 FeishuClient）：
  - 发送链接消息：send_message(to, text)
"""

import json
import logging
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional

import httpx

from agent.tools import tool
from agent.tools.models import ToolContext, ToolResult
from agent.tools.builtin.feishu.feishu_client import FeishuClient

_MINI_DIR = Path(__file__).resolve().parent.parent
_TOOLS_DIR = Path(__file__).resolve().parent
if str(_MINI_DIR) not in sys.path:
    sys.path.insert(0, str(_MINI_DIR))
if str(_TOOLS_DIR) not in sys.path:
    sys.path.insert(0, str(_TOOLS_DIR))

from config import (
    FEISHU_APP_ID,
    FEISHU_APP_SECRET,
    FEISHU_OPERATOR_CHAT_ID,
    REPORTS_DIR,
)

logger = logging.getLogger(__name__)

# ═══════════════════════════════════════════
# 常量
# ═══════════════════════════════════════════

FEISHU_BASE_URL = "https://open.feishu.cn/open-apis"
_HTTP_TIMEOUT = 30
_IMPORT_POLL_INTERVAL = 2  # 秒
_IMPORT_MAX_WAIT = 60  # 秒

# 全局 FeishuClient（复用框架能力发 IM 消息）
_feishu = FeishuClient(app_id=FEISHU_APP_ID, app_secret=FEISHU_APP_SECRET)

# token 缓存
_token_cache: Dict[str, object] = {}
_TOKEN_TTL = 1500  # 25 分钟（官方有效期 2 小时，留余量）


# ═══════════════════════════════════════════
# 内部方法：飞书 tenant_access_token
# ═══════════════════════════════════════════

def _get_tenant_token() -> str:
    """获取飞书 tenant_access_token（带缓存）"""
    cached = _token_cache.get("tenant")
    if cached and time.time() - cached["ts"] < _TOKEN_TTL:
        return cached["token"]

    resp = httpx.post(
        f"{FEISHU_BASE_URL}/auth/v3/tenant_access_token/internal",
        json={"app_id": FEISHU_APP_ID, "app_secret": FEISHU_APP_SECRET},
        timeout=_HTTP_TIMEOUT,
    )
    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"获取飞书 token 失败: {data}")

    token = data["tenant_access_token"]
    _token_cache["tenant"] = {"token": token, "ts": time.time()}
    logger.info("飞书 tenant_access_token 已刷新")
    return token


def _auth_headers(token: str) -> Dict[str, str]:
    return {"Authorization": f"Bearer {token}"}


# ═══════════════════════════════════════════
# 内部方法：Drive API
# ═══════════════════════════════════════════

def _upload_media(token: str, file_path: Path) -> str:
    """上传素材文件到飞书（用于后续导入）

    Returns:
        file_token: 上传后的文件标识
    """
    file_size = file_path.stat().st_size
    file_name = file_path.name

    with open(file_path, "rb") as f:
        resp = httpx.post(
            f"{FEISHU_BASE_URL}/drive/v1/medias/upload_all",
            headers=_auth_headers(token),
            data={
                "file_name": file_name,
                "parent_type": "explorer",
                "parent_node": "",
                "size": str(file_size),
            },
            files={"file": (file_name, f, "application/octet-stream")},
            timeout=60,
        )

    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"上传素材失败: {data.get('msg', data)}")

    file_token = data["data"]["file_token"]
    logger.info("素材上传成功: file_token=%s, size=%d", file_token, file_size)
    return file_token


def _create_import_task(token: str, file_token: str, file_name: str) -> str:
    """创建导入任务（xlsx → 飞书在线表格）

    Returns:
        ticket: 导入任务标识
    """
    title = file_name.replace(".xlsx", "").replace(".xls", "")
    body = {
        "file_extension": "xlsx",
        "file_token": file_token,
        "type": "sheet",
        "file_name": title,
        "point": {
            "mount_type": 1,
            "mount_key": "",
        },
    }

    resp = httpx.post(
        f"{FEISHU_BASE_URL}/drive/v1/import_tasks",
        headers={**_auth_headers(token), "Content-Type": "application/json"},
        json=body,
        timeout=_HTTP_TIMEOUT,
    )

    data = resp.json()
    if data.get("code") != 0:
        raise RuntimeError(f"创建导入任务失败: {data.get('msg', data)}")

    ticket = data["data"]["ticket"]
    logger.info("导入任务已创建: ticket=%s, title=%s", ticket, title)
    return ticket


def _wait_import_result(token: str, ticket: str) -> Dict:
    """轮询导入结果

    Returns:
        dict: 包含 token, type, url 的结果字典
    """
    start = time.time()
    while time.time() - start < _IMPORT_MAX_WAIT:
        resp = httpx.get(
            f"{FEISHU_BASE_URL}/drive/v1/import_tasks/{ticket}",
            headers=_auth_headers(token),
            timeout=_HTTP_TIMEOUT,
        )
        data = resp.json()

        if data.get("code") != 0:
            raise RuntimeError(f"查询导入结果失败: {data}")

        result = data.get("data", {}).get("result", {})
        job_status = result.get("job_status", -1)

        if job_status == 0:
            logger.info("导入成功: url=%s", result.get("url", ""))
            return result
        if job_status == 3:
            error_msg = result.get("job_error_msg", "未知错误")
            raise RuntimeError(f"导入失败: {error_msg}")

        time.sleep(_IMPORT_POLL_INTERVAL)

    raise RuntimeError(f"导入超时（等待 {_IMPORT_MAX_WAIT} 秒）")


def _set_permission(token: str, file_token: str, file_type: str = "sheet") -> None:
    """设置文档权限：任何人可编辑（最大权限）"""
    resp = httpx.patch(
        f"{FEISHU_BASE_URL}/drive/v1/permissions/{file_token}/public",
        headers={**_auth_headers(token), "Content-Type": "application/json"},
        params={"type": file_type},
        json={
            "external_access_entity": "open",  # 外部开放
            "link_share_entity": "anyone_editable",  # ✅ 修改：任何人可编辑
        },
        timeout=_HTTP_TIMEOUT,
    )

    data = resp.json()
    if data.get("code") != 0:
        logger.warning("设置权限失败（不影响主流程）: %s", data.get("msg", data))
    else:
        logger.info("文档权限已设置: anyone_editable（任何人获得链接可编辑 - 最大权限）")


def _send_link_message(chat_id: str, url: str, title: str) -> bool:
    """通过 IM 发送在线表格链接到群聊。

    文案：口语化，第二人称「您」，不用【】公文头。
    """
    try:
        now_label = datetime.now().strftime("%m-%d %H:%M")
        text = (
            f"这是 {now_label} 这批决策的详单，方便您对照查看：\n"
            f"{url}\n\n"
            f"审批请在聊天里直接回复我 ——\n"
            f"例：『通过』/『拒绝』/『广告 xxx 不动』/『只批准降价』"
        )
        _feishu.send_message(to=chat_id, text=text)
        logger.info("报告链接已发送到群聊: chat_id=%s", chat_id)
        return True
    except Exception as e:
        logger.warning("发送报告链接失败（不影响主流程）: %s", e)
        return False


# ═══════════════════════════════════════════
# 对外工具：import_to_feishu
# ═══════════════════════════════════════════

@tool(description="将本地 xlsx 报告导入为飞书在线表格，设置任何人获得链接可查看权限，并通过 IM 发送链接到运营群")
async def import_to_feishu(
    ctx: ToolContext,
    xlsx_path: str = "",
    send_im: bool = True,
    chat_id: str = "",
) -> ToolResult:
    """将 xlsx 文件导入飞书在线表格并分享

    完整流程：上传素材 → 导入为在线表格 → 设置权限 → IM 发送链接

    Args:
        xlsx_path: xlsx 文件路径。为空时自动使用 outputs/reports/ 下最新的 xlsx
        send_im: 是否通过 IM 发送链接（默认 True）
        chat_id: 目标群聊 ID。为空时使用配置中的 FEISHU_OPERATOR_CHAT_ID
    """
    try:
        # --- 1. 定位 xlsx 文件 ---
        if xlsx_path:
            file_path = Path(xlsx_path)
        else:
            # 自动找最新的 xlsx
            candidates = sorted(REPORTS_DIR.glob("*.xlsx"), reverse=True)
            if not candidates:
                return ToolResult(
                    title="未找到报告文件",
                    output=f"在 {REPORTS_DIR} 下未找到 xlsx 文件，请先运行 generate_report 生成报告",
                )
            file_path = candidates[0]

        if not file_path.exists():
            return ToolResult(
                title="文件不存在",
                output=f"文件不存在: {file_path}",
            )

        logger.info("开始导入飞书: %s (%d bytes)", file_path.name, file_path.stat().st_size)

        # --- 2. 获取 token ---
        token = _get_tenant_token()

        # --- 3. 上传素材 ---
        file_token = _upload_media(token, file_path)

        # --- 4. 创建导入任务 & 等待完成 ---
        ticket = _create_import_task(token, file_token, file_path.name)
        result = _wait_import_result(token, ticket)

        url = result.get("url", "")
        sheet_token = result.get("token", "")
        file_type = result.get("type", "sheet")

        # --- 5. 设置权限 ---
        if sheet_token:
            _set_permission(token, sheet_token, file_type)

        # --- 6. IM 发送链接 ---
        im_sent = False
        if send_im and url:
            target_chat = chat_id or FEISHU_OPERATOR_CHAT_ID
            if target_chat:
                title = file_path.stem
                im_sent = _send_link_message(target_chat, url, title)

        # --- 结果 ---
        output_lines = [
            f"文件: {file_path.name}",
            f"在线表格: {url}",
            f"表格 token: {sheet_token}",
            f"IM 发送: {'成功' if im_sent else '未发送' if not send_im else '失败'}",
        ]

        return ToolResult(
            title=f"飞书导入成功: {file_path.stem}",
            output="\n".join(output_lines),
            metadata={
                "url": url,
                "sheet_token": sheet_token,
                "file_type": file_type,
                "xlsx_path": str(file_path),
                "im_sent": im_sent,
            },
        )

    except Exception as e:
        logger.error("import_to_feishu 失败: %s", e, exc_info=True)
        return ToolResult(
            title="飞书导入失败",
            output=str(e),
        )