kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/models/plan_execution.py

@@ -0,0 +1,102 @@
+"""Database model for plan executions."""
+
+from sqlalchemy import Column, String, DateTime, Text, Integer, Numeric, Index
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from sqlalchemy.sql import func
+from datetime import datetime
+import uuid
+import enum
+
+from control_plane_api.app.database import Base
+
+
+class PlanExecutionStatus(str, enum.Enum):
+    """Plan execution status"""
+    PENDING = "pending"
+    RUNNING = "running"
+    PENDING_USER_INPUT = "pending_user_input"  # Waiting for user to provide input
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class PlanExecution(Base):
+    """Model for tracking multi-task plan executions."""
+
+    __tablename__ = "plan_executions"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+
+    # Execution identifier (Temporal workflow ID)
+    execution_id = Column(String(255), unique=True, nullable=False, index=True)
+
+    # Organization and agent
+    organization_id = Column(String(255), nullable=False, index=True)
+    agent_id = Column(UUID(as_uuid=True), nullable=True)  # Primary agent for the plan
+
+    # Link to plan generation (for traceability)
+    plan_generation_id = Column(String(255), nullable=True, index=True)  # Execution ID of the plan generation
+
+    # Plan metadata
+    title = Column(Text, nullable=False)
+    summary = Column(Text, nullable=True)
+    total_tasks = Column(Integer, nullable=False, default=0)
+    completed_tasks = Column(Integer, nullable=False, default=0)
+    failed_tasks = Column(Integer, nullable=True, default=0)
+
+    # Status tracking
+    status = Column(String(50), nullable=False, default=PlanExecutionStatus.RUNNING.value, index=True)
+
+    # DAG state (stores dependency graph and task status)
+    dag_state = Column(JSONB, default={}, nullable=True)
+
+    # Waiting tasks (for user input continuation)
+    # Format: [{"task_id": 1, "execution_id": "abc-123", "question": "What?", "waiting_since": "2025-12-18T10:00:00Z"}]
+    waiting_tasks = Column(JSONB, default=[], nullable=True)
+
+    # Aggregated metrics
+    total_tokens = Column(Integer, default=0, nullable=True)
+    total_execution_time_seconds = Column(Integer, default=0, nullable=True)
+    estimated_cost_usd = Column(Numeric(10, 4), nullable=True)
+    actual_cost_usd = Column(Numeric(10, 4), nullable=True)
+
+    # Plan JSON (store original plan for reference)
+    plan_json = Column(JSONB, nullable=True)
+
+    # Timestamps
+    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False)
+    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False)
+    started_at = Column(DateTime(timezone=True), nullable=True)
+    completed_at = Column(DateTime(timezone=True), nullable=True)
+
+    # Indexes for common queries
+    __table_args__ = (
+        Index('idx_plan_executions_org_status', 'organization_id', 'status'),
+        Index('idx_plan_executions_created', 'created_at'),
+    )
+
+    def to_dict(self):
+        """Convert to dictionary for API responses."""
+        return {
+            "id": str(self.id),
+            "execution_id": self.execution_id,
+            "organization_id": self.organization_id,
+            "agent_id": str(self.agent_id) if self.agent_id else None,
+            "title": self.title,
+            "summary": self.summary,
+            "total_tasks": self.total_tasks,
+            "completed_tasks": self.completed_tasks,
+            "failed_tasks": self.failed_tasks,
+            "status": self.status,
+            "dag_state": self.dag_state,
+            "waiting_tasks": self.waiting_tasks,
+            "total_tokens": self.total_tokens,
+            "total_execution_time_seconds": self.total_execution_time_seconds,
+            "estimated_cost_usd": float(self.estimated_cost_usd) if self.estimated_cost_usd else None,
+            "actual_cost_usd": float(self.actual_cost_usd) if self.actual_cost_usd else None,
+            "plan_json": self.plan_json,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+            "started_at": self.started_at.isoformat() if self.started_at else None,
+            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+        }

control_plane_api/app/models/presence.py

