Merge pull request #2140 from cheahjs/feat/model-config

feat: configurable model name, description and vision capability

backend/apps/litellm/main.py

@@ -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
+from config import SRC_LOG_LEVELS
 from constants import MESSAGES
 
 import os
@@ -77,7 +78,7 @@ 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 = Models.get_all_models()
 
 app.state.ENABLE = ENABLE_LITELLM
 app.state.CONFIG = litellm_config
@@ -241,6 +242,8 @@ async def get_models(user=Depends(get_current_user)):
                         )
                     )
 
+            for model in data["data"]:
+                add_custom_info_to_model(model)
             return data
         except Exception as e:
 
@@ -261,6 +264,14 @@ async def get_models(user=Depends(get_current_user)):
                         "object": "model",
                         "created": int(time.time()),
                         "owned_by": "openai",
+                        "custom_info": next(
+                            (
+                                item
+                                for item in app.state.MODEL_CONFIG
+                                if item.id == model["model_name"]
+                            ),
+                            None,
+                        ),
                     }
                     for model in app.state.CONFIG["model_list"]
                 ],
@@ -273,6 +284,12 @@ 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"]), None
+    )
+
+
 @app.get("/model/info")
 async def get_model_list(user=Depends(get_admin_user)):
     return {"data": app.state.CONFIG["model_list"]}

backend/apps/ollama/main.py

@@ -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 (
@@ -67,6 +67,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.MODEL_CONFIG = Models.get_all_models()
 
 
 app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
@@ -191,12 +192,21 @@ async def get_all_models():
 
     else:
         models = {"models": []}
+        
+    for model in models["models"]:
+        add_custom_info_to_model(model)
 
     app.state.MODELS = {model["model"]: model for model in models["models"]}
 
     return 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"]), None
+    )
+
+
 @app.get("/api/tags")
 @app.get("/api/tags/{url_idx}")
 async def get_ollama_tags(

backend/apps/openai/main.py

@@ -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 (
@@ -52,6 +52,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.MODEL_CONFIG = Models.get_all_models()
 
 
 app.state.config.ENABLE_OPENAI_API = ENABLE_OPENAI_API
@@ -249,10 +250,19 @@ async def get_all_models():
             )
         }
 
+        for model in models["data"]:
+            add_custom_info_to_model(model)
+
         log.info(f"models: {models}")
         app.state.MODELS = {model["id"]: model for model in models["data"]}
 
-        return models
+    return 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"]), None
+    )
 
 
 @app.get("/models")

backend/apps/web/internal/db.py

@@ -1,3 +1,5 @@
+import json
+
 from peewee import *
 from peewee_migrate import Router
 from playhouse.db_url import connect
@@ -8,6 +10,16 @@ import logging
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["DB"])
 
+
+class JSONField(TextField):
+    def db_value(self, value):
+        return json.dumps(value)
+
+    def python_value(self, value):
+        if value is not None:
+            return json.loads(value)
+
+
 # Check if the file exists
 if os.path.exists(f"{DATA_DIR}/ollama.db"):
     # Rename the file

backend/apps/web/internal/migrations/009_add_models.py

@@ -0,0 +1,55 @@
+"""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
+
+
+with suppress(ImportError):
+    import playhouse.postgres_ext as pw_pext
+
+
+def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
+    """Write your migrations here."""
+
+    @migrator.create_model
+    class Model(pw.Model):
+        id = pw.TextField(unique=True)
+        meta = pw.TextField()
+        base_model_id = pw.TextField(null=True)
+        name = pw.TextField()
+        params = pw.TextField()
+
+        class Meta:
+            table_name = "model"
+
+
+def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
+    """Write your rollback migrations here."""
+
+    migrator.remove_model("model")

backend/apps/web/models/models.py

@@ -0,0 +1,136 @@
+import json
+import logging
+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, JSONField
+
+from config import SRC_LOG_LEVELS
+
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["MODELS"])
+
+
+####################
+# 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):
+    pass
+
+
+# ModelMeta is a model for the data stored in the meta field of the Model table
+# It isn't currently used in the backend, but it's here as a reference
+class ModelMeta(BaseModel):
+    description: str
+    """
+        User-facing description of the model.
+    """
+
+    vision_capable: bool
+    """
+        A flag indicating if the model is capable of vision and thus image inputs
+    """
+
+
+class Model(pw.Model):
+    id = pw.TextField(unique=True)
+    """
+        The model's id as used in the API. If set to an existing model, it will override the model.
+    """
+
+    meta = JSONField()
+    """
+        Holds a JSON encoded blob of metadata, see `ModelMeta`.
+    """
+
+    base_model_id = 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 = JSONField()
+    """
+        Holds a JSON encoded blob of parameters, see `ModelParams`.
+    """
+
+    class Meta:
+        database = DB
+
+
+class ModelModel(BaseModel):
+    id: str
+    meta: ModelMeta
+    base_model_id: Optional[str] = None
+    name: str
+    params: ModelParams
+
+
+####################
+# Forms
+####################
+
+
+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 update_all_models(self, models: list[ModelModel]) -> bool:
+        try:
+            with self.db.atomic():
+                # Fetch current models from the database
+                current_models = self.get_all_models()
+                current_model_dict = {model.id: model for model in current_models}
+
+                # Create a set of model IDs from the current models and the new models
+                current_model_keys = set(current_model_dict.keys())
+                new_model_keys = set(model.id 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 not in current_model_keys
+                ]
+                models_to_update = [
+                    model for model in models if model.id 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.model_dump())
+
+                for model in models_to_update:
+                    Model.update(**model.model_dump()).where(
+                        Model.id == model.id
+                    ).execute()
+
+                for model_id, model_source in models_to_delete:
+                    Model.delete().where(Model.id == model_id).execute()
+
+            return True
+        except Exception as e:
+            log.exception(e)
+            return False
+
+
+Models = ModelsTable(DB)

backend/main.py

@@ -36,9 +36,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
 from utils.utils import get_admin_user
 from apps.rag.utils import rag_messages
 
@@ -113,6 +113,8 @@ 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 = Models.get_all_models()
+
 app.state.config.WEBHOOK_URL = WEBHOOK_URL
 
 origins = ["*"]
@@ -318,6 +320,33 @@ async def update_model_filter_config(
     }
 
 
+class SetModelConfigForm(BaseModel):
+    models: List[ModelModel]
+
+
+@app.post("/api/config/models")
+async def update_model_config(
+    form_data: SetModelConfigForm, user=Depends(get_admin_user)
+):
+    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 = form_data.models
+    openai_app.state.MODEL_CONFIG = form_data.models
+    litellm_app.state.MODEL_CONFIG = form_data.models
+    app.state.MODEL_CONFIG = form_data.models
+
+    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.MODEL_CONFIG}
+
+
 @app.get("/api/webhook")
 async def get_webhook_url(user=Depends(get_admin_user)):
     return {