feat: enhance model creation with UUID preservation option and implement Space model synchronization in ModelManager

This commit is contained in:
Junyan Qin
2025-12-31 22:25:07 +08:00
parent 197258ae91
commit 96e40eaf25
11 changed files with 219 additions and 132 deletions

View File

@@ -1,77 +1,14 @@
import sqlalchemy
from .. import migration
# this is a deprecated migration
@migration.migration_class(15)
class DBMigrateModelSourceTracking(migration.DBMigration):
"""Add source tracking fields to models tables for Space integration"""
async def upgrade(self):
"""Upgrade"""
# Add source column to llm_models table
llm_columns = await self._get_columns('llm_models')
if 'source' not in llm_columns:
if self.ap.persistence_mgr.db.name == 'postgresql':
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text("ALTER TABLE llm_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL")
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text("ALTER TABLE llm_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL")
)
if 'space_model_id' not in llm_columns:
if self.ap.persistence_mgr.db.name == 'postgresql':
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN space_model_id VARCHAR(255)')
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN space_model_id VARCHAR(255)')
)
# Add source column to embedding_models table
embedding_columns = await self._get_columns('embedding_models')
if 'source' not in embedding_columns:
if self.ap.persistence_mgr.db.name == 'postgresql':
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text(
"ALTER TABLE embedding_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL"
)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text(
"ALTER TABLE embedding_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL"
)
)
if 'space_model_id' not in embedding_columns:
if self.ap.persistence_mgr.db.name == 'postgresql':
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN space_model_id VARCHAR(255)')
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN space_model_id VARCHAR(255)')
)
async def _get_columns(self, table_name: str) -> list:
"""Get column names for a table"""
if self.ap.persistence_mgr.db.name == 'postgresql':
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.text(
f"SELECT column_name FROM information_schema.columns WHERE table_name = '{table_name}';"
)
)
all_result = result.fetchall()
return [row[0] for row in all_result]
else:
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.text(f'PRAGMA table_info({table_name});'))
all_result = result.fetchall()
return [row[1] for row in all_result]
pass
async def downgrade(self):
"""Downgrade"""

View File

@@ -61,6 +61,12 @@ class DBMigrateModelProviderRefactor(migration.DBMigration):
sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN provider_uuid VARCHAR(255)')
)
# Add prefered_ranking column if not exists
if 'prefered_ranking' not in llm_columns:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN prefered_ranking INTEGER NOT NULL DEFAULT 0')
)
# Only migrate if old columns exist
if 'requester' not in llm_columns:
return
@@ -152,6 +158,12 @@ class DBMigrateModelProviderRefactor(migration.DBMigration):
sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN provider_uuid VARCHAR(255)')
)
# Add prefered_ranking column if not exists
if 'prefered_ranking' not in embedding_columns:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN prefered_ranking INTEGER NOT NULL DEFAULT 0')
)
# Only migrate if old columns exist
if 'requester' not in embedding_columns:
return