# Copyright 2026 zhaoxi826 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from kilostar.core.postgres_database.model.individual import ( BaseIndividualModel, SpecialistIndividualModel, OrdinaryIndividualModel, SpecialIndividualModel, ) from sqlalchemy import select, and_ from typing import List, Optional from kilostar.core.postgres_database.database_exception import database_exception from kilostar.utils.error import BusinessError from ulid import ULID _AGENT_TYPE_MODEL_MAP = { "specialist": SpecialistIndividualModel, "ordinary": OrdinaryIndividualModel, "special": SpecialIndividualModel, } class IndividualDatabase: """Individual 表族（Base/Specialist/Ordinary/Special）的 DAO，按 agent_type 选择具体子表。""" def __init__(self, async_session_maker): self.async_session_maker = async_session_maker @staticmethod def _select_model(agent_type: str): return _AGENT_TYPE_MODEL_MAP.get(agent_type, BaseIndividualModel) @database_exception async def add_worker_individual(self, **kwargs): """新建一个 Worker Individual：自动生成 ULID，按 ``agent_type`` 选择对应子表写入。""" async with self.async_session_maker() as session: agent_id = str(ULID()) agent_type = kwargs.get("agent_type", "base") model_cls = self._select_model(agent_type) individual = model_cls(agent_id=agent_id, **kwargs) session.add(individual) await session.commit() await session.refresh(individual) return individual @database_exception async def get_worker_individual(self, agent_id: str): """按 agent_id 取单个 Individual；不存在返回 None。""" async with self.async_session_maker() as session: statement = select(BaseIndividualModel).where( BaseIndividualModel.agent_id == agent_id ) results = await session.execute(statement) return results.scalar_one_or_none() @database_exception async def get_worker_individual_list(self, owner_id: str): """读取某用户名下的所有 Individual。""" async with self.async_session_maker() as session: statement = select(BaseIndividualModel).where( BaseIndividualModel.owner_id == owner_id ) results = await session.execute(statement) return list(results.scalars().all()) @database_exception async def update_worker_individual(self, agent_id: str, **kwargs): """部分更新 Individual：只覆盖 kwargs 中非 None 的字段；找不到返回 None。""" async with self.async_session_maker() as session: statement = select(BaseIndividualModel).where( BaseIndividualModel.agent_id == agent_id ) results = await session.execute(statement) individual = results.scalar_one_or_none() if not individual: return None for key, value in kwargs.items(): if value is not None: setattr(individual, key, value) session.add(individual) await session.commit() await session.refresh(individual) return individual @database_exception async def delete_worker_individual(self, agent_id: str) -> bool: """删除 Individual；不存在返回 False，删除成功返回 True。 ``plugin_owned`` 不为空时拒绝删除（插件 agent 由插件生命周期管理， 要清理该插件目录后由启动期兜底逻辑收回）。 """ async with self.async_session_maker() as session: statement = select(BaseIndividualModel).where( BaseIndividualModel.agent_id == agent_id ) results = await session.execute(statement) individual = results.scalar_one_or_none() if not individual: return False if individual.plugin_owned: raise BusinessError( f"agent {agent_id} 由插件 {individual.plugin_owned} 拥有，" f"不可删除（请改卸载插件）" ) await session.delete(individual) await session.commit() return True @database_exception async def get_all_worker_individual(self): """返回数据库中全部 Individual。""" async with self.async_session_maker() as session: statement = select(BaseIndividualModel) results = await session.execute(statement) return list(results.scalars().all()) # ─── plugin_owned slot 专用 ───────────────────────────────────── @database_exception async def find_plugin_slot(self, plugin_name: str, slot_name: str): """按 ``(plugin_owned, agent_name)`` 查 slot；返回 ORM 对象或 None。""" async with self.async_session_maker() as session: stmt = select(BaseIndividualModel).where( and_( BaseIndividualModel.plugin_owned == plugin_name, BaseIndividualModel.agent_name == slot_name, ) ) return (await session.execute(stmt)).scalar_one_or_none() @database_exception async def upsert_plugin_slot( self, plugin_name: str, slot_name: str, description: str, owner_id: str = "system", node_affinity: str = "cpu", ): """插件安装期登记一个 agent slot。 slot 的 provider/model 留空，等用户在前端 Agent 页面装配。已存在则 只刷新 description/node_affinity（用户自己装配的 provider+model 不被覆盖）。 """ async with self.async_session_maker() as session: stmt = select(BaseIndividualModel).where( and_( BaseIndividualModel.plugin_owned == plugin_name, BaseIndividualModel.agent_name == slot_name, ) ) existing = (await session.execute(stmt)).scalar_one_or_none() if existing is not None: existing.description = description existing.node_affinity = node_affinity session.add(existing) await session.commit() await session.refresh(existing) return existing row = BaseIndividualModel( agent_id=str(ULID()), agent_name=slot_name, description=description, provider_title="", model_id="", owner_id=owner_id, agent_type="base", node_affinity=node_affinity, plugin_owned=plugin_name, ) session.add(row) await session.commit() await session.refresh(row) return row @database_exception async def list_plugin_owned_names(self) -> List[str]: """返回当前 DB 中所有出现过的 ``plugin_owned`` 值（去重）。""" async with self.async_session_maker() as session: stmt = select(BaseIndividualModel.plugin_owned).where( BaseIndividualModel.plugin_owned.is_not(None) ).distinct() return [r for (r,) in (await session.execute(stmt)).all() if r] @database_exception async def delete_plugin_slots(self, plugin_name: str) -> int: """删掉 ``plugin_owned == plugin_name`` 的所有 slot；返回被删条数。""" async with self.async_session_maker() as session: stmt = select(BaseIndividualModel).where( BaseIndividualModel.plugin_owned == plugin_name ) rows = list((await session.execute(stmt)).scalars().all()) for row in rows: await session.delete(row) await session.commit() return len(rows)