huoyan-enterprise/backend/services/excel_mcp_context.py

"""
为 MCP Excel 分析工具生成系统上下文（含可下载的 file_url）。
"""
from __future__ import annotations

from typing import Any, List, Optional

import asyncpg

from core.config import get_settings
from core.mcp_client import is_mcp_enabled
from logger.logging import get_logger
from services.oss_service import get_oss_service

logger = get_logger(__name__)

EXCEL_SUFFIXES = (".xlsx", ".xls", ".csv")


def _is_excel_file_name(file_name: str) -> bool:
    lower = (file_name or "").lower()
    return any(lower.endswith(ext) for ext in EXCEL_SUFFIXES)


def _resolve_download_url(
    file_path: str,
    *,
    kb_id: Optional[int] = None,
    thread_id: Optional[str] = None,
    expires: int = 3600,
) -> str:
    """将 OSS 存储路径转为 MCP 容器可下载的 URL（优先签名 URL）。"""
    url = (file_path or "").strip()
    if not url:
        return url

    oss = get_oss_service()
    if not oss.enabled or not url.startswith(("http://", "https://")):
        return url

    object_name = oss.extract_object_name_from_url(url, kb_id=kb_id, thread_id=thread_id)
    if not object_name:
        return url

    signed = oss.get_signed_url(object_name, expires=expires)
    return signed or url


async def _fetch_thread_excel_files(
    conn: asyncpg.Connection,
    thread_id: str,
    limit: int = 10,
) -> List[dict[str, Any]]:
    rows = await conn.fetch(
        """
        SELECT id, file_name, file_path
        FROM chat_thread_file
        WHERE thread_id = $1 AND is_deleted = FALSE AND status = 'completed'
        ORDER BY created_at DESC
        LIMIT $2
        """,
        thread_id,
        limit,
    )
    items: List[dict[str, Any]] = []
    for row in rows:
        if not _is_excel_file_name(row["file_name"]):
            continue
        items.append(
            {
                "source": "chat",
                "file_id": row["id"],
                "file_name": row["file_name"],
                "file_url": _resolve_download_url(
                    row["file_path"], thread_id=thread_id
                ),
            }
        )
    return items


async def _fetch_kb_excel_files(
    conn: asyncpg.Connection,
    knowledge_base_id: int,
    limit: int = 10,
) -> List[dict[str, Any]]:
    rows = await conn.fetch(
        """
        SELECT id, file_name, file_path
        FROM knowledge_base_file
        WHERE knowledge_base_id = $1 AND is_deleted = FALSE AND status = 'completed'
        ORDER BY created_at DESC
        LIMIT $2
        """,
        knowledge_base_id,
        limit,
    )
    items: List[dict[str, Any]] = []
    for row in rows:
        if not _is_excel_file_name(row["file_name"]):
            continue
        items.append(
            {
                "source": "knowledge_base",
                "file_id": row["id"],
                "file_name": row["file_name"],
                "file_url": _resolve_download_url(
                    row["file_path"], kb_id=knowledge_base_id
                ),
            }
        )
    return items


def format_excel_mcp_system_context(files: List[dict[str, Any]]) -> str:
    """生成注入 Agent 系统提示的 Excel MCP 使用说明。"""
    if not files:
        return ""

    lines = [
        "\n" + "=" * 60,
        "📊 **Excel/CSV 表格分析（MCP 工具，Docker 沙箱执行）**",
        "=" * 60,
        "用户会话中存在表格文件。涉及单元格、筛选、汇总、排序、统计等问题时：",
        "1. 先调用 MCP 工具 **excel_get_schema**(file_url, sheet_name?) 查看 sheet/列名/预览",
        "2. 再调用 **excel_run_pandas**(file_url, pandas_code, sheet_name?) 执行 pandas",
        "   - 代码中只能使用 `pd` 与 `df`，禁止 import",
        "   - 必须将最终答案赋给 `result`，例如: `result = df['销售额'].sum()`",
        "3. 根据工具返回的 JSON 结果用中文回答用户",
        "",
        "**可用表格文件（file_url 传给 MCP 工具）**：",
    ]
    for idx, f in enumerate(files, 1):
        src = "聊天上传" if f.get("source") == "chat" else "知识库"
        lines.append(
            f"{idx}. [{src}] `{f['file_name']}` (file_id={f['file_id']})\n"
            f"   file_url: {f['file_url']}"
        )
    lines.append("=" * 60)
    return "\n".join(lines)


async def build_excel_mcp_system_context(
    pool,
    thread_id: str,
    knowledge_base_id: Optional[int] = None,
) -> str:
    """
    汇总当前会话聊天文件 + 绑定知识库中的 Excel/CSV，生成 MCP 使用说明。
    未配置任何 MCP 服务时返回空字符串。
    """
    if not is_mcp_enabled():
        return ""

    files: List[dict[str, Any]] = []
    async with pool.acquire() as conn:
        files.extend(await _fetch_thread_excel_files(conn, thread_id))
        if knowledge_base_id:
            files.extend(await _fetch_kb_excel_files(conn, knowledge_base_id))

    if not files:
        return ""

    logger.info(f"Excel MCP 上下文: {len(files)} 个表格文件")
    return format_excel_mcp_system_context(files)