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

control_plane_api/alembic/versions/6a4d4dc3d8dc_generate_execution_transitions.py

@@ -0,0 +1,50 @@
+"""generate execution_transitions
+Revision ID: 6a4d4dc3d8dc
+Revises: f71305fb69b9
+Create Date: 2025-12-22 14:34:50.860339
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '6a4d4dc3d8dc'
+down_revision: Union[str, Sequence[str], None] = 'f71305fb69b9'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('execution_transitions',
+    sa.Column('id', sa.UUID(), nullable=False),
+    sa.Column('organization_id', sa.String(), nullable=False),
+    sa.Column('execution_id', sa.UUID(), nullable=False),
+    sa.Column('turn_number', sa.Integer(), nullable=False),
+    sa.Column('from_state', sa.String(length=50), nullable=False),
+    sa.Column('to_state', sa.String(length=50), nullable=False),
+    sa.Column('reasoning', sa.String(), nullable=False),
+    sa.Column('confidence', sa.String(length=20), nullable=False),
+    sa.Column('decision_factors', sa.JSON(), nullable=True),
+    sa.Column('ai_model', sa.String(length=100), nullable=True),
+    sa.Column('decision_time_ms', sa.Integer(), nullable=True),
+    sa.Column('is_manual_override', sa.Boolean(), nullable=False),
+    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
+    sa.ForeignKeyConstraint(['execution_id'], ['executions.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_execution_transitions_execution_id'), 'execution_transitions', ['execution_id'], unique=False)
+    op.create_index(op.f('ix_execution_transitions_organization_id'), 'execution_transitions', ['organization_id'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_execution_transitions_organization_id'), table_name='execution_transitions')
+    op.drop_index(op.f('ix_execution_transitions_execution_id'), table_name='execution_transitions')
+    op.drop_table('execution_transitions')
+    # ### end Alembic commands ###

control_plane_api/alembic/versions/87d11cf0a783_add_disconnected_status_to_worker_.py

@@ -0,0 +1,44 @@
+"""add_disconnected_status_to_worker_heartbeats
+
+Revision ID: 87d11cf0a783
+Revises: 43abf98d6a01
+Create Date: 2025-12-08 21:37:05.347985
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '87d11cf0a783'
+down_revision: Union[str, Sequence[str], None] = '43abf98d6a01'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Drop the old constraint
+    op.drop_constraint('worker_heartbeats_status_check', 'worker_heartbeats', type_='check')
+
+    # Create the new constraint with 'disconnected' added
+    op.create_check_constraint(
+        'worker_heartbeats_status_check',
+        'worker_heartbeats',
+        "status IN ('active', 'idle', 'busy', 'offline', 'disconnected')"
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # Drop the new constraint
+    op.drop_constraint('worker_heartbeats_status_check', 'worker_heartbeats', type_='check')
+
+    # Recreate the old constraint without 'disconnected'
+    op.create_check_constraint(
+        'worker_heartbeats_status_check',
+        'worker_heartbeats',
+        "status IN ('active', 'idle', 'busy', 'offline')"
+    )

control_plane_api/alembic/versions/add_ephemeral_queue_support.py

@@ -0,0 +1,85 @@
+"""add_ephemeral_queue_support
+
+Adds ephemeral queue support for on-demand workers.
+This enables automatic provisioning and cleanup of temporary worker queues.
+
+Revision ID: e9f2a3b4c5d6
+Revises: 87d11cf0a783
+Create Date: 2025-12-11 08:00:00.000000
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'e9f2a3b4c5d6'
+down_revision: Union[str, Sequence[str], None] = '87d11cf0a783'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """
+    Add ephemeral queue support columns to worker_queues table.
+
+    New columns:
+    - ephemeral: Mark queue as temporary/ephemeral
+    - single_execution_mode: Worker processes one task then exits
+    - auto_cleanup_after_seconds: TTL for automatic cleanup
+    - parent_execution_id: Track the execution that created this queue
+    """
+
+    # Add ephemeral flag (default false for backward compatibility)
+    op.add_column('worker_queues',
+        sa.Column('ephemeral', sa.Boolean(), server_default=sa.text('false'), nullable=False)
+    )
+
+    # Add single execution mode flag (default false)
+    op.add_column('worker_queues',
+        sa.Column('single_execution_mode', sa.Boolean(), server_default=sa.text('false'), nullable=False)
+    )
+
+    # Add auto cleanup TTL (nullable - only set for ephemeral queues)
+    op.add_column('worker_queues',
+        sa.Column('auto_cleanup_after_seconds', sa.Integer(), nullable=True)
+    )
+
+    # Add parent execution ID for tracking (nullable)
+    op.add_column('worker_queues',
+        sa.Column('parent_execution_id', sa.Text(), nullable=True)
+    )
+
+    # Create index for efficient ephemeral queue cleanup queries
+    op.create_index(
+        'idx_worker_queues_ephemeral_cleanup',
+        'worker_queues',
+        ['ephemeral', 'created_at'],
+        postgresql_where=sa.text("ephemeral = true")
+    )
+
+    # Create index for parent execution tracking
+    op.create_index(
+        'idx_worker_queues_parent_execution',
+        'worker_queues',
+        ['parent_execution_id'],
+        postgresql_where=sa.text("parent_execution_id IS NOT NULL")
+    )
+
+
+def downgrade() -> None:
+    """
+    Remove ephemeral queue support columns.
+    """
+
+    # Drop indexes first
+    op.drop_index('idx_worker_queues_parent_execution', table_name='worker_queues')
+    op.drop_index('idx_worker_queues_ephemeral_cleanup', table_name='worker_queues')
+
+    # Drop columns
+    op.drop_column('worker_queues', 'parent_execution_id')
+    op.drop_column('worker_queues', 'auto_cleanup_after_seconds')
+    op.drop_column('worker_queues', 'single_execution_mode')
+    op.drop_column('worker_queues', 'ephemeral')

control_plane_api/alembic/versions/add_model_type_to_llm_models.py

@@ -0,0 +1,31 @@
+"""add model_type column to llm_models
+
+Revision ID: add_model_type_to_llm_models
+Revises: f25de6ad895a
+Create Date: 2025-01-10 00:00:00.000000
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = 'add_model_type_to_llm_models'
+down_revision = 'f25de6ad895a'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Add model_type column with default value 'text-generation'
+    op.add_column('llm_models', sa.Column('model_type', sa.String(), nullable=False, server_default='text-generation'))
+
+    # Create index on model_type for efficient filtering
+    op.create_index(op.f('ix_llm_models_model_type'), 'llm_models', ['model_type'], unique=False)
+
+
+def downgrade():
+    # Remove index
+    op.drop_index(op.f('ix_llm_models_model_type'), table_name='llm_models')
+
+    # Remove column
+    op.drop_column('llm_models', 'model_type')

control_plane_api/alembic/versions/add_plan_executions_table.py

@@ -0,0 +1,114 @@
+"""add_plan_executions_table
+
+Creates the plan_executions table for tracking multi-task plan orchestration.
+Also adds plan_execution_id foreign key to executions table.
+
+Revision ID: add_plan_executions
+Revises: 43abf98d6a01
+Create Date: 2025-12-11 00:00:00.000000
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = 'add_plan_executions'
+down_revision: Union[str, Sequence[str], None] = '43abf98d6a01'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """
+    Create plan_executions table and add foreign key to executions.
+    """
+
+    # Create plan_executions table
+    op.create_table(
+        'plan_executions',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, nullable=False),
+        sa.Column('execution_id', sa.String(255), unique=True, nullable=False, index=True),
+        sa.Column('organization_id', sa.String(255), nullable=False, index=True),
+        sa.Column('agent_id', postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column('title', sa.Text(), nullable=False),
+        sa.Column('summary', sa.Text(), nullable=True),
+        sa.Column('total_tasks', sa.Integer(), nullable=False, server_default='0'),
+        sa.Column('completed_tasks', sa.Integer(), nullable=False, server_default='0'),
+        sa.Column('failed_tasks', sa.Integer(), nullable=True, server_default='0'),
+        sa.Column('status', sa.String(50), nullable=False, server_default='running', index=True),
+        sa.Column('dag_state', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='{}'),
+        sa.Column('total_tokens', sa.Integer(), nullable=True, server_default='0'),
+        sa.Column('total_execution_time_seconds', sa.Integer(), nullable=True, server_default='0'),
+        sa.Column('estimated_cost_usd', sa.Numeric(10, 4), nullable=True),
+        sa.Column('actual_cost_usd', sa.Numeric(10, 4), nullable=True),
+        sa.Column('plan_json', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
+        sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
+        sa.Column('started_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
+    )
+
+    # Create indexes
+    op.create_index(
+        'idx_plan_executions_org_status',
+        'plan_executions',
+        ['organization_id', 'status']
+    )
+    op.create_index(
+        'idx_plan_executions_created',
+        'plan_executions',
+        ['created_at']
+    )
+
+    # Add plan_execution_id to executions table
+    op.add_column(
+        'executions',
+        sa.Column('plan_execution_id', postgresql.UUID(as_uuid=True), nullable=True)
+    )
+
+    # Create foreign key constraint
+    op.create_foreign_key(
+        'fk_executions_plan_execution',
+        'executions',
+        'plan_executions',
+        ['plan_execution_id'],
+        ['id'],
+        ondelete='CASCADE'
+    )
+
+    # Create index on plan_execution_id
+    op.create_index(
+        'idx_executions_plan_execution_id',
+        'executions',
+        ['plan_execution_id']
+    )
+
+    print("✅ Created plan_executions table")
+    print("✅ Added plan_execution_id to executions table")
+
+
+def downgrade() -> None:
+    """
+    Remove plan_executions table and foreign key.
+    """
+
+    # Drop index
+    op.drop_index('idx_executions_plan_execution_id', table_name='executions')
+
+    # Drop foreign key
+    op.drop_constraint('fk_executions_plan_execution', 'executions', type_='foreignkey')
+
+    # Drop column from executions
+    op.drop_column('executions', 'plan_execution_id')
+
+    # Drop indexes
+    op.drop_index('idx_plan_executions_created', table_name='plan_executions')
+    op.drop_index('idx_plan_executions_org_status', table_name='plan_executions')
+
+    # Drop table
+    op.drop_table('plan_executions')
+
+    print("✅ Removed plan_executions table and related changes")

control_plane_api/alembic/versions/add_trace_span_tables.py

@@ -0,0 +1,154 @@
+"""add_trace_span_tables
+
+Creates traces and spans tables for OTEL observability with
+real-time streaming support.
+
+Revision ID: add_trace_span_tables
+Revises: f71305fb69b9
+Create Date: 2026-01-13
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = 'add_trace_span_tables'
+down_revision: Union[str, Sequence[str], None] = 'f71305fb69b9'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+# Pre-create enum types to reference in columns
+trace_status_enum = postgresql.ENUM('success', 'error', 'running', name='trace_status', create_type=False)
+span_kind_enum = postgresql.ENUM('INTERNAL', 'SERVER', 'CLIENT', 'PRODUCER', 'CONSUMER', name='span_kind', create_type=False)
+span_status_code_enum = postgresql.ENUM('UNSET', 'OK', 'ERROR', name='span_status_code', create_type=False)
+
+
+def upgrade() -> None:
+    """Create traces and spans tables with enums and indexes."""
+
+    # Create trace_status enum
+    op.execute("""
+        DO $$ BEGIN
+            CREATE TYPE trace_status AS ENUM ('success', 'error', 'running');
+        EXCEPTION
+            WHEN duplicate_object THEN null;
+        END $$;
+    """)
+
+    # Create span_kind enum (OTEL standard)
+    op.execute("""
+        DO $$ BEGIN
+            CREATE TYPE span_kind AS ENUM ('INTERNAL', 'SERVER', 'CLIENT', 'PRODUCER', 'CONSUMER');
+        EXCEPTION
+            WHEN duplicate_object THEN null;
+        END $$;
+    """)
+
+    # Create span_status_code enum (OTEL standard)
+    op.execute("""
+        DO $$ BEGIN
+            CREATE TYPE span_status_code AS ENUM ('UNSET', 'OK', 'ERROR');
+        EXCEPTION
+            WHEN duplicate_object THEN null;
+        END $$;
+    """)
+
+    # Create traces table
+    op.create_table(
+        'traces',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('gen_random_uuid()')),
+        sa.Column('trace_id', sa.String(64), nullable=False, unique=True),
+        sa.Column('organization_id', sa.String(255), nullable=False),
+        sa.Column('name', sa.String(512), nullable=False),
+        sa.Column('service_name', sa.String(255), nullable=True),
+        sa.Column('status', trace_status_enum, nullable=False, server_default='running'),
+        sa.Column('execution_id', sa.String(255), nullable=True),
+        sa.Column('execution_type', sa.String(50), nullable=True),
+        sa.Column('user_id', sa.String(255), nullable=True),
+        sa.Column('user_email', sa.String(255), nullable=True),
+        sa.Column('started_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.text('NOW()')),
+        sa.Column('ended_at', sa.DateTime(timezone=True), nullable=True),
+        sa.Column('duration_ms', sa.BigInteger(), nullable=True),
+        sa.Column('span_count', sa.Integer(), nullable=False, server_default='0'),
+        sa.Column('error_count', sa.Integer(), nullable=False, server_default='0'),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.text('NOW()')),
+        sa.Column('updated_at', sa.DateTime(timezone=True), nullable=True, server_default=sa.text('NOW()')),
+    )
+
+    # Create traces indexes
+    op.create_index('ix_traces_trace_id', 'traces', ['trace_id'])
+    op.create_index('ix_traces_org_id', 'traces', ['organization_id'])
+    op.create_index('ix_traces_org_started', 'traces', ['organization_id', 'started_at'])
+    op.create_index('ix_traces_org_status', 'traces', ['organization_id', 'status'])
+    op.create_index('ix_traces_org_service', 'traces', ['organization_id', 'service_name'])
+    op.create_index('ix_traces_org_user', 'traces', ['organization_id', 'user_id'])
+    op.create_index('ix_traces_user_id', 'traces', ['user_id'])
+
+    # Create partial index for execution_id
+    op.execute("""
+        CREATE INDEX IF NOT EXISTS ix_traces_execution
+        ON traces(execution_id)
+        WHERE execution_id IS NOT NULL
+    """)
+
+    # Create spans table
+    op.create_table(
+        'spans',
+        sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('gen_random_uuid()')),
+        sa.Column('trace_id', sa.String(64), sa.ForeignKey('traces.trace_id', ondelete='CASCADE'), nullable=False),
+        sa.Column('span_id', sa.String(32), nullable=False),
+        sa.Column('parent_span_id', sa.String(32), nullable=True),
+        sa.Column('organization_id', sa.String(255), nullable=False),
+        sa.Column('name', sa.String(512), nullable=False),
+        sa.Column('kind', span_kind_enum, nullable=False, server_default='INTERNAL'),
+        sa.Column('status_code', span_status_code_enum, nullable=False, server_default='UNSET'),
+        sa.Column('status_message', sa.Text(), nullable=True),
+        sa.Column('start_time_unix_nano', sa.BigInteger(), nullable=False),
+        sa.Column('end_time_unix_nano', sa.BigInteger(), nullable=True),
+        sa.Column('duration_ns', sa.BigInteger(), nullable=True),
+        sa.Column('attributes', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='{}'),
+        sa.Column('resource_attributes', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='{}'),
+        sa.Column('events', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='[]'),
+        sa.Column('links', postgresql.JSONB(astext_type=sa.Text()), nullable=True, server_default='[]'),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.text('NOW()')),
+    )
+
+    # Create spans indexes
+    op.create_index('ix_spans_trace_id', 'spans', ['trace_id'])
+    op.create_index('ix_spans_span_id', 'spans', ['span_id'])
+    op.create_index('ix_spans_parent_span_id', 'spans', ['parent_span_id'])
+    op.create_index('ix_spans_org_id', 'spans', ['organization_id'])
+    op.create_index('ix_spans_trace_parent', 'spans', ['trace_id', 'parent_span_id'])
+    op.create_index('ix_spans_trace_start', 'spans', ['trace_id', 'start_time_unix_nano'])
+    op.create_index('ix_spans_org_name', 'spans', ['organization_id', 'name'])
+    op.create_index('ix_spans_org_start', 'spans', ['organization_id', 'start_time_unix_nano'])
+
+    # Unique constraint on span_id within a trace
+    op.create_index('ix_spans_trace_span_unique', 'spans', ['trace_id', 'span_id'], unique=True)
+
+    # GIN indexes for JSONB columns
+    op.execute("CREATE INDEX IF NOT EXISTS ix_spans_attributes_gin ON spans USING GIN (attributes)")
+    op.execute("CREATE INDEX IF NOT EXISTS ix_spans_resource_attrs_gin ON spans USING GIN (resource_attributes)")
+
+    # Add table comments
+    op.execute("COMMENT ON TABLE traces IS 'Aggregated OTEL traces for observability'")
+    op.execute("COMMENT ON TABLE spans IS 'Individual OTEL spans within traces'")
+    op.execute("COMMENT ON COLUMN traces.trace_id IS 'OpenTelemetry trace ID (32-char hex)'")
+    op.execute("COMMENT ON COLUMN spans.span_id IS 'OpenTelemetry span ID (16-char hex)'")
+
+
+def downgrade() -> None:
+    """Drop traces and spans tables with enums."""
+
+    # Drop tables (cascades to indexes)
+    op.drop_table('spans')
+    op.drop_table('traces')
+
+    # Drop enums
+    op.execute("DROP TYPE IF EXISTS trace_status")
+    op.execute("DROP TYPE IF EXISTS span_kind")
+    op.execute("DROP TYPE IF EXISTS span_status_code")

