open-webui/backend/apps/webui/internal/db.py
Беклемишев Петр Алексеевич e59e1f5049 Fix rebase artifacts
2024-05-30 20:44:13 +07:00

45 lines
1.1 KiB
Python

import os
import logging
import json
from peewee import *
from peewee_migrate import Router
from playhouse.db_url import connect
from apps.webui.internal.wrappers import PeeweeConnectionState, register_peewee_databases
from config import SRC_LOG_LEVELS, DATA_DIR, DATABASE_URL, BACKEND_DIR
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)
register_peewee_databases()
# Check if the file exists
if os.path.exists(f"{DATA_DIR}/ollama.db"):
# Rename the file
os.rename(f"{DATA_DIR}/ollama.db", f"{DATA_DIR}/webui.db")
log.info("Database migrated from Ollama-WebUI successfully.")
else:
pass
DB = connect(DATABASE_URL)
DB._state = PeeweeConnectionState()
log.info(f"Connected to a {DB.__class__.__name__} database.")
router = Router(
DB,
migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations",
logger=log,
)
router.run()
DB.connect(reuse_if_open=True)