mirror of
https://github.com/open-webui/open-webui
synced 2025-06-26 18:26:48 +00:00
feat: store model configs in the database
This commit is contained in:
Jun Siang Cheah
@@ -1,4 +1,4 @@
|
||||
"""Peewee migrations -- 002_add_local_sharing.py.
|
||||
"""Peewee migrations -- 008_add_models.py.
|
||||
|
||||
Some examples (model - class or model name)::
|
||||
|
||||
@@ -37,43 +37,24 @@ with suppress(ImportError):
|
||||
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your migrations here."""
|
||||
|
||||
# Adding fields created_at and updated_at to the 'user' table
|
||||
migrator.add_fields(
|
||||
"user",
|
||||
created_at=pw.BigIntegerField(null=True), # Allow null for transition
|
||||
updated_at=pw.BigIntegerField(null=True), # Allow null for transition
|
||||
last_active_at=pw.BigIntegerField(null=True), # Allow null for transition
|
||||
)
|
||||
@migrator.create_model
|
||||
class Model(pw.Model):
|
||||
id = pw.TextField()
|
||||
source = pw.TextField()
|
||||
base_model = pw.TextField(null=True)
|
||||
name = pw.TextField()
|
||||
params = pw.TextField()
|
||||
|
||||
# Populate the new fields from an existing 'timestamp' field
|
||||
migrator.sql(
|
||||
'UPDATE "user" SET created_at = timestamp, updated_at = timestamp, last_active_at = timestamp WHERE timestamp IS NOT NULL'
|
||||
)
|
||||
class Meta:
|
||||
table_name = "model"
|
||||
|
||||
# Now that the data has been copied, remove the original 'timestamp' field
|
||||
migrator.remove_fields("user", "timestamp")
|
||||
|
||||
# Update the fields to be not null now that they are populated
|
||||
migrator.change_fields(
|
||||
"user",
|
||||
created_at=pw.BigIntegerField(null=False),
|
||||
updated_at=pw.BigIntegerField(null=False),
|
||||
last_active_at=pw.BigIntegerField(null=False),
|
||||
)
|
||||
indexes = (
|
||||
# Create a unique index on the id, source columns
|
||||
(("id", "source"), True),
|
||||
)
|
||||
|
||||
|
||||
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
|
||||
"""Write your rollback migrations here."""
|
||||
|
||||
# Recreate the timestamp field initially allowing null values for safe transition
|
||||
migrator.add_fields("user", timestamp=pw.BigIntegerField(null=True))
|
||||
|
||||
# Copy the earliest created_at date back into the new timestamp field
|
||||
# This assumes created_at was originally a copy of timestamp
|
||||
migrator.sql('UPDATE "user" SET timestamp = created_at')
|
||||
|
||||
# Remove the created_at and updated_at fields
|
||||
migrator.remove_fields("user", "created_at", "updated_at", "last_active_at")
|
||||
|
||||
# Finally, alter the timestamp field to not allow nulls if that was the original setting
|
||||
migrator.change_fields("user", timestamp=pw.BigIntegerField(null=False))
|
||||
migrator.remove_model("model")
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
import json
|
||||
from typing import Optional
|
||||
|
||||
import peewee as pw
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel
|
||||
|
||||
from apps.web.internal.db import DB
|
||||
|
||||
|
||||
####################
|
||||
# Models DB Schema
|
||||
####################
|
||||
|
||||
|
||||
# ModelParams is a model for the data stored in the params field of the Model table
|
||||
# It isn't currently used in the backend, but it's here as a reference
|
||||
class ModelParams(BaseModel):
|
||||
"""
|
||||
A Pydantic model that represents the parameters of a model.
|
||||
|
||||
Attributes:
|
||||
description (str): A description of the model.
|
||||
vision_capable (bool): A flag indicating if the model is capable of vision and thus image inputs.
|
||||
"""
|
||||
|
||||
description: str
|
||||
vision_capable: bool
|
||||
|
||||
|
||||
class Model(pw.Model):
|
||||
id = pw.TextField()
|
||||
"""
|
||||
The model's id as used in the API. If set to an existing model, it will override the model.
|
||||
"""
|
||||
|
||||
source = pw.TextField()
|
||||
"""
|
||||
The source of the model, e.g., ollama, openai, or litellm.
|
||||
"""
|
||||
|
||||
base_model = pw.TextField(null=True)
|
||||
"""
|
||||
An optional pointer to the actual model that should be used when proxying requests.
|
||||
Currently unused - but will be used to support Modelfile like behaviour in the future
|
||||
"""
|
||||
|
||||
name = pw.TextField()
|
||||
"""
|
||||
The human-readable display name of the model.
|
||||
"""
|
||||
|
||||
params = pw.TextField()
|
||||
"""
|
||||
Holds a JSON encoded blob of parameters, see `ModelParams`.
|
||||
"""
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
|
||||
indexes = (
|
||||
# Create a unique index on the id, source columns
|
||||
(("id", "source"), True),
|
||||
)
|
||||
|
||||
|
||||
class ModelModel(BaseModel):
|
||||
id: str
|
||||
source: str
|
||||
base_model: Optional[str] = None
|
||||
name: str
|
||||
params: str
|
||||
|
||||
def to_form(self) -> "ModelForm":
|
||||
return ModelForm(**{**self.model_dump(), "params": json.loads(self.params)})
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
####################
|
||||
|
||||
|
||||
class ModelForm(BaseModel):
|
||||
id: str
|
||||
source: str
|
||||
base_model: Optional[str] = None
|
||||
name: str
|
||||
params: dict
|
||||
|
||||
def to_db_model(self) -> ModelModel:
|
||||
return ModelModel(**{**self.model_dump(), "params": json.dumps(self.params)})
|
||||
|
||||
|
||||
class ModelsTable:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: pw.SqliteDatabase | pw.PostgresqlDatabase,
|
||||
):
|
||||
self.db = db
|
||||
self.db.create_tables([Model])
|
||||
|
||||
def get_all_models(self) -> list[ModelModel]:
|
||||
return [ModelModel(**model_to_dict(model)) for model in Model.select()]
|
||||
|
||||
def get_all_models_by_source(self, source: str) -> list[ModelModel]:
|
||||
return [
|
||||
ModelModel(**model_to_dict(model))
|
||||
for model in Model.select().where(Model.source == source)
|
||||
]
|
||||
|
||||
def update_all_models(self, models: list[ModelForm]) -> bool:
|
||||
try:
|
||||
with self.db.atomic():
|
||||
# Fetch current models from the database
|
||||
current_models = self.get_all_models()
|
||||
current_model_dict = {
|
||||
(model.id, model.source): model for model in current_models
|
||||
}
|
||||
|
||||
# Create a set of model IDs and sources from the current models and the new models
|
||||
current_model_keys = set(current_model_dict.keys())
|
||||
new_model_keys = set((model.id, model.source) for model in models)
|
||||
|
||||
# Determine which models need to be created, updated, or deleted
|
||||
models_to_create = [
|
||||
model
|
||||
for model in models
|
||||
if (model.id, model.source) not in current_model_keys
|
||||
]
|
||||
models_to_update = [
|
||||
model
|
||||
for model in models
|
||||
if (model.id, model.source) in current_model_keys
|
||||
]
|
||||
models_to_delete = current_model_keys - new_model_keys
|
||||
|
||||
# Perform the necessary database operations
|
||||
for model in models_to_create:
|
||||
Model.create(**model.to_db_model().model_dump())
|
||||
|
||||
for model in models_to_update:
|
||||
Model.update(**model.to_db_model().model_dump()).where(
|
||||
(Model.id == model.id) & (Model.source == model.source)
|
||||
).execute()
|
||||
|
||||
for model_id, model_source in models_to_delete:
|
||||
Model.delete().where(
|
||||
(Model.id == model_id) & (Model.source == model_source)
|
||||
).execute()
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
return False
|
||||
|
||||
|
||||
Models = ModelsTable(DB)
|
||||
|
||||
Reference in New Issue
Block a user