ForcePilot/backend/package/yuxi/agents/toolkits/buildin/crawl/tools.py

"""Web 抓取 buildin 工具注册：crawl_url、crawl_site、crawl_to_kb。

工具注册通过 @tool 装饰器完成，模块加载即注册。
- crawl_url：单页抓取，返回 Markdown
- crawl_site：站点递归抓取，异步任务，立即返回 task_id
- crawl_to_kb：将爬取内容入库到知识库
"""
import json

from langgraph.prebuilt.tool_node import ToolRuntime

from yuxi.agents.toolkits.registry import tool
from yuxi.config import config
from yuxi.services.task_service import TaskContext, tasker
from yuxi.utils.logging_config import logger
from yuxi.utils.paths import VIRTUAL_PATH_OUTPUTS

from .client import crawl4ai_client
from .ingester import check_kb_permission, ingest_crawl_result, resolve_user_dict
from .rate_limiter import crawl_rate_limiter
from .security import should_respect_robots, validate_crawl_url


CRAWL_URL_DESCRIPTION = """
抓取单个网页 URL，返回干净的 Markdown 内容。

使用场景：
1. 从 web_search 结果中选取感兴趣的 URL，抓取全文
2. 用户指定某个网页 URL，需要获取完整内容
3. 需要将网页内容入库到知识库前的抓取步骤

返回结果：
- markdown：网页正文 Markdown（已剥离导航栏、广告、页脚等噪声）
- url：最终 URL（可能经过重定向）
- title：页面标题
- status_code：HTTP 状态码
- crawled_at：抓取时间

注意事项：
- 结果长度限制 50000 字符（约 12k token），超出截断并提示
- 遵守 robots.txt，被禁止的 URL 返回错误
- SSRF 防护：内网地址、私有 IP 段会被拒绝
"""


CRAWL_SITE_DESCRIPTION = """
递归抓取一个站点的多个页面，返回每页 Markdown 摘要。

使用场景：
1. 把一个技术文档站（如 FastAPI 官网）整站入库
2. 批量抓取同域下的多个页面

输入参数：
- start_url：起始 URL
- max_depth：递归深度（1-3，默认 2）
- max_pages：最大抓取页数（1-100，默认 20）
- url_pattern：URL 通配符过滤（如 /docs/*）

返回结果：
- 摘要：成功页数、失败页数、总字符数
- 每页含 URL + Markdown 摘要（前 500 字符）
- 完整内容存到临时文件，通过 file_path 字段引用，可按需读取

执行方式：
- 异步任务，立即返回 task_id
- 通过 GET /api/tasks/{task_id} 查询进度
- 通过 POST /api/tasks/{task_id}/cancel 取消任务
"""


CRAWL_TO_KB_DESCRIPTION = """
将爬取的 Markdown 内容入库到指定知识库。

使用场景：
1. agent 已通过 crawl_url 抓取到全文，用户确认入库到某个知识库
2. 需要将网页内容持久化到知识库供后续检索

输入参数：
- content：Markdown 内容（通常来自 crawl_url 的返回）
- kb_id：目标知识库 ID
- source_url：来源 URL，记入文件元数据
- filename：可选，文件名，默认从 URL 推导

返回结果：
- file_id：入库后的文件 ID
- kb_id：知识库 ID
- status：入库状态（indexed 表示成功）

注意事项：
- 用户必须有目标知识库的写入权限
- 基于 content_hash 去重，重复内容返回提示
- 入库后自动触发解析与索引
"""


def _get_rate_limit_key(runtime: ToolRuntime) -> str:
    """获取限流 key：thread_id + uid。

    BaseContext 仅暴露 thread_id 与 uid（无 agent_id 字段），
    故限流维度为 (thread_id, uid)。
    """
    ctx = getattr(runtime, "context", None)
    if ctx is None:
        return "::"
    thread_id = getattr(ctx, "thread_id", "") or ""
    uid = getattr(ctx, "uid", "") or ""
    return f"{thread_id}:{uid}"


def _extract_runtime_identity(runtime: ToolRuntime) -> tuple[str, str]:
    """从 runtime 提取 (thread_id, uid)，供闭包捕获。

    必须在工具调用同步阶段提取，不能在异步任务协程内访问 runtime——
    工具调用返回后 runtime 可能失效。
    """
    ctx = getattr(runtime, "context", None)
    if ctx is None:
        return "", ""
    thread_id = getattr(ctx, "thread_id", "") or ""
    uid = getattr(ctx, "uid", "") or ""
    return thread_id, uid


