增加了文件上传和url的支持

api.py

@@ -1,8 +1,8 @@
 import asyncio
 import os
-import shutil
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional
+import urllib.parse
 
 import aiofiles
 import aiohttp
@@ -10,7 +10,7 @@ import uvicorn
 from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.responses import StreamingResponse
 from loguru import logger
-from models import Task, TaskCreateRequest, TaskStatus
+from models import Task, TaskStatus
 from processor import ProcessManager
 
 
@@ -18,7 +18,9 @@ from processor import ProcessManager
 class TaskManager:
     def __init__(self):
         self.tasks: Dict[str, Task] = {}
-        self.process_manager = ProcessManager(config_path="config/machines.json", task_manager=self)
+        self.process_manager = ProcessManager(
+            config_path="config/machines.json", task_manager=self
+        )
         self.callbacks: Dict[str, List[Callable[[Task], Any]]] = {}
         self._task_lock = asyncio.Lock()
         self._callback_lock = asyncio.Lock()
@@ -39,12 +41,19 @@ class TaskManager:
             logger.info("任务管理器已停止")
 
     async def add_task(
-        self, audio_file_path: str, video_file_path: str, output_dir: str, description: str, callback_url: Optional[str] = None
+        self,
+        audio_file_path: str,
+        video_file_path: str,
+        output_dir: str,
+        description: str,
+        callback_url: Optional[str] = None,
     ) -> Task:
         """添加新任务"""
         async with self._task_lock:
             task_id = f"task_{len(self.tasks) + 1}"
-            output_file_path = os.path.join(output_dir, f"{task_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4")
+            output_file_path = os.path.join(
+                output_dir, f"{task_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
+            )
             task = Task(
                 task_id=task_id,
                 audio_file_path=audio_file_path,
@@ -60,7 +69,9 @@ class TaskManager:
 
             # 如果有回调URL，添加回调
             if callback_url:
-                await self.add_callback(task_id, self._create_callback_handler(callback_url))
+                await self.add_callback(
+                    task_id, self._create_callback_handler(callback_url)
+                )
 
             return task
 
@@ -104,12 +115,22 @@ class TaskManager:
                 # 更新任务状态为已入队
                 return await self.update_task_status(task_id, TaskStatus.QUEUED)
             else:
-                return await self.update_task_status(task_id, TaskStatus.FAILED, error_message="Failed to queue video processing task")
+                return await self.update_task_status(
+                    task_id,
+                    TaskStatus.FAILED,
+                    error_message="Failed to queue video processing task",
+                )
         except Exception as e:
-            return await self.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
+            return await self.update_task_status(
+                task_id, TaskStatus.FAILED, error_message=str(e)
+            )
 
     async def update_task_status(
-        self, task_id: str, status: TaskStatus, machine_name: Optional[str] = None, error_message: Optional[str] = None
+        self,
+        task_id: str,
+        status: TaskStatus,
+        machine_name: Optional[str] = None,
+        error_message: Optional[str] = None,
     ) -> Task:
         """更新任务状态"""
         async with self._task_lock:
@@ -129,7 +150,9 @@ class TaskManager:
 
             # 触发回调
             if task_id in self.callbacks:
-                logger.info(f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}")
+                logger.info(
+                    f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}"
+                )
                 for callback in self.callbacks[task_id]:
                     try:
                         await callback(task)
@@ -154,7 +177,9 @@ class TaskManager:
                 self.callbacks[task_id] = []
             self.callbacks[task_id].append(callback)
 
-    async def remove_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None:
+    async def remove_callback(
+        self, task_id: str, callback: Callable[[Task], Any]
+    ) -> None:
         """移除任务回调"""
         async with self._callback_lock:
             if task_id in self.callbacks:
@@ -181,63 +206,32 @@ async def shutdown_event():
 # API路由
 @app.post("/tasks/", response_model=Task)
 async def create_task(
-    audio: UploadFile = File(...),
-    video: UploadFile = File(...),
+    audio: Optional[UploadFile] = File(None),
+    video: Optional[UploadFile] = File(None),
+    audio_url: Optional[str] = Form(None),
+    video_url: Optional[str] = Form(None),
     description: str = Form(...),
     callback_url: Optional[str] = Form(None),
     background_tasks: BackgroundTasks = BackgroundTasks(),
 ):
