@@ -37,85 +37,98 @@ async def get_db() -> AsyncSession:
3737 await session .close ()
3838
3939
40+ def _get_column_type_sql (column ) -> str :
41+ """Convert SQLAlchemy column type to SQLite type string."""
42+ from sqlalchemy import JSON , Boolean , DateTime , Float , Integer , String , Text
43+
44+ col_type = type (column .type )
45+
46+ if col_type == Integer or "Integer" in str (col_type ):
47+ return "INTEGER"
48+ elif col_type == String or "String" in str (col_type ):
49+ length = getattr (column .type , "length" , None )
50+ return f"VARCHAR({ length } if length else "VARCHAR(255)"
51+ elif col_type == Text or "Text" in str (col_type ):
52+ return "TEXT"
53+ elif col_type == Boolean or "Boolean" in str (col_type ):
54+ return "BOOLEAN"
55+ elif col_type == Float or "Float" in str (col_type ):
56+ return "FLOAT"
57+ elif col_type == DateTime or "DateTime" in str (col_type ):
58+ return "DATETIME"
59+ elif col_type == JSON or "JSON" in str (col_type ):
60+ return "JSON"
61+ else :
62+ # Default fallback
63+ return "TEXT"
64+
65+
4066async def _run_migrations (conn ):
41- """Run schema migrations for new columns (SQLite compatible) ."""
67+ """Auto-detect and add missing columns by comparing models with database schema ."""
4268 from sqlalchemy import text
4369
44- async def column_exists (table_name : str , column_name : str ) -> bool :
45- """Check if a column exists in a table."""
46- result = await conn .execute (text (f"PRAGMA table_info({ table_name } ))
47- columns = [row [1 ] for row in result .fetchall ()]
48- return column_name in columns
49-
50- # Migration: Add container_name to deployments (for Windows Docker compatibility)
51- if not await column_exists ("deployments" , "container_name" ):
52- logger .info ("Adding 'container_name' column to deployments table..." )
53- await conn .execute (text ("ALTER TABLE deployments ADD COLUMN container_name VARCHAR(255)" ))
54- logger .info ("'container_name' column added!" )
55-
56- # Migration: Add is_local to registration_tokens (for local worker detection)
57- if not await column_exists ("registration_tokens" , "is_local" ):
58- logger .info ("Adding 'is_local' column to registration_tokens table..." )
59- await conn .execute (
60- text ("ALTER TABLE registration_tokens ADD COLUMN is_local BOOLEAN DEFAULT 0" )
61- )
62- logger .info ("'is_local' column added!" )
63-
64- # Migration: Add conversation_type to conversations (for Agent chat support)
65- if not await column_exists ("conversations" , "conversation_type" ):
66- logger .info ("Adding 'conversation_type' column to conversations table..." )
67- await conn .execute (
68- text (
69- "ALTER TABLE conversations ADD COLUMN conversation_type VARCHAR(20) DEFAULT 'chat' NOT NULL"
70- )
70+ async def get_table_columns (table_name : str ) -> set [str ]:
71+ """Get all column names from a database table."""
72+ try :
73+ result = await conn .execute (text (f"PRAGMA table_info({ table_name } ))
74+ return {row [1 ] for row in result .fetchall ()}
75+ except Exception :
76+ return set ()
77+
78+ async def table_exists (table_name : str ) -> bool :
79+ """Check if a table exists in the database."""
80+ result = await conn .execute (
81+ text ("SELECT name FROM sqlite_master WHERE type='table' AND name=:name" ),
82+ {"name" : table_name },
7183 )
72- logger .info ("'conversation_type' column added!" )
73-
74- # Migration: Add agent_config to conversations (for Agent configuration)
75- if not await column_exists ("conversations" , "agent_config" ):
76- logger .info ("Adding 'agent_config' column to conversations table..." )
77- await conn .execute (text ("ALTER TABLE conversations ADD COLUMN agent_config JSON" ))
78- logger .info ("'agent_config' column added!" )
79-
80- # Migration: Add tool_calls to messages (for Agent tool calls)
81- if not await column_exists ("messages" , "tool_calls" ):
82- logger .info ("Adding 'tool_calls' column to messages table..." )
83- await conn .execute (text ("ALTER TABLE messages ADD COLUMN tool_calls JSON" ))
84- logger .info ("'tool_calls' column added!" )
85-
86- # Migration: Add tool_call_id to messages (for Agent tool results)
87- if not await column_exists ("messages" , "tool_call_id" ):
88- logger .info ("Adding 'tool_call_id' column to messages table..." )
89- await conn .execute (text ("ALTER TABLE messages ADD COLUMN tool_call_id VARCHAR(100)" ))
90- logger .info ("'tool_call_id' column added!" )
91-
92- # Migration: Add step_type to messages (for Agent execution steps)
93- if not await column_exists ("messages" , "step_type" ):
94- logger .info ("Adding 'step_type' column to messages table..." )
95- await conn .execute (text ("ALTER TABLE messages ADD COLUMN step_type VARCHAR(50)" ))
96- logger .info ("'step_type' column added!" )
97-
98- # Migration: Add execution_time_ms to messages (for tool execution timing)
99- if not await column_exists ("messages" , "execution_time_ms" ):
100- logger .info ("Adding 'execution_time_ms' column to messages table..." )
101- await conn .execute (text ("ALTER TABLE messages ADD COLUMN execution_time_ms FLOAT" ))
102- logger .info ("'execution_time_ms' column added!" )
103-
104- # Migration: Add tuning_config to tuning_jobs (for multi-framework testing)
105- if not await column_exists ("tuning_jobs" , "tuning_config" ):
106- logger .info ("Adding 'tuning_config' column to tuning_jobs table..." )
107- await conn .execute (text ("ALTER TABLE tuning_jobs ADD COLUMN tuning_config JSON" ))
108- logger .info ("'tuning_config' column added!" )
109-
110- # Migration: Add conversation_id to tuning_jobs (for Agent Chat integration)
111- if not await column_exists ("tuning_jobs" , "conversation_id" ):
112- logger .info ("Adding 'conversation_id' column to tuning_jobs table..." )
113- await conn .execute (text ("ALTER TABLE tuning_jobs ADD COLUMN conversation_id INTEGER" ))
114- logger .info ("'conversation_id' column added!" )
84+ return result .fetchone () is not None
85+
86+ # Iterate through all tables defined in models
87+ for table_name , table in Base .metadata .tables .items ():
88+ # Skip if table doesn't exist yet (will be created by create_all)
89+ if not await table_exists (table_name ):
90+ continue
91+
92+ # Get existing columns in database
93+ existing_columns = await get_table_columns (table_name )
94+
95+ # Check each column in the model
96+ for column in table .columns :
97+ if column .name not in existing_columns :
98+ # Build ALTER TABLE statement
99+ col_type = _get_column_type_sql (column )
100+
101+ # Handle default values
102+ default_clause = ""
103+ if column .default is not None :
104+ default_val = column .default .arg
105+ if callable (default_val ):
106+ default_val = default_val (None )
107+ if isinstance (default_val , str ):
108+ default_clause = f" DEFAULT '{ default_val }
109+ elif isinstance (default_val , bool ):
110+ default_clause = f" DEFAULT { 1 if default_val else 0 }
111+ elif default_val is not None :
112+ default_clause = f" DEFAULT { default_val }
113+
114+ sql = (
115+ f"ALTER TABLE { table_name } { column .name } { col_type } { default_clause }
116+ )
117+
118+ logger .info (f"Auto-migration: Adding '{ column .name } { table_name } )
119+ try :
120+ await conn .execute (text (sql ))
121+ logger .info (f"Column '{ column .name } { table_name } )
122+ except Exception as e :
123+ logger .warning (f"Failed to add column { column .name } { table_name } { e } )
115124
116125
117126async def init_db ():
118127 """Initialize database tables and run migrations"""
128+ # Import all models to register them with Base.metadata
129+ # This ensures all tables are created by create_all()
130+ import app .models # noqa: F401
131+
119132 try :
120133 async with engine .begin () as conn :
121134 await conn .run_sync (Base .metadata .create_all )
0 commit comments