my_bookmate/backend/app/api/models.py

import httpx
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import get_db
from app.models.ai_model import AiModel
from app.llm.hub import builtin_models, resolve_provider, _mask_key

router = APIRouter(prefix="/models", tags=["models"])


class AddModelRequest(BaseModel):
    label: str
    kind: str = "openai_compatible"  # openai_compatible | anthropic
    base_url: str = ""
    api_key: str = ""
    model_id: str


class UpdateModelRequest(BaseModel):
    label: str | None = None
    kind: str | None = None
    base_url: str | None = None
    api_key: str | None = None
    model_id: str | None = None


class ImportModelsRequest(BaseModel):
    api_key: str
    base_url: str = "https://openrouter.ai/api/v1"


class ImportOpenRouterRequest(BaseModel):
    api_key: str
    base_url: str = "https://openrouter.ai/api/v1"


OPENROUTER_BLOCKLIST = {
    # 一些非文本/量级过大的模型 id 特征，导入时快速排除
    "image", "tts", "whisper", "embed", "rerank", "moderation", "dall-e", "flux", "playground",
}


def _to_out(m: AiModel) -> dict:
    return {
        "name": m.name,
        "label": m.label,
        "kind": m.kind,
        "base_url": m.base_url,
        "api_key_masked": _mask_key(m.api_key),
        "model_id": m.model_id,
        "builtin": False,
    }


@router.get("")
async def list_models(db: AsyncSession = Depends(get_db)):
    """内置（.env）+ 用户自定义（DB）模型列表。api_key 只回显掩码。"""
    result = await db.execute(select(AiModel).order_by(AiModel.sort_order, AiModel.created_at))
    custom = [_to_out(m) for m in result.scalars().all()]
    return builtin_models() + custom


@router.post("")
async def add_model(req: AddModelRequest, db: AsyncSession = Depends(get_db)):
    label = (req.label or "").strip() or req.model_id
    # name = 请求模型 id；冲突时追加序号（-2, -3 ...）
    base_name = req.model_id.strip() or "custom-model"
    name = base_name
    suffix = 1
    while True:
        exists = (await db.execute(select(AiModel.id).where(AiModel.name == name))).scalar()
        if exists is None:
            break
        suffix += 1
        name = f"{base_name}-{suffix}"
    row = AiModel(
        name=name, label=label, kind=req.kind,
        base_url=(req.base_url or "").strip(), api_key=(req.api_key or "").strip(),
        model_id=req.model_id.strip(), sort_order=100,
    )
    db.add(row)
    await db.commit()
    await db.refresh(row)
    return _to_out(row)


@router.put("/{name}")
async def update_model(name: str, req: UpdateModelRequest, db: AsyncSession = Depends(get_db)):
    for m in builtin_models():
        if m["name"] == name:
            raise HTTPException(400, "内置模型不可修改，请在页面添加同名模型")
    row = (await db.execute(select(AiModel).where(AiModel.name == name))).scalar()
    if not row:
        raise HTTPException(404, "模型不存在")
    if req.label is not None:
        row.label = req.label.strip() or row.label
    if req.kind is not None:
        row.kind = req.kind
    if req.base_url is not None:
        row.base_url = req.base_url.strip()
    if req.api_key is not None:
        row.api_key = req.api_key.strip()
    if req.model_id is not None:
        row.model_id = req.model_id.strip()
    await db.commit()
    await db.refresh(row)
    return _to_out(row)


@router.delete("/{name}")
async def delete_model(name: str, db: AsyncSession = Depends(get_db)):
    for m in builtin_models():
        if m["name"] == name:
            raise HTTPException(400, "内置模型不可删除")
    row = (await db.execute(select(AiModel).where(AiModel.name == name))).scalar()
    if not row:
        raise HTTPException(404, "模型不存在")
    await db.delete(row)
    await db.commit()
    return {"deleted": name}


@router.post("/{name}/test")
async def test_model(name: str, db: AsyncSession = Depends(get_db)):
    """用该模型的配置发一条极短请求，验证 base_url/api_key/model_id 是否可用。"""
    provider = await resolve_provider(db, name)
    if provider is None:
        raise HTTPException(404, "模型不存在")
    try:
        resp = await provider.chat(
            [{"role": "user", "content": "hi"}],
            temperature=0,
        )
        return {"ok": True, "model": resp.model, "reply": resp.content[:80]}
    except Exception as e:
        return {"ok": False, "error": f"{type(e).__name__}: {e}"}


@router.post("/import")
async def import_models(req: ImportModelsRequest):
    """通用 OpenAI 兼容 /models 拉取（OpenRouter / 硅基流动 / SenseNova / 任取 base_url），
    供前端挑选后批量添加（不落库）。"""
    key = (req.api_key or "").strip()
    if not key:
        raise HTTPException(400, "需要 API Key")
    base = (req.base_url or "").strip().rstrip("/")
    if not base:
        raise HTTPException(400, "需要 Base URL")
    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.get(
                f"{base}/models",
                headers={"Authorization": f"Bearer {key}"},
            )
            if resp.status_code >= 400:
                raise HTTPException(resp.status_code, f"OpenRouter 返回 {resp.status_code}: {resp.text[:200]}")
            data = resp.json()
    except httpx.HTTPError as e:
        raise HTTPException(502, f"请求模型列表失败: {type(e).__name__}: {e}")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(502, f"解析模型列表失败: {e}")

    models = []
    for m in data.get("data", []):
        mid = (m.get("id") or "").strip()
        if not mid:
            continue
        low = mid.lower()
        if any(b in low for b in OPENROUTER_BLOCKLIST):
            continue
        models.append({"id": mid, "name": (m.get("name") or mid)})
    return {"models": models, "count": len(models)}


@router.post("/import-openrouter")
async def import_openrouter_models(req: ImportOpenRouterRequest, db: AsyncSession = Depends(get_db)):
    """兼容旧调用：等价于通用 /import（默认 base_url = OpenRouter）。"""
    return await import_models(ImportModelsRequest(api_key=req.api_key, base_url=req.base_url))