@@ -0,0 +1,49 @@
+from sqlalchemy import Column, String, DateTime, Boolean, Index
+from sqlalchemy.sql import func
+from datetime import datetime
+import uuid
+
+from control_plane_api.app.database import Base
+
+
+class UserPresence(Base):
+    """Model for tracking user presence on agents/tasks for real-time collaboration"""
+
+    __tablename__ = "user_presence"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+
+    # User information
+    user_id = Column(String, nullable=False, index=True)
+    user_email = Column(String, nullable=True)
+    user_name = Column(String, nullable=True)
+    user_avatar = Column(String, nullable=True)
+
+    # What the user is viewing/interacting with
+    agent_id = Column(String, nullable=True, index=True)
+    session_id = Column(String, nullable=True, index=True)
+    execution_id = Column(String, nullable=True, index=True)
+
+    # Presence state
+    is_active = Column(Boolean, default=True, nullable=False)
+    is_typing = Column(Boolean, default=False, nullable=False)
+
+    # Timestamps
+    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False)
+    last_active_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False)
+
+    # Create composite index for efficient queries
+    __table_args__ = (
+        Index('idx_presence_lookup', 'agent_id', 'is_active', 'last_active_at'),
+        Index('idx_presence_user', 'user_id', 'is_active'),
+    )
+
+    def __repr__(self):
+        return f"<UserPresence(user_id={self.user_id}, agent_id={self.agent_id}, is_active={self.is_active})>"
+
+    @property
+    def is_stale(self) -> bool:
+        """Check if presence is stale (no activity in last 5 minutes)"""
+        if not self.last_active_at:
+            return True
+        return (datetime.utcnow() - self.last_active_at).total_seconds() > 300  # 5 minutes

control_plane_api/app/models/project.py

@@ -0,0 +1,61 @@
+from sqlalchemy import Column, String, DateTime, Text, ForeignKey, Index, UniqueConstraint, CheckConstraint
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from sqlalchemy.orm import relationship
+from sqlalchemy.sql import text
+import enum
+
+from control_plane_api.app.database import Base
+
+
+class ProjectStatus(str, enum.Enum):
+    """Project status enumeration"""
+
+    ACTIVE = "active"
+    ARCHIVED = "archived"
+    DRAFT = "draft"
+
+
+class Project(Base):
+    """Project model for organizing agents, teams, and tasks"""
+
+    __tablename__ = "projects"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()"), nullable=False)
+    organization_id = Column(String(255), nullable=False)
+    name = Column(String(255), nullable=False)
+    key = Column(String(50), nullable=False)
+    description = Column(Text, nullable=True)
+    settings = Column(JSONB, server_default=text("'{}'::jsonb"), nullable=True)
+    status = Column(String(50), server_default=text("'active'::character varying"), nullable=True)
+    visibility = Column(String(20), server_default=text("'private'::character varying"), nullable=True)
+    owner_id = Column(String(255), nullable=True)
+    owner_email = Column(String(255), nullable=True)
+
+    # Foreign keys
+    environment_id = Column(UUID(as_uuid=True), ForeignKey("environments.id", ondelete="SET NULL"), nullable=True)
+
+    # Metadata
+    created_at = Column(DateTime(timezone=True), server_default=text("now()"), nullable=True)
+    updated_at = Column(DateTime(timezone=True), server_default=text("now()"), nullable=True)
+    archived_at = Column(DateTime(timezone=True), nullable=True)
+
+    # Relationships
+    environment = relationship("Environment", foreign_keys=[environment_id])
+
+    __table_args__ = (
+        UniqueConstraint('organization_id', 'key', name='unique_project_key_per_org'),
+        UniqueConstraint('organization_id', 'name', name='unique_project_name_per_org'),
+        CheckConstraint("status IN ('active', 'archived', 'draft')", name='projects_status_check'),
+        CheckConstraint("visibility IN ('private', 'org')", name='projects_visibility_check'),
+        Index('idx_projects_org_id', 'organization_id'),
+        Index('idx_projects_status', 'status', postgresql_where=text("status = 'active'")),
+        Index('idx_projects_visibility', 'organization_id', 'visibility'),
+        Index('idx_projects_owner', 'owner_id'),
+        Index('idx_projects_environment_id', 'environment_id'),
+    )
+
+    # Note: Project associations are managed through junction tables (project_agents, project_teams)
+    # in Supabase, not through direct foreign keys in the SQLAlchemy models
+
+    def __repr__(self):
+        return f"<Project(id={self.id}, name={self.name}, status={self.status})>"

control_plane_api/app/models/project_management.py

