PMFlow/backend/app/api/llm_engines.py

"""LLM 引擎配置 API。"""

from __future__ import annotations

from typing import Literal, Optional

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

from ..core.response import success_response
from ..database import get_db
from ..dependencies.auth import get_current_user
from ..dependencies.tenant import get_tenant_id
from ..models import User
from ..services.auth import has_any_role
from ..services.llm import llm_service

router = APIRouter(prefix="/ai-steward/llm-engines", tags=["AI Steward"])


class LLMEngineCreateRequest(BaseModel):
    name: str = Field(..., min_length=1, max_length=80)
    provider: Literal["deepseek", "glm", "openai", "ollama"]
    api_key: Optional[str] = Field(default="", max_length=500)
    model: Optional[str] = Field(default=None, max_length=120)
    base_url: Optional[str] = Field(default=None, max_length=255)
    set_active: bool = True


class LLMEngineUpdateRequest(BaseModel):
    name: Optional[str] = Field(default=None, min_length=1, max_length=80)
    provider: Optional[Literal["deepseek", "glm", "openai", "ollama"]] = None
    api_key: Optional[str] = Field(default=None, max_length=500)
    model: Optional[str] = Field(default=None, max_length=120)
    base_url: Optional[str] = Field(default=None, max_length=255)


class LLMEngineTestRequest(BaseModel):
    prompt: Optional[str] = Field(default=None, max_length=1000)


def _require_manager(user: User) -> None:
    if not has_any_role(str(user.role or ""), ["BOSS", "ADMIN"]):
        raise HTTPException(status_code=403, detail="仅管理员可操作 LLM 引擎配置")


@router.get("")
async def list_llm_engines(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _ = current_user
    data = await llm_service.list_engines(db, tenant_id=tenant_id)
    return success_response(data)


@router.post("")
async def create_llm_engine(
    body: LLMEngineCreateRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        engine = await llm_service.create_engine(
            db,
            tenant_id=tenant_id,
            name=body.name,
            provider=body.provider,
            api_key=body.api_key,
            model=body.model,
            base_url=body.base_url,
            set_active=body.set_active,
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({"engine": engine, **all_items})
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.put("/{engine_id}")
async def update_llm_engine(
    engine_id: str,
    body: LLMEngineUpdateRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        engine = await llm_service.update_engine(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            name=body.name,
            provider=body.provider,
            api_key=body.api_key,
            model=body.model,
            base_url=body.base_url,
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({"engine": engine, **all_items})
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/{engine_id}/activate")
async def activate_llm_engine(
    engine_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        result = await llm_service.activate_engine(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({**result, **all_items})
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/test-active")
async def test_active_llm_engine(
    body: LLMEngineTestRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        data = await llm_service.test_engine_connection(
            db,
            tenant_id=tenant_id,
            engine_id=None,
            prompt=body.prompt,
        )
        return success_response(data)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/{engine_id}/test")
async def test_llm_engine(
    engine_id: str,
    body: LLMEngineTestRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        data = await llm_service.test_engine_connection(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            prompt=body.prompt,
        )
        return success_response(data)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.delete("/{engine_id}")
async def delete_llm_engine(
    engine_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    _require_manager(current_user)
    try:
        result = await llm_service.delete_engine(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({**result, **all_items})
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc