firecrawl爬取代码备份

from pydantic import BaseModel, Field
import aiohttp
import asyncio
from typing import Callable, Awaitable, Any, Optional
import re


class Pipe:
    class Valves(BaseModel):
        firecrawl_api_token: str = Field(
            default="",
            description="Firecrawl API 的 Token",
        )

    def __init__(self):
        self.valves = self.Valves()

    async def pipe(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        user: Optional[dict] = None,
    ) -> str:
        # 1. 尝试从用户输入中提取 URL（优先级最高）
        user_provided_url = self._extract_user_input_url(body)
        if user_provided_url and self._is_valid_url(user_provided_url):
            method = "URL (用户输入)"
            firecrawl_api_url = "https://api.firecrawl.dev/v1/scrape"
            payload = {
                "formats": ["markdown"],
                "url": user_provided_url,
                "onlyMainContent": True,
                "removeBase64Images": True,
                "timeout": 70000,
            }
        else:
            return "未提供有效的 URL。请提供一个有效的网站链接。"

        # 2. 发送状态更新，告知用户正在进行爬取处理
        await __event_emitter__(
            {
                "type": "status",
                "data": {
                    "description": f"🔍 正在爬取 {method} 方式的网站内容，请稍候...",
                    "done": False,
                },
            }
        )

        try:
            headers = {
                "Authorization": f"Bearer {self.valves.firecrawl_api_token}",
                "Content-Type": "application/json",
            }
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    firecrawl_api_url, headers=headers, json=payload
                ) as response:
                    response.raise_for_status()  # 如果响应状态码是 4xx 或 5xx，将引发异常
                    firecrawl_result = await response.json()

            # 3. 处理完成，发送完成的状态更新
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": "✅ 网站内容爬取完成！",
                        "done": True,
                    },
                }
            )
            return self._format_firecrawl_result(firecrawl_result)

        except Exception as e:
            # 4. 发生错误，发送错误的状态更新
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": f"❌ Firecrawl API 请求失败：{e}",
                        "done": True,
                    },
                }
            )
            return f"Firecrawl API 请求失败：{e}"

    def _extract_user_input_url(self, body: dict) -> Optional[str]:
        """
        从用户输入中提取有效的 URL，仅检查最新的消息。
        """
        messages = body.get("messages", [])
        if not messages:
            return None

        # 获取最新的一条消息
        last_message = messages[-1]
        content = last_message.get("content", "")

        if isinstance(content, str):
            # 查找所有可能的 URL
            urls = self._find_urls_in_text(content)
            if urls:
                # 返回第一个有效的 URL
                return urls[0]

        return None

    def _find_urls_in_text(self, text: str) -> list:
        """
        使用正则表达式在文本中查找所有 URL。
        """
        # 简单的 URL 正则表达式
        url_pattern = re.compile(r"(https?://[^\s]+)", re.IGNORECASE)
        return url_pattern.findall(text)

    def _is_valid_url(self, url: str) -> bool:
        """
        检查 URL 是否有效。
        """
        # 简单的 URL 验证，可以根据需要扩展
        regex = re.compile(
            r"^(https?://)"
            r"("
            r"([A-Za-z0-9-]+\.)+[A-Za-z]{2,6}"  # 域名
            r"|"  # 或者
            r"localhost"  # 本地地址
            r")"
            r"(:\d+)?"  # 可选的端口
            r"(\/\S*)?$"  # 可选的路径
        )
        return re.match(regex, url) is not None

    def _format_firecrawl_result(self, firecrawl_result: dict) -> str:
        """
        格式化 Firecrawl API 的响应结果，仅返回 markdown 内容。
        """
        if firecrawl_result.get("success"):
            markdown_content = firecrawl_result.get("data", {}).get(
                "markdown", "未识别到 Markdown 内容。"
            )
            markdown_content = markdown_content.replace("\\", "")  # 去除多余的反斜杠
            markdown_content = markdown_content.encode("unicode_escape").decode(
                "unicode_escape"
            )  # 去除转义字符
            return f"{markdown_content}"
        else:
            error_message = firecrawl_result.get("error", "未知错误")
            return f"Firecrawl API 返回错误：{error_message}"