|
"""Peewee migrations -- 009_add_models.py. |
|
|
|
Some examples (model - class or model name):: |
|
|
|
> Model = migrator.orm['table_name'] # Return model in current state by name |
|
> Model = migrator.ModelClass # Return model in current state by name |
|
|
|
> migrator.sql(sql) # Run custom SQL |
|
> migrator.run(func, *args, **kwargs) # Run python function with the given args |
|
> migrator.create_model(Model) # Create a model (could be used as decorator) |
|
> migrator.remove_model(model, cascade=True) # Remove a model |
|
> migrator.add_fields(model, **fields) # Add fields to a model |
|
> migrator.change_fields(model, **fields) # Change fields |
|
> migrator.remove_fields(model, *field_names, cascade=True) |
|
> migrator.rename_field(model, old_field_name, new_field_name) |
|
> migrator.rename_table(model, new_table_name) |
|
> migrator.add_index(model, *col_names, unique=False) |
|
> migrator.add_not_null(model, *field_names) |
|
> migrator.add_default(model, field_name, default) |
|
> migrator.add_constraint(model, name, sql) |
|
> migrator.drop_index(model, *col_names) |
|
> migrator.drop_not_null(model, *field_names) |
|
> migrator.drop_constraints(model, *constraints) |
|
|
|
""" |
|
|
|
from contextlib import suppress |
|
|
|
import peewee as pw |
|
from peewee_migrate import Migrator |
|
import json |
|
|
|
from utils.misc import parse_ollama_modelfile |
|
|
|
with suppress(ImportError): |
|
import playhouse.postgres_ext as pw_pext |
|
|
|
|
|
def migrate(migrator: Migrator, database: pw.Database, *, fake=False): |
|
"""Write your migrations here.""" |
|
|
|
|
|
migrate_modelfile_to_model(migrator, database) |
|
|
|
migrator.remove_model("modelfile") |
|
|
|
|
|
def migrate_modelfile_to_model(migrator: Migrator, database: pw.Database): |
|
ModelFile = migrator.orm["modelfile"] |
|
Model = migrator.orm["model"] |
|
|
|
modelfiles = ModelFile.select() |
|
|
|
for modelfile in modelfiles: |
|
|
|
|
|
modelfile.modelfile = json.loads(modelfile.modelfile) |
|
meta = json.dumps( |
|
{ |
|
"description": modelfile.modelfile.get("desc"), |
|
"profile_image_url": modelfile.modelfile.get("imageUrl"), |
|
"ollama": {"modelfile": modelfile.modelfile.get("content")}, |
|
"suggestion_prompts": modelfile.modelfile.get("suggestionPrompts"), |
|
"categories": modelfile.modelfile.get("categories"), |
|
"user": {**modelfile.modelfile.get("user", {}), "community": True}, |
|
} |
|
) |
|
|
|
info = parse_ollama_modelfile(modelfile.modelfile.get("content")) |
|
|
|
|
|
Model.create( |
|
id=f"ollama-{modelfile.tag_name}", |
|
user_id=modelfile.user_id, |
|
base_model_id=info.get("base_model_id"), |
|
name=modelfile.modelfile.get("title"), |
|
meta=meta, |
|
params=json.dumps(info.get("params", {})), |
|
created_at=modelfile.timestamp, |
|
updated_at=modelfile.timestamp, |
|
) |
|
|
|
|
|
def rollback(migrator: Migrator, database: pw.Database, *, fake=False): |
|
"""Write your rollback migrations here.""" |
|
|
|
recreate_modelfile_table(migrator, database) |
|
move_data_back_to_modelfile(migrator, database) |
|
migrator.remove_model("model") |
|
|
|
|
|
def recreate_modelfile_table(migrator: Migrator, database: pw.Database): |
|
query = """ |
|
CREATE TABLE IF NOT EXISTS modelfile ( |
|
user_id TEXT, |
|
tag_name TEXT, |
|
modelfile JSON, |
|
timestamp BIGINT |
|
) |
|
""" |
|
migrator.sql(query) |
|
|
|
|
|
def move_data_back_to_modelfile(migrator: Migrator, database: pw.Database): |
|
Model = migrator.orm["model"] |
|
Modelfile = migrator.orm["modelfile"] |
|
|
|
models = Model.select() |
|
|
|
for model in models: |
|
|
|
meta = json.loads(model.meta) |
|
|
|
modelfile_data = { |
|
"title": model.name, |
|
"desc": meta.get("description"), |
|
"imageUrl": meta.get("profile_image_url"), |
|
"content": meta.get("ollama", {}).get("modelfile"), |
|
"suggestionPrompts": meta.get("suggestion_prompts"), |
|
"categories": meta.get("categories"), |
|
"user": {k: v for k, v in meta.get("user", {}).items() if k != "community"}, |
|
} |
|
|
|
|
|
Modelfile.create( |
|
user_id=model.user_id, |
|
tag_name=model.id, |
|
modelfile=modelfile_data, |
|
timestamp=model.created_at, |
|
) |
|
|