lightpdf-aipdf-mcp 0.1.31__py3-none-any.whl

lightpdf_aipdf_mcp/__init__.py

@@ -0,0 +1,8 @@
+from . import server
+
+def main():
+    """Main entry point for the package."""
+    server.cli_main()
+
+# 只导出main函数
+__all__ = ['main']

lightpdf_aipdf_mcp/common.py

@@ -0,0 +1,181 @@
+"""通用工具模块"""
+import asyncio
+import json
+import os
+import tempfile
+import time
+import urllib.parse
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any, Tuple
+
+import httpx
+from mcp.types import TextContent, LoggingMessageNotification, LoggingMessageNotificationParams, LoggingLevel
+
+@dataclass
+class BaseResult:
+    """基础结果数据类"""
+    success: bool
+    file_path: str
+    error_message: Optional[str] = None
+    download_url: Optional[str] = None
+
+class Logger:
+    """日志记录器类"""
+    def __init__(self, context, collect_info: bool = True):
+        self.context = context
+        self.collect_info = collect_info
+        self._info_log = []
+
+    async def log(self, level: str, message: str, add_to_result: bool = True):
+        """记录日志消息"""
+        if self.collect_info and add_to_result:
+            self._info_log.append(message)
+            
+        level_map = {
+            "debug": LoggingLevel.Debug,
+            "info": LoggingLevel.Info,
+            "warning": LoggingLevel.Warning,
+            "error": LoggingLevel.Error
+        }
+        
+        mcp_level = level_map.get(level.lower(), LoggingLevel.Info)
+        
+        # 使用更简洁的send_log_message方法
+        await self.context.session.send_log_message(mcp_level, message)
+    
+    async def error(self, message: str, error_class=RuntimeError):
+        """记录错误并引发异常"""
+        await self.log("error", message)
+        raise error_class(message)
+        
+    def get_result_info(self) -> List[str]:
+        """获取收集的信息日志"""
+        return self._info_log
+
+class FileHandler:
+    """文件处理工具类"""
+    def __init__(self, logger: Logger):
+        self.logger = logger
+
+    @staticmethod
+    def is_url(path: str) -> bool:
+        """检查路径是否为URL"""
+        return path.startswith(("http://", "https://"))
+
+    @staticmethod
+    def get_file_extension(file_path: str) -> str:
+        """获取文件扩展名（小写）"""
+        if "?" in file_path:  # 处理URL中的查询参数
+            file_path = file_path.split("?")[0]
+        return os.path.splitext(file_path)[1].lower()
+        
+    @staticmethod
+    def get_input_format(file_path: str):
+        """根据文件路径获取输入格式
+        
+        此方法需要导入InputFormat和INPUT_EXTENSIONS，
+        但为避免循环导入，由调用者提供转换逻辑
+        """
+        ext = FileHandler.get_file_extension(file_path)
+        return ext
+        
+    @staticmethod
+    def get_available_output_formats(input_format):
+        """获取指定输入格式支持的输出格式
+        
+        此方法需要导入FORMAT_CONVERSION_MAP，
+        但为避免循环导入，由调用者提供转换逻辑
+        """
+        # 实际实现在converter.py
+        return {}
+
+    async def validate_file_exists(self, file_path: str) -> Tuple[bool, bool]:
+        """验证文件是否存在
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            tuple: (文件是否存在, 是否为URL)
+        """
+        is_url = self.is_url(file_path)
+        if not is_url and not os.path.exists(file_path):
+            await self.logger.error(f"文件不存在：{file_path}", FileNotFoundError)
+            return False, is_url
+        return True, is_url
+
+class BaseApiClient:
+    """API客户端基类"""
+    def __init__(self, logger: Logger, file_handler: FileHandler):
+        self.logger = logger
+        self.file_handler = file_handler
+        self.api_key = os.getenv("API_KEY")
+        # 子类必须设置api_base_url
+        self.api_base_url = None
+
+    async def _wait_for_task(self, client: httpx.AsyncClient, task_id: str, operation_type: str = "处理") -> str:
+        """等待任务完成并返回下载链接
+        
+        Args:
+            client: HTTP客户端
+            task_id: 任务ID
+            operation_type: 操作类型描述，用于日志，默认为"处理"
+            
+        Returns:
+            str: 下载链接
+            
+        Raises:
+            RuntimeError: 如果任务失败或超时
+        """
+        headers = {"X-API-KEY": self.api_key}
+        MAX_ATTEMPTS = 100
+        
+        for attempt in range(MAX_ATTEMPTS):
+            await asyncio.sleep(3)
+            
+            status_response = await client.get(
+                f"{self.api_base_url}/{task_id}",
+                headers=headers
+            )
+            
+            if status_response.status_code != 200:
+                await self.logger.log("warning", f"获取任务状态失败。状态码: {status_response.status_code}")
+                continue
+            
+            status_result = status_response.json()
+            state = status_result.get("data", {}).get("state")
+            progress = status_result.get("data", {}).get("progress", 0)
+            
+            if state == 1:  # 完成
+                download_url = status_result.get("data", {}).get("file")
+                if not download_url:
+                    await self.logger.error(f"任务完成但未找到下载链接。任务状态：{json.dumps(status_result, ensure_ascii=False)}")
+                return download_url
+            elif state < 0:  # 失败
+                await self.logger.error(f"任务失败: {json.dumps(status_result, ensure_ascii=False)}")
+            else:  # 进行中
+                await self.logger.log("debug", f"{operation_type}进度: {progress}%", add_to_result=False)
+        
+        await self.logger.error(f"超过最大尝试次数（{MAX_ATTEMPTS}），任务未完成")
+
+    async def _handle_api_response(self, response: httpx.Response, error_prefix: str) -> str:
+        """处理API响应并提取任务ID
+        
+        Args:
+            response: API响应
+            error_prefix: 错误消息前缀
+            
+        Returns:
+            str: 任务ID
+            
+        Raises:
+            RuntimeError: 如果响应无效或任务ID缺失
+        """
+        if response.status_code != 200:
+            await self.logger.error(f"{error_prefix}失败。状态码: {response.status_code}\n响应: {response.text}")
+        
+        result = response.json()
+        if "data" not in result or "task_id" not in result["data"]:
+            await self.logger.error(f"无法获取任务ID。API响应：{json.dumps(result, ensure_ascii=False)}")
+        
+        return result["data"]["task_id"] 