@@ -0,0 +1,85 @@
+"""Project management and organization models"""
+from sqlalchemy import Column, String, DateTime, Text, JSON, UUID, CheckConstraint, ForeignKey, Index, UniqueConstraint, text
+from sqlalchemy.sql import func
+from sqlalchemy.orm import relationship
+from datetime import datetime
+
+from control_plane_api.app.database import Base
+
+
+class Profile(Base):
+    """Model for user profiles"""
+
+    __tablename__ = "profiles"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, nullable=False)
+    email = Column(Text, nullable=False, unique=True)
+    full_name = Column(Text, nullable=True)
+    avatar_url = Column(Text, nullable=True)
+
+    # Metadata
+    created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=True)
+    updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=True)
+
+    __table_args__ = (
+        Index('idx_profiles_email', 'email'),
+    )
+
+    def __repr__(self):
+        return f"<Profile(email={self.email}, full_name={self.full_name})>"
+
+
+class ProjectAgent(Base):
+    """Model for many-to-many relationship between projects and agents"""
+
+    __tablename__ = "project_agents"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()"), nullable=False)
+    project_id = Column(UUID(as_uuid=True), ForeignKey('projects.id', ondelete='CASCADE'), nullable=False)
+    agent_id = Column(UUID(as_uuid=True), ForeignKey('agents.id', ondelete='CASCADE'), nullable=False)
+    role = Column(String(50), nullable=True)
+    added_by = Column(String(255), nullable=True)
+
+    # Metadata
+    added_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=True)
+
+    # Relationships
+    project = relationship("Project", foreign_keys=[project_id], viewonly=True)
+    agent = relationship("Agent", foreign_keys=[agent_id], viewonly=True)
+
+    __table_args__ = (
+        Index('idx_project_agents_project', 'project_id'),
+        Index('idx_project_agents_agent', 'agent_id'),
+        UniqueConstraint('project_id', 'agent_id', name='unique_project_agent'),
+    )
+
+    def __repr__(self):
+        return f"<ProjectAgent(project_id={self.project_id}, agent_id={self.agent_id}, role={self.role})>"
+
+
+class ProjectTeam(Base):
+    """Model for many-to-many relationship between projects and teams"""
+
+    __tablename__ = "project_teams"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()"), nullable=False)
+    project_id = Column(UUID(as_uuid=True), ForeignKey('projects.id', ondelete='CASCADE'), nullable=False)
+    team_id = Column(UUID(as_uuid=True), ForeignKey('teams.id', ondelete='CASCADE'), nullable=False)
+    role = Column(String(50), nullable=True)
+    added_by = Column(String(255), nullable=True)
+
+    # Metadata
+    added_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=True)
+
+    # Relationships
+    project = relationship("Project", foreign_keys=[project_id], viewonly=True)
+    team = relationship("Team", foreign_keys=[team_id], viewonly=True)
+
+    __table_args__ = (
+        Index('idx_project_teams_project', 'project_id'),
+        Index('idx_project_teams_team', 'team_id'),
+        UniqueConstraint('project_id', 'team_id', name='unique_project_team'),
+    )
+
+    def __repr__(self):
+        return f"<ProjectTeam(project_id={self.project_id}, team_id={self.team_id}, role={self.role})>"

control_plane_api/app/models/session.py

@@ -0,0 +1,29 @@
+from sqlalchemy import Column, DateTime, Text, Index, text, desc
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.sql import func
+
+from control_plane_api.app.database import Base
+
+
+class Session(Base):
+    """Session model for storing agent session information"""
+
+    __tablename__ = "sessions"
+
+    execution_id = Column(Text, primary_key=True, nullable=False)
+    session_id = Column(Text, nullable=False)
+    organization_id = Column(Text, nullable=False)
+    user_id = Column(Text, nullable=True)
+    messages = Column(JSONB, nullable=False, server_default=text("'[]'::jsonb"))
+    metadata_ = Column("metadata", JSONB, nullable=True, server_default=text("'{}'::jsonb"))
+    created_at = Column(DateTime(timezone=True), nullable=True, server_default=func.now())
+    updated_at = Column(DateTime(timezone=True), nullable=True, server_default=func.now())
+
+    __table_args__ = (
+        Index("idx_sessions_organization_id", "organization_id"),
+        Index("idx_sessions_session_id", "session_id"),
+        Index("idx_sessions_updated_at", desc('updated_at')),
+    )
+
+    def __repr__(self):
+        return f"<Session(execution_id={self.execution_id}, session_id={self.session_id}, organization_id={self.organization_id})>"

