core/models/agent.py

"""
AI Agent Models

Database models for AI agent drafts and approvals.
"""

from datetime import datetime
from uuid import uuid4

from sqlalchemy import JSON, Column, DateTime, Float, String, Text

from core.models.base import Base


class AgentDraft(Base):
    """Model for AI-generated agent drafts"""

    __tablename__ = "agent_drafts"

    id = Column(String(36), primary_key=True, default=lambda: str(uuid4()))
    agent_name = Column(String(100), nullable=False)
    draft_type = Column(String(50), nullable=False)
    title = Column(Text, nullable=False)
    content = Column(Text, nullable=False)
    target_entity = Column(String(255))
    confidence = Column(Float, default=0.0)
    status = Column(String(50), default="draft")
    reviewer = Column(String(100))
    tags = Column(JSON, default=list)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    def to_dict(self):
        return {
            "id": self.id,
            "agentName": self.agent_name,
            "draftType": self.draft_type,
            "title": self.title,
            "content": self.content,
            "targetEntity": self.target_entity,
            "confidence": self.confidence,
            "status": self.status,
            "reviewer": self.reviewer,
            "tags": self.tags,
            "createdAt": self.created_at.isoformat() if self.created_at else None,
            "lastModified": self.updated_at.isoformat() if self.updated_at else None,
        }


class AgentApproval(Base):
    """Model for AI agent approval requests"""

    __tablename__ = "agent_approvals"

    id = Column(String(36), primary_key=True, default=lambda: str(uuid4()))
    agent_name = Column(String(100), nullable=False)
    action = Column(String(100), nullable=False)
    target = Column(String(255), nullable=False)
    confidence = Column(Float, nullable=False)
    status = Column(String(50), default="pending")
    risk = Column(String(20))
    details = Column(Text)
    approved_by = Column(String(100))
    approved_at = Column(DateTime)
    created_at = Column(DateTime, default=datetime.utcnow)

    def to_dict(self):
        return {
            "id": self.id,
            "agentName": self.agent_name,
            "action": self.action,
            "target": self.target,
            "confidence": self.confidence,
            "status": self.status,
            "risk": self.risk,
            "details": self.details,
            "approvedBy": self.approved_by,
            "approvedAt": self.approved_at.isoformat() if self.approved_at else None,
            "timestamp": self.created_at.isoformat() if self.created_at else None,
        }