pit-router/app/models/__init__.py

"""
PIT Router 数据模型
"""
from datetime import datetime
from typing import Optional, List
from sqlalchemy import String, DateTime, Integer, Text, JSON, ForeignKey, Boolean
from sqlalchemy.orm import Mapped, mapped_column, relationship
import uuid

from app.extensions import db


def generate_uuid() -> str:
    """生成 UUID"""
    return str(uuid.uuid4())


class User(db.Model):
    """用户模型"""
    __tablename__ = 'users'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    username: Mapped[str] = mapped_column(String(80), unique=True, nullable=False)
    password_hash: Mapped[str] = mapped_column(String(256), nullable=False)
    email: Mapped[str] = mapped_column(String(120), unique=True, nullable=False)
    nickname: Mapped[Optional[str]] = mapped_column(String(80), nullable=True)
    role: Mapped[str] = mapped_column(String(20), default='user')
    status: Mapped[str] = mapped_column(String(20), default='active')
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    last_login_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    
    # 关联
    sessions = relationship('Session', back_populates='user', cascade='all, delete-orphan')
    
    def __repr__(self):
        return f'<User {self.username}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'username': self.username,
            'email': self.email,
            'nickname': self.nickname,
            'role': self.role,
            'status': self.status,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'last_login_at': self.last_login_at.isoformat() if self.last_login_at else None,
        }


class Gateway(db.Model):
    """Gateway 模型"""
    __tablename__ = 'gateways'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    name: Mapped[str] = mapped_column(String(80), unique=True, nullable=False)
    url: Mapped[str] = mapped_column(String(256), nullable=False)
    token_hash: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)
    status: Mapped[str] = mapped_column(String(20), default='offline')
    agent_count: Mapped[int] = mapped_column(Integer, default=0)
    connection_limit: Mapped[int] = mapped_column(Integer, default=10)
    heartbeat_interval: Mapped[int] = mapped_column(Integer, default=60)
    allowed_ips: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
    last_heartbeat: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    
    # 关联
    agents = relationship('Agent', back_populates='gateway', cascade='all, delete-orphan')
    
    def __repr__(self):
        return f'<Gateway {self.name}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'url': self.url,
            'status': self.status,
            'agent_count': self.agent_count,
            'connection_limit': self.connection_limit,
            'heartbeat_interval': self.heartbeat_interval,
            'last_heartbeat': self.last_heartbeat.isoformat() if self.last_heartbeat else None,
            'created_at': self.created_at.isoformat() if self.created_at else None,
        }


class Agent(db.Model):
    """Agent 模型"""
    __tablename__ = 'agents'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    name: Mapped[str] = mapped_column(String(80), nullable=False)
    display_name: Mapped[Optional[str]] = mapped_column(String(80), nullable=True)
    gateway_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey('gateways.id'), nullable=True)
    socket_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    model: Mapped[Optional[str]] = mapped_column(String(80), nullable=True)
    capabilities: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
    status: Mapped[str] = mapped_column(String(20), default='offline')
    priority: Mapped[int] = mapped_column(Integer, default=5)
    weight: Mapped[int] = mapped_column(Integer, default=10)
    connection_limit: Mapped[int] = mapped_column(Integer, default=5)
    current_sessions: Mapped[int] = mapped_column(Integer, default=0)
    last_heartbeat: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    
    # 关联
    gateway = relationship('Gateway', back_populates='agents')
    sessions = relationship('Session', back_populates='agent')
    
    def __repr__(self):
        return f'<Agent {self.name}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'display_name': self.display_name,
            'gateway_id': self.gateway_id,
            'socket_id': self.socket_id,
            'model': self.model,
            'capabilities': self.capabilities,
            'status': self.status,
            'priority': self.priority,
            'weight': self.weight,
            'connection_limit': self.connection_limit,
            'current_sessions': self.current_sessions,
            'last_heartbeat': self.last_heartbeat.isoformat() if self.last_heartbeat else None,
            'created_at': self.created_at.isoformat() if self.created_at else None,
        }


class Bot(db.Model):
    """机器人模型 - Agent 的用户侧展示配置"""
    __tablename__ = 'bots'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    name: Mapped[str] = mapped_column(String(80), unique=True, nullable=False)
    display_name: Mapped[Optional[str]] = mapped_column(String(80), nullable=True)
    avatar: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)
    description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    owner_id: Mapped[str] = mapped_column(String(36), ForeignKey('users.id'), nullable=False)
    agent_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey('agents.id'), nullable=True)
    token_hash: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)
    is_system: Mapped[bool] = mapped_column(Boolean, default=False)
    status: Mapped[str] = mapped_column(String(20), default='offline')
    capabilities: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
    config: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    last_active_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    
    # 关联
    owner = relationship('User', foreign_keys=[owner_id])
    agent = relationship('Agent', foreign_keys=[agent_id])
    
    def __repr__(self):
        return f'<Bot {self.name}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'name': self.name,
            'display_name': self.display_name,
            'avatar': self.avatar,
            'description': self.description,
            'owner_id': self.owner_id,
            'agent_id': self.agent_id,
            'is_system': self.is_system,
            'status': self.status,
            'capabilities': self.capabilities,
            'config': self.config,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'last_active_at': self.last_active_at.isoformat() if self.last_active_at else None,
        }


