qihuaiflow/backend/python/app/api/admin/models.py

"""
大模型库管理 API（管理员）
"""
from fastapi import APIRouter, Depends, HTTPException, Query, Body
from sqlalchemy.orm import Session
from sqlalchemy import func
from typing import List, Optional
from pydantic import BaseModel
from ...database import get_db, SessionLocal
from ...models import AIModel, ModelProvider
from ...auth import get_current_admin
from ...middleware.log_service import LogService
import uuid
import json
import logging
import httpx
from datetime import datetime

router = APIRouter()

# ============================================================
# 协议预设列表
# ============================================================

PROTOCOLS = [
    {"id": "openai_compatible", "name": "OpenAI兼容", "base_path": "/v1", "description": "兼容OpenAI API格式的模型服务"},
    {"id": "new_api", "name": "New API", "base_path": "/v1", "description": "新一代API聚合平台"},
    {"id": "volcengine", "name": "火山引擎", "base_path": "/v1", "description": "字节跳动火山引擎模型服务"},
    {"id": "aliyun_bailian", "name": "阿里云百炼", "base_path": "/v1", "description": "阿里云百炼大模型平台"},
    {"id": "ollama", "name": "Ollama", "base_path": "/v1", "description": "本地Ollama模型服务"},
    {"id": "xiaoma", "name": "小马算力", "base_path": "/v1", "description": "小马算力模型服务"},
]

# ============================================================
# 模型列表（用于 fetch-models 的回退数据）
# ============================================================

FETCHABLE_MODELS = [
    {"id": "gpt-4o", "name": "GPT-4o", "capabilities": ["text", "vision", "tool"], "pricing": "paid"},
    {"id": "gpt-4o-mini", "name": "GPT-4o Mini", "capabilities": ["text", "vision"], "pricing": "paid"},
    {"id": "deepseek-v3", "name": "DeepSeek-V3", "capabilities": ["text", "thinking"], "pricing": "free"},
    {"id": "deepseek-r1", "name": "DeepSeek-R1", "capabilities": ["text", "thinking"], "pricing": "paid"},
    {"id": "qwen-max", "name": "Qwen-Max", "capabilities": ["text", "tool"], "pricing": "free"},
    {"id": "qwen-vl", "name": "Qwen-VL", "capabilities": ["text", "vision"], "pricing": "paid"},
    {"id": "whisper", "name": "Whisper", "capabilities": ["audio"], "pricing": "paid"},
    {"id": "dalle", "name": "DALL-E", "capabilities": ["image"], "pricing": "paid"},
]

# 支持通过 API 获取模型列表的协议类型
FETCHABLE_PROTOCOLS = {"openai_compatible", "new_api", "volcengine", "aliyun_bailian", "ollama", "xiaoma"}


# ============================================================
# 辅助函数
# ============================================================

def provider_to_dict(p: ModelProvider) -> dict:
    return {
        "id": p.id,
        "name": p.name,
        "protocol": p.protocol,
        "api_key": p.api_key,
        "api_base_url": p.api_base_url,
        "apiKey": p.api_key,
        "apiBaseUrl": p.api_base_url,
        "created_at": p.created_at.isoformat() if p.created_at else None,
        "updated_at": p.updated_at.isoformat() if p.updated_at else None,
        "createdAt": p.created_at.isoformat() if p.created_at else None,
        "updatedAt": p.updated_at.isoformat() if p.updated_at else None,
    }


def model_to_dict(m: AIModel) -> dict:
    return {
        "id": m.id,
        "provider_id": m.provider_id,
        "name": m.name,
        "api_key": m.api_key,
        "endpoint": m.endpoint,
        "capabilities": m.capabilities,
        "timeout": m.timeout,
        "status": m.status,
        "price_per_million_tokens": float(m.price_per_million_tokens) if m.price_per_million_tokens else 0,
        "is_enabled": m.is_enabled,
        "is_default": bool(m.is_default),
        "total_calls": m.total_calls or 0,
        "total_tokens": m.total_tokens or 0,
        "created_at": m.created_at.isoformat() if m.created_at else None,
        "updated_at": m.updated_at.isoformat() if m.updated_at else None,
    }


# ============================================================
# 协议列表
# ============================================================