control_plane_api/app/models/skill.py

@@ -0,0 +1,155 @@
+from sqlalchemy import Column, String, DateTime, Text, JSON, Boolean, Integer, CheckConstraint, ForeignKey, Index, UniqueConstraint, text
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from sqlalchemy.orm import relationship
+from datetime import datetime
+import uuid
+import enum
+
+from control_plane_api.app.database import Base
+
+
+class SkillType(str, enum.Enum):
+    """Skill type enumeration"""
+    FILE_SYSTEM = "file_system"
+    SHELL = "shell"
+    PYTHON = "python"
+    DOCKER = "docker"
+    SLEEP = "sleep"
+    FILE_GENERATION = "file_generation"
+    DATA_VISUALIZATION = "data_visualization"
+    WORKFLOW_EXECUTOR = "workflow_executor"
+    CUSTOM = "custom"
+
+
+class SkillEntityType(str, enum.Enum):
+    """Entity types that can be associated with skills"""
+    AGENT = "agent"
+    TEAM = "team"
+    ENVIRONMENT = "environment"
+
+
+class SlashCommandStatus(str, enum.Enum):
+    """Slash command execution status"""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class Skill(Base):
+    """Skill definitions with type-specific configurations"""
+
+    __tablename__ = "skills"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, nullable=False, server_default=text("gen_random_uuid()"))
+    organization_id = Column(String(255), nullable=False, index=True)
+    name = Column(String(255), nullable=False)
+    skill_type = Column(String(50), nullable=False)
+    description = Column(Text, nullable=True)
+    icon = Column(String(50), nullable=True, server_default=text("'Tool'::character varying"))
+    enabled = Column(Boolean, nullable=True, server_default=text("true"))
+    configuration = Column(JSONB, nullable=True, server_default=text("'{}'::jsonb"))
+    created_at = Column(DateTime, nullable=True, server_default=text("now()"))
+    updated_at = Column(DateTime, nullable=True, server_default=text("now()"))
+
+    # Relationships
+    associations = relationship("SkillAssociation", back_populates="skill", cascade="all, delete-orphan")
+
+    __table_args__ = (
+        UniqueConstraint('organization_id', 'name', name='unique_toolset_per_org'),
+        Index('idx_skills_type', 'skill_type'),
+        Index('idx_skills_enabled', 'enabled'),
+        CheckConstraint(
+            "skill_type IN ('file_system', 'shell', 'python', 'docker', 'sleep', 'file_generation', 'data_visualization', 'workflow_executor', 'custom')",
+            name="toolsets_type_check"
+        ),
+    )
+
+    def __repr__(self):
+        return f"<Skill {self.id} ({self.name}, type={self.skill_type})>"
+
+
+class SkillAssociation(Base):
+    """Associates skills with agents, teams, or environments"""
+
+    __tablename__ = "skill_associations"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+    organization_id = Column(String(255), nullable=False, index=True)
+    skill_id = Column(UUID(as_uuid=True), ForeignKey('skills.id', ondelete='CASCADE'), nullable=False)
+    entity_type = Column(String(50), nullable=False)
+    entity_id = Column(UUID(as_uuid=True), nullable=False)
+    configuration_override = Column(JSON, default=dict)  # Entity-specific configuration overrides
+    created_at = Column(DateTime, default=datetime.utcnow)
+
+    # Relationships
+    skill = relationship("Skill", back_populates="associations")
+
+    __table_args__ = (
+        UniqueConstraint('skill_id', 'entity_type', 'entity_id', name='unique_skill_entity'),
+        Index('ix_skill_associations_skill_id', 'skill_id'),
+        Index('ix_skill_associations_entity', 'entity_type', 'entity_id'),
+        Index('ix_skill_associations_org_skill', 'organization_id', 'skill_id'),
+        CheckConstraint(
+            "entity_type IN ('agent', 'team', 'environment')",
+            name="toolset_associations_entity_type_check"
+        ),
+    )
+
+    def __repr__(self):
+        return f"<SkillAssociation {self.id} ({self.entity_type}:{self.entity_id} -> {self.skill_id})>"
+
+
+class SlashCommand(Base):
+    """Custom slash commands for workflows"""
+
+    __tablename__ = "slash_commands"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+    command = Column(Text, nullable=False)
+    workflow_id = Column(Text, nullable=False)
+    workflow_name = Column(Text, nullable=False)
+    description = Column(Text, nullable=True)
+    enabled = Column(Boolean, default=True)
+    runner = Column(Text, nullable=True)
+    args = Column(JSON, default=list)
+    user_id = Column(Text, nullable=False, index=True)
+    project_id = Column(Text, nullable=False)
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(tz=datetime.timezone.utc))
+    updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(tz=datetime.timezone.utc))
+    last_used = Column(DateTime(timezone=True), nullable=True)
+
+    # Relationships
+    executions = relationship("SlashCommandExecution", back_populates="command", cascade="all, delete-orphan")
+
+    def __repr__(self):
+        return f"<SlashCommand {self.id} (/{self.command})>"
+
+
+class SlashCommandExecution(Base):
+    """Execution history for slash commands"""
+
+    __tablename__ = "slash_command_executions"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+    command_id = Column(UUID(as_uuid=True), ForeignKey('slash_commands.id', ondelete='CASCADE'), nullable=False)
+    command = Column(Text, nullable=False)
+    args = Column(JSON, default=dict)
+    workflow_id = Column(Text, nullable=False)
+    task_id = Column(UUID(as_uuid=True), nullable=False)
+    status = Column(Text, default="pending", nullable=False)
+    started_at = Column(DateTime(timezone=True), default=lambda: datetime.now(tz=datetime.timezone.utc))
+    completed_at = Column(DateTime(timezone=True), nullable=True)
+    error = Column(Text, nullable=True)
+    user_id = Column(Text, nullable=False, index=True)
+
+    # Relationships
+    command = relationship("SlashCommand", back_populates="executions")
+
+    __table_args__ = (
+        Index('ix_slash_command_executions_command_id', 'command_id'),
+        Index('ix_slash_command_executions_user_id', 'user_id'),
+    )
+
+    def __repr__(self):
+        return f"<SlashCommandExecution {self.id} ({self.command}, status={self.status})>"

