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 ..core.database import get_db
from ..dependencies.permissions import require_platform_ai_ops
from ..dependencies.tenant import get_tenant_id
from ..models import User
from ..services.llm_service 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", "dashscope", "openai", "zhipu", "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
    capabilities: dict[str, object] = Field(default_factory=dict)


class LLMEngineUpdateRequest(BaseModel):
    name: Optional[str] = Field(default=None, min_length=1, max_length=80)
    provider: Optional[Literal["deepseek", "dashscope", "openai", "zhipu", "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)
    capabilities: Optional[dict[str, object]] = None


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


class LLMEngineCapabilityProbeRequest(BaseModel):
    mode: Literal["quick", "full", "advanced"] = "full"
    targets: list[
        Literal[
            "basic",
            "chat",
            "basic_chat",
            "model_self_report",
            "self_report",
            "model_report",
            "structured",
            "structured_output",
            "json",
            "json_schema",
            "tool",
            "tool_calling",
            "function_calling",
            "stream",
            "streaming",
            "long_context",
            "long_context_32k",
            "long_context_128k",
            "vision",
            "vision_input",
            "audio",
            "audio_input",
        ]
    ] = Field(default_factory=list, max_length=12)
    allow_high_cost: bool = False


class LLMEngineCapabilityAttestationRequest(BaseModel):
    capability_key: Literal["long_context"] = "long_context"
    tested_tokens: int = Field(..., gt=0, le=10_000_000)
    source_type: Literal["official_doc", "provider_console", "contract", "manual_verification"] = "manual_verification"
    source_url: Optional[str] = Field(default=None, max_length=500)
    source_note: Optional[str] = Field(default=None, max_length=500)
    expires_at: Optional[str] = Field(default=None, max_length=80)


class LLMEngineSettingsUpdateRequest(BaseModel):
    enable_local_kb_qa: bool = False


class LLMEngineEnabledUpdateRequest(BaseModel):
    enabled: bool = True


@router.get("")
async def list_llm_engines(
    current_user: User = Depends(require_platform_ai_ops),
    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(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    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,
            capabilities=body.capabilities,
        )
        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("/settings")
async def update_llm_engine_settings(
    body: LLMEngineSettingsUpdateRequest,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        settings = await llm_service.update_engine_settings(
            db,
            tenant_id=tenant_id,
            enable_local_kb_qa=body.enable_local_kb_qa,
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({"settings": settings, **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(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    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,
            capabilities=body.capabilities,
        )
        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(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        result = await llm_service.set_default_fallback_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("/{engine_id}/default-fallback")
async def set_llm_engine_default_fallback(
    engine_id: str,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        result = await llm_service.set_default_fallback_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.put("/{engine_id}/enabled")
async def update_llm_engine_enabled(
    engine_id: str,
    body: LLMEngineEnabledUpdateRequest,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        result = await llm_service.set_engine_enabled(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            enabled=body.enabled,
        )
        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(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        data = await llm_service.test_engine_connection(
            db,
            tenant_id=tenant_id,
            engine_id=None,
            prompt=body.prompt,
            user_id=int(current_user.id),
        )
        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(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        data = await llm_service.test_engine_connection(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            prompt=body.prompt,
            user_id=int(current_user.id),
        )
        return success_response(data)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/{engine_id}/capability-probe")
async def probe_llm_engine_capabilities(
    engine_id: str,
    body: LLMEngineCapabilityProbeRequest,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        data = await llm_service.probe_engine_capabilities(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            targets=body.targets,
            mode=body.mode,
            allow_high_cost=body.allow_high_cost,
            user_id=int(current_user.id),
        )
        all_items = await llm_service.list_engines(db, tenant_id=tenant_id)
        return success_response({"probe": data, **all_items})
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@router.post("/{engine_id}/capability-attest")
async def attest_llm_engine_capability(
    engine_id: str,
    body: LLMEngineCapabilityAttestationRequest,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    try:
        engine = await llm_service.attest_engine_capability(
            db,
            tenant_id=tenant_id,
            engine_id=engine_id,
            capability_key=body.capability_key,
            tested_tokens=body.tested_tokens,
            source_type=body.source_type,
            source_url=body.source_url,
            source_note=body.source_note,
            expires_at=body.expires_at,
            user_id=int(current_user.id),
        )
        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.delete("/{engine_id}")
async def delete_llm_engine(
    engine_id: str,
    current_user: User = Depends(require_platform_ai_ops),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    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