ForcePilot/backend/package/yuxi/repositories/subagent_repository.py

"""SubAgent 数据访问层"""

from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from yuxi.storage.postgres.models_business import SubAgent
from yuxi.utils.datetime_utils import utc_now_naive


class SubAgentRepository:
    def __init__(self, db_session: AsyncSession):
        self.db = db_session

    async def list_all(self) -> list[SubAgent]:
        """获取所有 SubAgent，按 updated_at 降序"""
        result = await self.db.execute(select(SubAgent).order_by(SubAgent.updated_at.desc()))
        return list(result.scalars().all())

    async def list_enabled(self) -> list[SubAgent]:
        """获取已启用的 SubAgent。"""
        result = await self.db.execute(
            select(SubAgent).where(SubAgent.enabled.is_(True)).order_by(SubAgent.updated_at.desc())
        )
        return list(result.scalars().all())

    async def list_all_specs(self) -> list[dict[str, Any]]:
        """获取已启用的 SubAgent 运行规格，按 updated_at 降序。"""
        items = await self.list_enabled()
        return [item.to_subagent_spec() for item in items]

    async def get_by_name(self, name: str) -> SubAgent | None:
        """根据名称获取 SubAgent"""
        result = await self.db.execute(select(SubAgent).where(SubAgent.name == name))
        return result.scalar_one_or_none()

    async def exists_name(self, name: str) -> bool:
        """检查名称是否存在（仅查询计数，不获取完整数据）"""
        from sqlalchemy import func, select

        result = await self.db.execute(select(func.count()).select_from(SubAgent).where(SubAgent.name == name))
        return result.scalar() > 0

    async def create(
        self,
        *,
        name: str,
        description: str,
        system_prompt: str,
        tools: list[str] | None,
        model: str | None,
        is_builtin: bool,
        created_by: str | None,
    ) -> SubAgent:
        now = utc_now_naive()
        item = SubAgent(
            name=name,
            description=description,
            system_prompt=system_prompt,
            tools=tools or [],
            model=model,
            enabled=True,
            is_builtin=is_builtin,
            created_by=created_by,
            updated_by=created_by,
            created_at=now,
            updated_at=now,
        )
        self.db.add(item)
        await self.db.commit()
        await self.db.refresh(item)
        return item

    async def update(
        self,
        item: SubAgent,
        *,
        description: str | None,
        system_prompt: str | None,
        tools: list[str] | None,
        model: str | None,
        model_provided: bool = False,
        updated_by: str | None,
    ) -> SubAgent:
        # 批量更新非空字段
        updates = {
            "description": description,
            "system_prompt": system_prompt,
            "tools": tools,
        }
        for field, value in updates.items():
            if value is not None:
                setattr(item, field, value)
        # model_provided=True 时显式设置 model 值（包括 None），用于清空字段
        if model_provided:
            item.model = model
        item.updated_by = updated_by
        item.updated_at = utc_now_naive()
        await self.db.commit()
        await self.db.refresh(item)
        return item

    async def delete(self, item: SubAgent) -> None:
        """删除 SubAgent"""
        await self.db.delete(item)
        await self.db.commit()