control_plane_api/app/models/system_tables.py

@@ -0,0 +1,43 @@
+from sqlalchemy import Column, String, DateTime, Text, JSON, Boolean, Integer, BigInteger, LargeBinary, CheckConstraint, \
+    ForeignKey, Index, Table, MetaData, UniqueConstraint, text, desc
+from sqlalchemy.dialects.postgresql import UUID, INET
+from sqlalchemy.orm import relationship
+from datetime import datetime
+
+from control_plane_api.app.database import Base
+
+class PolicyAssociation(Base):
+    """Policy associations for resources"""
+
+    __tablename__ = "policy_associations"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()"))
+    organization_id = Column(String(255), nullable=False)
+    policy_id = Column(String(255), nullable=False)
+    policy_name = Column(String(255), nullable=False)
+    entity_type = Column(String(50), nullable=False)
+    entity_id = Column(UUID(as_uuid=True), nullable=False)
+    enabled = Column(Boolean, nullable=True, default=True, server_default=text("true"))
+    priority = Column(Integer, nullable=True, default=0, server_default=text("0"))
+    metadata_ = Column("metadata", JSON, nullable=True, default=dict, server_default=text("'{}'::jsonb"))
+    created_at = Column(DateTime(timezone=False), nullable=True, default=datetime.utcnow, server_default=text("now()"))
+    updated_at = Column(DateTime(timezone=False), nullable=True, default=datetime.utcnow, server_default=text("now()"), onupdate=datetime.utcnow)
+    created_by = Column(String(255), nullable=True)
+
+    __table_args__ = (
+        UniqueConstraint('policy_id', 'entity_type', 'entity_id', name='unique_policy_association'),
+        Index('idx_policy_assoc_org', 'organization_id'),
+        Index('idx_policy_assoc_entity', 'entity_type', 'entity_id'),
+        Index('idx_policy_assoc_policy', 'policy_id'),
+        Index('idx_policy_assoc_enabled', 'enabled'),
+        Index('idx_policy_assoc_priority', desc('priority')),
+        Index('idx_policy_assoc_entity_enabled', 'entity_type', 'entity_id', 'enabled', desc('priority')),
+        CheckConstraint(
+            "entity_type IN ('agent', 'team', 'environment')",
+            name='policy_associations_entity_type_check'
+        ),
+        {'schema': 'public'}
+    )
+
+    def __repr__(self):
+        return f"<PolicyAssociation {self.id} ({self.policy_id} on {self.entity_type}:{self.entity_id})>"