-    """创建新的视频处理任务"""
+    """创建新的视频处理任务，支持文件上传或URL下载"""
     try:
-        # 创建临时目录
-        temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
-        os.makedirs(temp_dir, exist_ok=True)
-        try:
-            # 确保文件名不为None
-            if not audio.filename or not video.filename:
-                raise ValueError("文件名不能为空")
-
-            # 保存上传的文件
-            audio_path = os.path.join(temp_dir, audio.filename)
-            video_path = os.path.join(temp_dir, video.filename)
-
-            # 保存音频文件
-            logger.info(f"开始保存音频文件: {audio_path}")
-            async with aiofiles.open(audio_path, "wb") as f:
-                content = await audio.read()
-                await f.write(content)
-            logger.info(f"音频文件保存完成: {audio_path}")
-
-            # 保存视频文件
-            logger.info(f"开始保存视频文件: {video_path}")
-            async with aiofiles.open(video_path, "wb") as f:
-                content = await video.read()
-                await f.write(content)
-            logger.info(f"视频文件保存完成: {video_path}")
-
-            # 验证文件是否保存成功
-            if not os.path.exists(audio_path) or not os.path.exists(video_path):
-                raise ValueError("文件保存失败")
-
-            # 创建输出目录
-            output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
-            os.makedirs(output_dir, exist_ok=True)
-
-            # 创建任务
-            task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
+        # 使用公共函数准备任务文件
+        task, temp_dir = await prepare_task_files(
+            audio, video, audio_url, video_url, description, callback_url
+        )
 
         # 异步处理任务
-            background_tasks.add_task(process_video_task, task.task_id, audio_path, video_path, task.output_file_path)
+        background_tasks.add_task(
+            process_video_task,
+            task.task_id,
+            task.audio_file_path,
+            task.video_file_path,
+            task.output_file_path,
+        )
 
         return task
 
