ForcePilot/backend/package/yuxi/channel/extensions/feishu/tools/doc.py

from __future__ import annotations

import logging

from yuxi.channel.extensions.feishu.client import get_client

logger = logging.getLogger(__name__)

FEISHU_DOC_TOOLS = [
    {
        "name": "feishu_doc_read",
        "description": "读取飞书文档内容",
        "parameters": {"doc_token": "文档 token"},
    },
    {
        "name": "feishu_doc_create",
        "description": "创建飞书文档",
        "parameters": {"title": "文档标题", "folder_token": "文件夹 token（可选）"},
    },
    {
        "name": "feishu_doc_append",
        "description": "追加内容到飞书文档",
        "parameters": {"doc_token": "文档 token", "content": "要追加的内容"},
    },
    {
        "name": "feishu_doc_list_blocks",
        "description": "列出文档的所有块",
        "parameters": {"doc_token": "文档 token"},
    },
    {
        "name": "feishu_doc_export",
        "description": "导出飞书文档为 PDF 或其他格式",
        "parameters": {"doc_token": "文档 token", "export_format": "导出格式（pdf/docx，默认 pdf）"},
    },
]


async def read_doc(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.GetDocumentRequest(document_id=doc_token)
        resp = await client.docx.v1.document.get_async(req)

        if resp.success() and resp.data:
            doc = resp.data.document
            return {
                "success": True,
                "title": getattr(doc, "title", ""),
                "document_id": getattr(doc, "document_id", ""),
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu read_doc error")
        return {"success": False, "error": "Internal error"}


async def write_doc(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    content: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        blocks = [
            lark_oapi.docx.v1.Block(
                block_type=2,
                text=lark_oapi.docx.v1.Text(
                    elements=[
                        lark_oapi.docx.v1.TextElement(
                            text_run=lark_oapi.docx.v1.TextRun(content=content),
                        )
                    ]
                ),
            )
        ]
        req = lark_oapi.docx.v1.UpdateDocumentRequest(
            document_id=doc_token,
            request_body=lark_oapi.docx.v1.UpdateDocumentRequestBody(
                document_id=doc_token,
                blocks=blocks,
            ),
        )
        resp = await client.docx.v1.document.update_async(req)

        if resp.success():
            return {"success": True}
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu write_doc error")
        return {"success": False, "error": "Internal error"}


async def append_doc(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    content: str,
    http_timeout_ms: int = 30000,
) -> dict:
    return await write_doc(app_id, app_secret, domain, doc_token, content, http_timeout_ms)


async def create_doc(
    app_id: str,
    app_secret: str,
    domain: str,
    title: str,
    folder_token: str = "",
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.CreateDocumentRequest(
            request_body=lark_oapi.docx.v1.CreateDocumentRequestBody(
                title=title,
                folder_token=folder_token if folder_token else None,
            ),
        )
        resp = await client.docx.v1.document.create_async(req)

        if resp.success() and resp.data:
            return {
                "success": True,
                "document_id": getattr(resp.data.document, "document_id", ""),
                "url": getattr(resp.data.document, "url", ""),
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu create_doc error")
        return {"success": False, "error": "Internal error"}


async def list_doc_blocks(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.ListDocumentBlockRequest(
            document_id=doc_token,
            page_size=500,
        )
        resp = await client.docx.v1.document_block.list_async(req)

        if resp.success() and resp.data:
            items = getattr(resp.data, "items", []) or []
            return {
                "success": True,
                "blocks": [
                    {
                        "block_id": getattr(b, "block_id", ""),
                        "block_type": getattr(b, "block_type", 0),
                        "text": getattr(getattr(b, "text", None), "elements", []),
                    }
                    for b in items
                ],
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu list_doc_blocks error")
        return {"success": False, "error": "Internal error"}


async def get_doc_block(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    block_id: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.GetDocumentBlockRequest(
            document_id=doc_token,
            block_id=block_id,
        )
        resp = await client.docx.v1.document_block.get_async(req)

        if resp.success() and resp.data:
            block = resp.data.block
            return {
                "success": True,
                "block": {
                    "block_id": getattr(block, "block_id", ""),
                    "block_type": getattr(block, "block_type", 0),
                },
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu get_doc_block error")
        return {"success": False, "error": "Internal error"}


async def update_doc_block(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    block_id: str,
    content: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.UpdateDocumentBlockRequest(
            document_id=doc_token,
            block_id=block_id,
            request_body=lark_oapi.docx.v1.UpdateDocumentBlockRequestBody(
                update_text_elements={
                    "elements": [
                        {"text_run": {"content": content}}
                    ]
                },
            ),
        )
        resp = await client.docx.v1.document_block.update_async(req)

        if resp.success():
            return {"success": True}
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu update_doc_block error")
        return {"success": False, "error": "Internal error"}


async def delete_doc_block(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    block_id: str,
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.DeleteDocumentBlockRequest(
            document_id=doc_token,
            block_id=block_id,
        )
        resp = await client.docx.v1.document_block.delete_async(req)

        if resp.success():
            return {"success": True}
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu delete_doc_block error")
        return {"success": False, "error": "Internal error"}


async def create_doc_table(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    row_size: int,
    column_size: int,
    http_timeout_ms: int = 30000,
) -> dict:
    header = "| " + " | ".join(f"列{i + 1}" for i in range(column_size)) + " |\n"
    sep = "| " + " | ".join("---" for _ in range(column_size)) + " |\n"
    rows = "".join("| " + " | ".join("" for _ in range(column_size)) + " |\n" for _ in range(row_size))
    table_md = header + sep + rows
    return await write_doc(app_id, app_secret, domain, doc_token, table_md, http_timeout_ms)


async def write_table_cells(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    table_block_id: str,
    values: list[list[str]],
    http_timeout_ms: int = 30000,
) -> dict:
    text = "".join("| " + " | ".join(row) + " |\n" for row in values)
    return await update_doc_block(app_id, app_secret, domain, doc_token, table_block_id, text, http_timeout_ms)


async def export_doc(
    app_id: str,
    app_secret: str,
    domain: str,
    doc_token: str,
    export_format: str = "pdf",
    http_timeout_ms: int = 30000,
) -> dict:
    try:
        client = get_client(app_id, app_secret, domain, http_timeout_ms)

        import lark_oapi
        req = lark_oapi.docx.v1.GetDocumentExportRequest(
            document_id=doc_token,
            export_format=export_format,
            export_type="download",
        )
        resp = await client.docx.v1.document.export_async(req)

        if resp.success():
            return {
                "success": True,
                "document_id": doc_token,
                "format": export_format,
            }
        else:
            return {"success": False, "error": getattr(resp, "msg", "")}
    except ImportError:
        return {"success": False, "error": "lark-oapi SDK not installed"}
    except Exception:
        logger.exception("Feishu export_doc error")
        return {"success": False, "error": "Internal error"}