@tool(
    category="buildin",
    tags=["抓取"],
    display_name="网页抓取",
    description=CRAWL_URL_DESCRIPTION,
)
async def crawl_url(
    url: str,
    output_format: str = "markdown",
    runtime: ToolRuntime = None,
) -> str:
    """抓取单个 URL，返回干净 Markdown。"""
    # 限流检查
    key = _get_rate_limit_key(runtime)
    if not await crawl_rate_limiter.check(key):
        return json.dumps(
            {"error": "抓取调用过于频繁，请稍后再试", "url": url},
            ensure_ascii=False,
        )

    # SSRF 校验
    ok, reason = await validate_crawl_url(url)
    if not ok:
        return json.dumps(
            {"error": f"URL 被 SSRF 防护拒绝: {reason}", "url": url},
            ensure_ascii=False,
        )

    respect_robots = should_respect_robots()
    try:
        result = await crawl4ai_client.crawl_single(
            url,
            output_format=output_format,
            check_robots_txt=respect_robots,
        )
    except Exception as exc:
        logger.warning(f"crawl_url_failed url={url} error={exc}")
        return json.dumps(
            {"error": f"抓取失败: {exc}", "url": url},
            ensure_ascii=False,
        )

    if not result.get("success", False):
        error_msg = result.get("error_message", "未知错误")
        status_code = result.get("status_code")
        return json.dumps(
            {
                "error": f"抓取失败: {error_msg}",
                "url": url,
                "status_code": status_code,
            },
            ensure_ascii=False,
        )

    markdown = result.get("markdown", "")
    # 截断到 50000 字符
    truncated = False
    if len(markdown) > 50000:
        markdown = markdown[:50000]
        truncated = True

    return json.dumps(
        {
            "markdown": markdown,
            "url": result.get("url", url),
            "title": result.get("metadata", {}).get("title", ""),
            "status_code": result.get("status_code"),
            "crawled_at": result.get("metadata", {}).get("crawl_date", ""),
            "truncated": truncated,
        },
        ensure_ascii=False,
    )