control_plane_api/alembic/versions/add_user_info_to_traces.py

@@ -0,0 +1,36 @@
+"""add_user_info_to_traces
+
+Add user_name and user_avatar columns to traces table for
+displaying user attribution in the observability UI.
+
+Revision ID: add_user_info_to_traces
+Revises: add_trace_span_tables
+Create Date: 2026-01-13
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision: str = 'add_user_info_to_traces'
+down_revision: Union[str, Sequence[str], None] = 'add_trace_span_tables'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Add user_name and user_avatar columns to traces table."""
+    op.add_column('traces', sa.Column('user_name', sa.String(255), nullable=True))
+    op.add_column('traces', sa.Column('user_avatar', sa.String(512), nullable=True))
+
+    # Add comment for documentation
+    op.execute("COMMENT ON COLUMN traces.user_name IS 'Display name of user who triggered the trace'")
+    op.execute("COMMENT ON COLUMN traces.user_avatar IS 'URL to user avatar image'")
+
+
+def downgrade() -> None:
+    """Remove user_name and user_avatar columns from traces table."""
+    op.drop_column('traces', 'user_avatar')
+    op.drop_column('traces', 'user_name')

control_plane_api/alembic/versions/adjusting_foreign_keys.py

@@ -0,0 +1,32 @@
+"""a
+
+Revision ID: 91520433aafc
+Revises: f25de6ad895a
+Create Date: 2025-11-29 10:57:24.435772
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '91520433aafc'
+down_revision: Union[str, Sequence[str], None] = 'f25de6ad895a'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.drop_constraint(op.f('executions_worker_queue_id_fkey'), 'executions', type_='foreignkey')
+    op.create_foreign_key(None, 'executions', 'worker_queues', ['worker_queue_id'], ['id'], ondelete='CASCADE')
+    op.create_unique_constraint('uq_team_org_name', 'teams', ['organization_id', 'name'])
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_constraint('uq_team_org_name', 'teams', type_='unique')
+    op.drop_constraint(None, 'executions', type_='foreignkey')
+    op.create_foreign_key(op.f('executions_worker_queue_id_fkey'), 'executions', 'worker_queues', ['worker_queue_id'], ['id'])