lightpdf_aipdf_mcp/converter.py

@@ -0,0 +1,298 @@
+"""PDF文档转换模块"""
+import os
+import httpx
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional, Set
+
+from .common import BaseResult, Logger, FileHandler, BaseApiClient
+
+class InputFormat(str, Enum):
+    """支持的输入文件格式"""
+    PDF = "pdf"
+    WORD = "docx"
+    EXCEL = "xlsx"
+    PPT = "pptx"
+    JPG = "jpg"
+    JPEG = "jpeg"
+    PNG = "png"
+    BMP = "bmp"
+    GIF = "gif"
+    CAD = "dwg"
+    CAJ = "caj"
+    OFD = "ofd"
+
+class OutputFormat(str, Enum):
+    """支持的输出文件格式"""
+    PDF = "pdf"
+    WORD = "docx"
+    EXCEL = "xlsx"
+    PPT = "pptx"
+    JPG = "jpg"
+    JPEG = "jpeg"
+    PNG = "png"
+    BMP = "bmp"
+    GIF = "gif"
+    HTML = "html"
+    TEXT = "txt"
+
+# 文件扩展名到输入格式的映射
+INPUT_EXTENSIONS = {
+    ".pdf": InputFormat.PDF,
+    ".docx": InputFormat.WORD,
+    ".xlsx": InputFormat.EXCEL,
+    ".pptx": InputFormat.PPT,
+    ".jpg": InputFormat.JPG,
+    ".jpeg": InputFormat.JPEG,
+    ".png": InputFormat.PNG,
+    ".bmp": InputFormat.BMP,
+    ".gif": InputFormat.GIF,
+    ".dwg": InputFormat.CAD,
+    ".caj": InputFormat.CAJ,
+    ".ofd": InputFormat.OFD,
+}
+
+# 输入格式到可用输出格式的映射
+FORMAT_CONVERSION_MAP = {
+    InputFormat.PDF: {
+        OutputFormat.WORD,   # PDF转Word
+        OutputFormat.EXCEL,  # PDF转Excel
+        OutputFormat.PPT,    # PDF转PPT
+        OutputFormat.JPG,    # PDF转JPG
+        OutputFormat.JPEG,   # PDF转JPEG
+        OutputFormat.PNG,    # PDF转PNG
+        OutputFormat.BMP,    # PDF转BMP
+        OutputFormat.GIF,    # PDF转GIF
+        OutputFormat.HTML,   # PDF转HTML
+        OutputFormat.TEXT,   # PDF转文本
+    },
+    InputFormat.WORD: {OutputFormat.PDF},    # Word转PDF
+    InputFormat.EXCEL: {OutputFormat.PDF},   # Excel转PDF
+    InputFormat.PPT: {OutputFormat.PDF},     # PPT转PDF
+    InputFormat.JPG: {OutputFormat.PDF},     # JPG转PDF
+    InputFormat.JPEG: {OutputFormat.PDF},    # JPEG转PDF
+    InputFormat.PNG: {OutputFormat.PDF},     # PNG转PDF
+    InputFormat.BMP: {OutputFormat.PDF},     # BMP转PDF
+    InputFormat.GIF: {OutputFormat.PDF},     # GIF转PDF
+    InputFormat.CAD: {OutputFormat.PDF},     # CAD转PDF
+    InputFormat.CAJ: {OutputFormat.PDF},     # CAJ转PDF
+    InputFormat.OFD: {OutputFormat.PDF},     # OFD转PDF
+}
+
+# 扩展FileHandler类的方法
+def get_input_format(file_path: str) -> Optional[InputFormat]:
+    """根据文件路径获取输入格式"""
+    ext = FileHandler.get_file_extension(file_path)
+    return INPUT_EXTENSIONS.get(ext)
+
+def get_available_output_formats(input_format: InputFormat) -> Set[OutputFormat]:
+    """获取指定输入格式支持的输出格式"""
+    return FORMAT_CONVERSION_MAP.get(input_format, set())
+
+# 为FileHandler类注入方法
+FileHandler.get_input_format = staticmethod(get_input_format)
+FileHandler.get_available_output_formats = staticmethod(get_available_output_formats)
+
+@dataclass
+class ConversionResult(BaseResult):
+    """转换结果数据类"""
+    pass
+
+class Converter(BaseApiClient):
+    """PDF文档转换器"""
+    def __init__(self, logger: Logger, file_handler: FileHandler):
+        super().__init__(logger, file_handler)
+        self.api_base_url = "https://techsz.aoscdn.com/api/tasks/document/conversion"
+
+    async def add_page_numbers(self, file_path: str, start_num: int = 1, position: str = "5", margin: int = 30, password: str = None) -> ConversionResult:
+        """为PDF文档添加页码
+        
+        Args:
+            file_path: 要添加页码的PDF文件路径
+            start_num: 起始页码，整数类型，默认为1
+            position: 页码显示位置，字符串类型，可选值1-6（左上/上中/右上/左下/下中/右下），默认为5（下中）
+            margin: 页码显示的边距，整数类型，可选值10/30/60，默认为30
+            password: 文档密码，如果文档受密码保护，则需要提供（可选）
+            
+        Returns:
+            ConversionResult: 转换结果
+        """
+        # 验证输入文件是否为PDF
+        input_format = self.file_handler.get_input_format(file_path)
+        if input_format != InputFormat.PDF:
+            await self.logger.error(f"添加页码功能仅支持PDF文件，当前文件格式为: {self.file_handler.get_file_extension(file_path)}")
+        
+        # 验证参数
+        valid_positions = {"1", "2", "3", "4", "5", "6"}
+        valid_margins = {10, 30, 60}
+        
+        # 验证position参数
+        if position not in valid_positions:
+            await self.logger.error(f"无效的页码位置值: {position}。有效值为: 1(左上), 2(上中), 3(右上), 4(左下), 5(下中), 6(右下)")
+        
+        # 验证margin参数
+        if margin not in valid_margins:
+            await self.logger.error(f"无效的页码边距值: {margin}。有效值为: 10, 30, 60")
+        
+        # 验证start_num是否为正整数
+        if not isinstance(start_num, int) or start_num < 1:
+            await self.logger.error(f"起始页码必须是正整数，当前值为: {start_num}")
+        
+        # 构建API参数
+        extra_params = {
+            "start_num": start_num,
+            "position": position,
+            "margin": margin
+        }
+        
+        # 记录操作描述
+        await self.logger.log("info", f"正在为PDF添加页码（起始页码: {start_num}, 位置: {position}, 边距: {margin}）...")
+        
+        # 调用convert_file方法处理API请求
+        return await self.convert_file(file_path, "number-pdf", extra_params, password)
+        
+    async def convert_file(self, file_path: str, format: str, extra_params: dict = None, password: str = None) -> ConversionResult:
+        """转换单个文件
+        
+        Args:
+            file_path: 要转换的文件路径
+            format: 目标格式
+            extra_params: 额外的API参数，例如去除水印
+            password: 文档密码，如果文档受密码保护，则需要提供（可选）
+            
+        Returns:
+            ConversionResult: 转换结果
+        """
+        if not self.api_key:
+            await self.logger.error("未找到API_KEY。请在客户端配置API_KEY环境变量。")
+
+        # 特殊格式：doc-repair用于去除水印，number-pdf用于添加页码，输出均为PDF
+        is_special_operation = format in ["doc-repair", "number-pdf"]
+        actual_output_format = "pdf" if is_special_operation else format
+
+        # 验证输入文件格式
+        input_format = self.file_handler.get_input_format(file_path)
+        if not input_format and not is_special_operation:
+            await self.logger.error(f"不支持的输入文件格式: {self.file_handler.get_file_extension(file_path)}")
+
+        # 如果是去除水印操作，检查是否PDF文件
+        if format == "doc-repair" and input_format != InputFormat.PDF:
+            await self.logger.error("去除水印功能仅支持PDF文件")
+            
+        # 如果是添加页码操作，检查是否PDF文件
+        if format == "number-pdf" and input_format != InputFormat.PDF:
+            await self.logger.error("添加页码功能仅支持PDF文件")
+
+        # 验证输出格式（除去特殊操作外）
+        if not is_special_operation:
+            try:
+                output_format = OutputFormat(format)
+            except ValueError:
+                await self.logger.error(f"不支持的输出格式: {format}")
+
+            # 验证格式转换是否支持
+            if input_format:  # 确保input_format有效
+                available_formats = self.file_handler.get_available_output_formats(input_format)
+                if output_format not in available_formats:
+                    await self.logger.error(
+                        f"不支持从 {input_format.value} 格式转换为 {output_format.value} 格式。"
+                        f"支持的输出格式: {', '.join(f.value for f in available_formats)}"
+                    )
+        else:
+            # 对于特殊操作，设置输出格式为PDF
+            output_format = OutputFormat.PDF
+
+        # 验证文件
+        file_exists_result = await self.file_handler.validate_file_exists(file_path)
+        if not file_exists_result[0]:
+            return ConversionResult(success=False, file_path=file_path, error_message="文件不存在")
+        is_url = file_exists_result[1]
+
+        # 操作描述
+        if format == "doc-repair":
+            operation_desc = "去除水印"
+        elif format == "number-pdf":
+            operation_desc = "添加页码"
+        else:
+            operation_desc = f"将 {input_format.value.upper()} 转换为 {output_format.value.upper()} 格式"
+        await self.logger.log("info", f"正在{operation_desc}...")
+
+        import httpx
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            try:
+                # 初始化extra_params（如果为None）
+                if extra_params is None:
+                    extra_params = {}
+                
+                # 如果提供了密码，将其添加到extra_params
+                if password:
+                    extra_params["password"] = password
+                
+                # 创建转换任务
+                task_id = await self._create_task(client, file_path, format, is_url, extra_params)
+                
+                # 等待任务完成
+                download_url = await self._wait_for_task(client, task_id, "转换")
+                
+                # 记录完成信息
+                await self.logger.log("info", "转换完成。可通过下载链接获取结果文件。")
+                
+                return ConversionResult(
+                    success=True,
+                    file_path=file_path,
+                    error_message=None,
+                    download_url=download_url
+                )
+
+            except Exception as e:
+                return ConversionResult(
+                    success=False,
+                    file_path=file_path,
+                    error_message=str(e),
+                    download_url=None
+                )
+
+    async def _create_task(self, client: httpx.AsyncClient, file_path: str, format: str, is_url: bool, extra_params: dict = None) -> str:
+        """创建转换任务
+        
+        Args:
+            client: HTTP客户端
+            file_path: 文件路径
+            format: 目标格式，特殊格式"doc-repair"用于去除水印，"number-pdf"用于添加页码
+            is_url: 是否URL路径
+            extra_params: 额外API参数(可选)
+        
+        Returns:
+            str: 任务ID
+        """
+        await self.logger.log("info", "正在提交转换任务...")
+        
+        headers = {"X-API-KEY": self.api_key}
+        data = {"format": format}
+        
+        # 添加额外参数
+        if extra_params:
+            data.update(extra_params)
+        
+        if is_url:
+            # 使用JSON方式时添加Content-Type
+            headers["Content-Type"] = "application/json"
+            response = await client.post(
+                self.api_base_url,
+                json=data,
+                headers=headers
+            )
+        else:
+            # 对于文件上传，使用表单方式，不需要添加Content-Type
+            with open(file_path, "rb") as f:
+                files = {"file": f}
+                response = await client.post(
+                    self.api_base_url,
+                    files=files,
+                    data=data,
+                    headers=headers
+                )
+        
+        # 使用基类的方法处理API响应
+        return await self._handle_api_response(response, "创建任务")