KiloStar/kilostar/core/postgres_database/model/individual.py

#  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 enum import Enum
from typing import List, Optional, Dict, Any
from sqlalchemy import String, Text, text, ForeignKey, Enum as SAEnum
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship

from .base import BaseDataModel


class ModalityType(str, Enum):
    TEXT = "text"
    VISION = "vision"
    AUDIO = "audio"
    MULTIMODAL = "multimodal"


# ==========================================
# 1. 通用基类表 (身份中心)
# ==========================================
class BaseIndividualModel(BaseDataModel):
    __tablename__ = "base_individual"

    agent_id: Mapped[str] = mapped_column(String(64), primary_key=True)
    agent_name: Mapped[str] = mapped_column(String(100), index=True, nullable=False)
    description: Mapped[str] = mapped_column(Text, nullable=False)
    system_prompt: Mapped[Optional[str]] = mapped_column(Text)
    provider_title: Mapped[str] = mapped_column(String(50))
    model_id: Mapped[str] = mapped_column(String(100))
    owner_id: Mapped[str] = mapped_column(String(64), index=True)

    agent_type: Mapped[str] = mapped_column(String(32))
    node_affinity: Mapped[str] = mapped_column(String(32), nullable=False, default="cpu")
    template_origin_id: Mapped[Optional[str]] = mapped_column(
        ForeignKey("persona_template.template_id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )

    __mapper_args__ = {"polymorphic_on": "agent_type", "polymorphic_identity": "base"}


# ==========================================
# 2. 专家子个体 (技能与复杂工作流)
# ==========================================
class SpecialistIndividualModel(BaseIndividualModel):
    __tablename__ = "specialist_individual"

    agent_id: Mapped[str] = mapped_column(
        ForeignKey("base_individual.agent_id", ondelete="CASCADE"), primary_key=True
    )
    bound_skill: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSONB)
    workspace: Mapped[Optional[List[str]]] = mapped_column(JSONB)
    tools: Mapped[Optional[List[str]]] = mapped_column(
        JSONB, default=list, server_default=text("'[]'::jsonb")
    )

    # 逻辑关联：作为管理者，管理下属个体
    sub_ordinary_agents: Mapped[List["OrdinaryIndividualModel"]] = relationship(
        back_populates="manager",
        cascade="all, delete-orphan",
        foreign_keys="[OrdinaryIndividualModel.manager_id]",
    )
    sub_special_agents: Mapped[List["SpecialIndividualModel"]] = relationship(
        back_populates="manager",
        cascade="all, delete-orphan",
        foreign_keys="[SpecialIndividualModel.manager_id]",
    )

    __mapper_args__ = {
        "polymorphic_identity": "specialist",
    }


# ==========================================
# 3. 基础子个体 (普通微调模型)
# ==========================================
class OrdinaryIndividualModel(BaseIndividualModel):
    __tablename__ = "ordinary_individual"

    agent_id: Mapped[str] = mapped_column(
        ForeignKey("base_individual.agent_id", ondelete="CASCADE"), primary_key=True
    )
    finetuned_from: Mapped[Optional[str]] = mapped_column(String(100))
    tools: Mapped[Optional[List[str]]] = mapped_column(
        JSONB, default=list, server_default=text("'[]'::jsonb")
    )

    # 【修复1】：必须显式定义物理外键
    manager_id: Mapped[Optional[str]] = mapped_column(
        ForeignKey("specialist_individual.agent_id", ondelete="SET NULL")
    )

    # 逻辑关联：指向上级专家
    manager: Mapped[Optional["SpecialistIndividualModel"]] = relationship(
        back_populates="sub_ordinary_agents",
        foreign_keys=[manager_id],  # 显式指定使用 manager_id 解析关系
    )

    __mapper_args__ = {
        "polymorphic_identity": "ordinary",
    }


# ==========================================
# 4. 特殊子个体 (多模态)
# ==========================================
class SpecialIndividualModel(BaseIndividualModel):
    __tablename__ = "special_individual"

    agent_id: Mapped[str] = mapped_column(
        ForeignKey("base_individual.agent_id", ondelete="CASCADE"), primary_key=True
    )
    modality_type: Mapped[ModalityType] = mapped_column(
        SAEnum(ModalityType),
        default=ModalityType.MULTIMODAL,
    )
    multimodal_config: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSONB)

    # 【修复1】：添加缺失的物理外键
    manager_id: Mapped[Optional[str]] = mapped_column(
        ForeignKey("specialist_individual.agent_id", ondelete="SET NULL")
    )

    # 【修复2】：修正 back_populates 指向正确的变量名
    manager: Mapped[Optional["SpecialistIndividualModel"]] = relationship(
        back_populates="sub_special_agents", foreign_keys=[manager_id]
    )

    __mapper_args__ = {
        "polymorphic_identity": "special",
    }