doc-forge/backend/routers/models.py

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

from database import get_db
from models.ai_model import AiModel
from schemas.schemas import AiModelCreate, AiModelUpdate, Response
from services.ai_service import call_ai
from services.security import decrypt_text, encrypt_text, mask_secret

router = APIRouter()


def _is_deepseek_model(model: AiModel) -> bool:
    provider = (model.provider or "").strip().lower()
    return provider == "deepseek"


def _serialize_model(model: AiModel) -> dict:
    api_key = decrypt_text(model.api_key_encrypted)
    return {
        "id": model.id,
        "name": model.name,
        "provider": model.provider,
        "api_format": model.api_format,
        "api_endpoint": model.api_endpoint,
        "api_key_preview": mask_secret(api_key),
        "supports_streaming": bool(model.supports_streaming),
        "enable_reasoning": bool(model.enable_reasoning),
        "status": model.status,
        "created_at": model.created_at,
    }


@router.get("")
async def list_models(db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(AiModel).order_by(AiModel.id.desc()))
    items = [_serialize_model(item) for item in result.scalars().all()]
    return Response(data=items)


@router.post("")
async def create_model(body: AiModelCreate, db: AsyncSession = Depends(get_db)):
    model = AiModel(
        name=body.name,
        provider=body.provider,
        api_format=body.api_format,
        api_endpoint=body.api_endpoint,
        api_key_encrypted=encrypt_text(body.api_key),
        supports_streaming=body.supports_streaming,
        enable_reasoning=body.enable_reasoning,
        status=body.status,
    )
    db.add(model)
    await db.commit()
    await db.refresh(model)
    return Response(data=_serialize_model(model))


@router.put("/{model_id}")
async def update_model(model_id: int, body: AiModelUpdate, db: AsyncSession = Depends(get_db)):
    model = await db.get(AiModel, model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型不存在")

    if body.name is not None:
        model.name = body.name
    if body.provider is not None:
        model.provider = body.provider
    if body.api_format is not None:
        model.api_format = body.api_format
    if body.api_endpoint is not None:
        model.api_endpoint = body.api_endpoint
    if body.supports_streaming is not None:
        model.supports_streaming = body.supports_streaming
    if body.enable_reasoning is not None:
        model.enable_reasoning = body.enable_reasoning
    if body.status is not None:
        model.status = body.status
    if body.api_key:
        model.api_key_encrypted = encrypt_text(body.api_key)

    await db.commit()
    await db.refresh(model)
    return Response(data=_serialize_model(model))


@router.delete("/{model_id}")
async def delete_model(model_id: int, db: AsyncSession = Depends(get_db)):
    model = await db.get(AiModel, model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型不存在")

    await db.delete(model)
    await db.commit()
    return Response(data={"id": model_id})


@router.post("/{model_id}/test")
async def test_model(model_id: int, db: AsyncSession = Depends(get_db)):
    model = await db.get(AiModel, model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型不存在")

    class FakeParagraph:
        title = "连接测试"
        content = "请返回一段非常简短的测试文本。"
        need_prompt = False
        prompt_text = ""
        output_format = "text"
        enable_reasoning = bool(model.enable_reasoning)

    try:
        result = await call_ai(FakeParagraph(), model)
        return Response(
            data={
                "id": model.id,
                "success": True,
                "message": f"模型 {model.name} 连接测试成功",
                "preview": result.content,
            }
        )
    except Exception as error:
        return Response(
            code=-1,
            message=str(error),
            data={"id": model.id, "success": False},
        )


@router.get("/{model_id}/balance")
async def get_model_balance(model_id: int, db: AsyncSession = Depends(get_db)):
    model = await db.get(AiModel, model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型不存在")
    if not _is_deepseek_model(model):
        raise HTTPException(status_code=400, detail="仅 DeepSeek 模型支持余额查询")

    api_key = decrypt_text(model.api_key_encrypted)
    if not api_key:
        raise HTTPException(status_code=400, detail="模型 API Key 不可用")

    try:
        async with httpx.AsyncClient(timeout=20, trust_env=False) as client:
            response = await client.get(
                "https://api.deepseek.com/user/balance",
                headers={"Authorization": f"Bearer {api_key}", "Accept": "application/json"},
            )
            response.raise_for_status()
            payload = response.json()
    except Exception as error:
        raise HTTPException(status_code=400, detail=f"查询余额失败：{error}")

    return Response(
        data={
            "id": model.id,
            "provider": model.provider,
            "is_available": payload.get("is_available", False),
            "balance_infos": payload.get("balance_infos", []),
        }
    )