zhang_crawl_bio/ui/backend/routers/worker.py

import json
from fastapi import APIRouter, HTTPException, Request, Depends
from fastapi.responses import JSONResponse, StreamingResponse
import httpx
import os
from pathlib import Path
from typing import Dict, List
from pydantic import BaseModel
from typing import Optional 
import pathlib
import os
from utils.config import WORKER_SERVICE_URL,config,Browser
from src.services.celery_worker import celery_worker
from utils.logu import logger

router = APIRouter()

class Endpoint(BaseModel):
    service_url: str = WORKER_SERVICE_URL
    token: Optional[str] = "temp_token_123"
    health: Optional[Dict] = {}


class PathValidationRequest(BaseModel):
    file_path: str

class PathValidationResponse(BaseModel):
    valid: bool
    message: str
    file_name: Optional[str] = None
    file_size: Optional[int] = None

class ResponseStatus(BaseModel):
    endpoint: Endpoint = Endpoint()
    browser_config: Optional[Browser] = None
    celery_status: Optional[Dict] = {}
    
def transform_workers_data(workers_data: list) -> list:
    """转换worker数据为前端需要的格式"""
    return [
        {
            "hostname": worker.get("hostname"),
            "worker-online": worker.get("worker-online"),
            "pid": worker.get("pid"),
            "active": worker.get("active"),
            "processed": worker.get("processed"),
            "status": worker.get("status")
        }
        for worker in workers_data
    ]

async def get_proxy_pool() -> Optional[list[str]]:
    """获取代理池数据"""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"http://{config.backend.host}:{config.backend.port}/api/proxy/proxies-pool",
                params={"force_refresh": False}
            )
            response.raise_for_status()
            data = response.json()
            return data.get("proxies", None)
    except Exception as e:
        logger.error(f"获取代理池失败: {str(e)}")
        return None

@router.post("/browser_config", tags=["worker"])
async def update_browser_config(new_config: Browser):
    """更新浏览器配置（包含自动路径校验）"""
    exe_path = new_config.exe_path
    # 自动处理路径格式
    if exe_path.startswith('"') and exe_path.endswith('"'):
        exe_path = exe_path.strip('"')
    
    path = Path(exe_path)
    try:
        if not path.exists():
            return JSONResponse(status_code=400, content={
                "status": "error",
                "message": "路径不存在",
                "detail": "请检查浏览器可执行文件路径是否正确"
            })
            
        if not path.is_file():
            return JSONResponse(status_code=400, content={
                "status": "error",
                "message": "路径不是文件",
                "detail": "请选择一个可执行文件而不是目录"
            })
            
        if not os.access(path, os.X_OK):
            return JSONResponse(status_code=400, content={
                "status": "error", 
                "message": "文件不可执行",
                "detail": "请检查文件权限或是否损坏"
            })
            
        if path.suffix.lower() != '.exe':
            return JSONResponse(status_code=400, content={
                "status": "error",
                "message": "仅支持可执行文件 (.exe)",
                "detail": "请选择扩展名为.exe的浏览器可执行文件"
            })
        new_config.exe_path = str(path)
        # 校验通过后保存配置
        config.browser = new_config
        config.save()
        return {"status": "success", "browser_config": config.browser}
        
    except Exception as e:
        logger.error(f"浏览器路径验证失败: {str(e)}")
        return JSONResponse(status_code=500, content={
            "status": "error",
            "message": f"路径验证异常: {str(e)}"
        })

@router.post("/reset_browser_config", tags=["worker"])
async def reset_browser_config():
    """重置浏览器配置为默认"""
    config.browser = Browser()
    config.save()
    return {"status": "success", "browser_config": config.browser}

@router.get("/status", tags=["worker"])
async def status() -> ResponseStatus:
    global config,celery_worker
    try:
        health = await health_check()
    except Exception as e:
        logger.error(e)
        health = {"err": 1, "msg": str(e)}
    celery_status = await celery_worker.check_status()
    queue_lengths = await celery_worker.get_queue_lengths()
    
    # 数据转换处理
    if celery_status.get("workers") and celery_status["workers"].get("data"):
        celery_status["workers"]["data"] = transform_workers_data(
            celery_status["workers"]["data"]
        )
    
    # 添加队列长度信息
    celery_status["queue_lengths"] = queue_lengths
    
    return ResponseStatus(
        endpoint=Endpoint(health=health),
        browser_config=config.browser,
        celery_status=celery_status
    )

class StartupRequest(BaseModel):
    worker_name: str
    action: str
    select_proxy: Optional[str] = ''
    data: Optional[Dict] = {}

@router.post("/ctrl", tags=["worker"])
async def ctrl_worker(request: StartupRequest):
    global celery_worker
    logger.info(f"{request.model_dump_json(indent=4)}")

    if request.action == "start":
        worker_model = await celery_worker.start_worker(request.worker_name, in_cmd_windows=True)
    elif request.action == "stop":
        worker_model = await celery_worker.stop_worker(request.worker_name)
    elif request.action == "clean":
        await celery_worker.clear_specific_queue(request.worker_name)
        logger.info(f"clean {request.data}")
        return {"status": "success"}
    elif request.action == "submit":
        res = await celery_worker.submit_tasks(
            request.worker_name,
            request.data,
            select_proxy=request.select_proxy
        )
        return {"status": "success", "res": res}
    elif request.action == "submit_all":
        res = await celery_worker.submit_all_tasks(
            request.worker_name,
            request.data,
            select_proxy=request.select_proxy
        )
        return {"status": "success", "res": res}
    else:
        logger.error(f"Invalid action: {request.action}")
        raise HTTPException(status_code=400, detail=f"Invalid action: {request.action}")
    return {"status": "success", "worker_model": worker_model}

async def health_check():
    """健康检查"""
    async with httpx.AsyncClient(base_url=WORKER_SERVICE_URL) as client:
        response = await client.get("/health")
        response.raise_for_status()
        ret = response.json()
        ret.update({"err": 0})
        return ret

@router.delete("/keywords/{keyword}", summary="删除关键词")
async def delete_keyword(keyword: str):
    '''
    DELETE: API_BASE_URL/worker/keywords/{keyword}
    '''
    async with httpx.AsyncClient(base_url=WORKER_SERVICE_URL) as client:
        # 发送 DELETE 请求
        logger.info(f"DELETE: {WORKER_SERVICE_URL}/keywords/{keyword}")
        response = await client.delete(f"/keywords/{keyword}")
        
        # 检查响应状态码，如果请求失败则抛出异常
        response.raise_for_status()
        ret = response.json()
        logger.info(f"response {ret}")
    if ret.get("err") == 0:
        return JSONResponse(
            status_code=200,
            content={
                "err": 0,
                "status": "success",
                "message": f"关键词 '{keyword}' 删除成功",
                "model": ret
            }
        )
    elif ret.get("err") == 1:
        return JSONResponse(
            status_code=404,
            content={
                "err": 1,
                "status": "error",
                "message": f"关键词 '{keyword}' 不存在"
            }
        )
    return JSONResponse(
        status_code=500,
        content={
            "err": ret.get("err", 2),
            "status": "error",
            "message": ret.get("message", "服务器内部错误")
        }
    )