@router.get("/protocols")
async def get_protocols(
    admin_user=Depends(get_current_admin),
):
    """获取支持的协议列表"""
    return PROTOCOLS


# ============================================================
# 模型供应商管理
# ============================================================

@router.get("/providers")
async def get_providers(
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """获取模型供应商列表"""
    providers = db.query(ModelProvider).order_by(
        ModelProvider.created_at.desc()
    ).all()
    return [provider_to_dict(p) for p in providers]


@router.post("/providers")
async def create_provider(
    name: str = Body(...),
    protocol: str = Body(...),
    api_key: str = Body(default=""),
    api_base_url: str = Body(default=None),
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """添加模型供应商"""
    # 验证协议
    valid_protocols = [p["id"] for p in PROTOCOLS]
    if protocol not in valid_protocols:
        raise HTTPException(status_code=400, detail=f"无效的协议类型，支持: {', '.join(valid_protocols)}")

    provider = ModelProvider(
        id=str(uuid.uuid4()),
        name=name,
        protocol=protocol,
        api_key=api_key,
        api_base_url=api_base_url,
    )
    db.add(provider)
    db.commit()
    db.refresh(provider)

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="create_provider",
        target_type="model_provider",
        target_id=provider.id,
        detail=f"添加模型供应商 {name} (协议: {protocol})",
    )

    return {"message": "供应商已创建", "id": provider.id, "provider": provider_to_dict(provider)}


@router.put("/providers/{id}")
async def update_provider(
    id: str,
    name: str = Body(default=None),
    protocol: str = Body(default=None),
    api_key: str = Body(default=None),
    api_base_url: str = Body(default=None),
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """更新模型供应商"""
    provider = db.query(ModelProvider).filter(ModelProvider.id == id).first()
    if not provider:
        raise HTTPException(status_code=404, detail="供应商不存在")

    if name is not None:
        provider.name = name
    if protocol is not None:
        valid_protocols = [p["id"] for p in PROTOCOLS]
        if protocol not in valid_protocols:
            raise HTTPException(status_code=400, detail=f"无效的协议类型")
        provider.protocol = protocol
    if api_key is not None:
        provider.api_key = api_key
    if api_base_url is not None:
        provider.api_base_url = api_base_url
    provider.updated_at = datetime.now()
    db.commit()
    db.refresh(provider)

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="update_provider",
        target_type="model_provider",
        target_id=id,
        detail=f"更新模型供应商 {provider.name}",
    )

    return {"message": "供应商已更新", "provider": provider_to_dict(provider)}


@router.delete("/providers/{id}")
async def delete_provider(
    id: str,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """删除模型供应商"""
    provider = db.query(ModelProvider).filter(ModelProvider.id == id).first()
    if not provider:
        raise HTTPException(status_code=404, detail="供应商不存在")

    # 检查是否有模型关联
    model_count = db.query(func.count(AIModel.id)).filter(
        AIModel.provider_id == id
    ).scalar()
    if model_count > 0:
        raise HTTPException(status_code=400, detail=f"该供应商下有 {model_count} 个模型，请先删除模型")

    db.delete(provider)
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="delete_provider",
        target_type="model_provider",
        target_id=id,
        detail=f"删除模型供应商 {provider.name}",
    )

    return {"message": "供应商已删除"}


@router.post("/providers/{id}/fetch-models")
async def fetch_models(
    id: str,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """获取供应商可用模型列表（优先从真实API获取）"""
    provider = db.query(ModelProvider).filter(ModelProvider.id == id).first()
    if not provider:
        raise HTTPException(status_code=404, detail="供应商不存在")

    logger = logging.getLogger(__name__)
    logger.info(f"Fetching models for provider {provider.name} (protocol: {provider.protocol})")

    api_base = (provider.api_base_url or "").rstrip("/")
    api_key = provider.api_key or ""

    # 尝试从真实API获取模型列表（只要有 api_base 就尝试调用，api_key 可选）
    if api_base and provider.protocol in FETCHABLE_PROTOCOLS:
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                headers = {}
                if api_key:
                    headers["Authorization"] = f"Bearer {api_key}"
                resp = await client.get(f"{api_base}/v1/models", headers=headers)
                if resp.status_code == 200:
                    data = resp.json()
                    models_list = data.get("data", [])
                    if models_list:
                        result = []
                        for m in models_list:
                            mid = m.get("id", "")
                            if not mid:
                                continue
                            result.append({
                                "id": mid,
                                "name": mid,
                                "capabilities": ["text"],
                                "pricing": "paid",
                            })
                        logger.info(f"Successfully fetched {len(result)} models from provider API")
                        return {"models": result, "source": "api"}
                    else:
                        logger.warning(f"Provider API returned empty model list")
                        return {"models": [], "source": "api", "message": "API返回的模型列表为空"}
                elif resp.status_code == 401 or resp.status_code == 403:
                    logger.warning(f"Provider API returned {resp.status_code}: auth error")
                    return {"models": [], "source": "error", "message": "API认证失败，请检查API Key是否正确"}
                else:
                    logger.warning(f"Provider API returned status {resp.status_code}: {resp.text[:200]}")
                    return {"models": [], "source": "error", "message": f"API请求失败(HTTP {resp.status_code})，请检查API地址和密钥"}
        except httpx.TimeoutException:
            logger.warning(f"Timeout fetching models from provider API: {api_base}")
            return {"models": [], "source": "error", "message": "请求API超时，请检查API地址是否正确"}
        except Exception as e:
            logger.warning(f"Failed to fetch models from provider API: {e}")
            return {"models": [], "source": "error", "message": f"请求API失败: {str(e)}"}

    # 回退：返回固定模型列表（供无 API 的离线环境使用）
    logger.info(f"Using fallback model list for provider {provider.name}")
    return {"models": FETCHABLE_MODELS, "source": "fallback"}


class AddModelsRequest(BaseModel):
    model_ids: List[str]


@router.post("/providers/{id}/add-models")
async def add_models(
    id: str,
    req: AddModelsRequest,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """批量添加模型到供应商"""
    provider = db.query(ModelProvider).filter(ModelProvider.id == id).first()
    if not provider:
        raise HTTPException(status_code=404, detail="供应商不存在")

    # 找到协议对应的 base_path
    protocol_info = next((p for p in PROTOCOLS if p["id"] == provider.protocol), None)
    base_path = protocol_info["base_path"] if protocol_info else "/v1"

    added = []
    for model_id in req.model_ids:
        # 在 FETCHABLE_MODELS 中查找额外信息
        model_info = next((m for m in FETCHABLE_MODELS if m["id"] == model_id), None)

        # 使用 model_id 作为模型名称
        model_name = model_info["name"] if model_info else model_id

        # 检查是否已存在同名模型
        existing = db.query(AIModel).filter(
            AIModel.provider_id == id,
            AIModel.name == model_name
        ).first()
        if existing:
            continue

        # 自动生成 endpoint
        if provider.api_base_url:
            endpoint = f"{provider.api_base_url.rstrip('/')}{base_path}/chat/completions"
        else:
            endpoint = ""

        model = AIModel(
            id=str(uuid.uuid4()),
            provider_id=id,
            name=model_name,
            api_key=provider.api_key or "",
            endpoint=endpoint,
            capabilities=model_info.get("capabilities") if model_info else ["text"],
            timeout=60,
            status="free" if (model_info and model_info.get("pricing") == "free") else "paid",
            price_per_million_tokens=0,
            is_enabled=True,
        )
        db.add(model)
        added.append(model_name)

    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="add_models",
        target_type="model_provider",
        target_id=id,
        detail=f"为供应商 {provider.name} 批量添加 {len(added)} 个模型",
    )

    return {"message": f"成功添加 {len(added)} 个模型", "added": added}


# ============================================================
# 模型管理
# ============================================================

@router.get("/")
async def get_models(
    provider_id: str = Query(None),
    status: str = Query(None, description="模型状态筛选: free/paid"),
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """获取模型列表（管理员）"""
    query = db.query(AIModel)

    if provider_id:
        query = query.filter(AIModel.provider_id == provider_id)
    if status:
        query = query.filter(AIModel.status == status)

    total = query.count()
    models = query.order_by(AIModel.created_at.desc()).offset(
        (page - 1) * page_size
    ).limit(page_size).all()

    return {
        "total": total,
        "page": page,
        "page_size": page_size,
        "items": [model_to_dict(m) for m in models],
    }


@router.post("/")
async def create_model(
    provider_id: str,
    name: str,
    api_key: str = None,
    endpoint: str = None,
    capabilities: List[str] = None,
    timeout: int = 60,
    status: str = "paid",
    price_per_million_tokens: float = 0,
    is_enabled: bool = True,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """添加模型"""
    # 验证供应商存在
    provider = db.query(ModelProvider).filter(ModelProvider.id == provider_id).first()
    if not provider:
        raise HTTPException(status_code=404, detail="供应商不存在")

    model = AIModel(
        id=str(uuid.uuid4()),
        provider_id=provider_id,
        name=name,
        api_key=api_key,
        endpoint=endpoint,
        capabilities=capabilities,
        timeout=timeout,
        status=status,
        price_per_million_tokens=price_per_million_tokens,
        is_enabled=is_enabled,
    )
    db.add(model)
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="create_model",
        target_type="model",
        target_id=model.id,
        detail=f"添加模型 {name}",
    )

    return {"message": "模型已创建", "id": model.id}


@router.put("/{id}")
async def update_model(
    id: str,
    provider_id: str = None,
    name: str = None,
    api_key: str = None,
    endpoint: str = None,
    capabilities: List[str] = None,
    timeout: int = None,
    status: str = None,
    price_per_million_tokens: float = None,
    is_enabled: bool = None,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """更新模型"""
    model = db.query(AIModel).filter(AIModel.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="模型不存在")

    if provider_id is not None:
        provider = db.query(ModelProvider).filter(ModelProvider.id == provider_id).first()
        if not provider:
            raise HTTPException(status_code=404, detail="供应商不存在")
        model.provider_id = provider_id
    if name is not None:
        model.name = name
    if api_key is not None:
        model.api_key = api_key
    if endpoint is not None:
        model.endpoint = endpoint
    if capabilities is not None:
        model.capabilities = capabilities
    if timeout is not None:
        model.timeout = timeout
    if status is not None:
        model.status = status
    if price_per_million_tokens is not None:
        model.price_per_million_tokens = price_per_million_tokens
    if is_enabled is not None:
        model.is_enabled = is_enabled
    model.updated_at = datetime.now()
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="update_model",
        target_type="model",
        target_id=id,
        detail=f"更新模型 {model.name}",
    )

    return {"message": "模型已更新"}


@router.delete("/{id}")
async def delete_model(
    id: str,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """删除模型"""
    model = db.query(AIModel).filter(AIModel.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="模型不存在")

    model_name = model.name
    db.delete(model)
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="delete_model",
        target_type="model",
        target_id=id,
        detail=f"删除模型 {model_name}",
    )

    return {"message": "模型已删除"}


@router.put("/{id}/toggle")
async def toggle_model(
    id: str,
    is_enabled: bool = Query(...),
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """开关模型状态"""
    model = db.query(AIModel).filter(AIModel.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="模型不存在")

    model.is_enabled = is_enabled
    model.updated_at = datetime.now()
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="toggle_model",
        target_type="model",
        target_id=id,
        detail=f"将模型 {model.name} {'启用' if is_enabled else '禁用'}",
    )

    return {"message": f"模型已{'启用' if is_enabled else '禁用'}"}


@router.put("/{id}/set-default")
async def set_default_model(
    id: str,
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """设为全局默认模型（取消其他模型的默认状态）"""
    model = db.query(AIModel).filter(AIModel.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="模型不存在")

    # 取消所有模型的默认状态
    db.query(AIModel).filter(AIModel.is_default == True).update({"is_default": False})
    # 设置当前模型为默认
    model.is_default = True
    model.updated_at = datetime.now()
    db.commit()

    LogService.log_audit(
        admin_id=admin_user.id,
        admin_name=admin_user.username,
        action="set_default_model",
        target_type="model",
        target_id=id,
        detail=f"设置模型 {model.name} 为全局默认",
    )

    return {"message": f"已将 {model.name} 设为全局默认模型"}


@router.get("/stats")
async def get_model_stats(
    admin_user=Depends(get_current_admin),
    db: Session = Depends(get_db),
):
    """获取模型调用统计"""
    total_calls = db.query(func.sum(AIModel.total_calls)).scalar() or 0
    total_tokens = db.query(func.sum(AIModel.total_tokens)).scalar() or 0

    return {
        "total_calls": int(total_calls),
        "total_tokens": int(total_tokens),
    }