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 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:
        text = f"**广告决策报告: {title}**\n\n报告已生成，点击查看: [打开在线表格]({url})"
        _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),
        )