OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: store model configs in the database

Browse Source
This commit is contained in:
Jun Siang Cheah
2024-05-19 18:46:24 +08:00
parent 1bacd5d93f
commit 4002ead6af
50 changed files with 434 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
backend/apps/litellm/main.py
View File

@@ -18,8 +18,9 @@ import requests
from pydantic import BaseModel, ConfigDict
from typing import Optional, List
from apps.web.models.models import Models
from utils.utils import get_verified_user, get_current_user, get_admin_user
from config import SRC_LOG_LEVELS, ENV, MODEL_CONFIG
from config import SRC_LOG_LEVELS
from constants import MESSAGES
import os
@@ -77,11 +78,12 @@ with open(LITELLM_CONFIG_DIR, "r") as file:
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER.value
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST.value
app.state.MODEL_CONFIG = [
model.to_form() for model in Models.get_all_models_by_source("litellm")
]
app.state.ENABLE = ENABLE_LITELLM
app.state.CONFIG = litellm_config
app.state.MODEL_CONFIG = MODEL_CONFIG.value.get("litellm", [])
# Global variable to store the subprocess reference
background_process = None
@@ -268,9 +270,9 @@ async def get_models(user=Depends(get_current_user)):
(
item
for item in app.state.MODEL_CONFIG
if item["name"] == model["model_name"]
if item.id == model["model_name"]
),
{},
None,
),
}
for model in app.state.CONFIG["model_list"]
@@ -286,7 +288,7 @@ async def get_models(user=Depends(get_current_user)):
def add_custom_info_to_model(model: dict):
model["custom_info"] = next(
(item for item in app.state.MODEL_CONFIG if item["id"] == model["id"]), {}
(item for item in app.state.MODEL_CONFIG if item.id == model["id"]), None
)

9
backend/apps/ollama/main.py
View File

@@ -29,7 +29,7 @@ import time
from urllib.parse import urlparse
from typing import Optional, List, Union
from apps.web.models.models import Models
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
from utils.utils import (
@@ -46,7 +46,6 @@ from config import (
ENABLE_MODEL_FILTER,
MODEL_FILTER_LIST,
UPLOAD_DIR,
MODEL_CONFIG,
AppConfig,
)
from utils.misc import calculate_sha256
@@ -67,7 +66,9 @@ app.state.config = AppConfig()
app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
app.state.MODEL_CONFIG = MODEL_CONFIG.value.get("ollama", [])
app.state.MODEL_CONFIG = [
model.to_form() for model in Models.get_all_models_by_source("ollama")
]
app.state.config.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
app.state.MODELS = {}
@@ -179,7 +180,7 @@ async def get_all_models():
def add_custom_info_to_model(model: dict):
model["custom_info"] = next(
(item for item in app.state.MODEL_CONFIG if item["id"] == model["model"]), {}
(item for item in app.state.MODEL_CONFIG if item.id == model["model"]), None
)

9
backend/apps/openai/main.py
View File

@@ -10,7 +10,7 @@ import logging
from pydantic import BaseModel
from apps.web.models.models import Models
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
from utils.utils import (
@@ -27,7 +27,6 @@ from config import (
CACHE_DIR,
ENABLE_MODEL_FILTER,
MODEL_FILTER_LIST,
MODEL_CONFIG,
AppConfig,
)
from typing import List, Optional
@@ -53,7 +52,9 @@ app.state.config = AppConfig()
app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
app.state.MODEL_CONFIG = MODEL_CONFIG.value.get("openai", [])
app.state.MODEL_CONFIG = [
model.to_form() for model in Models.get_all_models_by_source("openai")
]
app.state.config.ENABLE_OPENAI_API = ENABLE_OPENAI_API
@@ -262,7 +263,7 @@ async def get_all_models():
def add_custom_info_to_model(model: dict):
model["custom_info"] = next(
(item for item in app.state.MODEL_CONFIG if item["id"] == model["id"]), {}
(item for item in app.state.MODEL_CONFIG if item.id == model["id"]), None
)

49
backend/apps/web/internal/migrations/008_add_models.py
View File

@@ -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")

157
backend/apps/web/models/models.py
View File

@@ -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)

4
backend/config.py
View File

@@ -549,10 +549,6 @@ WEBHOOK_URL = PersistentConfig(
ENABLE_ADMIN_EXPORT = os.environ.get("ENABLE_ADMIN_EXPORT", "True").lower() == "true"
MODEL_CONFIG = PersistentConfig(
"CONFIG_DATA", "models", {"ollama": [], "litellm": [], "openai": []}
)
####################################
# WEBUI_SECRET_KEY
####################################

46
backend/main.py
View File

@@ -35,9 +35,9 @@ from apps.web.main import app as webui_app
import asyncio
from pydantic import BaseModel
from typing import List
from typing import List, Optional
from apps.web.models.models import Models, ModelModel, ModelForm
from utils.utils import get_admin_user
from apps.rag.utils import rag_messages
@@ -59,7 +59,6 @@ from config import (
SRC_LOG_LEVELS,
WEBHOOK_URL,
ENABLE_ADMIN_EXPORT,
MODEL_CONFIG,
AppConfig,
)
from constants import ERROR_MESSAGES
@@ -113,7 +112,7 @@ app.state.config = AppConfig()
app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
app.state.config.MODEL_CONFIG = MODEL_CONFIG
app.state.MODEL_CONFIG = [model.to_form() for model in Models.get_all_models()]
app.state.config.WEBHOOK_URL = WEBHOOK_URL
@@ -310,43 +309,40 @@ async def update_model_filter_config(
}
class ModelConfig(BaseModel):
id: str
name: str
description: str
vision_capable: bool
class SetModelConfigForm(BaseModel):
ollama: List[ModelConfig]
litellm: List[ModelConfig]
openai: List[ModelConfig]
models: List[ModelForm]
@app.post("/api/config/models")
async def update_model_config(
form_data: SetModelConfigForm, user=Depends(get_admin_user)
):
data = form_data.model_dump()
if not Models.update_all_models(form_data.models):
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=ERROR_MESSAGES.DEFAULT("Failed to update model config"),
)
ollama_app.state.MODEL_CONFIG = data.get("ollama", [])
ollama_app.state.MODEL_CONFIG = [
model for model in form_data.models if model.source == "ollama"
]
openai_app.state.MODEL_CONFIG = data.get("openai", [])
openai_app.state.MODEL_CONFIG = [
model for model in form_data.models if model.source == "openai"
]
litellm_app.state.MODEL_CONFIG = data.get("litellm", [])
litellm_app.state.MODEL_CONFIG = [
model for model in form_data.models if model.source == "litellm"
]
app.state.config.MODEL_CONFIG = {
"ollama": ollama_app.state.MODEL_CONFIG,
"openai": openai_app.state.MODEL_CONFIG,
"litellm": litellm_app.state.MODEL_CONFIG,
}
app.state.MODEL_CONFIG = [model for model in form_data.models]
return {"models": app.state.config.MODEL_CONFIG}
return {"models": app.state.MODEL_CONFIG}
@app.get("/api/config/models")
async def get_model_config(user=Depends(get_admin_user)):
return {"models": app.state.config.MODEL_CONFIG}
return {"models": app.state.MODEL_CONFIG}
@app.get("/api/webhook")