新推荐接口first commit

app/api/api_import_sys_sketch.py

@@ -0,0 +1,116 @@
+import logging
+import sys
+from typing import Optional
+from fastapi import APIRouter, HTTPException, Query
+from concurrent.futures import ThreadPoolExecutor
+import threading
+
+from app.schemas.response_template import ResponseModel
+from app.service.recommendation_system.import_sys_sketch_to_milvus import main as import_main
+
+logger = logging.getLogger()
+router = APIRouter()
+
+# 使用线程池执行器来运行长时间任务
+executor = ThreadPoolExecutor(max_workers=1)
+# 用于跟踪任务状态
+task_status = {"running": False}
+
+
+def run_import_task(batch_size: int, retry_times: int, limit: Optional[int], offset: int, skip_create_collection: bool):
+    """在后台线程中运行导入任务"""
+    original_argv = None
+    try:
+        task_status["running"] = True
+        # 保存原始 sys.argv
+        original_argv = sys.argv.copy()
+        
+        # 模拟命令行参数
+        sys.argv = [
+            "import_sys_sketch_to_milvus.py",
+            "--batch-size", str(batch_size),
+            "--retry-times", str(retry_times),
+        ]
+        if limit is not None:
+            sys.argv.extend(["--limit", str(limit)])
+        if offset > 0:
+            sys.argv.extend(["--offset", str(offset)])
+        if skip_create_collection:
+            sys.argv.append("--skip-create-collection")
+        
+        import_main()
+        task_status["running"] = False
+        logger.info("导入任务完成")
+    except Exception as e:
+        task_status["running"] = False
+        logger.error(f"导入任务失败: {e}", exc_info=True)
+        raise
+    finally:
+        # 恢复原始 sys.argv
+        if original_argv is not None:
+            sys.argv = original_argv
+
+
+@router.post("/import-sys-sketch", response_model=ResponseModel)
+async def import_sys_sketch(
+    batch_size: int = Query(1000, description="批量处理大小（默认：1000）"),
+    retry_times: int = Query(3, description="失败重试次数（默认：3）"),
+    limit: Optional[int] = Query(None, description="限制处理数量（用于测试，默认：不限制）"),
+    offset: int = Query(0, description="起始偏移量（默认：0）"),
+    skip_create_collection: bool = Query(False, description="跳过创建集合（如果集合已存在）"),
+):
+    """
+    从 t_sys_file 导入系统图向量到 Milvus
+    
+    该接口会异步执行导入任务，任务在后台运行。
+    """
+    try:
+        # 检查是否有任务正在运行
+        if task_status["running"]:
+            raise HTTPException(
+                status_code=409,
+                detail="已有导入任务正在运行，请等待完成后再试"
+            )
+        
+        # 在后台线程中执行任务
+        executor.submit(
+            run_import_task,
+            batch_size,
+            retry_times,
+            limit,
+            offset,
+            skip_create_collection
+        )
+        
+        return ResponseModel(
+            code=200,
+            msg="导入任务已启动，正在后台执行",
+            data={
+                "status": "started",
+                "batch_size": batch_size,
+                "retry_times": retry_times,
+                "limit": limit,
+                "offset": offset,
+                "skip_create_collection": skip_create_collection
+            }
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"启动导入任务失败: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"启动导入任务失败: {str(e)}")
+
+
+@router.get("/import-sys-sketch/status", response_model=ResponseModel)
+async def get_import_status():
+    """
+    获取导入任务状态
+    """
+    return ResponseModel(
+        code=200,
+        msg="OK",
+        data={
+            "running": task_status["running"]
+        }
+    )
+