ai-video-fullstack/backend/routes/assistants.py

"""Assistant CRUD backed by capability-to-model-resource bindings."""

import uuid

from db.models import Assistant, AssistantModelBinding, ModelResource
from db.session import get_session
from fastapi import APIRouter, Depends, HTTPException
from schemas import AssistantOut, AssistantUpsert
from services.masking import mask, resolve_incoming_key
from services.node_specs import validate_graph
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

router = APIRouter(prefix="/api/assistants", tags=["assistants"])
CAPABILITIES = ("LLM", "ASR", "TTS", "Realtime", "Embedding")


def _validate_workflow(body: AssistantUpsert) -> None:
    """workflow 类型:保存前校验图结构,不通过则 400。其他类型跳过。"""
    if body.type != "workflow":
        return
    errors = validate_graph(body.graph or {})
    if errors:
        raise HTTPException(400, "工作流校验失败:" + ";".join(errors))


async def _sync_bindings(
    session: AsyncSession, assistant_id: str, resource_ids: dict[str, str]
) -> None:
    for capability in CAPABILITIES:
        resource_id = resource_ids.get(capability)
        binding = await session.get(AssistantModelBinding, (assistant_id, capability))
        if not resource_id:
            if binding:
                await session.delete(binding)
            continue
        resource = await session.get(ModelResource, resource_id)
        if not resource or resource.capability != capability:
            raise HTTPException(400, f"{capability} 绑定必须引用同能力的模型资源")
        if binding:
            binding.model_resource_id = resource_id
        else:
            session.add(
                AssistantModelBinding(
                    assistant_id=assistant_id,
                    capability=capability,
                    model_resource_id=resource_id,
                    config={},
                )
            )


async def _resource_ids(session: AsyncSession, assistant_id: str) -> dict[str, str]:
    bindings = (
        await session.execute(
            select(AssistantModelBinding).where(
                AssistantModelBinding.assistant_id == assistant_id
            )
        )
    ).scalars().all()
    return {binding.capability: binding.model_resource_id for binding in bindings}


async def _to_out(session: AsyncSession, assistant: Assistant) -> AssistantOut:
    return AssistantOut(
        id=assistant.id,
        name=assistant.name,
        type=assistant.type,  # type: ignore[arg-type]
        runtime_mode=assistant.runtime_mode,  # type: ignore[arg-type]
        greeting=assistant.greeting,
        enable_interrupt=assistant.enable_interrupt,
        model_resource_ids=await _resource_ids(session, assistant.id),
        knowledge_base_id=assistant.knowledge_base_id,
        prompt=assistant.prompt,
        api_url=assistant.api_url,
        api_key=mask(assistant.api_key),
        app_id=assistant.app_id,
        graph=assistant.graph or {},
        updated_at=assistant.updated_at.isoformat() if assistant.updated_at else None,
    )


@router.get("", response_model=list[AssistantOut])
async def list_assistants(session: AsyncSession = Depends(get_session)):
    rows = (
        await session.execute(select(Assistant).order_by(Assistant.updated_at.desc()))
    ).scalars().all()
    return [await _to_out(session, assistant) for assistant in rows]


@router.post("", response_model=AssistantOut)
async def create_assistant(
    body: AssistantUpsert, session: AsyncSession = Depends(get_session)
):
    _validate_workflow(body)
    data = body.model_dump()
    resource_ids = data.pop("model_resource_ids")
    assistant = Assistant(id=f"asst_{uuid.uuid4().hex[:12]}", **data)
    session.add(assistant)
    await session.flush()
    await _sync_bindings(session, assistant.id, resource_ids)
    await session.commit()
    await session.refresh(assistant)
    return await _to_out(session, assistant)


@router.get("/{assistant_id}", response_model=AssistantOut)
async def get_assistant(
    assistant_id: str, session: AsyncSession = Depends(get_session)
):
    assistant = await session.get(Assistant, assistant_id)
    if not assistant:
        raise HTTPException(404, "助手不存在")
    return await _to_out(session, assistant)


@router.post("/{assistant_id}/duplicate", response_model=AssistantOut)
async def duplicate_assistant(
    assistant_id: str, session: AsyncSession = Depends(get_session)
):
    source = await session.get(Assistant, assistant_id)
    if not source:
        raise HTTPException(404, "助手不存在")
    assistant = Assistant(
        id=f"asst_{uuid.uuid4().hex[:12]}",
        name=f"{source.name} 副本",
        type=source.type,
        runtime_mode=source.runtime_mode,
        greeting=source.greeting,
        enable_interrupt=source.enable_interrupt,
        knowledge_base_id=source.knowledge_base_id,
        prompt=source.prompt,
        api_url=source.api_url,
        api_key=source.api_key,
        app_id=source.app_id,
        graph=dict(source.graph or {}),
    )
    session.add(assistant)
    await session.flush()
    await _sync_bindings(session, assistant.id, await _resource_ids(session, source.id))
    await session.commit()
    await session.refresh(assistant)
    return await _to_out(session, assistant)


@router.put("/{assistant_id}", response_model=AssistantOut)
async def update_assistant(
    assistant_id: str,
    body: AssistantUpsert,
    session: AsyncSession = Depends(get_session),
):
    assistant = await session.get(Assistant, assistant_id)
    if not assistant:
        raise HTTPException(404, "助手不存在")
    _validate_workflow(body)
    data = body.model_dump()
    resource_ids = data.pop("model_resource_ids")
    data["api_key"] = resolve_incoming_key(data["api_key"], assistant.api_key)
    for key, value in data.items():
        setattr(assistant, key, value)
    await _sync_bindings(session, assistant.id, resource_ids)
    await session.commit()
    await session.refresh(assistant)
    return await _to_out(session, assistant)


@router.delete("/{assistant_id}")
async def delete_assistant(
    assistant_id: str, session: AsyncSession = Depends(get_session)
):
    assistant = await session.get(Assistant, assistant_id)
    if not assistant:
        raise HTTPException(404, "助手不存在")
    await session.delete(assistant)
    await session.commit()
    return {"ok": True}