1  """Add system_prompt and context columns to output table."""
 2  import sqlalchemy as sa
 3  from alembic import op
 4  
 5  
 6  revision = '033'
 7  down_revision = '032'
 8  branch_labels = None
 9  depends_on = None
10  
11  
12  def upgrade():
13      try:
14          op.add_column('output', sa.Column('system_prompt', sa.Text(), nullable=True))
15      except Exception as e:
16          print(f"Column system_prompt may already exist: {e}")
17  
18      try:
19          op.add_column('output', sa.Column('context', sa.Text(), nullable=True))
20      except Exception as e:
21          print(f"Column context may already exist: {e}")
22  
23  
24  def downgrade():
25      try:
26          op.drop_column('output', 'context')
27      except Exception as e:
28          print(f"Error dropping context column: {e}")
29  
30      try:
31          op.drop_column('output', 'system_prompt')
32      except Exception as e:
33          print(f"Error dropping system_prompt column: {e}")