@tool(
    category="buildin",
    tags=["抓取"],
    display_name="站点递归抓取",
    description=CRAWL_SITE_DESCRIPTION,
)
async def crawl_site(
    start_url: str,
    max_depth: int = 2,
    max_pages: int = 20,
    url_pattern: str = "",
    runtime: ToolRuntime = None,
) -> str:
    """递归抓取站点，异步执行，立即返回 task_id。"""
    # 限流检查
    key = _get_rate_limit_key(runtime)
    if not await crawl_rate_limiter.check(key):
        return json.dumps(
            {"error": "抓取调用过于频繁，请稍后再试", "url": start_url},
            ensure_ascii=False,
        )

    # 参数 clamp
    max_depth = max(1, min(3, max_depth))
    max_pages = max(1, min(100, max_pages))

    # SSRF 校验起始 URL
    ok, reason = await validate_crawl_url(start_url)
    if not ok:
        return json.dumps(
            {"error": f"起始 URL 被 SSRF 防护拒绝: {reason}", "url": start_url},
            ensure_ascii=False,
        )

    # 在工具调用同步阶段提取 (thread_id, uid)，供闭包捕获。
    # 不能在异步任务协程内访问 runtime——工具调用返回后 runtime 可能失效。
    thread_id, uid = _extract_runtime_identity(runtime)

    # 闭包捕获参数，TaskContext 仅持有 task_id，不持有 Task 引用。
    # 闭包内不访问 runtime，仅使用已提取的 thread_id / uid。
    async def _run_crawl_site(context: TaskContext) -> dict:
        """crawl_site 异步任务协程，通过闭包捕获 start_url 等参数。"""
        respect_robots = should_respect_robots()

        await context.set_progress(5.0, f"开始抓取 {start_url}")
        await context.raise_if_cancelled()

        try:
            results = await crawl4ai_client.crawl_many(
                [start_url],
                max_depth=max_depth,
                max_pages=max_pages,
                url_pattern=url_pattern,
                check_robots_txt=respect_robots,
                semaphore_count=config.web_crawl_default_concurrency,
            )
        except Exception as exc:
            await context.set_result({"error": str(exc), "start_url": start_url})
            raise

        await context.set_progress(80.0, f"抓取完成，共 {len(results)} 页")
        await context.raise_if_cancelled()

        # 聚合结果：每页含 URL + 摘要（前 500 字符）
        pages = []
        full_content_parts = []
        total_chars = 0
        success_count = 0
        failed_count = 0
        for item in results:
            if item.get("success", False):
                markdown = item.get("markdown", "")
                total_chars += len(markdown)
                success_count += 1
                pages.append(
                    {
                        "url": item.get("url", ""),
                        "title": item.get("metadata", {}).get("title", ""),
                        "snippet": markdown[:500],
                        "status_code": item.get("status_code"),
                    }
                )
                full_content_parts.append(f"# {item.get('url', '')}\n\n{markdown}\n\n---\n")
            else:
                failed_count += 1
                pages.append(
                    {
                        "url": item.get("url", ""),
                        "error": item.get("error_message", "未知错误"),
                        "status_code": item.get("status_code"),
                    }
                )

        # 完整内容存到临时文件，agent 可按需读取
        file_path = ""
        if full_content_parts and thread_id and uid:
            from yuxi.agents.backends.sandbox.paths import (
                ensure_thread_dirs,
                sandbox_outputs_dir,
            )

            ensure_thread_dirs(thread_id, uid)
            outputs_dir = sandbox_outputs_dir(thread_id)
            file_name = f"crawl_site_{start_url.replace('/', '_').replace(':', '_')[:50]}_{success_count}pages.md"
            full_path = outputs_dir / file_name
            full_path.write_text("\n".join(full_content_parts), encoding="utf-8")
            file_path = f"{VIRTUAL_PATH_OUTPUTS}/{file_name}"

        summary = {
            "start_url": start_url,
            "success_count": success_count,
            "failed_count": failed_count,
            "total_chars": total_chars,
            "pages": pages,
            "file_path": file_path,  # 完整内容文件路径，agent 可按需读取
        }
        await context.set_result(summary)
        await context.set_progress(100.0, f"抓取完成：成功 {success_count} 页，失败 {failed_count} 页")
        return summary

    # 异步任务去重：基于 start_url 去重
    task, created = await tasker.enqueue_unique_by_payload(
        name=f"站点抓取 ({start_url})",
        task_type="web_crawl_site",
        payload={
            "start_url": start_url,
            "max_depth": max_depth,
            "max_pages": max_pages,
            "url_pattern": url_pattern,
        },
        coroutine=_run_crawl_site,
        payload_match={"start_url": start_url},
        statuses={"pending", "running"},
    )

    if not created:
        return json.dumps(
            {
                "message": "已有相同起始 URL 的抓取任务正在执行",
                "task_id": task.id,
                "status": task.status,
            },
            ensure_ascii=False,
        )

    return json.dumps(
        {
            "message": "抓取任务已提交",
            "task_id": task.id,
            "status": task.status,
            "progress_query": f"GET /api/tasks/{task.id}",
            "cancel_query": f"POST /api/tasks/{task.id}/cancel",
        },
        ensure_ascii=False,
    )


@tool(
    category="buildin",
    tags=["抓取", "知识库"],
    display_name="爬取内容入库",
    description=CRAWL_TO_KB_DESCRIPTION,
)
async def crawl_to_kb(
    content: str,
    kb_id: str,
    source_url: str,
    filename: str | None = None,
    runtime: ToolRuntime = None,
) -> str:
    """将爬取的 Markdown 内容入库到指定知识库。"""
    if not content or not content.strip():
        return json.dumps(
            {"error": "内容不能为空", "kb_id": kb_id},
            ensure_ascii=False,
        )

    # 权限校验：BaseContext 仅暴露 uid，需查询 UserRepository 补全 role/department_id，
    # 复用 KnowledgeBaseManager.check_accessible(user, kb_id)。
    _, uid = _extract_runtime_identity(runtime)
    user = await resolve_user_dict(uid)
    has_permission = await check_kb_permission(kb_id, user)
    if not has_permission:
        return json.dumps(
            {"error": f"无知识库 {kb_id} 的写入权限", "kb_id": kb_id},
            ensure_ascii=False,
        )

    try:
        result = await ingest_crawl_result(
            content=content,
            kb_id=kb_id,
            source_url=source_url,
            filename=filename,
            operator_id=uid,
        )
        return json.dumps(result, ensure_ascii=False)
    except ValueError as exc:
        return json.dumps(
            {"error": str(exc), "kb_id": kb_id, "source_url": source_url},
            ensure_ascii=False,
        )
    except Exception as exc:
        logger.error(f"crawl_to_kb_failed kb_id={kb_id} source_url={source_url} error={exc}")
        return json.dumps(
            {"error": f"入库失败: {exc}", "kb_id": kb_id, "source_url": source_url},
            ensure_ascii=False,
        )