mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-10 07:46:02 +00:00
103 lines
4.0 KiB
Python
103 lines
4.0 KiB
Python
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},
|
|
)
|