class Session(db.Model):
    """会话模型"""
    __tablename__ = 'sessions'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    user_id: Mapped[str] = mapped_column(String(36), ForeignKey('users.id'), nullable=False)
    bot_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey('bots.id'), nullable=True)
    primary_agent_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey('agents.id'), nullable=True)
    participating_agent_ids: Mapped[Optional[str]] = mapped_column(JSON, nullable=True)
    user_socket_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    title: Mapped[Optional[str]] = mapped_column(String(200), nullable=True)
    channel_type: Mapped[str] = mapped_column(String(20), default='web')
    status: Mapped[str] = mapped_column(String(20), default='active')
    message_count: Mapped[int] = mapped_column(Integer, default=0)
    unread_count: Mapped[int] = mapped_column(Integer, default=0)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    last_active_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    
    # 关联
    user = relationship('User', back_populates='sessions')
    bot = relationship('Bot', foreign_keys=[bot_id])
    agent = relationship('Agent', back_populates='sessions')
    messages = relationship('Message', back_populates='session', cascade='all, delete-orphan')
    
    def __repr__(self):
        return f'<Session {self.id}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'user_id': self.user_id,
            'bot_id': self.bot_id,
            'primary_agent_id': self.primary_agent_id,
            'participating_agent_ids': self.participating_agent_ids,
            'user_socket_id': self.user_socket_id,
            'title': self.title,
            'channel_type': self.channel_type,
            'status': self.status,
            'message_count': self.message_count,
            'unread_count': self.unread_count,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
            'last_active_at': self.last_active_at.isoformat() if self.last_active_at else None,
        }


class Message(db.Model):
    """消息模型"""
    __tablename__ = 'messages'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    session_id: Mapped[str] = mapped_column(String(36), ForeignKey('sessions.id'), nullable=False)
    sender_type: Mapped[str] = mapped_column(String(20), nullable=False)
    sender_id: Mapped[str] = mapped_column(String(36), nullable=False)
    sender_name: Mapped[Optional[str]] = mapped_column(String(80), nullable=True)
    bot_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey('bots.id'), nullable=True)
    message_type: Mapped[str] = mapped_column(String(20), default='text')
    content: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    content_type: Mapped[str] = mapped_column(String(20), default='markdown')
    reply_to: Mapped[Optional[str]] = mapped_column(String(36), nullable=True)
    status: Mapped[str] = mapped_column(String(20), default='sent')
    ack_status: Mapped[str] = mapped_column(String(20), default='pending')
    retry_count: Mapped[int] = mapped_column(Integer, default=0)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    delivered_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    
    # 关联
    session = relationship('Session', back_populates='messages')
    bot = relationship('Bot', foreign_keys=[bot_id])
    
    def __repr__(self):
        return f'<Message {self.id}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'session_id': self.session_id,
            'sender_type': self.sender_type,
            'sender_id': self.sender_id,
            'sender_name': self.sender_name,
            'bot_id': self.bot_id,
            'message_type': self.message_type,
            'content': self.content,
            'content_type': self.content_type,
            'reply_to': self.reply_to,
            'status': self.status,
            'ack_status': self.ack_status,
            'retry_count': self.retry_count,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'delivered_at': self.delivered_at.isoformat() if self.delivered_at else None,
        }


class Connection(db.Model):
    """连接模型"""
    __tablename__ = 'connections'
    
    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=generate_uuid)
    socket_id: Mapped[str] = mapped_column(String(100), unique=True, nullable=False)
    connection_type: Mapped[str] = mapped_column(String(20), nullable=False)
    entity_id: Mapped[str] = mapped_column(String(36), nullable=False)
    entity_type: Mapped[str] = mapped_column(String(20), nullable=False)
    ip_address: Mapped[Optional[str]] = mapped_column(String(45), nullable=True)
    user_agent: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    status: Mapped[str] = mapped_column(String(20), default='connected')
    auth_token: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    connected_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    last_activity: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    disconnected_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    
    def __repr__(self):
        return f'<Connection {self.socket_id}>'
    
    def to_dict(self):
        return {
            'id': self.id,
            'socket_id': self.socket_id,
            'connection_type': self.connection_type,
            'entity_id': self.entity_id,
            'entity_type': self.entity_type,
            'ip_address': self.ip_address,
            'user_agent': self.user_agent,
            'status': self.status,
            'connected_at': self.connected_at.isoformat() if self.connected_at else None,
            'last_activity': self.last_activity.isoformat() if self.last_activity else None,
        }