mirror of
https://git.mirrors.martin98.com/https://github.com/open-webui/open-webui
synced 2025-08-17 11:15:56 +08:00
Merge pull request #3595 from open-webui/dev-migration
feat: db migration
This commit is contained in:
commit
8f6f76682b
26
.github/workflows/integration-test.yml
vendored
26
.github/workflows/integration-test.yml
vendored
@ -67,6 +67,28 @@ jobs:
|
||||
path: compose-logs.txt
|
||||
if-no-files-found: ignore
|
||||
|
||||
pytest:
|
||||
name: Run backend tests
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r backend/requirements.txt
|
||||
|
||||
- name: pytest run
|
||||
run: |
|
||||
ls -al
|
||||
cd backend
|
||||
PYTHONPATH=. pytest . -o log_cli=true -o log_cli_level=INFO
|
||||
|
||||
migration_test:
|
||||
name: Run Migration Tests
|
||||
runs-on: ubuntu-latest
|
||||
@ -171,7 +193,7 @@ jobs:
|
||||
fi
|
||||
|
||||
# Check that service will reconnect to postgres when connection will be closed
|
||||
status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health)
|
||||
status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db)
|
||||
if [[ "$status_code" -ne 200 ]] ; then
|
||||
echo "Server has failed before postgres reconnect check"
|
||||
exit 1
|
||||
@ -183,7 +205,7 @@ jobs:
|
||||
cur = conn.cursor(); \
|
||||
cur.execute('SELECT pg_terminate_backend(psa.pid) FROM pg_stat_activity psa WHERE datname = current_database() AND pid <> pg_backend_pid();')"
|
||||
|
||||
status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health)
|
||||
status_code=$(curl --write-out %{http_code} -s --output /dev/null http://localhost:8081/health/db)
|
||||
if [[ "$status_code" -ne 200 ]] ; then
|
||||
echo "Server has not reconnected to postgres after connection was closed: returned status $status_code"
|
||||
exit 1
|
||||
|
114
backend/alembic.ini
Normal file
114
backend/alembic.ini
Normal file
@ -0,0 +1,114 @@
|
||||
# A generic, single database configuration.
|
||||
|
||||
[alembic]
|
||||
# path to migration scripts
|
||||
script_location = migrations
|
||||
|
||||
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
|
||||
# Uncomment the line below if you want the files to be prepended with date and time
|
||||
# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||
|
||||
# sys.path path, will be prepended to sys.path if present.
|
||||
# defaults to the current working directory.
|
||||
prepend_sys_path = .
|
||||
|
||||
# timezone to use when rendering the date within the migration file
|
||||
# as well as the filename.
|
||||
# If specified, requires the python>=3.9 or backports.zoneinfo library.
|
||||
# Any required deps can installed by adding `alembic[tz]` to the pip requirements
|
||||
# string value is passed to ZoneInfo()
|
||||
# leave blank for localtime
|
||||
# timezone =
|
||||
|
||||
# max length of characters to apply to the
|
||||
# "slug" field
|
||||
# truncate_slug_length = 40
|
||||
|
||||
# set to 'true' to run the environment during
|
||||
# the 'revision' command, regardless of autogenerate
|
||||
# revision_environment = false
|
||||
|
||||
# set to 'true' to allow .pyc and .pyo files without
|
||||
# a source .py file to be detected as revisions in the
|
||||
# versions/ directory
|
||||
# sourceless = false
|
||||
|
||||
# version location specification; This defaults
|
||||
# to migrations/versions. When using multiple version
|
||||
# directories, initial revisions must be specified with --version-path.
|
||||
# The path separator used here should be the separator specified by "version_path_separator" below.
|
||||
# version_locations = %(here)s/bar:%(here)s/bat:migrations/versions
|
||||
|
||||
# version path separator; As mentioned above, this is the character used to split
|
||||
# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
|
||||
# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
|
||||
# Valid values for version_path_separator are:
|
||||
#
|
||||
# version_path_separator = :
|
||||
# version_path_separator = ;
|
||||
# version_path_separator = space
|
||||
version_path_separator = os # Use os.pathsep. Default configuration used for new projects.
|
||||
|
||||
# set to 'true' to search source files recursively
|
||||
# in each "version_locations" directory
|
||||
# new in Alembic version 1.10
|
||||
# recursive_version_locations = false
|
||||
|
||||
# the output encoding used when revision files
|
||||
# are written from script.py.mako
|
||||
# output_encoding = utf-8
|
||||
|
||||
# sqlalchemy.url = REPLACE_WITH_DATABASE_URL
|
||||
|
||||
|
||||
[post_write_hooks]
|
||||
# post_write_hooks defines scripts or Python functions that are run
|
||||
# on newly generated revision scripts. See the documentation for further
|
||||
# detail and examples
|
||||
|
||||
# format using "black" - use the console_scripts runner, against the "black" entrypoint
|
||||
# hooks = black
|
||||
# black.type = console_scripts
|
||||
# black.entrypoint = black
|
||||
# black.options = -l 79 REVISION_SCRIPT_FILENAME
|
||||
|
||||
# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
|
||||
# hooks = ruff
|
||||
# ruff.type = exec
|
||||
# ruff.executable = %(here)s/.venv/bin/ruff
|
||||
# ruff.options = --fix REVISION_SCRIPT_FILENAME
|
||||
|
||||
# Logging configuration
|
||||
[loggers]
|
||||
keys = root,sqlalchemy,alembic
|
||||
|
||||
[handlers]
|
||||
keys = console
|
||||
|
||||
[formatters]
|
||||
keys = generic
|
||||
|
||||
[logger_root]
|
||||
level = WARN
|
||||
handlers = console
|
||||
qualname =
|
||||
|
||||
[logger_sqlalchemy]
|
||||
level = WARN
|
||||
handlers =
|
||||
qualname = sqlalchemy.engine
|
||||
|
||||
[logger_alembic]
|
||||
level = INFO
|
||||
handlers =
|
||||
qualname = alembic
|
||||
|
||||
[handler_console]
|
||||
class = StreamHandler
|
||||
args = (sys.stderr,)
|
||||
level = NOTSET
|
||||
formatter = generic
|
||||
|
||||
[formatter_generic]
|
||||
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||
datefmt = %H:%M:%S
|
@ -1,18 +1,39 @@
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
|
||||
from peewee import *
|
||||
from peewee_migrate import Router
|
||||
|
||||
from apps.webui.internal.wrappers import register_connection
|
||||
|
||||
from typing import Optional, Any
|
||||
from typing_extensions import Self
|
||||
|
||||
from sqlalchemy import create_engine, types, Dialect
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker, scoped_session
|
||||
from sqlalchemy.sql.type_api import _T
|
||||
|
||||
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):
|
||||
class JSONField(types.TypeDecorator):
|
||||
impl = types.Text
|
||||
cache_ok = True
|
||||
|
||||
def process_bind_param(self, value: Optional[_T], dialect: Dialect) -> Any:
|
||||
return json.dumps(value)
|
||||
|
||||
def process_result_value(self, value: Optional[_T], dialect: Dialect) -> Any:
|
||||
if value is not None:
|
||||
return json.loads(value)
|
||||
|
||||
def copy(self, **kw: Any) -> Self:
|
||||
return JSONField(self.impl.length)
|
||||
|
||||
def db_value(self, value):
|
||||
return json.dumps(value)
|
||||
|
||||
@ -30,25 +51,57 @@ else:
|
||||
pass
|
||||
|
||||
|
||||
# The `register_connection` function encapsulates the logic for setting up
|
||||
# the database connection based on the connection string, while `connect`
|
||||
# is a Peewee-specific method to manage the connection state and avoid errors
|
||||
# when a connection is already open.
|
||||
try:
|
||||
DB = register_connection(DATABASE_URL)
|
||||
log.info(f"Connected to a {DB.__class__.__name__} database.")
|
||||
except Exception as e:
|
||||
log.error(f"Failed to initialize the database connection: {e}")
|
||||
raise
|
||||
# Workaround to handle the peewee migration
|
||||
# This is required to ensure the peewee migration is handled before the alembic migration
|
||||
def handle_peewee_migration():
|
||||
try:
|
||||
db = register_connection(DATABASE_URL)
|
||||
migrate_dir = BACKEND_DIR / "apps" / "webui" / "internal" / "migrations"
|
||||
router = Router(db, logger=log, migrate_dir=migrate_dir)
|
||||
router.run()
|
||||
db.close()
|
||||
|
||||
router = Router(
|
||||
DB,
|
||||
migrate_dir=BACKEND_DIR / "apps" / "webui" / "internal" / "migrations",
|
||||
logger=log,
|
||||
# check if db connection has been closed
|
||||
|
||||
except Exception as e:
|
||||
log.error(f"Failed to initialize the database connection: {e}")
|
||||
raise
|
||||
|
||||
finally:
|
||||
# Properly closing the database connection
|
||||
if db and not db.is_closed():
|
||||
db.close()
|
||||
|
||||
# Assert if db connection has been closed
|
||||
assert db.is_closed(), "Database connection is still open."
|
||||
|
||||
|
||||
handle_peewee_migration()
|
||||
|
||||
|
||||
SQLALCHEMY_DATABASE_URL = DATABASE_URL
|
||||
if "sqlite" in SQLALCHEMY_DATABASE_URL:
|
||||
engine = create_engine(
|
||||
SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
|
||||
)
|
||||
else:
|
||||
engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
|
||||
|
||||
|
||||
SessionLocal = sessionmaker(
|
||||
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
||||
)
|
||||
router.run()
|
||||
try:
|
||||
DB.connect(reuse_if_open=True)
|
||||
except OperationalError as e:
|
||||
log.info(f"Failed to connect to database again due to: {e}")
|
||||
pass
|
||||
Base = declarative_base()
|
||||
Session = scoped_session(SessionLocal)
|
||||
|
||||
|
||||
# Dependency
|
||||
def get_session():
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
get_db = contextmanager(get_session)
|
||||
|
@ -1,10 +1,7 @@
|
||||
"""Peewee migrations -- 017_add_user_oauth_sub.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)
|
||||
@ -21,7 +18,6 @@ Some examples (model - class or model name)::
|
||||
> migrator.drop_index(model, *col_names)
|
||||
> migrator.drop_not_null(model, *field_names)
|
||||
> migrator.drop_constraints(model, *constraints)
|
||||
|
||||
"""
|
||||
|
||||
from contextlib import suppress
|
||||
|
@ -1,21 +0,0 @@
|
||||
# Database Migrations
|
||||
|
||||
This directory contains all the database migrations for the web app.
|
||||
Migrations are done using the [`peewee-migrate`](https://github.com/klen/peewee_migrate) library.
|
||||
|
||||
Migrations are automatically ran at app startup.
|
||||
|
||||
## Creating a migration
|
||||
|
||||
Have you made a change to the schema of an existing model?
|
||||
You will need to create a migration file to ensure that existing databases are updated for backwards compatibility.
|
||||
|
||||
1. Have a database file (`webui.db`) that has the old schema prior to any of your changes.
|
||||
2. Make your changes to the models.
|
||||
3. From the `backend` directory, run the following command:
|
||||
```bash
|
||||
pw_migrate create --auto --auto-source apps.webui.models --database sqlite:///${SQLITE_DB} --directory apps/web/internal/migrations ${MIGRATION_NAME}
|
||||
```
|
||||
- `$SQLITE_DB` should be the path to the database file.
|
||||
- `$MIGRATION_NAME` should be a descriptive name for the migration.
|
||||
4. The migration file will be created in the `apps/web/internal/migrations` directory.
|
@ -3,7 +3,7 @@ from fastapi.routing import APIRoute
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from starlette.middleware.sessions import SessionMiddleware
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
from apps.webui.routers import (
|
||||
auths,
|
||||
users,
|
||||
|
@ -1,14 +1,13 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
from typing import Optional
|
||||
import uuid
|
||||
import logging
|
||||
from peewee import *
|
||||
from sqlalchemy import String, Column, Boolean, Text
|
||||
|
||||
from apps.webui.models.users import UserModel, Users
|
||||
from utils.utils import verify_password
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
@ -20,14 +19,13 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class Auth(Model):
|
||||
id = CharField(unique=True)
|
||||
email = CharField()
|
||||
password = TextField()
|
||||
active = BooleanField()
|
||||
class Auth(Base):
|
||||
__tablename__ = "auth"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
email = Column(String)
|
||||
password = Column(Text)
|
||||
active = Column(Boolean)
|
||||
|
||||
|
||||
class AuthModel(BaseModel):
|
||||
@ -94,9 +92,6 @@ class AddUserForm(SignupForm):
|
||||
|
||||
|
||||
class AuthsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Auth])
|
||||
|
||||
def insert_new_auth(
|
||||
self,
|
||||
@ -107,36 +102,44 @@ class AuthsTable:
|
||||
role: str = "pending",
|
||||
oauth_sub: Optional[str] = None,
|
||||
) -> Optional[UserModel]:
|
||||
log.info("insert_new_auth")
|
||||
with get_db() as db:
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
log.info("insert_new_auth")
|
||||
|
||||
auth = AuthModel(
|
||||
**{"id": id, "email": email, "password": password, "active": True}
|
||||
)
|
||||
result = Auth.create(**auth.model_dump())
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
user = Users.insert_new_user(
|
||||
id, name, email, profile_image_url, role, oauth_sub
|
||||
)
|
||||
auth = AuthModel(
|
||||
**{"id": id, "email": email, "password": password, "active": True}
|
||||
)
|
||||
result = Auth(**auth.model_dump())
|
||||
db.add(result)
|
||||
|
||||
if result and user:
|
||||
return user
|
||||
else:
|
||||
return None
|
||||
user = Users.insert_new_user(
|
||||
id, name, email, profile_image_url, role, oauth_sub
|
||||
)
|
||||
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
|
||||
if result and user:
|
||||
return user
|
||||
else:
|
||||
return None
|
||||
|
||||
def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
|
||||
log.info(f"authenticate_user: {email}")
|
||||
try:
|
||||
auth = Auth.get(Auth.email == email, Auth.active == True)
|
||||
if auth:
|
||||
if verify_password(password, auth.password):
|
||||
user = Users.get_user_by_id(auth.id)
|
||||
return user
|
||||
with get_db() as db:
|
||||
|
||||
auth = db.query(Auth).filter_by(email=email, active=True).first()
|
||||
if auth:
|
||||
if verify_password(password, auth.password):
|
||||
user = Users.get_user_by_id(auth.id)
|
||||
return user
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
|
||||
@ -155,46 +158,49 @@ class AuthsTable:
|
||||
def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]:
|
||||
log.info(f"authenticate_user_by_trusted_header: {email}")
|
||||
try:
|
||||
auth = Auth.get(Auth.email == email, Auth.active == True)
|
||||
if auth:
|
||||
user = Users.get_user_by_id(auth.id)
|
||||
return user
|
||||
with get_db() as db:
|
||||
auth = db.query(Auth).filter(email=email, active=True).first()
|
||||
if auth:
|
||||
user = Users.get_user_by_id(auth.id)
|
||||
return user
|
||||
except:
|
||||
return None
|
||||
|
||||
def update_user_password_by_id(self, id: str, new_password: str) -> bool:
|
||||
try:
|
||||
query = Auth.update(password=new_password).where(Auth.id == id)
|
||||
result = query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
return True if result == 1 else False
|
||||
result = (
|
||||
db.query(Auth).filter_by(id=id).update({"password": new_password})
|
||||
)
|
||||
return True if result == 1 else False
|
||||
except:
|
||||
return False
|
||||
|
||||
def update_email_by_id(self, id: str, email: str) -> bool:
|
||||
try:
|
||||
query = Auth.update(email=email).where(Auth.id == id)
|
||||
result = query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
return True if result == 1 else False
|
||||
result = db.query(Auth).filter_by(id=id).update({"email": email})
|
||||
return True if result == 1 else False
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_auth_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
# Delete User
|
||||
result = Users.delete_user_by_id(id)
|
||||
with get_db() as db:
|
||||
|
||||
if result:
|
||||
# Delete Auth
|
||||
query = Auth.delete().where(Auth.id == id)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
# Delete User
|
||||
result = Users.delete_user_by_id(id)
|
||||
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
if result:
|
||||
db.query(Auth).filter_by(id=id).delete()
|
||||
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Auths = AuthsTable(DB)
|
||||
Auths = AuthsTable()
|
||||
|
@ -1,36 +1,38 @@
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Union, Optional
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
|
||||
import json
|
||||
import uuid
|
||||
import time
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from sqlalchemy import Column, String, BigInteger, Boolean, Text
|
||||
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
|
||||
####################
|
||||
# Chat DB Schema
|
||||
####################
|
||||
|
||||
|
||||
class Chat(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
title = TextField()
|
||||
chat = TextField() # Save Chat JSON as Text
|
||||
class Chat(Base):
|
||||
__tablename__ = "chat"
|
||||
|
||||
created_at = BigIntegerField()
|
||||
updated_at = BigIntegerField()
|
||||
id = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
title = Column(Text)
|
||||
chat = Column(Text) # Save Chat JSON as Text
|
||||
|
||||
share_id = CharField(null=True, unique=True)
|
||||
archived = BooleanField(default=False)
|
||||
created_at = Column(BigInteger)
|
||||
updated_at = Column(BigInteger)
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
share_id = Column(Text, unique=True, nullable=True)
|
||||
archived = Column(Boolean, default=False)
|
||||
|
||||
|
||||
class ChatModel(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
user_id: str
|
||||
title: str
|
||||
@ -75,91 +77,101 @@ class ChatTitleIdResponse(BaseModel):
|
||||
|
||||
|
||||
class ChatTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
db.create_tables([Chat])
|
||||
|
||||
def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
|
||||
id = str(uuid.uuid4())
|
||||
chat = ChatModel(
|
||||
**{
|
||||
"id": id,
|
||||
"user_id": user_id,
|
||||
"title": (
|
||||
form_data.chat["title"] if "title" in form_data.chat else "New Chat"
|
||||
),
|
||||
"chat": json.dumps(form_data.chat),
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
result = Chat.create(**chat.model_dump())
|
||||
return chat if result else None
|
||||
id = str(uuid.uuid4())
|
||||
chat = ChatModel(
|
||||
**{
|
||||
"id": id,
|
||||
"user_id": user_id,
|
||||
"title": (
|
||||
form_data.chat["title"]
|
||||
if "title" in form_data.chat
|
||||
else "New Chat"
|
||||
),
|
||||
"chat": json.dumps(form_data.chat),
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
result = Chat(**chat.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
return ChatModel.model_validate(result) if result else None
|
||||
|
||||
def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
|
||||
try:
|
||||
query = Chat.update(
|
||||
chat=json.dumps(chat),
|
||||
title=chat["title"] if "title" in chat else "New Chat",
|
||||
updated_at=int(time.time()),
|
||||
).where(Chat.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
chat = Chat.get(Chat.id == id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
except:
|
||||
chat_obj = db.get(Chat, id)
|
||||
chat_obj.chat = json.dumps(chat)
|
||||
chat_obj.title = chat["title"] if "title" in chat else "New Chat"
|
||||
chat_obj.updated_at = int(time.time())
|
||||
db.commit()
|
||||
db.refresh(chat_obj)
|
||||
|
||||
return ChatModel.model_validate(chat_obj)
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
|
||||
# Get the existing chat to share
|
||||
chat = Chat.get(Chat.id == chat_id)
|
||||
# Check if the chat is already shared
|
||||
if chat.share_id:
|
||||
return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
|
||||
# Create a new chat with the same data, but with a new ID
|
||||
shared_chat = ChatModel(
|
||||
**{
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_id": f"shared-{chat_id}",
|
||||
"title": chat.title,
|
||||
"chat": chat.chat,
|
||||
"created_at": chat.created_at,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
shared_result = Chat.create(**shared_chat.model_dump())
|
||||
# Update the original chat with the share_id
|
||||
result = (
|
||||
Chat.update(share_id=shared_chat.id).where(Chat.id == chat_id).execute()
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
return shared_chat if (shared_result and result) else None
|
||||
# Get the existing chat to share
|
||||
chat = db.get(Chat, chat_id)
|
||||
# Check if the chat is already shared
|
||||
if chat.share_id:
|
||||
return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
|
||||
# Create a new chat with the same data, but with a new ID
|
||||
shared_chat = ChatModel(
|
||||
**{
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_id": f"shared-{chat_id}",
|
||||
"title": chat.title,
|
||||
"chat": chat.chat,
|
||||
"created_at": chat.created_at,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
shared_result = Chat(**shared_chat.model_dump())
|
||||
db.add(shared_result)
|
||||
db.commit()
|
||||
db.refresh(shared_result)
|
||||
# Update the original chat with the share_id
|
||||
result = (
|
||||
db.query(Chat)
|
||||
.filter_by(id=chat_id)
|
||||
.update({"share_id": shared_chat.id})
|
||||
)
|
||||
|
||||
return shared_chat if (shared_result and result) else None
|
||||
|
||||
def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
|
||||
try:
|
||||
print("update_shared_chat_by_id")
|
||||
chat = Chat.get(Chat.id == chat_id)
|
||||
print(chat)
|
||||
with get_db() as db:
|
||||
|
||||
query = Chat.update(
|
||||
title=chat.title,
|
||||
chat=chat.chat,
|
||||
).where(Chat.id == chat.share_id)
|
||||
print("update_shared_chat_by_id")
|
||||
chat = db.get(Chat, chat_id)
|
||||
print(chat)
|
||||
chat.title = chat.title
|
||||
chat.chat = chat.chat
|
||||
db.commit()
|
||||
db.refresh(chat)
|
||||
|
||||
query.execute()
|
||||
|
||||
chat = Chat.get(Chat.id == chat.share_id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
return self.get_chat_by_id(chat.share_id)
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool:
|
||||
try:
|
||||
query = Chat.delete().where(Chat.user_id == f"shared-{chat_id}")
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
@ -167,56 +179,50 @@ class ChatTable:
|
||||
self, id: str, share_id: Optional[str]
|
||||
) -> Optional[ChatModel]:
|
||||
try:
|
||||
query = Chat.update(
|
||||
share_id=share_id,
|
||||
).where(Chat.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
chat = Chat.get(Chat.id == id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
chat = db.get(Chat, id)
|
||||
chat.share_id = share_id
|
||||
db.commit()
|
||||
db.refresh(chat)
|
||||
return ChatModel.model_validate(chat)
|
||||
except:
|
||||
return None
|
||||
|
||||
def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
|
||||
try:
|
||||
chat = self.get_chat_by_id(id)
|
||||
query = Chat.update(
|
||||
archived=(not chat.archived),
|
||||
).where(Chat.id == id)
|
||||
with get_db() as db:
|
||||
|
||||
query.execute()
|
||||
|
||||
chat = Chat.get(Chat.id == id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
chat = db.get(Chat, id)
|
||||
chat.archived = not chat.archived
|
||||
db.commit()
|
||||
db.refresh(chat)
|
||||
return ChatModel.model_validate(chat)
|
||||
except:
|
||||
return None
|
||||
|
||||
def archive_all_chats_by_user_id(self, user_id: str) -> bool:
|
||||
try:
|
||||
chats = self.get_chats_by_user_id(user_id)
|
||||
for chat in chats:
|
||||
query = Chat.update(
|
||||
archived=True,
|
||||
).where(Chat.id == chat.id)
|
||||
with get_db() as db:
|
||||
|
||||
query.execute()
|
||||
|
||||
return True
|
||||
db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def get_archived_chat_list_by_user_id(
|
||||
self, user_id: str, skip: int = 0, limit: int = 50
|
||||
) -> List[ChatModel]:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.archived == True)
|
||||
.where(Chat.user_id == user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit)
|
||||
# .offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
all_chats = (
|
||||
db.query(Chat)
|
||||
.filter_by(user_id=user_id, archived=True)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit).offset(skip)
|
||||
.all()
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def get_chat_list_by_user_id(
|
||||
self,
|
||||
@ -225,131 +231,138 @@ class ChatTable:
|
||||
skip: int = 0,
|
||||
limit: int = 50,
|
||||
) -> List[ChatModel]:
|
||||
if include_archived:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.user_id == user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit)
|
||||
# .offset(skip)
|
||||
]
|
||||
else:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.archived == False)
|
||||
.where(Chat.user_id == user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit)
|
||||
# .offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
query = db.query(Chat).filter_by(user_id=user_id)
|
||||
if not include_archived:
|
||||
query = query.filter_by(archived=False)
|
||||
all_chats = (
|
||||
query.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit).offset(skip)
|
||||
.all()
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def get_chat_list_by_chat_ids(
|
||||
self, chat_ids: List[str], skip: int = 0, limit: int = 50
|
||||
) -> List[ChatModel]:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.archived == False)
|
||||
.where(Chat.id.in_(chat_ids))
|
||||
.order_by(Chat.updated_at.desc())
|
||||
]
|
||||
|
||||
with get_db() as db:
|
||||
|
||||
all_chats = (
|
||||
db.query(Chat)
|
||||
.filter(Chat.id.in_(chat_ids))
|
||||
.filter_by(archived=False)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
.all()
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
|
||||
try:
|
||||
chat = Chat.get(Chat.id == id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
with get_db() as db:
|
||||
|
||||
chat = db.get(Chat, id)
|
||||
return ChatModel.model_validate(chat)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
|
||||
try:
|
||||
chat = Chat.get(Chat.share_id == id)
|
||||
with get_db() as db:
|
||||
|
||||
if chat:
|
||||
chat = Chat.get(Chat.id == id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
chat = db.query(Chat).filter_by(share_id=id).first()
|
||||
|
||||
if chat:
|
||||
return self.get_chat_by_id(id)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
|
||||
try:
|
||||
chat = Chat.get(Chat.id == id, Chat.user_id == user_id)
|
||||
return ChatModel(**model_to_dict(chat))
|
||||
with get_db() as db:
|
||||
|
||||
chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
|
||||
return ChatModel.model_validate(chat)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select().order_by(Chat.updated_at.desc())
|
||||
# .limit(limit).offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
all_chats = (
|
||||
db.query(Chat)
|
||||
# .limit(limit).offset(skip)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.user_id == user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
# .limit(limit).offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
all_chats = (
|
||||
db.query(Chat)
|
||||
.filter_by(user_id=user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
|
||||
return [
|
||||
ChatModel(**model_to_dict(chat))
|
||||
for chat in Chat.select()
|
||||
.where(Chat.archived == True)
|
||||
.where(Chat.user_id == user_id)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
all_chats = (
|
||||
db.query(Chat)
|
||||
.filter_by(user_id=user_id, archived=True)
|
||||
.order_by(Chat.updated_at.desc())
|
||||
)
|
||||
return [ChatModel.model_validate(chat) for chat in all_chats]
|
||||
|
||||
def delete_chat_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Chat.delete().where((Chat.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True and self.delete_shared_chat_by_chat_id(id)
|
||||
db.query(Chat).filter_by(id=id).delete()
|
||||
|
||||
return True and self.delete_shared_chat_by_chat_id(id)
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
|
||||
try:
|
||||
query = Chat.delete().where((Chat.id == id) & (Chat.user_id == user_id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True and self.delete_shared_chat_by_chat_id(id)
|
||||
db.query(Chat).filter_by(id=id, user_id=user_id).delete()
|
||||
|
||||
return True and self.delete_shared_chat_by_chat_id(id)
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_chats_by_user_id(self, user_id: str) -> bool:
|
||||
try:
|
||||
|
||||
self.delete_shared_chats_by_user_id(user_id)
|
||||
with get_db() as db:
|
||||
|
||||
query = Chat.delete().where(Chat.user_id == user_id)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
self.delete_shared_chats_by_user_id(user_id)
|
||||
|
||||
return True
|
||||
db.query(Chat).filter_by(user_id=user_id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
|
||||
try:
|
||||
shared_chat_ids = [
|
||||
f"shared-{chat.id}"
|
||||
for chat in Chat.select().where(Chat.user_id == user_id)
|
||||
]
|
||||
|
||||
query = Chat.delete().where(Chat.user_id << shared_chat_ids)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
|
||||
shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
|
||||
|
||||
db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Chats = ChatTable(DB)
|
||||
Chats = ChatTable()
|
||||
|
@ -1,14 +1,11 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from typing import List, Union, Optional
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Optional
|
||||
import time
|
||||
import logging
|
||||
|
||||
from utils.utils import decode_token
|
||||
from utils.misc import get_gravatar_url
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
import json
|
||||
|
||||
@ -22,20 +19,21 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class Document(Model):
|
||||
collection_name = CharField(unique=True)
|
||||
name = CharField(unique=True)
|
||||
title = TextField()
|
||||
filename = TextField()
|
||||
content = TextField(null=True)
|
||||
user_id = CharField()
|
||||
timestamp = BigIntegerField()
|
||||
class Document(Base):
|
||||
__tablename__ = "document"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
collection_name = Column(String, primary_key=True)
|
||||
name = Column(String, unique=True)
|
||||
title = Column(Text)
|
||||
filename = Column(Text)
|
||||
content = Column(Text, nullable=True)
|
||||
user_id = Column(String)
|
||||
timestamp = Column(BigInteger)
|
||||
|
||||
|
||||
class DocumentModel(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
collection_name: str
|
||||
name: str
|
||||
title: str
|
||||
@ -72,57 +70,63 @@ class DocumentForm(DocumentUpdateForm):
|
||||
|
||||
|
||||
class DocumentsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Document])
|
||||
|
||||
def insert_new_doc(
|
||||
self, user_id: str, form_data: DocumentForm
|
||||
) -> Optional[DocumentModel]:
|
||||
document = DocumentModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"user_id": user_id,
|
||||
"timestamp": int(time.time()),
|
||||
}
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
result = Document.create(**document.model_dump())
|
||||
if result:
|
||||
return document
|
||||
else:
|
||||
document = DocumentModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"user_id": user_id,
|
||||
"timestamp": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Document(**document.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return DocumentModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_doc_by_name(self, name: str) -> Optional[DocumentModel]:
|
||||
try:
|
||||
document = Document.get(Document.name == name)
|
||||
return DocumentModel(**model_to_dict(document))
|
||||
with get_db() as db:
|
||||
|
||||
document = db.query(Document).filter_by(name=name).first()
|
||||
return DocumentModel.model_validate(document) if document else None
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_docs(self) -> List[DocumentModel]:
|
||||
return [
|
||||
DocumentModel(**model_to_dict(doc))
|
||||
for doc in Document.select()
|
||||
# .limit(limit).offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
return [
|
||||
DocumentModel.model_validate(doc) for doc in db.query(Document).all()
|
||||
]
|
||||
|
||||
def update_doc_by_name(
|
||||
self, name: str, form_data: DocumentUpdateForm
|
||||
) -> Optional[DocumentModel]:
|
||||
try:
|
||||
query = Document.update(
|
||||
title=form_data.title,
|
||||
name=form_data.name,
|
||||
timestamp=int(time.time()),
|
||||
).where(Document.name == name)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
doc = Document.get(Document.name == form_data.name)
|
||||
return DocumentModel(**model_to_dict(doc))
|
||||
db.query(Document).filter_by(name=name).update(
|
||||
{
|
||||
"title": form_data.title,
|
||||
"name": form_data.name,
|
||||
"timestamp": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
return self.get_doc_by_name(form_data.name)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
return None
|
||||
@ -135,26 +139,28 @@ class DocumentsTable:
|
||||
doc_content = json.loads(doc.content if doc.content else "{}")
|
||||
doc_content = {**doc_content, **updated}
|
||||
|
||||
query = Document.update(
|
||||
content=json.dumps(doc_content),
|
||||
timestamp=int(time.time()),
|
||||
).where(Document.name == name)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
doc = Document.get(Document.name == name)
|
||||
return DocumentModel(**model_to_dict(doc))
|
||||
db.query(Document).filter_by(name=name).update(
|
||||
{
|
||||
"content": json.dumps(doc_content),
|
||||
"timestamp": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
return self.get_doc_by_name(name)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
return None
|
||||
|
||||
def delete_doc_by_name(self, name: str) -> bool:
|
||||
try:
|
||||
query = Document.delete().where((Document.name == name))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
db.query(Document).filter_by(name=name).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Documents = DocumentsTable(DB)
|
||||
Documents = DocumentsTable()
|
||||
|
@ -1,10 +1,11 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
import logging
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
|
||||
from sqlalchemy import Column, String, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import JSONField, Base, get_db
|
||||
|
||||
import json
|
||||
|
||||
@ -18,15 +19,14 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class File(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
filename = TextField()
|
||||
meta = JSONField()
|
||||
created_at = BigIntegerField()
|
||||
class File(Base):
|
||||
__tablename__ = "file"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
filename = Column(Text)
|
||||
meta = Column(JSONField)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
|
||||
class FileModel(BaseModel):
|
||||
@ -36,6 +36,8 @@ class FileModel(BaseModel):
|
||||
meta: dict
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -57,56 +59,64 @@ class FileForm(BaseModel):
|
||||
|
||||
|
||||
class FilesTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([File])
|
||||
|
||||
def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
|
||||
file = FileModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"user_id": user_id,
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
result = File.create(**file.model_dump())
|
||||
if result:
|
||||
return file
|
||||
else:
|
||||
file = FileModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"user_id": user_id,
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = File(**file.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return FileModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
|
||||
def get_file_by_id(self, id: str) -> Optional[FileModel]:
|
||||
try:
|
||||
file = File.get(File.id == id)
|
||||
return FileModel(**model_to_dict(file))
|
||||
except:
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
file = db.get(File, id)
|
||||
return FileModel.model_validate(file)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_files(self) -> List[FileModel]:
|
||||
return [FileModel(**model_to_dict(file)) for file in File.select()]
|
||||
with get_db() as db:
|
||||
|
||||
return [FileModel.model_validate(file) for file in db.query(File).all()]
|
||||
|
||||
def delete_file_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = File.delete().where((File.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
db.query(File).filter_by(id=id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_all_files(self) -> bool:
|
||||
try:
|
||||
query = File.delete()
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
db.query(File).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Files = FilesTable(DB)
|
||||
Files = FilesTable()
|
||||
|
@ -1,10 +1,11 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
import logging
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
|
||||
from sqlalchemy import Column, String, Text, BigInteger, Boolean
|
||||
|
||||
from apps.webui.internal.db import JSONField, Base, get_db
|
||||
from apps.webui.models.users import Users
|
||||
|
||||
import json
|
||||
@ -21,21 +22,20 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class Function(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
name = TextField()
|
||||
type = TextField()
|
||||
content = TextField()
|
||||
meta = JSONField()
|
||||
valves = JSONField()
|
||||
is_active = BooleanField(default=False)
|
||||
is_global = BooleanField(default=False)
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
class Function(Base):
|
||||
__tablename__ = "function"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
name = Column(Text)
|
||||
type = Column(Text)
|
||||
content = Column(Text)
|
||||
meta = Column(JSONField)
|
||||
valves = Column(JSONField)
|
||||
is_active = Column(Boolean)
|
||||
is_global = Column(Boolean)
|
||||
updated_at = Column(BigInteger)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
|
||||
class FunctionMeta(BaseModel):
|
||||
@ -55,6 +55,8 @@ class FunctionModel(BaseModel):
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -85,13 +87,11 @@ class FunctionValves(BaseModel):
|
||||
|
||||
|
||||
class FunctionsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Function])
|
||||
|
||||
def insert_new_function(
|
||||
self, user_id: str, type: str, form_data: FunctionForm
|
||||
) -> Optional[FunctionModel]:
|
||||
|
||||
function = FunctionModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
@ -103,86 +103,99 @@ class FunctionsTable:
|
||||
)
|
||||
|
||||
try:
|
||||
result = Function.create(**function.model_dump())
|
||||
if result:
|
||||
return function
|
||||
else:
|
||||
return None
|
||||
with get_db() as db:
|
||||
result = Function(**function.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return FunctionModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
|
||||
def get_function_by_id(self, id: str) -> Optional[FunctionModel]:
|
||||
try:
|
||||
function = Function.get(Function.id == id)
|
||||
return FunctionModel(**model_to_dict(function))
|
||||
with get_db() as db:
|
||||
|
||||
function = db.get(Function, id)
|
||||
return FunctionModel.model_validate(function)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_functions(self, active_only=False) -> List[FunctionModel]:
|
||||
if active_only:
|
||||
return [
|
||||
FunctionModel(**model_to_dict(function))
|
||||
for function in Function.select().where(Function.is_active == True)
|
||||
]
|
||||
else:
|
||||
return [
|
||||
FunctionModel(**model_to_dict(function))
|
||||
for function in Function.select()
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
if active_only:
|
||||
return [
|
||||
FunctionModel.model_validate(function)
|
||||
for function in db.query(Function).filter_by(is_active=True).all()
|
||||
]
|
||||
else:
|
||||
return [
|
||||
FunctionModel.model_validate(function)
|
||||
for function in db.query(Function).all()
|
||||
]
|
||||
|
||||
def get_functions_by_type(
|
||||
self, type: str, active_only=False
|
||||
) -> List[FunctionModel]:
|
||||
if active_only:
|
||||
return [
|
||||
FunctionModel(**model_to_dict(function))
|
||||
for function in Function.select().where(
|
||||
Function.type == type, Function.is_active == True
|
||||
)
|
||||
]
|
||||
else:
|
||||
return [
|
||||
FunctionModel(**model_to_dict(function))
|
||||
for function in Function.select().where(Function.type == type)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
if active_only:
|
||||
return [
|
||||
FunctionModel.model_validate(function)
|
||||
for function in db.query(Function)
|
||||
.filter_by(type=type, is_active=True)
|
||||
.all()
|
||||
]
|
||||
else:
|
||||
return [
|
||||
FunctionModel.model_validate(function)
|
||||
for function in db.query(Function).filter_by(type=type).all()
|
||||
]
|
||||
|
||||
def get_global_filter_functions(self) -> List[FunctionModel]:
|
||||
return [
|
||||
FunctionModel(**model_to_dict(function))
|
||||
for function in Function.select().where(
|
||||
Function.type == "filter",
|
||||
Function.is_active == True,
|
||||
Function.is_global == True,
|
||||
)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
return [
|
||||
FunctionModel.model_validate(function)
|
||||
for function in db.query(Function)
|
||||
.filter_by(type="filter", is_active=True, is_global=True)
|
||||
.all()
|
||||
]
|
||||
|
||||
def get_function_valves_by_id(self, id: str) -> Optional[dict]:
|
||||
try:
|
||||
function = Function.get(Function.id == id)
|
||||
return function.valves if function.valves else {}
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
function = db.get(Function, id)
|
||||
return function.valves if function.valves else {}
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
return None
|
||||
|
||||
def update_function_valves_by_id(
|
||||
self, id: str, valves: dict
|
||||
) -> Optional[FunctionValves]:
|
||||
try:
|
||||
query = Function.update(
|
||||
**{"valves": valves},
|
||||
updated_at=int(time.time()),
|
||||
).where(Function.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
function = Function.get(Function.id == id)
|
||||
return FunctionValves(**model_to_dict(function))
|
||||
except:
|
||||
return None
|
||||
try:
|
||||
function = db.get(Function, id)
|
||||
function.valves = valves
|
||||
function.updated_at = int(time.time())
|
||||
db.commit()
|
||||
db.refresh(function)
|
||||
return self.get_function_by_id(id)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_user_valves_by_id_and_user_id(
|
||||
self, id: str, user_id: str
|
||||
) -> Optional[dict]:
|
||||
|
||||
try:
|
||||
user = Users.get_user_by_id(user_id)
|
||||
user_settings = user.settings.model_dump() if user.settings else {}
|
||||
@ -201,6 +214,7 @@ class FunctionsTable:
|
||||
def update_user_valves_by_id_and_user_id(
|
||||
self, id: str, user_id: str, valves: dict
|
||||
) -> Optional[dict]:
|
||||
|
||||
try:
|
||||
user = Users.get_user_by_id(user_id)
|
||||
user_settings = user.settings.model_dump() if user.settings else {}
|
||||
@ -222,39 +236,43 @@ class FunctionsTable:
|
||||
return None
|
||||
|
||||
def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]:
|
||||
try:
|
||||
query = Function.update(
|
||||
**updated,
|
||||
updated_at=int(time.time()),
|
||||
).where(Function.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
function = Function.get(Function.id == id)
|
||||
return FunctionModel(**model_to_dict(function))
|
||||
except:
|
||||
return None
|
||||
try:
|
||||
db.query(Function).filter_by(id=id).update(
|
||||
{
|
||||
**updated,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
return self.get_function_by_id(id)
|
||||
except:
|
||||
return None
|
||||
|
||||
def deactivate_all_functions(self) -> Optional[bool]:
|
||||
try:
|
||||
query = Function.update(
|
||||
**{"is_active": False},
|
||||
updated_at=int(time.time()),
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
query.execute()
|
||||
|
||||
return True
|
||||
except:
|
||||
return None
|
||||
try:
|
||||
db.query(Function).update(
|
||||
{
|
||||
"is_active": False,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
return True
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_function_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Function.delete().where((Function.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
try:
|
||||
db.query(Function).filter_by(id=id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Functions = FunctionsTable(DB)
|
||||
Functions = FunctionsTable()
|
||||
|
@ -1,10 +1,9 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Union, Optional
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from apps.webui.models.chats import Chats
|
||||
from sqlalchemy import Column, String, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
import time
|
||||
import uuid
|
||||
@ -14,15 +13,14 @@ import uuid
|
||||
####################
|
||||
|
||||
|
||||
class Memory(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
content = TextField()
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
class Memory(Base):
|
||||
__tablename__ = "memory"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
content = Column(Text)
|
||||
updated_at = Column(BigInteger)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
|
||||
class MemoryModel(BaseModel):
|
||||
@ -32,6 +30,8 @@ class MemoryModel(BaseModel):
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -39,94 +39,104 @@ class MemoryModel(BaseModel):
|
||||
|
||||
|
||||
class MemoriesTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Memory])
|
||||
|
||||
def insert_new_memory(
|
||||
self,
|
||||
user_id: str,
|
||||
content: str,
|
||||
) -> Optional[MemoryModel]:
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
memory = MemoryModel(
|
||||
**{
|
||||
"id": id,
|
||||
"user_id": user_id,
|
||||
"content": content,
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
result = Memory.create(**memory.model_dump())
|
||||
if result:
|
||||
return memory
|
||||
else:
|
||||
return None
|
||||
with get_db() as db:
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
memory = MemoryModel(
|
||||
**{
|
||||
"id": id,
|
||||
"user_id": user_id,
|
||||
"content": content,
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
result = Memory(**memory.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return MemoryModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
|
||||
def update_memory_by_id(
|
||||
self,
|
||||
id: str,
|
||||
content: str,
|
||||
) -> Optional[MemoryModel]:
|
||||
try:
|
||||
memory = Memory.get(Memory.id == id)
|
||||
memory.content = content
|
||||
memory.updated_at = int(time.time())
|
||||
memory.save()
|
||||
return MemoryModel(**model_to_dict(memory))
|
||||
except:
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
db.query(Memory).filter_by(id=id).update(
|
||||
{"content": content, "updated_at": int(time.time())}
|
||||
)
|
||||
db.commit()
|
||||
return self.get_memory_by_id(id)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_memories(self) -> List[MemoryModel]:
|
||||
try:
|
||||
memories = Memory.select()
|
||||
return [MemoryModel(**model_to_dict(memory)) for memory in memories]
|
||||
except:
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
memories = db.query(Memory).all()
|
||||
return [MemoryModel.model_validate(memory) for memory in memories]
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_memories_by_user_id(self, user_id: str) -> List[MemoryModel]:
|
||||
try:
|
||||
memories = Memory.select().where(Memory.user_id == user_id)
|
||||
return [MemoryModel(**model_to_dict(memory)) for memory in memories]
|
||||
except:
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
def get_memory_by_id(self, id) -> Optional[MemoryModel]:
|
||||
try:
|
||||
memory = Memory.get(Memory.id == id)
|
||||
return MemoryModel(**model_to_dict(memory))
|
||||
except:
|
||||
return None
|
||||
try:
|
||||
memories = db.query(Memory).filter_by(user_id=user_id).all()
|
||||
return [MemoryModel.model_validate(memory) for memory in memories]
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_memory_by_id(self, id: str) -> Optional[MemoryModel]:
|
||||
with get_db() as db:
|
||||
|
||||
try:
|
||||
memory = db.get(Memory, id)
|
||||
return MemoryModel.model_validate(memory)
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_memory_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Memory.delete().where(Memory.id == id)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
try:
|
||||
db.query(Memory).filter_by(id=id).delete()
|
||||
return True
|
||||
|
||||
except:
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_memories_by_user_id(self, user_id: str) -> bool:
|
||||
try:
|
||||
query = Memory.delete().where(Memory.user_id == user_id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
try:
|
||||
db.query(Memory).filter_by(user_id=user_id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool:
|
||||
try:
|
||||
query = Memory.delete().where(Memory.id == id, Memory.user_id == user_id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
try:
|
||||
db.query(Memory).filter_by(id=id, user_id=user_id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Memories = MemoriesTable(DB)
|
||||
Memories = MemoriesTable()
|
||||
|
@ -2,13 +2,10 @@ import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import peewee as pw
|
||||
from peewee import *
|
||||
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
from apps.webui.internal.db import Base, JSONField, get_db
|
||||
|
||||
from typing import List, Union, Optional
|
||||
from config import SRC_LOG_LEVELS
|
||||
@ -46,38 +43,37 @@ class ModelMeta(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class Model(pw.Model):
|
||||
id = pw.TextField(unique=True)
|
||||
class Model(Base):
|
||||
__tablename__ = "model"
|
||||
|
||||
id = Column(Text, primary_key=True)
|
||||
"""
|
||||
The model's id as used in the API. If set to an existing model, it will override the model.
|
||||
"""
|
||||
user_id = pw.TextField()
|
||||
user_id = Column(Text)
|
||||
|
||||
base_model_id = pw.TextField(null=True)
|
||||
base_model_id = Column(Text, nullable=True)
|
||||
"""
|
||||
An optional pointer to the actual model that should be used when proxying requests.
|
||||
"""
|
||||
|
||||
name = pw.TextField()
|
||||
name = Column(Text)
|
||||
"""
|
||||
The human-readable display name of the model.
|
||||
"""
|
||||
|
||||
params = JSONField()
|
||||
params = Column(JSONField)
|
||||
"""
|
||||
Holds a JSON encoded blob of parameters, see `ModelParams`.
|
||||
"""
|
||||
|
||||
meta = JSONField()
|
||||
meta = Column(JSONField)
|
||||
"""
|
||||
Holds a JSON encoded blob of metadata, see `ModelMeta`.
|
||||
"""
|
||||
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
updated_at = Column(BigInteger)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
|
||||
class ModelModel(BaseModel):
|
||||
@ -92,6 +88,8 @@ class ModelModel(BaseModel):
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -115,12 +113,6 @@ class ModelForm(BaseModel):
|
||||
|
||||
|
||||
class ModelsTable:
|
||||
def __init__(
|
||||
self,
|
||||
db: pw.SqliteDatabase | pw.PostgresqlDatabase,
|
||||
):
|
||||
self.db = db
|
||||
self.db.create_tables([Model])
|
||||
|
||||
def insert_new_model(
|
||||
self, form_data: ModelForm, user_id: str
|
||||
@ -134,34 +126,46 @@ class ModelsTable:
|
||||
}
|
||||
)
|
||||
try:
|
||||
result = Model.create(**model.model_dump())
|
||||
|
||||
if result:
|
||||
return model
|
||||
else:
|
||||
return None
|
||||
with get_db() as db:
|
||||
|
||||
result = Model(**model.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
|
||||
if result:
|
||||
return ModelModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return None
|
||||
|
||||
def get_all_models(self) -> List[ModelModel]:
|
||||
return [ModelModel(**model_to_dict(model)) for model in Model.select()]
|
||||
with get_db() as db:
|
||||
|
||||
return [ModelModel.model_validate(model) for model in db.query(Model).all()]
|
||||
|
||||
def get_model_by_id(self, id: str) -> Optional[ModelModel]:
|
||||
try:
|
||||
model = Model.get(Model.id == id)
|
||||
return ModelModel(**model_to_dict(model))
|
||||
with get_db() as db:
|
||||
|
||||
model = db.get(Model, id)
|
||||
return ModelModel.model_validate(model)
|
||||
except:
|
||||
return None
|
||||
|
||||
def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
|
||||
try:
|
||||
# update only the fields that are present in the model
|
||||
query = Model.update(**model.model_dump()).where(Model.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
model = Model.get(Model.id == id)
|
||||
return ModelModel(**model_to_dict(model))
|
||||
# update only the fields that are present in the model
|
||||
model = db.query(Model).get(id)
|
||||
model.update(**model.model_dump())
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
return ModelModel.model_validate(model)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
@ -169,11 +173,12 @@ class ModelsTable:
|
||||
|
||||
def delete_model_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Model.delete().where(Model.id == id)
|
||||
query.execute()
|
||||
return True
|
||||
with get_db() as db:
|
||||
|
||||
db.query(Model).filter_by(id=id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Models = ModelsTable(DB)
|
||||
Models = ModelsTable()
|
||||
|
@ -1,13 +1,10 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from typing import List, Union, Optional
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Optional
|
||||
import time
|
||||
|
||||
from utils.utils import decode_token
|
||||
from utils.misc import get_gravatar_url
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
import json
|
||||
|
||||
@ -16,15 +13,14 @@ import json
|
||||
####################
|
||||
|
||||
|
||||
class Prompt(Model):
|
||||
command = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
title = TextField()
|
||||
content = TextField()
|
||||
timestamp = BigIntegerField()
|
||||
class Prompt(Base):
|
||||
__tablename__ = "prompt"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
command = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
title = Column(Text)
|
||||
content = Column(Text)
|
||||
timestamp = Column(BigInteger)
|
||||
|
||||
|
||||
class PromptModel(BaseModel):
|
||||
@ -34,6 +30,8 @@ class PromptModel(BaseModel):
|
||||
content: str
|
||||
timestamp: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -48,10 +46,6 @@ class PromptForm(BaseModel):
|
||||
|
||||
class PromptsTable:
|
||||
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Prompt])
|
||||
|
||||
def insert_new_prompt(
|
||||
self, user_id: str, form_data: PromptForm
|
||||
) -> Optional[PromptModel]:
|
||||
@ -66,53 +60,58 @@ class PromptsTable:
|
||||
)
|
||||
|
||||
try:
|
||||
result = Prompt.create(**prompt.model_dump())
|
||||
if result:
|
||||
return prompt
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
with get_db() as db:
|
||||
|
||||
result = Prompt(**prompt.dict())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return PromptModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def get_prompt_by_command(self, command: str) -> Optional[PromptModel]:
|
||||
try:
|
||||
prompt = Prompt.get(Prompt.command == command)
|
||||
return PromptModel(**model_to_dict(prompt))
|
||||
with get_db() as db:
|
||||
|
||||
prompt = db.query(Prompt).filter_by(command=command).first()
|
||||
return PromptModel.model_validate(prompt)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_prompts(self) -> List[PromptModel]:
|
||||
return [
|
||||
PromptModel(**model_to_dict(prompt))
|
||||
for prompt in Prompt.select()
|
||||
# .limit(limit).offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
return [
|
||||
PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()
|
||||
]
|
||||
|
||||
def update_prompt_by_command(
|
||||
self, command: str, form_data: PromptForm
|
||||
) -> Optional[PromptModel]:
|
||||
try:
|
||||
query = Prompt.update(
|
||||
title=form_data.title,
|
||||
content=form_data.content,
|
||||
timestamp=int(time.time()),
|
||||
).where(Prompt.command == command)
|
||||
with get_db() as db:
|
||||
|
||||
query.execute()
|
||||
|
||||
prompt = Prompt.get(Prompt.command == command)
|
||||
return PromptModel(**model_to_dict(prompt))
|
||||
prompt = db.query(Prompt).filter_by(command=command).first()
|
||||
prompt.title = form_data.title
|
||||
prompt.content = form_data.content
|
||||
prompt.timestamp = int(time.time())
|
||||
db.commit()
|
||||
return PromptModel.model_validate(prompt)
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_prompt_by_command(self, command: str) -> bool:
|
||||
try:
|
||||
query = Prompt.delete().where((Prompt.command == command))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
|
||||
return True
|
||||
db.query(Prompt).filter_by(command=command).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Prompts = PromptsTable(DB)
|
||||
Prompts = PromptsTable()
|
||||
|
@ -1,14 +1,14 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Union, Optional
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Optional
|
||||
|
||||
import json
|
||||
import uuid
|
||||
import time
|
||||
import logging
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import Base, get_db
|
||||
|
||||
from config import SRC_LOG_LEVELS
|
||||
|
||||
@ -20,25 +20,23 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class Tag(Model):
|
||||
id = CharField(unique=True)
|
||||
name = CharField()
|
||||
user_id = CharField()
|
||||
data = TextField(null=True)
|
||||
class Tag(Base):
|
||||
__tablename__ = "tag"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
name = Column(String)
|
||||
user_id = Column(String)
|
||||
data = Column(Text, nullable=True)
|
||||
|
||||
|
||||
class ChatIdTag(Model):
|
||||
id = CharField(unique=True)
|
||||
tag_name = CharField()
|
||||
chat_id = CharField()
|
||||
user_id = CharField()
|
||||
timestamp = BigIntegerField()
|
||||
class ChatIdTag(Base):
|
||||
__tablename__ = "chatidtag"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
tag_name = Column(String)
|
||||
chat_id = Column(String)
|
||||
user_id = Column(String)
|
||||
timestamp = Column(BigInteger)
|
||||
|
||||
|
||||
class TagModel(BaseModel):
|
||||
@ -47,6 +45,8 @@ class TagModel(BaseModel):
|
||||
user_id: str
|
||||
data: Optional[str] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class ChatIdTagModel(BaseModel):
|
||||
id: str
|
||||
@ -55,6 +55,8 @@ class ChatIdTagModel(BaseModel):
|
||||
user_id: str
|
||||
timestamp: int
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -75,28 +77,31 @@ class ChatTagsResponse(BaseModel):
|
||||
|
||||
|
||||
class TagTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
db.create_tables([Tag, ChatIdTag])
|
||||
|
||||
def insert_new_tag(self, name: str, user_id: str) -> Optional[TagModel]:
|
||||
id = str(uuid.uuid4())
|
||||
tag = TagModel(**{"id": id, "user_id": user_id, "name": name})
|
||||
try:
|
||||
result = Tag.create(**tag.model_dump())
|
||||
if result:
|
||||
return tag
|
||||
else:
|
||||
with get_db() as db:
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
tag = TagModel(**{"id": id, "user_id": user_id, "name": name})
|
||||
try:
|
||||
result = Tag(**tag.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return TagModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
return None
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def get_tag_by_name_and_user_id(
|
||||
self, name: str, user_id: str
|
||||
) -> Optional[TagModel]:
|
||||
try:
|
||||
tag = Tag.get(Tag.name == name, Tag.user_id == user_id)
|
||||
return TagModel(**model_to_dict(tag))
|
||||
with get_db() as db:
|
||||
tag = db.query(Tag).filter(name=name, user_id=user_id).first()
|
||||
return TagModel.model_validate(tag)
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
@ -118,82 +123,109 @@ class TagTable:
|
||||
}
|
||||
)
|
||||
try:
|
||||
result = ChatIdTag.create(**chatIdTag.model_dump())
|
||||
if result:
|
||||
return chatIdTag
|
||||
else:
|
||||
return None
|
||||
with get_db() as db:
|
||||
result = ChatIdTag(**chatIdTag.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return ChatIdTagModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_tags_by_user_id(self, user_id: str) -> List[TagModel]:
|
||||
tag_names = [
|
||||
ChatIdTagModel(**model_to_dict(chat_id_tag)).tag_name
|
||||
for chat_id_tag in ChatIdTag.select()
|
||||
.where(ChatIdTag.user_id == user_id)
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
]
|
||||
with get_db() as db:
|
||||
tag_names = [
|
||||
chat_id_tag.tag_name
|
||||
for chat_id_tag in (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(user_id=user_id)
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
.all()
|
||||
)
|
||||
]
|
||||
|
||||
return [
|
||||
TagModel(**model_to_dict(tag))
|
||||
for tag in Tag.select()
|
||||
.where(Tag.user_id == user_id)
|
||||
.where(Tag.name.in_(tag_names))
|
||||
]
|
||||
return [
|
||||
TagModel.model_validate(tag)
|
||||
for tag in (
|
||||
db.query(Tag)
|
||||
.filter_by(user_id=user_id)
|
||||
.filter(Tag.name.in_(tag_names))
|
||||
.all()
|
||||
)
|
||||
]
|
||||
|
||||
def get_tags_by_chat_id_and_user_id(
|
||||
self, chat_id: str, user_id: str
|
||||
) -> List[TagModel]:
|
||||
tag_names = [
|
||||
ChatIdTagModel(**model_to_dict(chat_id_tag)).tag_name
|
||||
for chat_id_tag in ChatIdTag.select()
|
||||
.where((ChatIdTag.user_id == user_id) & (ChatIdTag.chat_id == chat_id))
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
]
|
||||
with get_db() as db:
|
||||
|
||||
return [
|
||||
TagModel(**model_to_dict(tag))
|
||||
for tag in Tag.select()
|
||||
.where(Tag.user_id == user_id)
|
||||
.where(Tag.name.in_(tag_names))
|
||||
]
|
||||
tag_names = [
|
||||
chat_id_tag.tag_name
|
||||
for chat_id_tag in (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(user_id=user_id, chat_id=chat_id)
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
.all()
|
||||
)
|
||||
]
|
||||
|
||||
return [
|
||||
TagModel.model_validate(tag)
|
||||
for tag in (
|
||||
db.query(Tag)
|
||||
.filter_by(user_id=user_id)
|
||||
.filter(Tag.name.in_(tag_names))
|
||||
.all()
|
||||
)
|
||||
]
|
||||
|
||||
def get_chat_ids_by_tag_name_and_user_id(
|
||||
self, tag_name: str, user_id: str
|
||||
) -> Optional[ChatIdTagModel]:
|
||||
return [
|
||||
ChatIdTagModel(**model_to_dict(chat_id_tag))
|
||||
for chat_id_tag in ChatIdTag.select()
|
||||
.where((ChatIdTag.user_id == user_id) & (ChatIdTag.tag_name == tag_name))
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
]
|
||||
) -> List[ChatIdTagModel]:
|
||||
with get_db() as db:
|
||||
|
||||
return [
|
||||
ChatIdTagModel.model_validate(chat_id_tag)
|
||||
for chat_id_tag in (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(user_id=user_id, tag_name=tag_name)
|
||||
.order_by(ChatIdTag.timestamp.desc())
|
||||
.all()
|
||||
)
|
||||
]
|
||||
|
||||
def count_chat_ids_by_tag_name_and_user_id(
|
||||
self, tag_name: str, user_id: str
|
||||
) -> int:
|
||||
return (
|
||||
ChatIdTag.select()
|
||||
.where((ChatIdTag.tag_name == tag_name) & (ChatIdTag.user_id == user_id))
|
||||
.count()
|
||||
)
|
||||
with get_db() as db:
|
||||
|
||||
return (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(tag_name=tag_name, user_id=user_id)
|
||||
.count()
|
||||
)
|
||||
|
||||
def delete_tag_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> bool:
|
||||
try:
|
||||
query = ChatIdTag.delete().where(
|
||||
(ChatIdTag.tag_name == tag_name) & (ChatIdTag.user_id == user_id)
|
||||
)
|
||||
res = query.execute() # Remove the rows, return number of rows removed.
|
||||
log.debug(f"res: {res}")
|
||||
|
||||
tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id)
|
||||
if tag_count == 0:
|
||||
# Remove tag item from Tag col as well
|
||||
query = Tag.delete().where(
|
||||
(Tag.name == tag_name) & (Tag.user_id == user_id)
|
||||
with get_db() as db:
|
||||
res = (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(tag_name=tag_name, user_id=user_id)
|
||||
.delete()
|
||||
)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
log.debug(f"res: {res}")
|
||||
db.commit()
|
||||
|
||||
return True
|
||||
tag_count = self.count_chat_ids_by_tag_name_and_user_id(
|
||||
tag_name, user_id
|
||||
)
|
||||
if tag_count == 0:
|
||||
# Remove tag item from Tag col as well
|
||||
db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete()
|
||||
return True
|
||||
except Exception as e:
|
||||
log.error(f"delete_tag: {e}")
|
||||
return False
|
||||
@ -202,23 +234,24 @@ class TagTable:
|
||||
self, tag_name: str, chat_id: str, user_id: str
|
||||
) -> bool:
|
||||
try:
|
||||
query = ChatIdTag.delete().where(
|
||||
(ChatIdTag.tag_name == tag_name)
|
||||
& (ChatIdTag.chat_id == chat_id)
|
||||
& (ChatIdTag.user_id == user_id)
|
||||
)
|
||||
res = query.execute() # Remove the rows, return number of rows removed.
|
||||
log.debug(f"res: {res}")
|
||||
with get_db() as db:
|
||||
|
||||
tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id)
|
||||
if tag_count == 0:
|
||||
# Remove tag item from Tag col as well
|
||||
query = Tag.delete().where(
|
||||
(Tag.name == tag_name) & (Tag.user_id == user_id)
|
||||
res = (
|
||||
db.query(ChatIdTag)
|
||||
.filter_by(tag_name=tag_name, chat_id=chat_id, user_id=user_id)
|
||||
.delete()
|
||||
)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
log.debug(f"res: {res}")
|
||||
db.commit()
|
||||
|
||||
return True
|
||||
tag_count = self.count_chat_ids_by_tag_name_and_user_id(
|
||||
tag_name, user_id
|
||||
)
|
||||
if tag_count == 0:
|
||||
# Remove tag item from Tag col as well
|
||||
db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete()
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
log.error(f"delete_tag: {e}")
|
||||
return False
|
||||
@ -234,4 +267,4 @@ class TagTable:
|
||||
return True
|
||||
|
||||
|
||||
Tags = TagTable(DB)
|
||||
Tags = TagTable()
|
||||
|
@ -1,10 +1,10 @@
|
||||
from pydantic import BaseModel
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from typing import List, Union, Optional
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing import List, Optional
|
||||
import time
|
||||
import logging
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from apps.webui.internal.db import Base, JSONField, get_db
|
||||
from apps.webui.models.users import Users
|
||||
|
||||
import json
|
||||
@ -21,19 +21,18 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
|
||||
####################
|
||||
|
||||
|
||||
class Tool(Model):
|
||||
id = CharField(unique=True)
|
||||
user_id = CharField()
|
||||
name = TextField()
|
||||
content = TextField()
|
||||
specs = JSONField()
|
||||
meta = JSONField()
|
||||
valves = JSONField()
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
class Tool(Base):
|
||||
__tablename__ = "tool"
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
id = Column(String, primary_key=True)
|
||||
user_id = Column(String)
|
||||
name = Column(Text)
|
||||
content = Column(Text)
|
||||
specs = Column(JSONField)
|
||||
meta = Column(JSONField)
|
||||
valves = Column(JSONField)
|
||||
updated_at = Column(BigInteger)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
|
||||
class ToolMeta(BaseModel):
|
||||
@ -51,6 +50,8 @@ class ToolModel(BaseModel):
|
||||
updated_at: int # timestamp in epoch
|
||||
created_at: int # timestamp in epoch
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -78,61 +79,68 @@ class ToolValves(BaseModel):
|
||||
|
||||
|
||||
class ToolsTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([Tool])
|
||||
|
||||
def insert_new_tool(
|
||||
self, user_id: str, form_data: ToolForm, specs: List[dict]
|
||||
) -> Optional[ToolModel]:
|
||||
tool = ToolModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"specs": specs,
|
||||
"user_id": user_id,
|
||||
"updated_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Tool.create(**tool.model_dump())
|
||||
if result:
|
||||
return tool
|
||||
else:
|
||||
with get_db() as db:
|
||||
|
||||
tool = ToolModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
"specs": specs,
|
||||
"user_id": user_id,
|
||||
"updated_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
result = Tool(**tool.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return ToolModel.model_validate(result)
|
||||
else:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error creating tool: {e}")
|
||||
return None
|
||||
|
||||
def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
|
||||
try:
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
with get_db() as db:
|
||||
|
||||
tool = db.get(Tool, id)
|
||||
return ToolModel.model_validate(tool)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_tools(self) -> List[ToolModel]:
|
||||
return [ToolModel(**model_to_dict(tool)) for tool in Tool.select()]
|
||||
with get_db() as db:
|
||||
return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
|
||||
|
||||
def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
|
||||
try:
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return tool.valves if tool.valves else {}
|
||||
with get_db() as db:
|
||||
|
||||
tool = db.get(Tool, id)
|
||||
return tool.valves if tool.valves else {}
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
return None
|
||||
|
||||
def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
|
||||
try:
|
||||
query = Tool.update(
|
||||
**{"valves": valves},
|
||||
updated_at=int(time.time()),
|
||||
).where(Tool.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolValves(**model_to_dict(tool))
|
||||
db.query(Tool).filter_by(id=id).update(
|
||||
{"valves": valves, "updated_at": int(time.time())}
|
||||
)
|
||||
db.commit()
|
||||
return self.get_tool_by_id(id)
|
||||
except:
|
||||
return None
|
||||
|
||||
@ -179,25 +187,23 @@ class ToolsTable:
|
||||
|
||||
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
|
||||
try:
|
||||
query = Tool.update(
|
||||
**updated,
|
||||
updated_at=int(time.time()),
|
||||
).where(Tool.id == id)
|
||||
query.execute()
|
||||
|
||||
tool = Tool.get(Tool.id == id)
|
||||
return ToolModel(**model_to_dict(tool))
|
||||
with get_db() as db:
|
||||
tool = db.get(Tool, id)
|
||||
tool.update(**updated)
|
||||
tool.updated_at = int(time.time())
|
||||
db.commit()
|
||||
db.refresh(tool)
|
||||
return ToolModel.model_validate(tool)
|
||||
except:
|
||||
return None
|
||||
|
||||
def delete_tool_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
query = Tool.delete().where((Tool.id == id))
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
|
||||
return True
|
||||
with get_db() as db:
|
||||
db.query(Tool).filter_by(id=id).delete()
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
Tools = ToolsTable(DB)
|
||||
Tools = ToolsTable()
|
||||
|
@ -1,11 +1,12 @@
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from peewee import *
|
||||
from playhouse.shortcuts import model_to_dict
|
||||
from pydantic import BaseModel, ConfigDict, parse_obj_as
|
||||
from typing import List, Union, Optional
|
||||
import time
|
||||
|
||||
from sqlalchemy import String, Column, BigInteger, Text
|
||||
|
||||
from utils.misc import get_gravatar_url
|
||||
|
||||
from apps.webui.internal.db import DB, JSONField
|
||||
from apps.webui.internal.db import Base, JSONField, Session, get_db
|
||||
from apps.webui.models.chats import Chats
|
||||
|
||||
####################
|
||||
@ -13,25 +14,24 @@ from apps.webui.models.chats import Chats
|
||||
####################
|
||||
|
||||
|
||||
class User(Model):
|
||||
id = CharField(unique=True)
|
||||
name = CharField()
|
||||
email = CharField()
|
||||
role = CharField()
|
||||
profile_image_url = TextField()
|
||||
class User(Base):
|
||||
__tablename__ = "user"
|
||||
|
||||
last_active_at = BigIntegerField()
|
||||
updated_at = BigIntegerField()
|
||||
created_at = BigIntegerField()
|
||||
id = Column(String, primary_key=True)
|
||||
name = Column(String)
|
||||
email = Column(String)
|
||||
role = Column(String)
|
||||
profile_image_url = Column(Text)
|
||||
|
||||
api_key = CharField(null=True, unique=True)
|
||||
settings = JSONField(null=True)
|
||||
info = JSONField(null=True)
|
||||
last_active_at = Column(BigInteger)
|
||||
updated_at = Column(BigInteger)
|
||||
created_at = Column(BigInteger)
|
||||
|
||||
oauth_sub = TextField(null=True, unique=True)
|
||||
api_key = Column(String, nullable=True, unique=True)
|
||||
settings = Column(JSONField, nullable=True)
|
||||
info = Column(JSONField, nullable=True)
|
||||
|
||||
class Meta:
|
||||
database = DB
|
||||
oauth_sub = Column(Text, unique=True)
|
||||
|
||||
|
||||
class UserSettings(BaseModel):
|
||||
@ -57,6 +57,8 @@ class UserModel(BaseModel):
|
||||
|
||||
oauth_sub: Optional[str] = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
####################
|
||||
# Forms
|
||||
@ -76,9 +78,6 @@ class UserUpdateForm(BaseModel):
|
||||
|
||||
|
||||
class UsersTable:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
self.db.create_tables([User])
|
||||
|
||||
def insert_new_user(
|
||||
self,
|
||||
@ -89,77 +88,92 @@ class UsersTable:
|
||||
role: str = "pending",
|
||||
oauth_sub: Optional[str] = None,
|
||||
) -> Optional[UserModel]:
|
||||
user = UserModel(
|
||||
**{
|
||||
"id": id,
|
||||
"name": name,
|
||||
"email": email,
|
||||
"role": role,
|
||||
"profile_image_url": profile_image_url,
|
||||
"last_active_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
"oauth_sub": oauth_sub,
|
||||
}
|
||||
)
|
||||
result = User.create(**user.model_dump())
|
||||
if result:
|
||||
return user
|
||||
else:
|
||||
return None
|
||||
with get_db() as db:
|
||||
user = UserModel(
|
||||
**{
|
||||
"id": id,
|
||||
"name": name,
|
||||
"email": email,
|
||||
"role": role,
|
||||
"profile_image_url": profile_image_url,
|
||||
"last_active_at": int(time.time()),
|
||||
"created_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
"oauth_sub": oauth_sub,
|
||||
}
|
||||
)
|
||||
result = User(**user.model_dump())
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
if result:
|
||||
return user
|
||||
else:
|
||||
return None
|
||||
|
||||
def get_user_by_id(self, id: str) -> Optional[UserModel]:
|
||||
try:
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
except:
|
||||
with get_db() as db:
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
|
||||
try:
|
||||
user = User.get(User.api_key == api_key)
|
||||
return UserModel(**model_to_dict(user))
|
||||
with get_db() as db:
|
||||
|
||||
user = db.query(User).filter_by(api_key=api_key).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_user_by_email(self, email: str) -> Optional[UserModel]:
|
||||
try:
|
||||
user = User.get(User.email == email)
|
||||
return UserModel(**model_to_dict(user))
|
||||
with get_db() as db:
|
||||
|
||||
user = db.query(User).filter_by(email=email).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]:
|
||||
try:
|
||||
user = User.get(User.oauth_sub == sub)
|
||||
return UserModel(**model_to_dict(user))
|
||||
with get_db() as db:
|
||||
|
||||
user = db.query(User).filter_by(oauth_sub=sub).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
|
||||
return [
|
||||
UserModel(**model_to_dict(user))
|
||||
for user in User.select()
|
||||
# .limit(limit).offset(skip)
|
||||
]
|
||||
with get_db() as db:
|
||||
users = (
|
||||
db.query(User)
|
||||
# .offset(skip).limit(limit)
|
||||
.all()
|
||||
)
|
||||
return [UserModel.model_validate(user) for user in users]
|
||||
|
||||
def get_num_users(self) -> Optional[int]:
|
||||
return User.select().count()
|
||||
with get_db() as db:
|
||||
return db.query(User).count()
|
||||
|
||||
def get_first_user(self) -> UserModel:
|
||||
try:
|
||||
user = User.select().order_by(User.created_at).first()
|
||||
return UserModel(**model_to_dict(user))
|
||||
with get_db() as db:
|
||||
user = db.query(User).order_by(User.created_at).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
|
||||
try:
|
||||
query = User.update(role=role).where(User.id == id)
|
||||
query.execute()
|
||||
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
with get_db() as db:
|
||||
db.query(User).filter_by(id=id).update({"role": role})
|
||||
db.commit()
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
@ -167,23 +181,28 @@ class UsersTable:
|
||||
self, id: str, profile_image_url: str
|
||||
) -> Optional[UserModel]:
|
||||
try:
|
||||
query = User.update(profile_image_url=profile_image_url).where(
|
||||
User.id == id
|
||||
)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
db.query(User).filter_by(id=id).update(
|
||||
{"profile_image_url": profile_image_url}
|
||||
)
|
||||
db.commit()
|
||||
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
|
||||
try:
|
||||
query = User.update(last_active_at=int(time.time())).where(User.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
db.query(User).filter_by(id=id).update(
|
||||
{"last_active_at": int(time.time())}
|
||||
)
|
||||
db.commit()
|
||||
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
@ -191,22 +210,24 @@ class UsersTable:
|
||||
self, id: str, oauth_sub: str
|
||||
) -> Optional[UserModel]:
|
||||
try:
|
||||
query = User.update(oauth_sub=oauth_sub).where(User.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
|
||||
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
except:
|
||||
return None
|
||||
|
||||
def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
|
||||
try:
|
||||
query = User.update(**updated).where(User.id == id)
|
||||
query.execute()
|
||||
with get_db() as db:
|
||||
db.query(User).filter_by(id=id).update(updated)
|
||||
db.commit()
|
||||
|
||||
user = User.get(User.id == id)
|
||||
return UserModel(**model_to_dict(user))
|
||||
except:
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return UserModel.model_validate(user)
|
||||
# return UserModel(**user.dict())
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
def delete_user_by_id(self, id: str) -> bool:
|
||||
@ -215,9 +236,10 @@ class UsersTable:
|
||||
result = Chats.delete_chats_by_user_id(id)
|
||||
|
||||
if result:
|
||||
# Delete User
|
||||
query = User.delete().where(User.id == id)
|
||||
query.execute() # Remove the rows, return number of rows removed.
|
||||
with get_db() as db:
|
||||
# Delete User
|
||||
db.query(User).filter_by(id=id).delete()
|
||||
db.commit()
|
||||
|
||||
return True
|
||||
else:
|
||||
@ -227,19 +249,20 @@ class UsersTable:
|
||||
|
||||
def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
|
||||
try:
|
||||
query = User.update(api_key=api_key).where(User.id == id)
|
||||
result = query.execute()
|
||||
|
||||
return True if result == 1 else False
|
||||
with get_db() as db:
|
||||
result = db.query(User).filter_by(id=id).update({"api_key": api_key})
|
||||
db.commit()
|
||||
return True if result == 1 else False
|
||||
except:
|
||||
return False
|
||||
|
||||
def get_user_api_key_by_id(self, id: str) -> Optional[str]:
|
||||
try:
|
||||
user = User.get(User.id == id)
|
||||
return user.api_key
|
||||
except:
|
||||
with get_db() as db:
|
||||
user = db.query(User).filter_by(id=id).first()
|
||||
return user.api_key
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
|
||||
Users = UsersTable(DB)
|
||||
Users = UsersTable()
|
||||
|
@ -76,7 +76,10 @@ async def delete_all_user_chats(request: Request, user=Depends(get_verified_user
|
||||
|
||||
@router.get("/list/user/{user_id}", response_model=List[ChatTitleIdResponse])
|
||||
async def get_user_chat_list_by_user_id(
|
||||
user_id: str, user=Depends(get_admin_user), skip: int = 0, limit: int = 50
|
||||
user_id: str,
|
||||
user=Depends(get_admin_user),
|
||||
skip: int = 0,
|
||||
limit: int = 50,
|
||||
):
|
||||
return Chats.get_chat_list_by_user_id(
|
||||
user_id, include_archived=True, skip=skip, limit=limit
|
||||
@ -119,7 +122,7 @@ async def get_user_chats(user=Depends(get_verified_user)):
|
||||
|
||||
|
||||
@router.get("/all/archived", response_model=List[ChatResponse])
|
||||
async def get_user_chats(user=Depends(get_verified_user)):
|
||||
async def get_user_archived_chats(user=Depends(get_verified_user)):
|
||||
return [
|
||||
ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
|
||||
for chat in Chats.get_archived_chats_by_user_id(user.id)
|
||||
|
@ -130,7 +130,9 @@ async def tag_doc_by_name(form_data: TagDocumentForm, user=Depends(get_verified_
|
||||
|
||||
@router.post("/doc/update", response_model=Optional[DocumentResponse])
|
||||
async def update_doc_by_name(
|
||||
name: str, form_data: DocumentUpdateForm, user=Depends(get_admin_user)
|
||||
name: str,
|
||||
form_data: DocumentUpdateForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
doc = Documents.update_doc_by_name(name, form_data)
|
||||
if doc:
|
||||
|
@ -50,10 +50,7 @@ router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/")
|
||||
def upload_file(
|
||||
file: UploadFile = File(...),
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
def upload_file(file: UploadFile = File(...), user=Depends(get_verified_user)):
|
||||
log.info(f"file.content_type: {file.content_type}")
|
||||
try:
|
||||
unsanitized_filename = file.filename
|
||||
|
@ -50,7 +50,9 @@ class MemoryUpdateModel(BaseModel):
|
||||
|
||||
@router.post("/add", response_model=Optional[MemoryModel])
|
||||
async def add_memory(
|
||||
request: Request, form_data: AddMemoryForm, user=Depends(get_verified_user)
|
||||
request: Request,
|
||||
form_data: AddMemoryForm,
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
memory = Memories.insert_new_memory(user.id, form_data.content)
|
||||
memory_embedding = request.app.state.EMBEDDING_FUNCTION(memory.content)
|
||||
|
@ -5,6 +5,7 @@ from typing import List, Union, Optional
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
import json
|
||||
|
||||
from apps.webui.models.models import Models, ModelModel, ModelForm, ModelResponse
|
||||
|
||||
from utils.utils import get_verified_user, get_admin_user
|
||||
@ -29,7 +30,9 @@ async def get_models(user=Depends(get_verified_user)):
|
||||
|
||||
@router.post("/add", response_model=Optional[ModelModel])
|
||||
async def add_new_model(
|
||||
request: Request, form_data: ModelForm, user=Depends(get_admin_user)
|
||||
request: Request,
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
if form_data.id in request.app.state.MODELS:
|
||||
raise HTTPException(
|
||||
@ -73,7 +76,10 @@ async def get_model_by_id(id: str, user=Depends(get_verified_user)):
|
||||
|
||||
@router.post("/update", response_model=Optional[ModelModel])
|
||||
async def update_model_by_id(
|
||||
request: Request, id: str, form_data: ModelForm, user=Depends(get_admin_user)
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
model = Models.get_model_by_id(id)
|
||||
if model:
|
||||
|
@ -71,7 +71,9 @@ async def get_prompt_by_command(command: str, user=Depends(get_verified_user)):
|
||||
|
||||
@router.post("/command/{command}/update", response_model=Optional[PromptModel])
|
||||
async def update_prompt_by_command(
|
||||
command: str, form_data: PromptForm, user=Depends(get_admin_user)
|
||||
command: str,
|
||||
form_data: PromptForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
prompt = Prompts.update_prompt_by_command(f"/{command}", form_data)
|
||||
if prompt:
|
||||
|
@ -6,7 +6,6 @@ from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
import json
|
||||
|
||||
|
||||
from apps.webui.models.users import Users
|
||||
from apps.webui.models.tools import Tools, ToolForm, ToolModel, ToolResponse
|
||||
from apps.webui.utils import load_toolkit_module_by_id
|
||||
@ -57,7 +56,9 @@ async def get_toolkits(user=Depends(get_admin_user)):
|
||||
|
||||
@router.post("/create", response_model=Optional[ToolResponse])
|
||||
async def create_new_toolkit(
|
||||
request: Request, form_data: ToolForm, user=Depends(get_admin_user)
|
||||
request: Request,
|
||||
form_data: ToolForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
if not form_data.id.isidentifier():
|
||||
raise HTTPException(
|
||||
@ -131,7 +132,10 @@ async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)):
|
||||
|
||||
@router.post("/id/{id}/update", response_model=Optional[ToolModel])
|
||||
async def update_toolkit_by_id(
|
||||
request: Request, id: str, form_data: ToolForm, user=Depends(get_admin_user)
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: ToolForm,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
|
||||
|
||||
|
@ -138,7 +138,7 @@ async def get_user_info_by_session_user(user=Depends(get_verified_user)):
|
||||
|
||||
|
||||
@router.post("/user/info/update", response_model=Optional[dict])
|
||||
async def update_user_settings_by_session_user(
|
||||
async def update_user_info_by_session_user(
|
||||
form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
user = Users.get_user_by_id(user.id)
|
||||
@ -205,7 +205,9 @@ async def get_user_by_id(user_id: str, user=Depends(get_verified_user)):
|
||||
|
||||
@router.post("/{user_id}/update", response_model=Optional[UserModel])
|
||||
async def update_user_by_id(
|
||||
user_id: str, form_data: UserUpdateForm, session_user=Depends(get_admin_user)
|
||||
user_id: str,
|
||||
form_data: UserUpdateForm,
|
||||
session_user=Depends(get_admin_user),
|
||||
):
|
||||
user = Users.get_user_by_id(user_id)
|
||||
|
||||
|
@ -1,6 +1,5 @@
|
||||
from fastapi import APIRouter, UploadFile, File, Response
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from peewee import SqliteDatabase
|
||||
from starlette.responses import StreamingResponse, FileResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
@ -10,7 +9,6 @@ import markdown
|
||||
import black
|
||||
|
||||
|
||||
from apps.webui.internal.db import DB
|
||||
from utils.utils import get_admin_user
|
||||
from utils.misc import calculate_sha256, get_gravatar_url
|
||||
|
||||
@ -114,13 +112,15 @@ async def download_db(user=Depends(get_admin_user)):
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
if not isinstance(DB, SqliteDatabase):
|
||||
from apps.webui.internal.db import engine
|
||||
|
||||
if engine.name != "sqlite":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DB_NOT_SQLITE,
|
||||
)
|
||||
return FileResponse(
|
||||
DB.database,
|
||||
engine.url.database,
|
||||
media_type="application/octet-stream",
|
||||
filename="webui.db",
|
||||
)
|
||||
|
@ -778,6 +778,7 @@ class BannerModel(BaseModel):
|
||||
dismissible: bool
|
||||
timestamp: int
|
||||
|
||||
|
||||
try:
|
||||
banners = json.loads(os.environ.get("WEBUI_BANNERS", "[]"))
|
||||
banners = [BannerModel(**banner) for banner in banners]
|
||||
|
@ -1,5 +1,6 @@
|
||||
import base64
|
||||
import uuid
|
||||
import subprocess
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from authlib.integrations.starlette_client import OAuth
|
||||
@ -27,6 +28,7 @@ from fastapi.responses import JSONResponse
|
||||
from fastapi import HTTPException
|
||||
from fastapi.middleware.wsgi import WSGIMiddleware
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from sqlalchemy import text
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.middleware.sessions import SessionMiddleware
|
||||
@ -54,6 +56,7 @@ from apps.webui.main import (
|
||||
get_pipe_models,
|
||||
generate_function_chat_completion,
|
||||
)
|
||||
from apps.webui.internal.db import Session, SessionLocal
|
||||
|
||||
|
||||
from pydantic import BaseModel
|
||||
@ -125,6 +128,8 @@ from config import (
|
||||
WEBUI_SESSION_COOKIE_SAME_SITE,
|
||||
WEBUI_SESSION_COOKIE_SECURE,
|
||||
AppConfig,
|
||||
BACKEND_DIR,
|
||||
DATABASE_URL,
|
||||
)
|
||||
from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES, TASKS
|
||||
from utils.webhook import post_webhook
|
||||
@ -167,8 +172,19 @@ https://github.com/open-webui/open-webui
|
||||
)
|
||||
|
||||
|
||||
def run_migrations():
|
||||
env = os.environ.copy()
|
||||
env["DATABASE_URL"] = DATABASE_URL
|
||||
migration_task = subprocess.run(
|
||||
["alembic", f"-c{BACKEND_DIR}/alembic.ini", "upgrade", "head"], env=env
|
||||
)
|
||||
if migration_task.returncode > 0:
|
||||
raise ValueError("Error running migrations")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
run_migrations()
|
||||
yield
|
||||
|
||||
|
||||
@ -899,6 +915,14 @@ app.add_middleware(
|
||||
)
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def commit_session_after_request(request: Request, call_next):
|
||||
response = await call_next(request)
|
||||
log.debug("Commit session after request")
|
||||
Session.commit()
|
||||
return response
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def check_url(request: Request, call_next):
|
||||
if len(app.state.MODELS) == 0:
|
||||
@ -1744,7 +1768,9 @@ async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_use
|
||||
|
||||
@app.get("/api/pipelines/{pipeline_id}/valves")
|
||||
async def get_pipeline_valves(
|
||||
urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
|
||||
urlIdx: Optional[int],
|
||||
pipeline_id: str,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
models = await get_all_models()
|
||||
r = None
|
||||
@ -1782,7 +1808,9 @@ async def get_pipeline_valves(
|
||||
|
||||
@app.get("/api/pipelines/{pipeline_id}/valves/spec")
|
||||
async def get_pipeline_valves_spec(
|
||||
urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
|
||||
urlIdx: Optional[int],
|
||||
pipeline_id: str,
|
||||
user=Depends(get_admin_user),
|
||||
):
|
||||
models = await get_all_models()
|
||||
|
||||
@ -2171,6 +2199,12 @@ async def healthcheck():
|
||||
return {"status": True}
|
||||
|
||||
|
||||
@app.get("/health/db")
|
||||
async def healthcheck_with_db():
|
||||
Session.execute(text("SELECT 1;")).all()
|
||||
return {"status": True}
|
||||
|
||||
|
||||
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
|
||||
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
|
||||
|
||||
|
4
backend/migrations/README
Normal file
4
backend/migrations/README
Normal file
@ -0,0 +1,4 @@
|
||||
Generic single-database configuration.
|
||||
|
||||
Create new migrations with
|
||||
DATABASE_URL=<replace with actual url> alembic revision --autogenerate -m "a description"
|
93
backend/migrations/env.py
Normal file
93
backend/migrations/env.py
Normal file
@ -0,0 +1,93 @@
|
||||
import os
|
||||
from logging.config import fileConfig
|
||||
|
||||
from sqlalchemy import engine_from_config
|
||||
from sqlalchemy import pool
|
||||
|
||||
from alembic import context
|
||||
|
||||
from apps.webui.models.auths import Auth
|
||||
from apps.webui.models.chats import Chat
|
||||
from apps.webui.models.documents import Document
|
||||
from apps.webui.models.memories import Memory
|
||||
from apps.webui.models.models import Model
|
||||
from apps.webui.models.prompts import Prompt
|
||||
from apps.webui.models.tags import Tag, ChatIdTag
|
||||
from apps.webui.models.tools import Tool
|
||||
from apps.webui.models.users import User
|
||||
from apps.webui.models.files import File
|
||||
from apps.webui.models.functions import Function
|
||||
|
||||
# this is the Alembic Config object, which provides
|
||||
# access to the values within the .ini file in use.
|
||||
config = context.config
|
||||
|
||||
# Interpret the config file for Python logging.
|
||||
# This line sets up loggers basically.
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# add your model's MetaData object here
|
||||
# for 'autogenerate' support
|
||||
# from myapp import mymodel
|
||||
# target_metadata = mymodel.Base.metadata
|
||||
target_metadata = Auth.metadata
|
||||
|
||||
# other values from the config, defined by the needs of env.py,
|
||||
# can be acquired:
|
||||
# my_important_option = config.get_main_option("my_important_option")
|
||||
# ... etc.
|
||||
|
||||
database_url = os.getenv("DATABASE_URL", None)
|
||||
if database_url:
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""Run migrations in 'offline' mode.
|
||||
|
||||
This configures the context with just a URL
|
||||
and not an Engine, though an Engine is acceptable
|
||||
here as well. By skipping the Engine creation
|
||||
we don't even need a DBAPI to be available.
|
||||
|
||||
Calls to context.execute() here emit the given string to the
|
||||
script output.
|
||||
|
||||
"""
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine
|
||||
and associate a connection with the context.
|
||||
|
||||
"""
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
27
backend/migrations/script.py.mako
Normal file
27
backend/migrations/script.py.mako
Normal file
@ -0,0 +1,27 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import apps.webui.internal.db
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
${downgrades if downgrades else "pass"}
|
9
backend/migrations/util.py
Normal file
9
backend/migrations/util.py
Normal file
@ -0,0 +1,9 @@
|
||||
from alembic import op
|
||||
from sqlalchemy import Inspector
|
||||
|
||||
|
||||
def get_existing_tables():
|
||||
con = op.get_bind()
|
||||
inspector = Inspector.from_engine(con)
|
||||
tables = set(inspector.get_table_names())
|
||||
return tables
|
202
backend/migrations/versions/7e5b5dc7342b_init.py
Normal file
202
backend/migrations/versions/7e5b5dc7342b_init.py
Normal file
@ -0,0 +1,202 @@
|
||||
"""init
|
||||
|
||||
Revision ID: 7e5b5dc7342b
|
||||
Revises:
|
||||
Create Date: 2024-06-24 13:15:33.808998
|
||||
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import apps.webui.internal.db
|
||||
from migrations.util import get_existing_tables
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "7e5b5dc7342b"
|
||||
down_revision: Union[str, None] = None
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
existing_tables = set(get_existing_tables())
|
||||
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
if "auth" not in existing_tables:
|
||||
op.create_table(
|
||||
"auth",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("email", sa.String(), nullable=True),
|
||||
sa.Column("password", sa.Text(), nullable=True),
|
||||
sa.Column("active", sa.Boolean(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "chat" not in existing_tables:
|
||||
op.create_table(
|
||||
"chat",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("title", sa.Text(), nullable=True),
|
||||
sa.Column("chat", sa.Text(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("share_id", sa.Text(), nullable=True),
|
||||
sa.Column("archived", sa.Boolean(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("share_id"),
|
||||
)
|
||||
|
||||
if "chatidtag" not in existing_tables:
|
||||
op.create_table(
|
||||
"chatidtag",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("tag_name", sa.String(), nullable=True),
|
||||
sa.Column("chat_id", sa.String(), nullable=True),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("timestamp", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "document" not in existing_tables:
|
||||
op.create_table(
|
||||
"document",
|
||||
sa.Column("collection_name", sa.String(), nullable=False),
|
||||
sa.Column("name", sa.String(), nullable=True),
|
||||
sa.Column("title", sa.Text(), nullable=True),
|
||||
sa.Column("filename", sa.Text(), nullable=True),
|
||||
sa.Column("content", sa.Text(), nullable=True),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("timestamp", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("collection_name"),
|
||||
sa.UniqueConstraint("name"),
|
||||
)
|
||||
|
||||
if "file" not in existing_tables:
|
||||
op.create_table(
|
||||
"file",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("filename", sa.Text(), nullable=True),
|
||||
sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "function" not in existing_tables:
|
||||
op.create_table(
|
||||
"function",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("name", sa.Text(), nullable=True),
|
||||
sa.Column("type", sa.Text(), nullable=True),
|
||||
sa.Column("content", sa.Text(), nullable=True),
|
||||
sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("valves", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("is_active", sa.Boolean(), nullable=True),
|
||||
sa.Column("is_global", sa.Boolean(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "memory" not in existing_tables:
|
||||
op.create_table(
|
||||
"memory",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("content", sa.Text(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "model" not in existing_tables:
|
||||
op.create_table(
|
||||
"model",
|
||||
sa.Column("id", sa.Text(), nullable=False),
|
||||
sa.Column("user_id", sa.Text(), nullable=True),
|
||||
sa.Column("base_model_id", sa.Text(), nullable=True),
|
||||
sa.Column("name", sa.Text(), nullable=True),
|
||||
sa.Column("params", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "prompt" not in existing_tables:
|
||||
op.create_table(
|
||||
"prompt",
|
||||
sa.Column("command", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("title", sa.Text(), nullable=True),
|
||||
sa.Column("content", sa.Text(), nullable=True),
|
||||
sa.Column("timestamp", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("command"),
|
||||
)
|
||||
|
||||
if "tag" not in existing_tables:
|
||||
op.create_table(
|
||||
"tag",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("name", sa.String(), nullable=True),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("data", sa.Text(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "tool" not in existing_tables:
|
||||
op.create_table(
|
||||
"tool",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("user_id", sa.String(), nullable=True),
|
||||
sa.Column("name", sa.Text(), nullable=True),
|
||||
sa.Column("content", sa.Text(), nullable=True),
|
||||
sa.Column("specs", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("meta", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("valves", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
if "user" not in existing_tables:
|
||||
op.create_table(
|
||||
"user",
|
||||
sa.Column("id", sa.String(), nullable=False),
|
||||
sa.Column("name", sa.String(), nullable=True),
|
||||
sa.Column("email", sa.String(), nullable=True),
|
||||
sa.Column("role", sa.String(), nullable=True),
|
||||
sa.Column("profile_image_url", sa.Text(), nullable=True),
|
||||
sa.Column("last_active_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("api_key", sa.String(), nullable=True),
|
||||
sa.Column("settings", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("info", apps.webui.internal.db.JSONField(), nullable=True),
|
||||
sa.Column("oauth_sub", sa.Text(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint("api_key"),
|
||||
sa.UniqueConstraint("oauth_sub"),
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("user")
|
||||
op.drop_table("tool")
|
||||
op.drop_table("tag")
|
||||
op.drop_table("prompt")
|
||||
op.drop_table("model")
|
||||
op.drop_table("memory")
|
||||
op.drop_table("function")
|
||||
op.drop_table("file")
|
||||
op.drop_table("document")
|
||||
op.drop_table("chatidtag")
|
||||
op.drop_table("chat")
|
||||
op.drop_table("auth")
|
||||
# ### end Alembic commands ###
|
@ -12,6 +12,8 @@ passlib[bcrypt]==1.7.4
|
||||
|
||||
requests==2.32.3
|
||||
aiohttp==3.9.5
|
||||
sqlalchemy==2.0.30
|
||||
alembic==1.13.1
|
||||
peewee==3.17.5
|
||||
peewee-migrate==1.12.2
|
||||
psycopg2-binary==2.9.9
|
||||
@ -67,4 +69,9 @@ pytube==15.0.0
|
||||
|
||||
extract_msg
|
||||
pydub
|
||||
duckduckgo-search~=6.1.7
|
||||
duckduckgo-search~=6.1.7
|
||||
|
||||
## Tests
|
||||
docker~=7.1.0
|
||||
pytest~=8.2.1
|
||||
pytest-docker~=3.1.1
|
||||
|
0
backend/test/__init__.py
Normal file
0
backend/test/__init__.py
Normal file
202
backend/test/apps/webui/routers/test_auths.py
Normal file
202
backend/test/apps/webui/routers/test_auths.py
Normal file
@ -0,0 +1,202 @@
|
||||
import pytest
|
||||
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestAuths(AbstractPostgresTest):
|
||||
BASE_PATH = "/api/v1/auths"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from apps.webui.models.users import Users
|
||||
from apps.webui.models.auths import Auths
|
||||
|
||||
cls.users = Users
|
||||
cls.auths = Auths
|
||||
|
||||
def test_get_session_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"id": "1",
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"role": "user",
|
||||
"profile_image_url": "/user.png",
|
||||
}
|
||||
|
||||
def test_update_profile(self):
|
||||
from utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("old_password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/profile"),
|
||||
json={"name": "John Doe 2", "profile_image_url": "/user2.png"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.name == "John Doe 2"
|
||||
assert db_user.profile_image_url == "/user2.png"
|
||||
|
||||
def test_update_password(self):
|
||||
from utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("old_password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/password"),
|
||||
json={"password": "old_password", "new_password": "new_password"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
old_auth = self.auths.authenticate_user(
|
||||
"john.doe@openwebui.com", "old_password"
|
||||
)
|
||||
assert old_auth is None
|
||||
new_auth = self.auths.authenticate_user(
|
||||
"john.doe@openwebui.com", "new_password"
|
||||
)
|
||||
assert new_auth is not None
|
||||
|
||||
def test_signin(self):
|
||||
from utils.utils import get_password_hash
|
||||
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password=get_password_hash("password"),
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="user",
|
||||
)
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/signin"),
|
||||
json={"email": "john.doe@openwebui.com", "password": "password"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == user.id
|
||||
assert data["name"] == "John Doe"
|
||||
assert data["email"] == "john.doe@openwebui.com"
|
||||
assert data["role"] == "user"
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_signup(self):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/signup"),
|
||||
json={
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"password": "password",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] is not None and len(data["id"]) > 0
|
||||
assert data["name"] == "John Doe"
|
||||
assert data["email"] == "john.doe@openwebui.com"
|
||||
assert data["role"] in ["admin", "user", "pending"]
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_add_user(self):
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/add"),
|
||||
json={
|
||||
"name": "John Doe 2",
|
||||
"email": "john.doe2@openwebui.com",
|
||||
"password": "password2",
|
||||
"role": "admin",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] is not None and len(data["id"]) > 0
|
||||
assert data["name"] == "John Doe 2"
|
||||
assert data["email"] == "john.doe2@openwebui.com"
|
||||
assert data["role"] == "admin"
|
||||
assert data["profile_image_url"] == "/user.png"
|
||||
assert data["token"] is not None and len(data["token"]) > 0
|
||||
assert data["token_type"] == "Bearer"
|
||||
|
||||
def test_get_admin_details(self):
|
||||
self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
with mock_webui_user():
|
||||
response = self.fast_api_client.get(self.create_url("/admin/details"))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
}
|
||||
|
||||
def test_create_api_key_(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.post(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["api_key"] is not None
|
||||
assert len(data["api_key"]) > 0
|
||||
|
||||
def test_delete_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, "abc")
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.delete(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == True
|
||||
db_user = self.users.get_user_by_id(user.id)
|
||||
assert db_user.api_key is None
|
||||
|
||||
def test_get_api_key(self):
|
||||
user = self.auths.insert_new_auth(
|
||||
email="john.doe@openwebui.com",
|
||||
password="password",
|
||||
name="John Doe",
|
||||
profile_image_url="/user.png",
|
||||
role="admin",
|
||||
)
|
||||
self.users.update_user_api_key_by_id(user.id, "abc")
|
||||
with mock_webui_user(id=user.id):
|
||||
response = self.fast_api_client.get(self.create_url("/api_key"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"api_key": "abc"}
|
238
backend/test/apps/webui/routers/test_chats.py
Normal file
238
backend/test/apps/webui/routers/test_chats.py
Normal file
@ -0,0 +1,238 @@
|
||||
import uuid
|
||||
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestChats(AbstractPostgresTest):
|
||||
|
||||
BASE_PATH = "/api/v1/chats"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
from apps.webui.models.chats import ChatForm
|
||||
from apps.webui.models.chats import Chats
|
||||
|
||||
self.chats = Chats
|
||||
self.chats.insert_new_chat(
|
||||
"2",
|
||||
ChatForm(
|
||||
**{
|
||||
"chat": {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def test_get_session_user_chat_list(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_delete_all_user_chats(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(self.chats.get_chats()) == 0
|
||||
|
||||
def test_get_user_chat_list_by_user_id(self):
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url("/list/user/2"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_create_new_chat(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/new"),
|
||||
json={
|
||||
"chat": {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["archived"] is False
|
||||
assert data["chat"] == {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
}
|
||||
assert data["user_id"] == "2"
|
||||
assert data["id"] is not None
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "New Chat"
|
||||
assert data["updated_at"] is not None
|
||||
assert data["created_at"] is not None
|
||||
assert len(self.chats.get_chats()) == 2
|
||||
|
||||
def test_get_user_chats(self):
|
||||
self.test_get_session_user_chat_list()
|
||||
|
||||
def test_get_user_archived_chats(self):
|
||||
self.chats.archive_all_chats_by_user_id("2")
|
||||
from apps.webui.internal.db import Session
|
||||
|
||||
Session.commit()
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/all/archived"))
|
||||
assert response.status_code == 200
|
||||
first_chat = response.json()[0]
|
||||
assert first_chat["id"] is not None
|
||||
assert first_chat["title"] == "New Chat"
|
||||
assert first_chat["created_at"] is not None
|
||||
assert first_chat["updated_at"] is not None
|
||||
|
||||
def test_get_all_user_chats_in_db(self):
|
||||
with mock_webui_user(id="4"):
|
||||
response = self.fast_api_client.get(self.create_url("/all/db"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
|
||||
def test_get_archived_session_user_chat_list(self):
|
||||
self.test_get_user_archived_chats()
|
||||
|
||||
def test_archive_all_chats(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(self.create_url("/archive/all"))
|
||||
assert response.status_code == 200
|
||||
assert len(self.chats.get_archived_chats_by_user_id("2")) == 1
|
||||
|
||||
def test_get_shared_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
self.chats.update_chat_share_id_by_id(chat_id, chat_id)
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/share/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["id"] == chat_id
|
||||
assert data["share_id"] == chat_id
|
||||
assert data["title"] == "New Chat"
|
||||
|
||||
def test_get_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat1",
|
||||
"description": "chat1 description",
|
||||
"tags": ["tag1", "tag2"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "New Chat"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_update_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url(f"/{chat_id}"),
|
||||
json={
|
||||
"chat": {
|
||||
"name": "chat2",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag2", "tag4"],
|
||||
"title": "Just another title",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == chat_id
|
||||
assert data["chat"] == {
|
||||
"name": "chat2",
|
||||
"title": "Just another title",
|
||||
"description": "chat2 description",
|
||||
"tags": ["tag2", "tag4"],
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "Just another title"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_delete_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url(f"/{chat_id}"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is True
|
||||
|
||||
def test_clone_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}/clone"))
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] != chat_id
|
||||
assert data["chat"] == {
|
||||
"branchPointMessageId": "1",
|
||||
"description": "chat1 description",
|
||||
"history": {"currentId": "1", "messages": []},
|
||||
"name": "chat1",
|
||||
"originalChatId": chat_id,
|
||||
"tags": ["tag1", "tag2"],
|
||||
"title": "Clone of New Chat",
|
||||
}
|
||||
assert data["share_id"] is None
|
||||
assert data["title"] == "Clone of New Chat"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
def test_archive_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url(f"/{chat_id}/archive"))
|
||||
assert response.status_code == 200
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.archived is True
|
||||
|
||||
def test_share_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(self.create_url(f"/{chat_id}/share"))
|
||||
assert response.status_code == 200
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.share_id is not None
|
||||
|
||||
def test_delete_shared_chat_by_id(self):
|
||||
chat_id = self.chats.get_chats()[0].id
|
||||
share_id = str(uuid.uuid4())
|
||||
self.chats.update_chat_share_id_by_id(chat_id, share_id)
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(self.create_url(f"/{chat_id}/share"))
|
||||
assert response.status_code
|
||||
|
||||
chat = self.chats.get_chat_by_id(chat_id)
|
||||
assert chat.share_id is None
|
106
backend/test/apps/webui/routers/test_documents.py
Normal file
106
backend/test/apps/webui/routers/test_documents.py
Normal file
@ -0,0 +1,106 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestDocuments(AbstractPostgresTest):
|
||||
|
||||
BASE_PATH = "/api/v1/documents"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from apps.webui.models.documents import Documents
|
||||
|
||||
cls.documents = Documents
|
||||
|
||||
def test_documents(self):
|
||||
# Empty database
|
||||
assert len(self.documents.get_docs()) == 0
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
# Create a new document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"name": "doc_name",
|
||||
"title": "doc title",
|
||||
"collection_name": "custom collection",
|
||||
"filename": "doc_name.pdf",
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["name"] == "doc_name"
|
||||
assert len(self.documents.get_docs()) == 1
|
||||
|
||||
# Get the document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/doc?name=doc_name"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["collection_name"] == "custom collection"
|
||||
assert data["name"] == "doc_name"
|
||||
assert data["title"] == "doc title"
|
||||
assert data["filename"] == "doc_name.pdf"
|
||||
assert data["content"] == {}
|
||||
|
||||
# Create another document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"name": "doc_name 2",
|
||||
"title": "doc title 2",
|
||||
"collection_name": "custom collection 2",
|
||||
"filename": "doc_name2.pdf",
|
||||
"content": "",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["name"] == "doc_name 2"
|
||||
assert len(self.documents.get_docs()) == 2
|
||||
|
||||
# Get all documents
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
# Update the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/doc/update?name=doc_name"),
|
||||
json={"name": "doc_name rework", "title": "updated title"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "doc_name rework"
|
||||
assert data["title"] == "updated title"
|
||||
|
||||
# Tag the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/doc/tags"),
|
||||
json={
|
||||
"name": "doc_name rework",
|
||||
"tags": [{"name": "testing-tag"}, {"name": "another-tag"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "doc_name rework"
|
||||
assert data["content"] == {
|
||||
"tags": [{"name": "testing-tag"}, {"name": "another-tag"}]
|
||||
}
|
||||
assert len(self.documents.get_docs()) == 2
|
||||
|
||||
# Delete the first document
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/doc/delete?name=doc_name rework")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert len(self.documents.get_docs()) == 1
|
62
backend/test/apps/webui/routers/test_models.py
Normal file
62
backend/test/apps/webui/routers/test_models.py
Normal file
@ -0,0 +1,62 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestModels(AbstractPostgresTest):
|
||||
|
||||
BASE_PATH = "/api/v1/models"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from apps.webui.models.models import Model
|
||||
|
||||
cls.models = Model
|
||||
|
||||
def test_models(self):
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/add"),
|
||||
json={
|
||||
"id": "my-model",
|
||||
"base_model_id": "base-model-id",
|
||||
"name": "Hello World",
|
||||
"meta": {
|
||||
"profile_image_url": "/favicon.png",
|
||||
"description": "description",
|
||||
"capabilities": None,
|
||||
"model_config": {},
|
||||
},
|
||||
"params": {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(
|
||||
self.create_url(query_params={"id": "my-model"})
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()[0]
|
||||
assert data["id"] == "my-model"
|
||||
assert data["name"] == "Hello World"
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/delete?id=my-model")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
92
backend/test/apps/webui/routers/test_prompts.py
Normal file
92
backend/test/apps/webui/routers/test_prompts.py
Normal file
@ -0,0 +1,92 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
class TestPrompts(AbstractPostgresTest):
|
||||
|
||||
BASE_PATH = "/api/v1/prompts"
|
||||
|
||||
def test_prompts(self):
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 0
|
||||
|
||||
# Create a two new prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"command": "/my-command",
|
||||
"title": "Hello World",
|
||||
"content": "description",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/create"),
|
||||
json={
|
||||
"command": "/my-command2",
|
||||
"title": "Hello World 2",
|
||||
"content": "description 2",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
|
||||
# Get prompt by command
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/command/my-command"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command"
|
||||
assert data["title"] == "Hello World"
|
||||
assert data["content"] == "description"
|
||||
assert data["user_id"] == "2"
|
||||
|
||||
# Update prompt
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/command/my-command2/update"),
|
||||
json={
|
||||
"command": "irrelevant for request",
|
||||
"title": "Hello World Updated",
|
||||
"content": "description Updated",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command2"
|
||||
assert data["title"] == "Hello World Updated"
|
||||
assert data["content"] == "description Updated"
|
||||
assert data["user_id"] == "3"
|
||||
|
||||
# Get prompt by command
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/command/my-command2"))
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["command"] == "/my-command2"
|
||||
assert data["title"] == "Hello World Updated"
|
||||
assert data["content"] == "description Updated"
|
||||
assert data["user_id"] == "3"
|
||||
|
||||
# Delete prompt
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.delete(
|
||||
self.create_url("/command/my-command/delete")
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all prompts
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/"))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
168
backend/test/apps/webui/routers/test_users.py
Normal file
168
backend/test/apps/webui/routers/test_users.py
Normal file
@ -0,0 +1,168 @@
|
||||
from test.util.abstract_integration_test import AbstractPostgresTest
|
||||
from test.util.mock_user import mock_webui_user
|
||||
|
||||
|
||||
def _get_user_by_id(data, param):
|
||||
return next((item for item in data if item["id"] == param), None)
|
||||
|
||||
|
||||
def _assert_user(data, id, **kwargs):
|
||||
user = _get_user_by_id(data, id)
|
||||
assert user is not None
|
||||
comparison_data = {
|
||||
"name": f"user {id}",
|
||||
"email": f"user{id}@openwebui.com",
|
||||
"profile_image_url": f"/user{id}.png",
|
||||
"role": "user",
|
||||
**kwargs,
|
||||
}
|
||||
for key, value in comparison_data.items():
|
||||
assert user[key] == value
|
||||
|
||||
|
||||
class TestUsers(AbstractPostgresTest):
|
||||
|
||||
BASE_PATH = "/api/v1/users"
|
||||
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
from apps.webui.models.users import Users
|
||||
|
||||
cls.users = Users
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
self.users.insert_new_user(
|
||||
id="1",
|
||||
name="user 1",
|
||||
email="user1@openwebui.com",
|
||||
profile_image_url="/user1.png",
|
||||
role="user",
|
||||
)
|
||||
self.users.insert_new_user(
|
||||
id="2",
|
||||
name="user 2",
|
||||
email="user2@openwebui.com",
|
||||
profile_image_url="/user2.png",
|
||||
role="user",
|
||||
)
|
||||
|
||||
def test_users(self):
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(data, "2")
|
||||
|
||||
# update role
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/update/role"), json={"id": "2", "role": "admin"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
_assert_user([response.json()], "2", role="admin")
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(data, "2", role="admin")
|
||||
|
||||
# Get (empty) user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/settings"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/user/settings/update"),
|
||||
json={
|
||||
"ui": {"attr1": "value1", "attr2": "value2"},
|
||||
"model_config": {"attr3": "value3", "attr4": "value4"},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user settings
|
||||
with mock_webui_user(id="2"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/settings"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"ui": {"attr1": "value1", "attr2": "value2"},
|
||||
"model_config": {"attr3": "value3", "attr4": "value4"},
|
||||
}
|
||||
|
||||
# Get (empty) user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/info"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() is None
|
||||
|
||||
# Update user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/user/info/update"),
|
||||
json={"attr1": "value1", "attr2": "value2"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get user info
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/user/info"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"attr1": "value1", "attr2": "value2"}
|
||||
|
||||
# Get user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.get(self.create_url("/2"))
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"name": "user 2", "profile_image_url": "/user2.png"}
|
||||
|
||||
# Update user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.post(
|
||||
self.create_url("/2/update"),
|
||||
json={
|
||||
"name": "user 2 updated",
|
||||
"email": "user2-updated@openwebui.com",
|
||||
"profile_image_url": "/user2-updated.png",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 2
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
||||
_assert_user(
|
||||
data,
|
||||
"2",
|
||||
role="admin",
|
||||
name="user 2 updated",
|
||||
email="user2-updated@openwebui.com",
|
||||
profile_image_url="/user2-updated.png",
|
||||
)
|
||||
|
||||
# Delete user by id
|
||||
with mock_webui_user(id="1"):
|
||||
response = self.fast_api_client.delete(self.create_url("/2"))
|
||||
assert response.status_code == 200
|
||||
|
||||
# Get all users
|
||||
with mock_webui_user(id="3"):
|
||||
response = self.fast_api_client.get(self.create_url(""))
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == 1
|
||||
data = response.json()
|
||||
_assert_user(data, "1")
|
161
backend/test/util/abstract_integration_test.py
Normal file
161
backend/test/util/abstract_integration_test.py
Normal file
@ -0,0 +1,161 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import docker
|
||||
import pytest
|
||||
from docker import DockerClient
|
||||
from pytest_docker.plugin import get_docker_ip
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import text, create_engine
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_fast_api_client():
|
||||
from main import app
|
||||
|
||||
with TestClient(app) as c:
|
||||
return c
|
||||
|
||||
|
||||
class AbstractIntegrationTest:
|
||||
BASE_PATH = None
|
||||
|
||||
def create_url(self, path="", query_params=None):
|
||||
if self.BASE_PATH is None:
|
||||
raise Exception("BASE_PATH is not set")
|
||||
parts = self.BASE_PATH.split("/")
|
||||
parts = [part.strip() for part in parts if part.strip() != ""]
|
||||
path_parts = path.split("/")
|
||||
path_parts = [part.strip() for part in path_parts if part.strip() != ""]
|
||||
query_parts = ""
|
||||
if query_params:
|
||||
query_parts = "&".join(
|
||||
[f"{key}={value}" for key, value in query_params.items()]
|
||||
)
|
||||
query_parts = f"?{query_parts}"
|
||||
return "/".join(parts + path_parts) + query_parts
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
pass
|
||||
|
||||
def setup_method(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls):
|
||||
pass
|
||||
|
||||
def teardown_method(self):
|
||||
pass
|
||||
|
||||
|
||||
class AbstractPostgresTest(AbstractIntegrationTest):
|
||||
DOCKER_CONTAINER_NAME = "postgres-test-container-will-get-deleted"
|
||||
docker_client: DockerClient
|
||||
|
||||
@classmethod
|
||||
def _create_db_url(cls, env_vars_postgres: dict) -> str:
|
||||
host = get_docker_ip()
|
||||
user = env_vars_postgres["POSTGRES_USER"]
|
||||
pw = env_vars_postgres["POSTGRES_PASSWORD"]
|
||||
port = 8081
|
||||
db = env_vars_postgres["POSTGRES_DB"]
|
||||
return f"postgresql://{user}:{pw}@{host}:{port}/{db}"
|
||||
|
||||
@classmethod
|
||||
def setup_class(cls):
|
||||
super().setup_class()
|
||||
try:
|
||||
env_vars_postgres = {
|
||||
"POSTGRES_USER": "user",
|
||||
"POSTGRES_PASSWORD": "example",
|
||||
"POSTGRES_DB": "openwebui",
|
||||
}
|
||||
cls.docker_client = docker.from_env()
|
||||
cls.docker_client.containers.run(
|
||||
"postgres:16.2",
|
||||
detach=True,
|
||||
environment=env_vars_postgres,
|
||||
name=cls.DOCKER_CONTAINER_NAME,
|
||||
ports={5432: ("0.0.0.0", 8081)},
|
||||
command="postgres -c log_statement=all",
|
||||
)
|
||||
time.sleep(0.5)
|
||||
|
||||
database_url = cls._create_db_url(env_vars_postgres)
|
||||
os.environ["DATABASE_URL"] = database_url
|
||||
retries = 10
|
||||
db = None
|
||||
while retries > 0:
|
||||
try:
|
||||
from config import BACKEND_DIR
|
||||
|
||||
db = create_engine(database_url, pool_pre_ping=True)
|
||||
db = db.connect()
|
||||
log.info("postgres is ready!")
|
||||
break
|
||||
except Exception as e:
|
||||
log.warning(e)
|
||||
time.sleep(3)
|
||||
retries -= 1
|
||||
|
||||
if db:
|
||||
# import must be after setting env!
|
||||
cls.fast_api_client = get_fast_api_client()
|
||||
db.close()
|
||||
else:
|
||||
raise Exception("Could not connect to Postgres")
|
||||
except Exception as ex:
|
||||
log.error(ex)
|
||||
cls.teardown_class()
|
||||
pytest.fail(f"Could not setup test environment: {ex}")
|
||||
|
||||
def _check_db_connection(self):
|
||||
from apps.webui.internal.db import Session
|
||||
|
||||
retries = 10
|
||||
while retries > 0:
|
||||
try:
|
||||
Session.execute(text("SELECT 1"))
|
||||
Session.commit()
|
||||
break
|
||||
except Exception as e:
|
||||
Session.rollback()
|
||||
log.warning(e)
|
||||
time.sleep(3)
|
||||
retries -= 1
|
||||
|
||||
def setup_method(self):
|
||||
super().setup_method()
|
||||
self._check_db_connection()
|
||||
|
||||
@classmethod
|
||||
def teardown_class(cls) -> None:
|
||||
super().teardown_class()
|
||||
cls.docker_client.containers.get(cls.DOCKER_CONTAINER_NAME).remove(force=True)
|
||||
|
||||
def teardown_method(self):
|
||||
from apps.webui.internal.db import Session
|
||||
|
||||
# rollback everything not yet committed
|
||||
Session.commit()
|
||||
|
||||
# truncate all tables
|
||||
tables = [
|
||||
"auth",
|
||||
"chat",
|
||||
"chatidtag",
|
||||
"document",
|
||||
"memory",
|
||||
"model",
|
||||
"prompt",
|
||||
"tag",
|
||||
'"user"',
|
||||
]
|
||||
for table in tables:
|
||||
Session.execute(text(f"TRUNCATE TABLE {table}"))
|
||||
Session.commit()
|
45
backend/test/util/mock_user.py
Normal file
45
backend/test/util/mock_user.py
Normal file
@ -0,0 +1,45 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mock_webui_user(**kwargs):
|
||||
from apps.webui.main import app
|
||||
|
||||
with mock_user(app, **kwargs):
|
||||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mock_user(app: FastAPI, **kwargs):
|
||||
from utils.utils import (
|
||||
get_current_user,
|
||||
get_verified_user,
|
||||
get_admin_user,
|
||||
get_current_user_by_api_key,
|
||||
)
|
||||
from apps.webui.models.users import User
|
||||
|
||||
def create_user():
|
||||
user_parameters = {
|
||||
"id": "1",
|
||||
"name": "John Doe",
|
||||
"email": "john.doe@openwebui.com",
|
||||
"role": "user",
|
||||
"profile_image_url": "/user.png",
|
||||
"last_active_at": 1627351200,
|
||||
"updated_at": 1627351200,
|
||||
"created_at": 162735120,
|
||||
**kwargs,
|
||||
}
|
||||
return User(**user_parameters)
|
||||
|
||||
app.dependency_overrides = {
|
||||
get_current_user: create_user,
|
||||
get_verified_user: create_user,
|
||||
get_admin_user: create_user,
|
||||
get_current_user_by_api_key: create_user,
|
||||
}
|
||||
yield
|
||||
app.dependency_overrides = {}
|
@ -1,5 +1,6 @@
|
||||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
|
||||
from fastapi import HTTPException, status, Depends, Request
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from apps.webui.models.users import Users
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user