feat: model fallback chain (#2017) (#2018)

src/langbot/pkg/persistence/migrations/dbm023_model_fallback_config.py

@@ -0,0 +1,102 @@
+from .. import migration
+
+import sqlalchemy
+import json
+
+
+@migration.migration_class(23)
+class DBMigrateModelFallbackConfig(migration.DBMigration):
+    """Convert model field from plain UUID string to object with primary/fallbacks"""
+
+    async def upgrade(self):
+        """Upgrade"""
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.text('SELECT uuid, config FROM legacy_pipelines')
+        )
+        pipelines = result.fetchall()
+
+        current_version = self.ap.ver_mgr.get_current_version()
+
+        for pipeline_row in pipelines:
+            uuid = pipeline_row[0]
+            config = json.loads(pipeline_row[1]) if isinstance(pipeline_row[1], str) else pipeline_row[1]
+
+            if 'ai' not in config or 'local-agent' not in config['ai']:
+                continue
+
+            local_agent = config['ai']['local-agent']
+            changed = False
+
+            # Convert model from string to object
+            model_value = local_agent.get('model', '')
+            if isinstance(model_value, str):
+                local_agent['model'] = {
+                    'primary': model_value,
+                    'fallbacks': [],
+                }
+                changed = True
+
+            # Remove leftover fallback-models field if present
+            if 'fallback-models' in local_agent:
+                del local_agent['fallback-models']
+                changed = True
+
+            if not changed:
+                continue
+
+            # Update using raw SQL with compatibility for both SQLite and PostgreSQL
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        'UPDATE legacy_pipelines SET config = :config::jsonb, for_version = :for_version WHERE uuid = :uuid'
+                    ),
+                    {'config': json.dumps(config), 'for_version': current_version, 'uuid': uuid},
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        'UPDATE legacy_pipelines SET config = :config, for_version = :for_version WHERE uuid = :uuid'
+                    ),
+                    {'config': json.dumps(config), 'for_version': current_version, 'uuid': uuid},
+                )
+
+    async def downgrade(self):
+        """Downgrade"""
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.text('SELECT uuid, config FROM legacy_pipelines')
+        )
+        pipelines = result.fetchall()
+
+        current_version = self.ap.ver_mgr.get_current_version()
+
+        for pipeline_row in pipelines:
+            uuid = pipeline_row[0]
+            config = json.loads(pipeline_row[1]) if isinstance(pipeline_row[1], str) else pipeline_row[1]
+
+            if 'ai' not in config or 'local-agent' not in config['ai']:
+                continue
+
+            local_agent = config['ai']['local-agent']
+
+            # Convert model from object back to string
+            model_value = local_agent.get('model', '')
+            if isinstance(model_value, dict):
+                local_agent['model'] = model_value.get('primary', '')
+            else:
+                continue
+
+            # Update using raw SQL with compatibility for both SQLite and PostgreSQL
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        'UPDATE legacy_pipelines SET config = :config::jsonb, for_version = :for_version WHERE uuid = :uuid'
+                    ),
+                    {'config': json.dumps(config), 'for_version': current_version, 'uuid': uuid},
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        'UPDATE legacy_pipelines SET config = :config, for_version = :for_version WHERE uuid = :uuid'
+                    ),
+                    {'config': json.dumps(config), 'for_version': current_version, 'uuid': uuid},
+                )