69 lines
3.0 KiB
Python
69 lines
3.0 KiB
Python
|
"""Add agent checkpoint and tree node tables
|
||
|
|
|
||
|
|
Revision ID: 007_add_agent_checkpoint_tables
|
||
|
|
Revises: 006_add_agent_tables
|
||
|
|
Create Date: 2024-12-12
|
||
|
|
|
||
|
|
"""
|
||
|
|
from alembic import op
|
||
|
|
import sqlalchemy as sa
|
||
|
|
|
||
|
|
|
||
|
|
# revision identifiers, used by Alembic.
|
||
|
|
revision = '007_add_agent_checkpoint_tables'
|
||
|
|
down_revision = '006_add_agent_tables'
|
||
|
|
branch_labels = None
|
||
|
|
depends_on = None
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
# Create agent_checkpoints table
|
||
|
|
op.create_table(
|
||
|
|
'agent_checkpoints',
|
||
|
|
sa.Column('id', sa.String(36), primary_key=True),
|
||
|
|
sa.Column('task_id', sa.String(36), sa.ForeignKey('agent_tasks.id', ondelete='CASCADE'), nullable=False, index=True),
|
||
|
|
sa.Column('agent_id', sa.String(50), nullable=False, index=True),
|
||
|
|
sa.Column('agent_name', sa.String(255), nullable=False),
|
||
|
|
sa.Column('agent_type', sa.String(50), nullable=False),
|
||
|
|
sa.Column('parent_agent_id', sa.String(50), nullable=True),
|
||
|
|
sa.Column('state_data', sa.Text, nullable=False),
|
||
|
|
sa.Column('iteration', sa.Integer, default=0),
|
||
|
|
sa.Column('status', sa.String(30), nullable=False),
|
||
|
|
sa.Column('total_tokens', sa.Integer, default=0),
|
||
|
|
sa.Column('tool_calls', sa.Integer, default=0),
|
||
|
|
sa.Column('findings_count', sa.Integer, default=0),
|
||
|
|
sa.Column('checkpoint_type', sa.String(30), default='auto'),
|
||
|
|
sa.Column('checkpoint_name', sa.String(255), nullable=True),
|
||
|
|
sa.Column('checkpoint_metadata', sa.JSON, nullable=True),
|
||
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), index=True),
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create agent_tree_nodes table
|
||
|
|
op.create_table(
|
||
|
|
'agent_tree_nodes',
|
||
|
|
sa.Column('id', sa.String(36), primary_key=True),
|
||
|
|
sa.Column('task_id', sa.String(36), sa.ForeignKey('agent_tasks.id', ondelete='CASCADE'), nullable=False, index=True),
|
||
|
|
sa.Column('agent_id', sa.String(50), nullable=False, unique=True, index=True),
|
||
|
|
sa.Column('agent_name', sa.String(255), nullable=False),
|
||
|
|
sa.Column('agent_type', sa.String(50), nullable=False),
|
||
|
|
sa.Column('parent_agent_id', sa.String(50), nullable=True, index=True),
|
||
|
|
sa.Column('depth', sa.Integer, default=0),
|
||
|
|
sa.Column('task_description', sa.Text, nullable=True),
|
||
|
|
sa.Column('knowledge_modules', sa.JSON, nullable=True),
|
||
|
|
sa.Column('status', sa.String(30), default='created'),
|
||
|
|
sa.Column('result_summary', sa.Text, nullable=True),
|
||
|
|
sa.Column('findings_count', sa.Integer, default=0),
|
||
|
|
sa.Column('iterations', sa.Integer, default=0),
|
||
|
|
sa.Column('tokens_used', sa.Integer, default=0),
|
||
|
|
sa.Column('tool_calls', sa.Integer, default=0),
|
||
|
|
sa.Column('duration_ms', sa.Integer, nullable=True),
|
||
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now()),
|
||
|
|
sa.Column('started_at', sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.Column('finished_at', sa.DateTime(timezone=True), nullable=True),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
op.drop_table('agent_tree_nodes')
|
||
|
|
op.drop_table('agent_checkpoints')
|