-        except Exception as e:
-            # 如果发生错误，立即清理临时文件
-            # if os.path.exists(temp_dir):
-            #     shutil.rmtree(temp_dir)
-            logger.error(f"创建任务失败: {str(e)}")
-            raise HTTPException(status_code=500, detail=str(e))
-
     except Exception as e:
         logger.error(f"创建任务失败: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
@@ -245,61 +239,36 @@ async def create_task(
 
 @app.post("/tasks/sync")
 async def create_task_sync(
-    audio: UploadFile = File(...),
-    video: UploadFile = File(...),
+    audio: Optional[UploadFile] = File(None),
+    video: Optional[UploadFile] = File(None),
+    audio_url: Optional[str] = Form(None),
+    video_url: Optional[str] = Form(None),
     description: str = Form(...),
     callback_url: Optional[str] = Form(None),
-    background_tasks: BackgroundTasks = BackgroundTasks(),
 ):
-    """创建新的视频处理任务并返回处理结果"""
+    """创建新的视频处理任务并返回处理结果，支持文件上传或URL下载"""
     try:
         logger.info("创建新的视频处理任务并返回处理结果")
-        # 创建临时目录
-        temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
-        logger.info(f"临时目录: {temp_dir}")
-        os.makedirs(temp_dir, exist_ok=True)
-        try:
-            # 确保文件名不为None
-            if not audio.filename or not video.filename:
-                raise ValueError("文件名不能为空")
-            logger.info(f"文件名: {audio.filename}, {video.filename}")
-            # 保存上传的文件
-            audio_path = os.path.join(temp_dir, audio.filename)
-            video_path = os.path.join(temp_dir, video.filename)
 
-            # 保存音频文件
-            logger.info(f"开始保存音频文件: {audio_path}")
-            async with aiofiles.open(audio_path, "wb") as f:
-                content = await audio.read()
-                await f.write(content)
-            logger.info(f"音频文件保存完成: {audio_path}")
-
-            # 保存视频文件
-            logger.info(f"开始保存视频文件: {video_path}")
-            async with aiofiles.open(video_path, "wb") as f:
-                content = await video.read()
-                await f.write(content)
-            logger.info(f"视频文件保存完成: {video_path}")
-
-            # 验证文件是否保存成功
-            if not os.path.exists(audio_path) or not os.path.exists(video_path):
-                raise ValueError("文件保存失败")
-
-            # 创建输出目录
-            output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
-            os.makedirs(output_dir, exist_ok=True)
-
-            # 创建任务
-            task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
+        # 使用公共函数准备任务文件
+        task, temp_dir = await prepare_task_files(
+            audio, video, audio_url, video_url, description, callback_url
+        )
+        logger.info(f"任务创建成功: {task.task_id}")
+        logger.info(f"音频文件路径: {task.audio_file_path}")
+        logger.info(f"视频文件路径: {task.video_file_path}")
 
         # 创建一个事件来等待任务完成
         task_completed = asyncio.Event()
         logger.info(f"创建任务完成事件: {task.task_id}")
 
         async def task_completion_callback(completed_task: Task):
-                logger.info(f"任务完成回调被触发: {completed_task.task_id}, 状态: {completed_task.status}")
+            logger.info(
+                f"任务完成回调被触发: {completed_task.task_id}, 状态: {completed_task.status}"
+            )
             if completed_task.task_id == task.task_id and (
-                    completed_task.status == TaskStatus.COMPLETED or completed_task.status == TaskStatus.FAILED
+                completed_task.status == TaskStatus.COMPLETED
+                or completed_task.status == TaskStatus.FAILED
             ):
                 logger.info(f"设置任务完成事件: {task.task_id}")
                 task_completed.set()
@@ -344,17 +313,14 @@ async def create_task_sync(
                         yield chunk
 
             return StreamingResponse(
-                    file_iterator(), media_type="video/mp4", headers={"Content-Disposition": f'attachment; filename="{filename}"'}
+                file_iterator(),
+                media_type="video/mp4",
+                headers={"Content-Disposition": f'attachment; filename="{filename}"'},
             )
         else:
-                raise HTTPException(status_code=500, detail=result.error_message or "处理失败")
-
-        except Exception as e:
-            # 如果发生错误，立即清理临时文件
-            # if os.path.exists(temp_dir):
-            #     shutil.rmtree(temp_dir)
-            logger.error(f"创建任务失败: {str(e)}", exc_info=True)
-            raise HTTPException(status_code=500, detail=str(e))
+            raise HTTPException(
+                status_code=500, detail=result.error_message or "处理失败"
+            )
 
     except Exception as e:
         logger.error(f"创建任务失败: {str(e)}", exc_info=True)
@@ -377,7 +343,9 @@ async def get_all_tasks():
 
 
 @app.post("/tasks/{task_id}/callback")
-async def add_task_callback(task_id: str, callback_url: str, background_tasks: BackgroundTasks):
+async def add_task_callback(
+    task_id: str, callback_url: str, background_tasks: BackgroundTasks
+):
     """添加任务完成回调"""
     if not await task_manager.get_task(task_id):
         raise HTTPException(status_code=404, detail="Task not found")
@@ -394,7 +362,9 @@ async def add_task_callback(task_id: str, callback_url: str, background_tasks: B
 
 
 # 后台任务处理函数
-async def process_video_task(task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str):
+async def process_video_task(
+    task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str
+):
     """处理视频任务的后台函数"""
     try:
         # 更新任务状态为处理中
@@ -404,21 +374,161 @@ async def process_video_task(task_id: str, audio_file_path: str, video_file_path
         file_id = f"{task_id}_{int(datetime.now().timestamp())}"
 
         # 处理视频
-        success = await task_manager.process_manager.process_video(task_id, audio_file_path, video_file_path, out_file_path)
+        success = await task_manager.process_manager.process_video(
+            task_id, audio_file_path, video_file_path, out_file_path
+        )
 
         if success:
             await task_manager.update_task_status(task_id, TaskStatus.COMPLETED)
         else:
-            await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message="Video processing failed")
+            await task_manager.update_task_status(
+                task_id, TaskStatus.FAILED, error_message="Video processing failed"
+            )
     except Exception as e:
-        await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
+        await task_manager.update_task_status(
+            task_id, TaskStatus.FAILED, error_message=str(e)
+        )
         # 清理临时文件
         # if os.path.exists(temp_dir):
         #     shutil.rmtree(temp_dir)
         #     logger.info(f"已清理临时目录: {temp_dir}")
 
 
+# 添加文件下载辅助函数
+async def download_file_from_url(url: str, destination_path: str) -> bool:
+    """从URL下载文件到指定路径"""
+    try:
+        logger.info(f"开始下载文件: {url} -> {destination_path}")
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                url, timeout=aiohttp.ClientTimeout(total=300)
+            ) as response:
+                if response.status == 200:
+                    # 获取Content-Length用于进度追踪
+                    content_length = response.headers.get("Content-Length")
+                    if content_length:
+                        logger.info(f"文件大小: {int(content_length)} bytes")
+
+                    async with aiofiles.open(destination_path, "wb") as f:
+                        downloaded = 0
+                        async for chunk in response.content.iter_chunked(8192):
+                            await f.write(chunk)
+                            downloaded += len(chunk)
+
+                    logger.info(
+                        f"文件下载完成: {destination_path} ({downloaded} bytes)"
+                    )
+                    return True
+                else:
+                    logger.error(f"下载失败，HTTP状态码: {response.status}")
+                    return False
+    except Exception as e:
+        logger.error(f"下载文件时发生错误: {str(e)}")
+        return False
+
+
+async def process_media_input(
+    media_file: Optional[UploadFile],
+    media_url: Optional[str],
+    media_type: str,
+    temp_dir: str,
+) -> str:
+    """处理媒体输入（文件或URL），返回本地文件路径"""
+    if media_file and media_url:
+        raise ValueError(f"{media_type}不能同时提供文件和URL")
+
+    if not media_file and not media_url:
+        raise ValueError(f"必须提供{media_type}文件或URL")
+
+    if media_file:
+        # 处理上传文件
+        if not media_file.filename:
+            raise ValueError(f"{media_type}文件名不能为空")
+
+        file_path = os.path.join(temp_dir, media_file.filename)
+        logger.info(f"开始保存{media_type}文件: {file_path}")
+        async with aiofiles.open(file_path, "wb") as f:
+            content = await media_file.read()
+            await f.write(content)
+        logger.info(f"{media_type}文件保存完成: {file_path}")
+        return file_path
+
+    else:
+        # 处理URL下载
+        parsed_url = urllib.parse.urlparse(media_url)
+        filename = os.path.basename(parsed_url.path)
+        if not filename or "." not in filename:
+            # 如果URL没有文件名或没有扩展名，根据媒体类型生成
+            if media_type == "音频":
+                extension = ".mp3"  # 默认音频扩展名
+            elif media_type == "视频":
+                extension = ".mp4"  # 默认视频扩展名
+            else:
+                extension = ""
+
+            if not filename:
+                filename = f"{media_type}_{int(datetime.now().timestamp())}{extension}"
+            elif "." not in filename:
+                filename = f"{filename}{extension}"
+
+        file_path = os.path.join(temp_dir, filename)
+        success = await download_file_from_url(media_url, file_path)
+        if not success:
+            raise ValueError(f"下载{media_type}文件失败: {media_url}")
+        return file_path
+
+
+async def prepare_task_files(
+    audio: Optional[UploadFile],
+    video: Optional[UploadFile],
+    audio_url: Optional[str],
+    video_url: Optional[str],
+    description: str,
+    callback_url: Optional[str],
+) -> tuple[Task, str]:
+    """
+    准备任务文件和创建任务的公共逻辑
+    返回创建的任务对象和临时目录路径
+    """
+    # 创建临时目录
+    temp_dir = os.path.join(
+        os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S")
+    )
+    os.makedirs(temp_dir, exist_ok=True)
+
+    try:
+        # 处理音频输入（文件或URL）
+        audio_path = await process_media_input(audio, audio_url, "音频", temp_dir)
+
+        # 处理视频输入（文件或URL）
+        video_path = await process_media_input(video, video_url, "视频", temp_dir)
+
+        # 验证文件是否存在
+        if not os.path.exists(audio_path) or not os.path.exists(video_path):
+            raise ValueError("文件处理失败")
+
+        # 创建输出目录
+        output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
+        os.makedirs(output_dir, exist_ok=True)
+
+        # 创建任务
+        task = await task_manager.add_task(
+            audio_path, video_path, output_dir, description, callback_url
+        )
+
+        return task, temp_dir
+
+    except Exception as e:
+        # 如果发生错误，立即清理临时文件
+        # if os.path.exists(temp_dir):
+        #     shutil.rmtree(temp_dir)
+        logger.error(f"准备任务文件失败: {str(e)}")
+        raise
+
+
 if __name__ == "__main__":
     import uvicorn
 
-    uvicorn.run("api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info")  # 开发模式下启用热重载  # 单进程模式
+    uvicorn.run(
+        "api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info"
+    )  # 开发模式下启用热重载  # 单进程模式