Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
open-webui
Commits
df09d083
Commit
df09d083
authored
Jun 18, 2024
by
Jonathan Rohde
Browse files
feat(sqlalchemy): Replace peewee with sqlalchemy
parent
8dac2a21
Changes
47
Show whitespace changes
Inline
Side-by-side
Showing
20 changed files
with
860 additions
and
804 deletions
+860
-804
.github/workflows/integration-test.yml
.github/workflows/integration-test.yml
+2
-2
backend/alembic.ini
backend/alembic.ini
+114
-0
backend/apps/ollama/main.py
backend/apps/ollama/main.py
+5
-2
backend/apps/openai/main.py
backend/apps/openai/main.py
+3
-1
backend/apps/socket/main.py
backend/apps/socket/main.py
+3
-1
backend/apps/webui/internal/db.py
backend/apps/webui/internal/db.py
+40
-26
backend/apps/webui/internal/wrappers.py
backend/apps/webui/internal/wrappers.py
+0
-72
backend/apps/webui/main.py
backend/apps/webui/main.py
+3
-3
backend/apps/webui/models/auths.py
backend/apps/webui/models/auths.py
+40
-37
backend/apps/webui/models/chats.py
backend/apps/webui/models/chats.py
+139
-162
backend/apps/webui/models/documents.py
backend/apps/webui/models/documents.py
+48
-56
backend/apps/webui/models/files.py
backend/apps/webui/models/files.py
+30
-32
backend/apps/webui/models/functions.py
backend/apps/webui/models/functions.py
+36
-38
backend/apps/webui/models/memories.py
backend/apps/webui/models/memories.py
+43
-44
backend/apps/webui/models/models.py
backend/apps/webui/models/models.py
+38
-41
backend/apps/webui/models/prompts.py
backend/apps/webui/models/prompts.py
+38
-48
backend/apps/webui/models/tags.py
backend/apps/webui/models/tags.py
+109
-89
backend/apps/webui/models/tools.py
backend/apps/webui/models/tools.py
+37
-41
backend/apps/webui/models/users.py
backend/apps/webui/models/users.py
+92
-83
backend/apps/webui/routers/auths.py
backend/apps/webui/routers/auths.py
+40
-26
No files found.
.github/workflows/integration-test.yml
View file @
df09d083
...
@@ -171,7 +171,7 @@ jobs:
...
@@ -171,7 +171,7 @@ jobs:
fi
fi
# Check that service will reconnect to postgres when connection will be closed
# 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
if [[ "$status_code" -ne 200 ]] ; then
echo "Server has failed before postgres reconnect check"
echo "Server has failed before postgres reconnect check"
exit 1
exit 1
...
@@ -183,7 +183,7 @@ jobs:
...
@@ -183,7 +183,7 @@ jobs:
cur = conn.cursor(); \
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();')"
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
if [[ "$status_code" -ne 200 ]] ; then
echo "Server has not reconnected to postgres after connection was closed: returned status $status_code"
echo "Server has not reconnected to postgres after connection was closed: returned status $status_code"
exit 1
exit 1
...
...
backend/alembic.ini
0 → 100644
View file @
df09d083
# 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
backend/apps/ollama/main.py
View file @
df09d083
...
@@ -31,6 +31,7 @@ from typing import Optional, List, Union
...
@@ -31,6 +31,7 @@ from typing import Optional, List, Union
from
starlette.background
import
BackgroundTask
from
starlette.background
import
BackgroundTask
from
apps.webui.internal.db
import
get_db
from
apps.webui.models.models
import
Models
from
apps.webui.models.models
import
Models
from
apps.webui.models.users
import
Users
from
apps.webui.models.users
import
Users
from
constants
import
ERROR_MESSAGES
from
constants
import
ERROR_MESSAGES
...
@@ -711,6 +712,7 @@ async def generate_chat_completion(
...
@@ -711,6 +712,7 @@ async def generate_chat_completion(
form_data
:
GenerateChatCompletionForm
,
form_data
:
GenerateChatCompletionForm
,
url_idx
:
Optional
[
int
]
=
None
,
url_idx
:
Optional
[
int
]
=
None
,
user
=
Depends
(
get_verified_user
),
user
=
Depends
(
get_verified_user
),
db
=
Depends
(
get_db
),
):
):
log
.
debug
(
log
.
debug
(
...
@@ -724,7 +726,7 @@ async def generate_chat_completion(
...
@@ -724,7 +726,7 @@ async def generate_chat_completion(
}
}
model_id
=
form_data
.
model
model_id
=
form_data
.
model
model_info
=
Models
.
get_model_by_id
(
model_id
)
model_info
=
Models
.
get_model_by_id
(
db
,
model_id
)
if
model_info
:
if
model_info
:
if
model_info
.
base_model_id
:
if
model_info
.
base_model_id
:
...
@@ -883,6 +885,7 @@ async def generate_openai_chat_completion(
...
@@ -883,6 +885,7 @@ async def generate_openai_chat_completion(
form_data
:
dict
,
form_data
:
dict
,
url_idx
:
Optional
[
int
]
=
None
,
url_idx
:
Optional
[
int
]
=
None
,
user
=
Depends
(
get_verified_user
),
user
=
Depends
(
get_verified_user
),
db
=
Depends
(
get_db
),
):
):
form_data
=
OpenAIChatCompletionForm
(
**
form_data
)
form_data
=
OpenAIChatCompletionForm
(
**
form_data
)
...
@@ -891,7 +894,7 @@ async def generate_openai_chat_completion(
...
@@ -891,7 +894,7 @@ async def generate_openai_chat_completion(
}
}
model_id
=
form_data
.
model
model_id
=
form_data
.
model
model_info
=
Models
.
get_model_by_id
(
model_id
)
model_info
=
Models
.
get_model_by_id
(
db
,
model_id
)
if
model_info
:
if
model_info
:
if
model_info
.
base_model_id
:
if
model_info
.
base_model_id
:
...
...
backend/apps/openai/main.py
View file @
df09d083
...
@@ -11,6 +11,7 @@ import logging
...
@@ -11,6 +11,7 @@ import logging
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
from
starlette.background
import
BackgroundTask
from
starlette.background
import
BackgroundTask
from
apps.webui.internal.db
import
get_db
from
apps.webui.models.models
import
Models
from
apps.webui.models.models
import
Models
from
apps.webui.models.users
import
Users
from
apps.webui.models.users
import
Users
from
constants
import
ERROR_MESSAGES
from
constants
import
ERROR_MESSAGES
...
@@ -353,12 +354,13 @@ async def generate_chat_completion(
...
@@ -353,12 +354,13 @@ async def generate_chat_completion(
form_data
:
dict
,
form_data
:
dict
,
url_idx
:
Optional
[
int
]
=
None
,
url_idx
:
Optional
[
int
]
=
None
,
user
=
Depends
(
get_verified_user
),
user
=
Depends
(
get_verified_user
),
db
=
Depends
(
get_db
),
):
):
idx
=
0
idx
=
0
payload
=
{
**
form_data
}
payload
=
{
**
form_data
}
model_id
=
form_data
.
get
(
"model"
)
model_id
=
form_data
.
get
(
"model"
)
model_info
=
Models
.
get_model_by_id
(
model_id
)
model_info
=
Models
.
get_model_by_id
(
db
,
model_id
)
if
model_info
:
if
model_info
:
if
model_info
.
base_model_id
:
if
model_info
.
base_model_id
:
...
...
backend/apps/socket/main.py
View file @
df09d083
...
@@ -24,7 +24,9 @@ async def connect(sid, environ, auth):
...
@@ -24,7 +24,9 @@ async def connect(sid, environ, auth):
data
=
decode_token
(
auth
[
"token"
])
data
=
decode_token
(
auth
[
"token"
])
if
data
is
not
None
and
"id"
in
data
:
if
data
is
not
None
and
"id"
in
data
:
user
=
Users
.
get_user_by_id
(
data
[
"id"
])
from
apps.webui.internal.db
import
SessionLocal
user
=
Users
.
get_user_by_id
(
SessionLocal
(),
data
[
"id"
])
if
user
:
if
user
:
SESSION_POOL
[
sid
]
=
user
.
id
SESSION_POOL
[
sid
]
=
user
.
id
...
...
backend/apps/webui/internal/db.py
View file @
df09d083
import
os
import
os
import
logging
import
logging
import
json
import
json
from
typing
import
Optional
,
Any
from
typing_extensions
import
Self
from
peewee
import
*
from
sqlalchemy
import
create_engine
,
types
,
Dialect
from
peewee_migrate
import
Router
from
sqlalchemy.ext.declarative
import
declarative_base
from
sqlalchemy.orm
import
sessionmaker
from
sqlalchemy.sql.type_api
import
_T
from
apps.webui.internal.wrappers
import
register_connection
from
config
import
SRC_LOG_LEVELS
,
DATA_DIR
,
DATABASE_URL
,
BACKEND_DIR
from
config
import
SRC_LOG_LEVELS
,
DATA_DIR
,
DATABASE_URL
,
BACKEND_DIR
log
=
logging
.
getLogger
(
__name__
)
log
=
logging
.
getLogger
(
__name__
)
log
.
setLevel
(
SRC_LOG_LEVELS
[
"DB"
])
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
):
def
db_value
(
self
,
value
):
return
json
.
dumps
(
value
)
return
json
.
dumps
(
value
)
...
@@ -29,26 +45,24 @@ if os.path.exists(f"{DATA_DIR}/ollama.db"):
...
@@ -29,26 +45,24 @@ if os.path.exists(f"{DATA_DIR}/ollama.db"):
else
:
else
:
pass
pass
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
)
Base
=
declarative_base
()
# The `register_connection` function encapsulates the logic for setting up
# the database connection based on the connection string, while `connect`
def
get_db
():
# is a Peewee-specific method to manage the connection state and avoid errors
db
=
SessionLocal
()
# when a connection is already open.
try
:
try
:
yield
db
DB
=
register_connection
(
DATABASE_URL
)
db
.
commit
()
log
.
info
(
f
"Connected to a
{
DB
.
__class__
.
__name__
}
database."
)
except
Exception
as
e
:
except
Exception
as
e
:
db
.
rollback
()
log
.
error
(
f
"Failed to initialize the database connection:
{
e
}
"
)
raise
e
raise
finally
:
db
.
close
()
router
=
Router
(
DB
,
migrate_dir
=
BACKEND_DIR
/
"apps"
/
"webui"
/
"internal"
/
"migrations"
,
logger
=
log
,
)
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
backend/apps/webui/internal/wrappers.py
deleted
100644 → 0
View file @
8dac2a21
from
contextvars
import
ContextVar
from
peewee
import
*
from
peewee
import
PostgresqlDatabase
,
InterfaceError
as
PeeWeeInterfaceError
import
logging
from
playhouse.db_url
import
connect
,
parse
from
playhouse.shortcuts
import
ReconnectMixin
from
config
import
SRC_LOG_LEVELS
log
=
logging
.
getLogger
(
__name__
)
log
.
setLevel
(
SRC_LOG_LEVELS
[
"DB"
])
db_state_default
=
{
"closed"
:
None
,
"conn"
:
None
,
"ctx"
:
None
,
"transactions"
:
None
}
db_state
=
ContextVar
(
"db_state"
,
default
=
db_state_default
.
copy
())
class
PeeweeConnectionState
(
object
):
def
__init__
(
self
,
**
kwargs
):
super
().
__setattr__
(
"_state"
,
db_state
)
super
().
__init__
(
**
kwargs
)
def
__setattr__
(
self
,
name
,
value
):
self
.
_state
.
get
()[
name
]
=
value
def
__getattr__
(
self
,
name
):
value
=
self
.
_state
.
get
()[
name
]
return
value
class
CustomReconnectMixin
(
ReconnectMixin
):
reconnect_errors
=
(
# psycopg2
(
OperationalError
,
"termin"
),
(
InterfaceError
,
"closed"
),
# peewee
(
PeeWeeInterfaceError
,
"closed"
),
)
class
ReconnectingPostgresqlDatabase
(
CustomReconnectMixin
,
PostgresqlDatabase
):
pass
def
register_connection
(
db_url
):
db
=
connect
(
db_url
)
if
isinstance
(
db
,
PostgresqlDatabase
):
# Enable autoconnect for SQLite databases, managed by Peewee
db
.
autoconnect
=
True
db
.
reuse_if_open
=
True
log
.
info
(
"Connected to PostgreSQL database"
)
# Get the connection details
connection
=
parse
(
db_url
)
# Use our custom database class that supports reconnection
db
=
ReconnectingPostgresqlDatabase
(
connection
[
"database"
],
user
=
connection
[
"user"
],
password
=
connection
[
"password"
],
host
=
connection
[
"host"
],
port
=
connection
[
"port"
],
)
db
.
connect
(
reuse_if_open
=
True
)
elif
isinstance
(
db
,
SqliteDatabase
):
# Enable autoconnect for SQLite databases, managed by Peewee
db
.
autoconnect
=
True
db
.
reuse_if_open
=
True
log
.
info
(
"Connected to SQLite database"
)
else
:
raise
ValueError
(
"Unsupported database connection"
)
return
db
backend/apps/webui/main.py
View file @
df09d083
...
@@ -3,7 +3,7 @@ from fastapi.routing import APIRoute
...
@@ -3,7 +3,7 @@ from fastapi.routing import APIRoute
from
fastapi.responses
import
StreamingResponse
from
fastapi.responses
import
StreamingResponse
from
fastapi.middleware.cors
import
CORSMiddleware
from
fastapi.middleware.cors
import
CORSMiddleware
from
starlette.middleware.sessions
import
SessionMiddleware
from
starlette.middleware.sessions
import
SessionMiddleware
from
sqlalchemy.orm
import
Session
from
apps.webui.routers
import
(
from
apps.webui.routers
import
(
auths
,
auths
,
users
,
users
,
...
@@ -114,8 +114,8 @@ async def get_status():
...
@@ -114,8 +114,8 @@ async def get_status():
}
}
async
def
get_pipe_models
():
async
def
get_pipe_models
(
db
:
Session
):
pipes
=
Functions
.
get_functions_by_type
(
"pipe"
,
active_only
=
True
)
pipes
=
Functions
.
get_functions_by_type
(
db
,
"pipe"
,
active_only
=
True
)
pipe_models
=
[]
pipe_models
=
[]
for
pipe
in
pipes
:
for
pipe
in
pipes
:
...
...
backend/apps/webui/models/auths.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
from
typing
import
List
,
Union
,
Optional
from
typing
import
Optional
import
time
import
uuid
import
uuid
import
logging
import
logging
from
peewee
import
*
from
sqlalchemy
import
String
,
Column
,
Boolean
from
sqlalchemy.orm
import
Session
from
apps.webui.models.users
import
UserModel
,
Users
from
apps.webui.models.users
import
UserModel
,
Users
from
utils.utils
import
verify_password
from
utils.utils
import
verify_password
from
apps.webui.internal.db
import
D
B
from
apps.webui.internal.db
import
B
ase
from
config
import
SRC_LOG_LEVELS
from
config
import
SRC_LOG_LEVELS
...
@@ -20,14 +20,13 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -20,14 +20,13 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
Auth
(
Model
):
class
Auth
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"auth"
email
=
CharField
()
password
=
TextField
()
active
=
BooleanField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
email
=
Column
(
String
)
password
=
Column
(
String
)
active
=
Column
(
Boolean
)
class
AuthModel
(
BaseModel
):
class
AuthModel
(
BaseModel
):
...
@@ -94,12 +93,10 @@ class AddUserForm(SignupForm):
...
@@ -94,12 +93,10 @@ class AddUserForm(SignupForm):
class
AuthsTable
:
class
AuthsTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Auth
])
def
insert_new_auth
(
def
insert_new_auth
(
self
,
self
,
db
:
Session
,
email
:
str
,
email
:
str
,
password
:
str
,
password
:
str
,
name
:
str
,
name
:
str
,
...
@@ -114,24 +111,30 @@ class AuthsTable:
...
@@ -114,24 +111,30 @@ class AuthsTable:
auth
=
AuthModel
(
auth
=
AuthModel
(
**
{
"id"
:
id
,
"email"
:
email
,
"password"
:
password
,
"active"
:
True
}
**
{
"id"
:
id
,
"email"
:
email
,
"password"
:
password
,
"active"
:
True
}
)
)
result
=
Auth
.
create
(
**
auth
.
model_dump
())
result
=
Auth
(
**
auth
.
model_dump
())
db
.
add
(
result
)
user
=
Users
.
insert_new_user
(
user
=
Users
.
insert_new_user
(
id
,
name
,
email
,
profile_image_url
,
role
,
oauth_sub
db
,
id
,
name
,
email
,
profile_image_url
,
role
,
oauth_sub
)
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
and
user
:
if
result
and
user
:
return
user
return
user
else
:
else
:
return
None
return
None
def
authenticate_user
(
self
,
email
:
str
,
password
:
str
)
->
Optional
[
UserModel
]:
def
authenticate_user
(
self
,
db
:
Session
,
email
:
str
,
password
:
str
)
->
Optional
[
UserModel
]:
log
.
info
(
f
"authenticate_user:
{
email
}
"
)
log
.
info
(
f
"authenticate_user:
{
email
}
"
)
try
:
try
:
auth
=
Auth
.
get
(
Auth
.
email
==
email
,
Auth
.
active
==
True
)
auth
=
db
.
query
(
Auth
).
filter_by
(
email
=
email
,
active
=
True
)
.
first
()
if
auth
:
if
auth
:
if
verify_password
(
password
,
auth
.
password
):
if
verify_password
(
password
,
auth
.
password
):
user
=
Users
.
get_user_by_id
(
auth
.
id
)
user
=
Users
.
get_user_by_id
(
db
,
auth
.
id
)
return
user
return
user
else
:
else
:
return
None
return
None
...
@@ -140,55 +143,55 @@ class AuthsTable:
...
@@ -140,55 +143,55 @@ class AuthsTable:
except
:
except
:
return
None
return
None
def
authenticate_user_by_api_key
(
self
,
api_key
:
str
)
->
Optional
[
UserModel
]:
def
authenticate_user_by_api_key
(
self
,
db
:
Session
,
api_key
:
str
)
->
Optional
[
UserModel
]:
log
.
info
(
f
"authenticate_user_by_api_key:
{
api_key
}
"
)
log
.
info
(
f
"authenticate_user_by_api_key:
{
api_key
}
"
)
# if no api_key, return None
# if no api_key, return None
if
not
api_key
:
if
not
api_key
:
return
None
return
None
try
:
try
:
user
=
Users
.
get_user_by_api_key
(
api_key
)
user
=
Users
.
get_user_by_api_key
(
db
,
api_key
)
return
user
if
user
else
None
return
user
if
user
else
None
except
:
except
:
return
False
return
False
def
authenticate_user_by_trusted_header
(
self
,
email
:
str
)
->
Optional
[
UserModel
]:
def
authenticate_user_by_trusted_header
(
self
,
db
:
Session
,
email
:
str
)
->
Optional
[
UserModel
]:
log
.
info
(
f
"authenticate_user_by_trusted_header:
{
email
}
"
)
log
.
info
(
f
"authenticate_user_by_trusted_header:
{
email
}
"
)
try
:
try
:
auth
=
Auth
.
get
(
Auth
.
email
==
email
,
Auth
.
active
==
True
)
auth
=
db
.
query
(
Auth
).
filter
(
email
=
email
,
active
=
True
)
.
first
()
if
auth
:
if
auth
:
user
=
Users
.
get_user_by_id
(
auth
.
id
)
user
=
Users
.
get_user_by_id
(
auth
.
id
)
return
user
return
user
except
:
except
:
return
None
return
None
def
update_user_password_by_id
(
self
,
id
:
str
,
new_password
:
str
)
->
bool
:
def
update_user_password_by_id
(
self
,
db
:
Session
,
id
:
str
,
new_password
:
str
)
->
bool
:
try
:
try
:
query
=
Auth
.
update
(
password
=
new_password
).
where
(
Auth
.
id
==
id
)
result
=
db
.
query
(
Auth
).
filter_by
(
id
=
id
).
update
({
"password"
:
new_password
})
result
=
query
.
execute
()
return
True
if
result
==
1
else
False
return
True
if
result
==
1
else
False
except
:
except
:
return
False
return
False
def
update_email_by_id
(
self
,
id
:
str
,
email
:
str
)
->
bool
:
def
update_email_by_id
(
self
,
db
:
Session
,
id
:
str
,
email
:
str
)
->
bool
:
try
:
try
:
query
=
Auth
.
update
(
email
=
email
).
where
(
Auth
.
id
==
id
)
result
=
db
.
query
(
Auth
).
filter_by
(
id
=
id
).
update
({
"email"
:
email
})
result
=
query
.
execute
()
return
True
if
result
==
1
else
False
return
True
if
result
==
1
else
False
except
:
except
:
return
False
return
False
def
delete_auth_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_auth_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
# Delete User
# Delete User
result
=
Users
.
delete_user_by_id
(
id
)
result
=
Users
.
delete_user_by_id
(
db
,
id
)
if
result
:
if
result
:
# Delete Auth
db
.
query
(
Auth
).
filter_by
(
id
=
id
).
delete
()
query
=
Auth
.
delete
().
where
(
Auth
.
id
==
id
)
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
else
:
else
:
...
@@ -197,4 +200,4 @@ class AuthsTable:
...
@@ -197,4 +200,4 @@ class AuthsTable:
return
False
return
False
Auths
=
AuthsTable
(
DB
)
Auths
=
AuthsTable
()
backend/apps/webui/models/chats.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
import
json
import
json
import
uuid
import
uuid
import
time
import
time
from
apps.webui.internal.db
import
DB
from
sqlalchemy
import
Column
,
String
,
BigInteger
,
Boolean
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
Base
####################
####################
# Chat DB Schema
# Chat DB Schema
####################
####################
class
Chat
(
Model
):
class
Chat
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"chat"
user_id
=
CharField
()
title
=
TextField
()
chat
=
TextField
()
# Save Chat JSON as Text
created_at
=
BigIntegerField
()
id
=
Column
(
String
,
primary_key
=
True
)
updated_at
=
BigIntegerField
()
user_id
=
Column
(
String
)
title
=
Column
(
String
)
chat
=
Column
(
String
)
# Save Chat JSON as Text
share_id
=
CharField
(
null
=
True
,
unique
=
True
)
created_at
=
Column
(
BigInteger
)
archived
=
BooleanField
(
default
=
False
)
updated_at
=
Column
(
BigInteger
)
class
Meta
:
share_id
=
Column
(
String
,
unique
=
True
,
nullable
=
True
)
database
=
DB
archived
=
Column
(
Boolean
,
default
=
False
)
class
ChatModel
(
BaseModel
):
class
ChatModel
(
BaseModel
):
model_config
=
ConfigDict
(
from_attributes
=
True
)
id
:
str
id
:
str
user_id
:
str
user_id
:
str
title
:
str
title
:
str
...
@@ -75,11 +78,10 @@ class ChatTitleIdResponse(BaseModel):
...
@@ -75,11 +78,10 @@ class ChatTitleIdResponse(BaseModel):
class
ChatTable
:
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
]:
def
insert_new_chat
(
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
ChatForm
)
->
Optional
[
ChatModel
]:
id
=
str
(
uuid
.
uuid4
())
id
=
str
(
uuid
.
uuid4
())
chat
=
ChatModel
(
chat
=
ChatModel
(
**
{
**
{
...
@@ -94,29 +96,36 @@ class ChatTable:
...
@@ -94,29 +96,36 @@ class ChatTable:
}
}
)
)
result
=
Chat
.
create
(
**
chat
.
model_dump
())
result
=
Chat
(
**
chat
.
model_dump
())
return
chat
if
result
else
None
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
]:
def
update_chat_by_id
(
self
,
db
:
Session
,
id
:
str
,
chat
:
dict
)
->
Optional
[
ChatModel
]:
try
:
try
:
query
=
Chat
.
update
(
db
.
query
(
Chat
).
filter_by
(
id
=
id
)
.
update
(
chat
=
json
.
dumps
(
chat
),
{
title
=
chat
[
"title"
]
if
"title"
in
chat
else
"New C
hat
"
,
"chat"
:
json
.
dumps
(
c
hat
)
,
updated_at
=
int
(
time
.
time
())
,
"title"
:
chat
[
"title"
]
if
"title"
in
chat
else
"New Chat"
,
).
where
(
Chat
.
id
==
id
)
"updated_at"
:
int
(
time
.
time
()),
query
.
execute
()
}
)
chat
=
Chat
.
get
(
Chat
.
id
==
id
)
return
ChatModel
(
**
model_to_dict
(
chat
)
)
return
self
.
get_chat_by_id
(
db
,
id
)
except
:
except
:
return
None
return
None
def
insert_shared_chat_by_chat_id
(
self
,
chat_id
:
str
)
->
Optional
[
ChatModel
]:
def
insert_shared_chat_by_chat_id
(
self
,
db
:
Session
,
chat_id
:
str
)
->
Optional
[
ChatModel
]:
# Get the existing chat to share
# Get the existing chat to share
chat
=
Chat
.
get
(
Chat
.
id
==
chat_id
)
chat
=
db
.
get
(
Chat
,
chat_id
)
# Check if the chat is already shared
# Check if the chat is already shared
if
chat
.
share_id
:
if
chat
.
share_id
:
return
self
.
get_chat_by_id_and_user_id
(
chat
.
share_id
,
"shared"
)
return
self
.
get_chat_by_id_and_user_id
(
db
,
chat
.
share_id
,
"shared"
)
# Create a new chat with the same data, but with a new ID
# Create a new chat with the same data, but with a new ID
shared_chat
=
ChatModel
(
shared_chat
=
ChatModel
(
**
{
**
{
...
@@ -128,228 +137,196 @@ class ChatTable:
...
@@ -128,228 +137,196 @@ class ChatTable:
"updated_at"
:
int
(
time
.
time
()),
"updated_at"
:
int
(
time
.
time
()),
}
}
)
)
shared_result
=
Chat
.
create
(
**
shared_chat
.
model_dump
())
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
# Update the original chat with the share_id
result
=
(
result
=
(
Chat
.
update
(
share_id
=
shared_chat
.
id
).
where
(
Chat
.
id
==
chat_id
).
execute
(
)
db
.
query
(
Chat
).
filter_by
(
id
=
chat_id
)
.
update
(
{
"
share_id
"
:
shared_chat
.
id
}
)
)
)
return
shared_chat
if
(
shared_result
and
result
)
else
None
return
shared_chat
if
(
shared_result
and
result
)
else
None
def
update_shared_chat_by_chat_id
(
self
,
chat_id
:
str
)
->
Optional
[
ChatModel
]:
def
update_shared_chat_by_chat_id
(
self
,
db
:
Session
,
chat_id
:
str
)
->
Optional
[
ChatModel
]:
try
:
try
:
print
(
"update_shared_chat_by_id"
)
print
(
"update_shared_chat_by_id"
)
chat
=
Chat
.
get
(
Chat
.
id
==
chat_id
)
chat
=
db
.
get
(
Chat
,
chat_id
)
print
(
chat
)
print
(
chat
)
query
=
Chat
.
update
(
db
.
query
(
Chat
).
filter_by
(
id
=
chat
.
share_id
).
update
(
title
=
chat
.
title
,
{
"title"
:
chat
.
title
,
"chat"
:
chat
.
chat
}
chat
=
chat
.
chat
,
)
).
where
(
Chat
.
id
==
chat
.
share_id
)
query
.
execute
()
chat
=
Chat
.
get
(
Chat
.
id
==
chat
.
share_id
)
return
self
.
get_chat_by_id
(
db
,
chat
.
share_id
)
return
ChatModel
(
**
model_to_dict
(
chat
))
except
:
except
:
return
None
return
None
def
delete_shared_chat_by_chat_id
(
self
,
chat_id
:
str
)
->
bool
:
def
delete_shared_chat_by_chat_id
(
self
,
db
:
Session
,
chat_id
:
str
)
->
bool
:
try
:
try
:
query
=
Chat
.
delete
().
where
(
Chat
.
user_id
==
f
"shared-
{
chat_id
}
"
)
db
.
query
(
Chat
).
filter_by
(
user_id
=
f
"shared-
{
chat_id
}
"
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
def
update_chat_share_id_by_id
(
def
update_chat_share_id_by_id
(
self
,
id
:
str
,
share_id
:
Optional
[
str
]
self
,
db
:
Session
,
id
:
str
,
share_id
:
Optional
[
str
]
)
->
Optional
[
ChatModel
]:
)
->
Optional
[
ChatModel
]:
try
:
try
:
query
=
Chat
.
update
(
db
.
query
(
Chat
).
filter_by
(
id
=
id
).
update
({
"share_id"
:
share_id
})
share_id
=
share_id
,
).
where
(
Chat
.
id
==
id
)
query
.
execute
()
chat
=
Chat
.
get
(
Chat
.
id
==
id
)
return
self
.
get_chat_by_id
(
db
,
id
)
return
ChatModel
(
**
model_to_dict
(
chat
))
except
:
except
:
return
None
return
None
def
toggle_chat_archive_by_id
(
self
,
id
:
str
)
->
Optional
[
ChatModel
]:
def
toggle_chat_archive_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
ChatModel
]:
try
:
try
:
chat
=
self
.
get_chat_by_id
(
id
)
chat
=
self
.
get_chat_by_id
(
db
,
id
)
query
=
Chat
.
update
(
db
.
query
(
Chat
).
filter_by
(
id
=
id
).
update
({
"archived"
:
not
chat
.
archived
})
archived
=
(
not
chat
.
archived
),
).
where
(
Chat
.
id
==
id
)
query
.
execute
()
return
self
.
get_chat_by_id
(
db
,
id
)
chat
=
Chat
.
get
(
Chat
.
id
==
id
)
return
ChatModel
(
**
model_to_dict
(
chat
))
except
:
except
:
return
None
return
None
def
archive_all_chats_by_user_id
(
self
,
user_id
:
str
)
->
bool
:
def
archive_all_chats_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
bool
:
try
:
try
:
chats
=
self
.
get_chats_by_user_id
(
user_id
)
db
.
query
(
Chat
).
filter_by
(
user_id
=
user_id
).
update
({
"archived"
:
True
})
for
chat
in
chats
:
query
=
Chat
.
update
(
archived
=
True
,
).
where
(
Chat
.
id
==
chat
.
id
)
query
.
execute
()
return
True
return
True
except
:
except
:
return
False
return
False
def
get_archived_chat_list_by_user_id
(
def
get_archived_chat_list_by_user_id
(
self
,
user_id
:
str
,
skip
:
int
=
0
,
limit
:
int
=
50
self
,
db
:
Session
,
user_id
:
str
,
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
ChatModel
]:
)
->
List
[
ChatModel
]:
return
[
all_chats
=
(
ChatModel
(
**
model_to_dict
(
chat
))
db
.
query
(
Chat
)
for
chat
in
Chat
.
select
()
.
filter_by
(
user_id
=
user_id
,
archived
=
True
)
.
where
(
Chat
.
archived
==
True
)
.
where
(
Chat
.
user_id
==
user_id
)
.
order_by
(
Chat
.
updated_at
.
desc
())
.
order_by
(
Chat
.
updated_at
.
desc
())
# .limit(limit)
# .limit(limit).offset(skip)
# .offset(skip)
.
all
()
]
)
return
[
ChatModel
.
model_validate
(
chat
)
for
chat
in
all_chats
]
def
get_chat_list_by_user_id
(
def
get_chat_list_by_user_id
(
self
,
self
,
db
:
Session
,
user_id
:
str
,
user_id
:
str
,
include_archived
:
bool
=
False
,
include_archived
:
bool
=
False
,
skip
:
int
=
0
,
skip
:
int
=
0
,
limit
:
int
=
50
,
limit
:
int
=
50
,
)
->
List
[
ChatModel
]:
)
->
List
[
ChatModel
]:
if
include_archived
:
query
=
db
.
query
(
Chat
).
filter_by
(
user_id
=
user_id
)
return
[
if
not
include_archived
:
ChatModel
(
**
model_to_dict
(
chat
))
query
=
query
.
filter_by
(
archived
=
False
)
for
chat
in
Chat
.
select
()
all_chats
=
(
.
where
(
Chat
.
user_id
==
user_id
)
query
.
order_by
(
Chat
.
updated_at
.
desc
())
.
order_by
(
Chat
.
updated_at
.
desc
())
# .limit(limit).offset(skip)
# .limit(limit)
.
all
()
# .offset(skip)
)
]
return
[
ChatModel
.
model_validate
(
chat
)
for
chat
in
all_chats
]
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)
]
def
get_chat_list_by_chat_ids
(
def
get_chat_list_by_chat_ids
(
self
,
chat_ids
:
List
[
str
],
skip
:
int
=
0
,
limit
:
int
=
50
self
,
db
:
Session
,
chat_ids
:
List
[
str
],
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
ChatModel
]:
)
->
List
[
ChatModel
]:
return
[
all_chats
=
(
ChatModel
(
**
model_to_dict
(
chat
))
db
.
query
(
Chat
)
for
chat
in
Chat
.
select
()
.
filter
(
Chat
.
id
.
in_
(
chat_ids
))
.
where
(
Chat
.
archived
==
False
)
.
filter_by
(
archived
=
False
)
.
where
(
Chat
.
id
.
in_
(
chat_ids
))
.
order_by
(
Chat
.
updated_at
.
desc
())
.
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
]:
def
get_chat_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
ChatModel
]:
try
:
try
:
chat
=
Chat
.
get
(
Chat
.
id
==
id
)
chat
=
db
.
get
(
Chat
,
id
)
return
ChatModel
(
**
model_
to_dict
(
chat
)
)
return
ChatModel
.
model_
validate
(
chat
)
except
:
except
:
return
None
return
None
def
get_chat_by_share_id
(
self
,
id
:
str
)
->
Optional
[
ChatModel
]:
def
get_chat_by_share_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
ChatModel
]:
try
:
try
:
chat
=
Chat
.
get
(
Chat
.
share_id
==
id
)
chat
=
db
.
query
(
Chat
).
filter_by
(
share_id
=
id
).
first
(
)
if
chat
:
if
chat
:
chat
=
Chat
.
get
(
Chat
.
id
==
id
)
return
self
.
get_chat_by_id
(
db
,
id
)
return
ChatModel
(
**
model_to_dict
(
chat
))
else
:
else
:
return
None
return
None
except
:
except
Exception
as
e
:
return
None
return
None
def
get_chat_by_id_and_user_id
(
self
,
id
:
str
,
user_id
:
str
)
->
Optional
[
ChatModel
]:
def
get_chat_by_id_and_user_id
(
self
,
db
:
Session
,
id
:
str
,
user_id
:
str
)
->
Optional
[
ChatModel
]:
try
:
try
:
chat
=
Chat
.
get
(
Chat
.
id
==
id
,
Chat
.
user_id
==
user_id
)
chat
=
db
.
query
(
Chat
).
filter_by
(
id
=
id
,
user_id
=
user_id
)
.
first
()
return
ChatModel
(
**
model_
to_dict
(
chat
)
)
return
ChatModel
.
model_
validate
(
chat
)
except
:
except
:
return
None
return
None
def
get_chats
(
self
,
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
ChatModel
]:
def
get_chats
(
self
,
db
:
Session
,
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
ChatModel
]:
return
[
all_chats
=
(
ChatModel
(
**
model_to_dict
(
chat
))
db
.
query
(
Chat
)
for
chat
in
Chat
.
select
().
order_by
(
Chat
.
updated_at
.
desc
())
# .limit(limit).offset(skip)
# .limit(limit).offset(skip)
]
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
())
.
order_by
(
Chat
.
updated_at
.
desc
())
# .limit(limit).offset(skip)
)
]
return
[
ChatModel
.
model_validate
(
chat
)
for
chat
in
all_chats
]
def
get_archived_chats_by_user_id
(
self
,
user_id
:
str
)
->
List
[
ChatModel
]:
def
get_chats_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
List
[
ChatModel
]:
return
[
all_chats
=
(
ChatModel
(
**
model_to_dict
(
chat
))
db
.
query
(
Chat
).
filter_by
(
user_id
=
user_id
).
order_by
(
Chat
.
updated_at
.
desc
())
for
chat
in
Chat
.
select
()
)
.
where
(
Chat
.
archived
==
True
)
return
[
ChatModel
.
model_validate
(
chat
)
for
chat
in
all_chats
]
.
where
(
Chat
.
user_id
==
user_id
)
def
get_archived_chats_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
List
[
ChatModel
]:
all_chats
=
(
db
.
query
(
Chat
)
.
filter_by
(
user_id
=
user_id
,
archived
=
True
)
.
order_by
(
Chat
.
updated_at
.
desc
())
.
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
:
def
delete_chat_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
Chat
.
delete
().
where
((
Chat
.
id
==
id
))
db
.
query
(
Chat
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
and
self
.
delete_shared_chat_by_chat_id
(
id
)
return
True
and
self
.
delete_shared_chat_by_chat_id
(
db
,
id
)
except
:
except
:
return
False
return
False
def
delete_chat_by_id_and_user_id
(
self
,
id
:
str
,
user_id
:
str
)
->
bool
:
def
delete_chat_by_id_and_user_id
(
self
,
db
:
Session
,
id
:
str
,
user_id
:
str
)
->
bool
:
try
:
try
:
query
=
Chat
.
delete
().
where
((
Chat
.
id
==
id
)
&
(
Chat
.
user_id
==
user_id
))
db
.
query
(
Chat
).
filter_by
(
id
=
id
,
user_id
=
user_id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
and
self
.
delete_shared_chat_by_chat_id
(
id
)
return
True
and
self
.
delete_shared_chat_by_chat_id
(
db
,
id
)
except
:
except
:
return
False
return
False
def
delete_chats_by_user_id
(
self
,
user_id
:
str
)
->
bool
:
def
delete_chats_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
bool
:
try
:
try
:
self
.
delete_shared_chats_by_user_id
(
user_id
)
self
.
delete_shared_chats_by_user_id
(
db
,
user_id
)
query
=
Chat
.
delete
().
where
(
Chat
.
user_id
==
user_id
)
query
.
execute
()
# Remove the rows, return number of rows removed.
db
.
query
(
Chat
).
filter_by
(
user_id
=
user_id
).
delete
()
return
True
return
True
except
:
except
:
return
False
return
False
def
delete_shared_chats_by_user_id
(
self
,
user_id
:
str
)
->
bool
:
def
delete_shared_chats_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
bool
:
try
:
try
:
shared_chat_ids
=
[
chats_by_user
=
db
.
query
(
Chat
).
filter_by
(
user_id
=
user_id
).
all
()
f
"shared-
{
chat
.
id
}
"
shared_chat_ids
=
[
f
"shared-
{
chat
.
id
}
"
for
chat
in
chats_by_user
]
for
chat
in
Chat
.
select
().
where
(
Chat
.
user_id
==
user_id
)
]
query
=
Chat
.
delete
().
where
(
Chat
.
user_id
<<
shared_chat_ids
)
db
.
query
(
Chat
).
filter
(
Chat
.
user_id
.
in_
(
shared_chat_ids
)).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Chats
=
ChatTable
(
DB
)
Chats
=
ChatTable
()
backend/apps/webui/models/documents.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
typing
import
List
,
Optional
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
import
time
import
time
import
logging
import
logging
from
utils.utils
import
decode_token
from
sqlalchemy
import
String
,
Column
,
BigInteger
from
utils.misc
import
get_gravatar_url
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
D
B
from
apps.webui.internal.db
import
B
ase
import
json
import
json
...
@@ -22,20 +20,21 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -22,20 +20,21 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
Document
(
Model
):
class
Document
(
Base
):
collection_name
=
CharField
(
unique
=
True
)
__tablename__
=
"document"
name
=
CharField
(
unique
=
True
)
title
=
TextField
()
filename
=
TextField
()
content
=
TextField
(
null
=
True
)
user_id
=
CharField
()
timestamp
=
BigIntegerField
()
class
Meta
:
collection_name
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
name
=
Column
(
String
,
unique
=
True
)
title
=
Column
(
String
)
filename
=
Column
(
String
)
content
=
Column
(
String
,
nullable
=
True
)
user_id
=
Column
(
String
)
timestamp
=
Column
(
BigInteger
)
class
DocumentModel
(
BaseModel
):
class
DocumentModel
(
BaseModel
):
model_config
=
ConfigDict
(
from_attributes
=
True
)
collection_name
:
str
collection_name
:
str
name
:
str
name
:
str
title
:
str
title
:
str
...
@@ -72,12 +71,9 @@ class DocumentForm(DocumentUpdateForm):
...
@@ -72,12 +71,9 @@ class DocumentForm(DocumentUpdateForm):
class
DocumentsTable
:
class
DocumentsTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Document
])
def
insert_new_doc
(
def
insert_new_doc
(
self
,
user_id
:
str
,
form_data
:
DocumentForm
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
DocumentForm
)
->
Optional
[
DocumentModel
]:
)
->
Optional
[
DocumentModel
]:
document
=
DocumentModel
(
document
=
DocumentModel
(
**
{
**
{
...
@@ -88,73 +84,69 @@ class DocumentsTable:
...
@@ -88,73 +84,69 @@ class DocumentsTable:
)
)
try
:
try
:
result
=
Document
.
create
(
**
document
.
model_dump
())
result
=
Document
(
**
document
.
model_dump
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
d
ocument
return
D
ocument
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
:
except
:
return
None
return
None
def
get_doc_by_name
(
self
,
name
:
str
)
->
Optional
[
DocumentModel
]:
def
get_doc_by_name
(
self
,
db
:
Session
,
name
:
str
)
->
Optional
[
DocumentModel
]:
try
:
try
:
document
=
Document
.
get
(
Document
.
name
==
name
)
document
=
db
.
query
(
Document
).
filter_by
(
name
=
name
).
first
(
)
return
DocumentModel
(
**
model_
to_dict
(
document
)
)
return
DocumentModel
.
model_
validate
(
document
)
if
document
else
None
except
:
except
:
return
None
return
None
def
get_docs
(
self
)
->
List
[
DocumentModel
]:
def
get_docs
(
self
,
db
:
Session
)
->
List
[
DocumentModel
]:
return
[
return
[
DocumentModel
.
model_validate
(
doc
)
for
doc
in
db
.
query
(
Document
).
all
()]
DocumentModel
(
**
model_to_dict
(
doc
))
for
doc
in
Document
.
select
()
# .limit(limit).offset(skip)
]
def
update_doc_by_name
(
def
update_doc_by_name
(
self
,
name
:
str
,
form_data
:
DocumentUpdateForm
self
,
db
:
Session
,
name
:
str
,
form_data
:
DocumentUpdateForm
)
->
Optional
[
DocumentModel
]:
)
->
Optional
[
DocumentModel
]:
try
:
try
:
query
=
Document
.
update
(
db
.
query
(
Document
).
filter_by
(
name
=
name
).
update
(
title
=
form_data
.
title
,
{
name
=
form_data
.
name
,
"title"
:
form_data
.
title
,
timestamp
=
int
(
time
.
time
()),
"name"
:
form_data
.
name
,
).
where
(
Document
.
name
==
name
)
"timestamp"
:
int
(
time
.
time
()),
query
.
execute
()
}
)
doc
=
Document
.
get
(
Document
.
name
==
form_data
.
name
)
return
self
.
get_doc_by_name
(
db
,
form_data
.
name
)
return
DocumentModel
(
**
model_to_dict
(
doc
))
except
Exception
as
e
:
except
Exception
as
e
:
log
.
exception
(
e
)
log
.
exception
(
e
)
return
None
return
None
def
update_doc_content_by_name
(
def
update_doc_content_by_name
(
self
,
name
:
str
,
updated
:
dict
self
,
db
:
Session
,
name
:
str
,
updated
:
dict
)
->
Optional
[
DocumentModel
]:
)
->
Optional
[
DocumentModel
]:
try
:
try
:
doc
=
self
.
get_doc_by_name
(
name
)
doc
=
self
.
get_doc_by_name
(
db
,
name
)
doc_content
=
json
.
loads
(
doc
.
content
if
doc
.
content
else
"{}"
)
doc_content
=
json
.
loads
(
doc
.
content
if
doc
.
content
else
"{}"
)
doc_content
=
{
**
doc_content
,
**
updated
}
doc_content
=
{
**
doc_content
,
**
updated
}
query
=
Document
.
update
(
db
.
query
(
Document
).
filter_by
(
name
=
name
).
update
(
content
=
json
.
dumps
(
doc_content
),
{
timestamp
=
int
(
time
.
time
()),
"content"
:
json
.
dumps
(
doc_content
),
).
where
(
Document
.
name
==
name
)
"timestamp"
:
int
(
time
.
time
()),
query
.
execute
()
}
)
doc
=
Document
.
get
(
Document
.
name
==
name
)
return
self
.
get_doc_by_name
(
db
,
name
)
return
DocumentModel
(
**
model_to_dict
(
doc
))
except
Exception
as
e
:
except
Exception
as
e
:
log
.
exception
(
e
)
log
.
exception
(
e
)
return
None
return
None
def
delete_doc_by_name
(
self
,
name
:
str
)
->
bool
:
def
delete_doc_by_name
(
self
,
db
:
Session
,
name
:
str
)
->
bool
:
try
:
try
:
query
=
Document
.
delete
().
where
((
Document
.
name
==
name
))
db
.
query
(
Document
).
filter_by
(
name
=
name
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Documents
=
DocumentsTable
(
DB
)
Documents
=
DocumentsTable
()
backend/apps/webui/models/files.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
import
time
import
time
import
logging
import
logging
from
apps.webui.internal.db
import
DB
,
JSONField
from
sqlalchemy
import
Column
,
String
,
BigInteger
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
JSONField
,
Base
import
json
import
json
...
@@ -18,15 +20,14 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -18,15 +20,14 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
File
(
Model
):
class
File
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"file"
user_id
=
CharField
()
filename
=
TextField
()
meta
=
JSONField
()
created_at
=
BigIntegerField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
user_id
=
Column
(
String
)
filename
=
Column
(
String
)
meta
=
Column
(
JSONField
)
created_at
=
Column
(
BigInteger
)
class
FileModel
(
BaseModel
):
class
FileModel
(
BaseModel
):
...
@@ -36,6 +37,7 @@ class FileModel(BaseModel):
...
@@ -36,6 +37,7 @@ class FileModel(BaseModel):
meta
:
dict
meta
:
dict
created_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -57,11 +59,8 @@ class FileForm(BaseModel):
...
@@ -57,11 +59,8 @@ class FileForm(BaseModel):
class
FilesTable
:
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
]:
def
insert_new_file
(
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
FileForm
)
->
Optional
[
FileModel
]:
file
=
FileModel
(
file
=
FileModel
(
**
{
**
{
**
form_data
.
model_dump
(),
**
form_data
.
model_dump
(),
...
@@ -71,42 +70,41 @@ class FilesTable:
...
@@ -71,42 +70,41 @@ class FilesTable:
)
)
try
:
try
:
result
=
File
.
create
(
**
file
.
model_dump
())
result
=
File
(
**
file
.
model_dump
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
f
ile
return
F
ile
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
Exception
as
e
:
except
Exception
as
e
:
print
(
f
"Error creating tool:
{
e
}
"
)
print
(
f
"Error creating tool:
{
e
}
"
)
return
None
return
None
def
get_file_by_id
(
self
,
id
:
str
)
->
Optional
[
FileModel
]:
def
get_file_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
FileModel
]:
try
:
try
:
file
=
File
.
get
(
File
.
id
==
id
)
file
=
db
.
get
(
File
,
id
)
return
FileModel
(
**
model_
to_dict
(
file
)
)
return
FileModel
.
model_
validate
(
file
)
except
:
except
:
return
None
return
None
def
get_files
(
self
)
->
List
[
FileModel
]:
def
get_files
(
self
,
db
:
Session
)
->
List
[
FileModel
]:
return
[
FileModel
(
**
model_
to_dict
(
file
)
)
for
file
in
File
.
select
()]
return
[
FileModel
.
model_
validate
(
file
)
for
file
in
db
.
query
(
File
).
all
()]
def
delete_file_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_file_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
File
.
delete
().
where
((
File
.
id
==
id
))
db
.
query
(
File
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
def
delete_all_files
(
self
)
->
bool
:
def
delete_all_files
(
self
,
db
:
Session
)
->
bool
:
try
:
try
:
query
=
File
.
delete
()
db
.
query
(
File
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Files
=
FilesTable
(
DB
)
Files
=
FilesTable
()
backend/apps/webui/models/functions.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
import
time
import
time
import
logging
import
logging
from
apps.webui.internal.db
import
DB
,
JSONField
from
sqlalchemy
import
Column
,
String
,
Text
,
BigInteger
,
Boolean
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
JSONField
,
Base
from
apps.webui.models.users
import
Users
from
apps.webui.models.users
import
Users
import
json
import
json
...
@@ -21,20 +23,19 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -21,20 +23,19 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
Function
(
Model
):
class
Function
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"function"
user_id
=
CharField
()
name
=
TextField
()
type
=
TextField
()
content
=
TextField
()
meta
=
JSONField
()
valves
=
JSONField
()
is_active
=
BooleanField
(
default
=
False
)
updated_at
=
BigIntegerField
()
created_at
=
BigIntegerField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
user_id
=
Column
(
String
)
name
=
Column
(
Text
)
type
=
Column
(
Text
)
content
=
Column
(
Text
)
meta
=
Column
(
JSONField
)
valves
=
Column
(
JSONField
)
is_active
=
Column
(
Boolean
)
updated_at
=
Column
(
BigInteger
)
created_at
=
Column
(
BigInteger
)
class
FunctionMeta
(
BaseModel
):
class
FunctionMeta
(
BaseModel
):
...
@@ -53,6 +54,8 @@ class FunctionModel(BaseModel):
...
@@ -53,6 +54,8 @@ class FunctionModel(BaseModel):
updated_at
:
int
# timestamp in epoch
updated_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -82,12 +85,9 @@ class FunctionValves(BaseModel):
...
@@ -82,12 +85,9 @@ class FunctionValves(BaseModel):
class
FunctionsTable
:
class
FunctionsTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Function
])
def
insert_new_function
(
def
insert_new_function
(
self
,
user_id
:
str
,
type
:
str
,
form_data
:
FunctionForm
self
,
db
:
Session
,
user_id
:
str
,
type
:
str
,
form_data
:
FunctionForm
)
->
Optional
[
FunctionModel
]:
)
->
Optional
[
FunctionModel
]:
function
=
FunctionModel
(
function
=
FunctionModel
(
**
{
**
{
...
@@ -100,19 +100,22 @@ class FunctionsTable:
...
@@ -100,19 +100,22 @@ class FunctionsTable:
)
)
try
:
try
:
result
=
Function
.
create
(
**
function
.
model_dump
())
result
=
Function
(
**
function
.
model_dump
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
f
unction
return
F
unction
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
Exception
as
e
:
except
Exception
as
e
:
print
(
f
"Error creating tool:
{
e
}
"
)
print
(
f
"Error creating tool:
{
e
}
"
)
return
None
return
None
def
get_function_by_id
(
self
,
id
:
str
)
->
Optional
[
FunctionModel
]:
def
get_function_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
FunctionModel
]:
try
:
try
:
function
=
Function
.
get
(
Function
.
id
==
id
)
function
=
db
.
get
(
Function
,
id
)
return
FunctionModel
(
**
model_
to_dict
(
function
)
)
return
FunctionModel
.
model_
validate
(
function
)
except
:
except
:
return
None
return
None
...
@@ -211,14 +214,11 @@ class FunctionsTable:
...
@@ -211,14 +214,11 @@ class FunctionsTable:
def
update_function_by_id
(
self
,
id
:
str
,
updated
:
dict
)
->
Optional
[
FunctionModel
]:
def
update_function_by_id
(
self
,
id
:
str
,
updated
:
dict
)
->
Optional
[
FunctionModel
]:
try
:
try
:
query
=
Function
.
update
(
db
.
query
(
Function
).
filter_by
(
id
=
id
)
.
update
(
{
**
updated
,
**
updated
,
updated_at
=
int
(
time
.
time
()),
"updated_at"
:
int
(
time
.
time
()),
).
where
(
Function
.
id
==
id
)
})
query
.
execute
()
return
self
.
get_function_by_id
(
db
,
id
)
function
=
Function
.
get
(
Function
.
id
==
id
)
return
FunctionModel
(
**
model_to_dict
(
function
))
except
:
except
:
return
None
return
None
...
@@ -235,14 +235,12 @@ class FunctionsTable:
...
@@ -235,14 +235,12 @@ class FunctionsTable:
except
:
except
:
return
None
return
None
def
delete_function_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_function_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
Function
.
delete
().
where
((
Function
.
id
==
id
))
db
.
query
(
Function
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Functions
=
FunctionsTable
(
DB
)
Functions
=
FunctionsTable
()
backend/apps/webui/models/memories.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
from
apps.webui.internal.db
import
DB
from
sqlalchemy
import
Column
,
String
,
BigInteger
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
Base
from
apps.webui.models.chats
import
Chats
from
apps.webui.models.chats
import
Chats
import
time
import
time
...
@@ -14,15 +15,14 @@ import uuid
...
@@ -14,15 +15,14 @@ import uuid
####################
####################
class
Memory
(
Model
):
class
Memory
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"memory"
user_id
=
CharField
()
content
=
TextField
()
updated_at
=
BigIntegerField
()
created_at
=
BigIntegerField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
user_id
=
Column
(
String
)
content
=
Column
(
String
)
updated_at
=
Column
(
BigInteger
)
created_at
=
Column
(
BigInteger
)
class
MemoryModel
(
BaseModel
):
class
MemoryModel
(
BaseModel
):
...
@@ -32,6 +32,8 @@ class MemoryModel(BaseModel):
...
@@ -32,6 +32,8 @@ class MemoryModel(BaseModel):
updated_at
:
int
# timestamp in epoch
updated_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -39,12 +41,10 @@ class MemoryModel(BaseModel):
...
@@ -39,12 +41,10 @@ class MemoryModel(BaseModel):
class
MemoriesTable
:
class
MemoriesTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Memory
])
def
insert_new_memory
(
def
insert_new_memory
(
self
,
self
,
db
:
Session
,
user_id
:
str
,
user_id
:
str
,
content
:
str
,
content
:
str
,
)
->
Optional
[
MemoryModel
]:
)
->
Optional
[
MemoryModel
]:
...
@@ -59,74 +59,73 @@ class MemoriesTable:
...
@@ -59,74 +59,73 @@ class MemoriesTable:
"updated_at"
:
int
(
time
.
time
()),
"updated_at"
:
int
(
time
.
time
()),
}
}
)
)
result
=
Memory
.
create
(
**
memory
.
model_dump
())
result
=
Memory
(
**
memory
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
m
emory
return
M
emory
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
def
update_memory_by_id
(
def
update_memory_by_id
(
self
,
self
,
db
:
Session
,
id
:
str
,
id
:
str
,
content
:
str
,
content
:
str
,
)
->
Optional
[
MemoryModel
]:
)
->
Optional
[
MemoryModel
]:
try
:
try
:
memory
=
Memory
.
get
(
Memory
.
id
==
id
)
db
.
query
(
Memory
).
filter_by
(
id
=
id
).
update
(
memory
.
content
=
content
{
"content"
:
content
,
"updated_at"
:
int
(
time
.
time
())}
memory
.
updated_at
=
int
(
time
.
time
())
)
memory
.
save
()
return
self
.
get_memory_by_id
(
db
,
id
)
return
MemoryModel
(
**
model_to_dict
(
memory
))
except
:
except
:
return
None
return
None
def
get_memories
(
self
)
->
List
[
MemoryModel
]:
def
get_memories
(
self
,
db
:
Session
)
->
List
[
MemoryModel
]:
try
:
try
:
memories
=
Memory
.
select
()
memories
=
db
.
query
(
Memory
).
all
()
return
[
MemoryModel
(
**
model_
to_dict
(
memory
)
)
for
memory
in
memories
]
return
[
MemoryModel
.
model_
validate
(
memory
)
for
memory
in
memories
]
except
:
except
:
return
None
return
None
def
get_memories_by_user_id
(
self
,
user_id
:
str
)
->
List
[
MemoryModel
]:
def
get_memories_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
List
[
MemoryModel
]:
try
:
try
:
memories
=
Memory
.
select
().
where
(
Memory
.
user_id
==
user_id
)
memories
=
db
.
query
(
Memory
).
filter_by
(
user_id
=
user_id
)
.
all
()
return
[
MemoryModel
(
**
model_
to_dict
(
memory
)
)
for
memory
in
memories
]
return
[
MemoryModel
.
model_
validate
(
memory
)
for
memory
in
memories
]
except
:
except
:
return
None
return
None
def
get_memory_by_id
(
self
,
i
d
)
->
Optional
[
MemoryModel
]:
def
get_memory_by_id
(
self
,
d
b
:
Session
,
id
:
str
)
->
Optional
[
MemoryModel
]:
try
:
try
:
memory
=
Memory
.
get
(
Memory
.
id
==
id
)
memory
=
db
.
get
(
Memory
,
id
)
return
MemoryModel
(
**
model_
to_dict
(
memory
)
)
return
MemoryModel
.
model_
validate
(
memory
)
except
:
except
:
return
None
return
None
def
delete_memory_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_memory_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
Memory
.
delete
().
where
(
Memory
.
id
==
id
)
db
.
query
(
Memory
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
def
delete_memories_by_user_id
(
self
,
user_id
:
str
)
->
bool
:
def
delete_memories_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
bool
:
try
:
try
:
query
=
Memory
.
delete
().
where
(
Memory
.
user_id
==
user_id
)
db
.
query
(
Memory
).
filter_by
(
user_id
=
user_id
).
delete
()
query
.
execute
()
return
True
return
True
except
:
except
:
return
False
return
False
def
delete_memory_by_id_and_user_id
(
self
,
id
:
str
,
user_id
:
str
)
->
bool
:
def
delete_memory_by_id_and_user_id
(
self
,
db
:
Session
,
id
:
str
,
user_id
:
str
)
->
bool
:
try
:
try
:
query
=
Memory
.
delete
().
where
(
Memory
.
id
==
id
,
Memory
.
user_id
==
user_id
)
db
.
query
(
Memory
).
filter_by
(
id
=
id
,
user_id
=
user_id
).
delete
()
query
.
execute
()
return
True
return
True
except
:
except
:
return
False
return
False
Memories
=
MemoriesTable
(
DB
)
Memories
=
MemoriesTable
()
backend/apps/webui/models/models.py
View file @
df09d083
...
@@ -2,13 +2,11 @@ import json
...
@@ -2,13 +2,11 @@ import json
import
logging
import
logging
from
typing
import
Optional
from
typing
import
Optional
import
peewee
as
pw
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
from
pydantic
import
BaseModel
,
ConfigDict
from
pydantic
import
BaseModel
,
ConfigDict
from
sqlalchemy
import
String
,
Column
,
BigInteger
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
D
B
,
JSONField
from
apps.webui.internal.db
import
B
ase
,
JSONField
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
from
config
import
SRC_LOG_LEVELS
from
config
import
SRC_LOG_LEVELS
...
@@ -46,41 +44,42 @@ class ModelMeta(BaseModel):
...
@@ -46,41 +44,42 @@ class ModelMeta(BaseModel):
pass
pass
class
Model
(
pw
.
Model
):
class
Model
(
Base
):
id
=
pw
.
TextField
(
unique
=
True
)
__tablename__
=
"model"
id
=
Column
(
String
,
primary_key
=
True
)
"""
"""
The model's id as used in the API. If set to an existing model, it will override the model.
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
(
String
)
base_model_id
=
pw
.
TextField
(
null
=
True
)
base_model_id
=
Column
(
String
,
nullable
=
True
)
"""
"""
An optional pointer to the actual model that should be used when proxying requests.
An optional pointer to the actual model that should be used when proxying requests.
"""
"""
name
=
pw
.
TextField
(
)
name
=
Column
(
String
)
"""
"""
The human-readable display name of the model.
The human-readable display name of the model.
"""
"""
params
=
JSONField
(
)
params
=
Column
(
JSONField
)
"""
"""
Holds a JSON encoded blob of parameters, see `ModelParams`.
Holds a JSON encoded blob of parameters, see `ModelParams`.
"""
"""
meta
=
JSONField
(
)
meta
=
Column
(
JSONField
)
"""
"""
Holds a JSON encoded blob of metadata, see `ModelMeta`.
Holds a JSON encoded blob of metadata, see `ModelMeta`.
"""
"""
updated_at
=
BigIntegerField
()
updated_at
=
Column
(
BigInteger
)
created_at
=
BigIntegerField
()
created_at
=
Column
(
BigInteger
)
class
Meta
:
database
=
DB
class
ModelModel
(
BaseModel
):
class
ModelModel
(
BaseModel
):
model_config
=
ConfigDict
(
from_attributes
=
True
)
id
:
str
id
:
str
user_id
:
str
user_id
:
str
base_model_id
:
Optional
[
str
]
=
None
base_model_id
:
Optional
[
str
]
=
None
...
@@ -115,15 +114,9 @@ class ModelForm(BaseModel):
...
@@ -115,15 +114,9 @@ class ModelForm(BaseModel):
class
ModelsTable
:
class
ModelsTable
:
def
__init__
(
self
,
db
:
pw
.
SqliteDatabase
|
pw
.
PostgresqlDatabase
,
):
self
.
db
=
db
self
.
db
.
create_tables
([
Model
])
def
insert_new_model
(
def
insert_new_model
(
self
,
form_data
:
ModelForm
,
user_id
:
str
self
,
db
:
Session
,
form_data
:
ModelForm
,
user_id
:
str
)
->
Optional
[
ModelModel
]:
)
->
Optional
[
ModelModel
]:
model
=
ModelModel
(
model
=
ModelModel
(
**
{
**
{
...
@@ -134,46 +127,50 @@ class ModelsTable:
...
@@ -134,46 +127,50 @@ class ModelsTable:
}
}
)
)
try
:
try
:
result
=
Model
.
create
(
**
model
.
model_dump
())
result
=
Model
(
**
model
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
model
return
ModelModel
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
Exception
as
e
:
except
Exception
as
e
:
print
(
e
)
print
(
e
)
return
None
return
None
def
get_all_models
(
self
)
->
List
[
ModelModel
]:
def
get_all_models
(
self
,
db
:
Session
)
->
List
[
ModelModel
]:
return
[
ModelModel
(
**
model_
to_dict
(
model
)
)
for
model
in
Model
.
select
()]
return
[
ModelModel
.
model_
validate
(
model
)
for
model
in
db
.
query
(
Model
).
all
()]
def
get_model_by_id
(
self
,
id
:
str
)
->
Optional
[
ModelModel
]:
def
get_model_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
ModelModel
]:
try
:
try
:
model
=
Model
.
get
(
Model
.
id
==
id
)
model
=
db
.
get
(
Model
,
id
)
return
ModelModel
(
**
model_
to_dict
(
model
)
)
return
ModelModel
.
model_
validate
(
model
)
except
:
except
:
return
None
return
None
def
update_model_by_id
(
self
,
id
:
str
,
model
:
ModelForm
)
->
Optional
[
ModelModel
]:
def
update_model_by_id
(
self
,
db
:
Session
,
id
:
str
,
model
:
ModelForm
)
->
Optional
[
ModelModel
]:
try
:
try
:
# update only the fields that are present in the model
# update only the fields that are present in the model
query
=
Model
.
update
(
**
model
.
model_dump
()).
where
(
Model
.
id
==
id
)
model
=
db
.
query
(
Model
).
get
(
id
)
query
.
execute
()
model
.
update
(
**
model
.
model_dump
()
)
db
.
commit
()
model
=
Model
.
get
(
Model
.
id
==
id
)
db
.
refresh
(
model
)
return
ModelModel
(
**
model_
to_dict
(
model
)
)
return
ModelModel
.
model_
validate
(
model
)
except
Exception
as
e
:
except
Exception
as
e
:
print
(
e
)
print
(
e
)
return
None
return
None
def
delete_model_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_model_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
Model
.
delete
().
where
(
Model
.
id
==
id
)
db
.
query
(
Model
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
return
True
return
True
except
:
except
:
return
False
return
False
Models
=
ModelsTable
(
DB
)
Models
=
ModelsTable
()
backend/apps/webui/models/prompts.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
typing
import
List
,
Optional
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
import
time
import
time
from
utils.utils
import
decode_token
from
sqlalchemy
import
String
,
Column
,
BigInteger
from
utils.misc
import
get_gravatar_url
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
D
B
from
apps.webui.internal.db
import
B
ase
import
json
import
json
...
@@ -16,15 +14,14 @@ import json
...
@@ -16,15 +14,14 @@ import json
####################
####################
class
Prompt
(
Model
):
class
Prompt
(
Base
):
command
=
CharField
(
unique
=
True
)
__tablename__
=
"prompt"
user_id
=
CharField
()
title
=
TextField
()
content
=
TextField
()
timestamp
=
BigIntegerField
()
class
Meta
:
command
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
user_id
=
Column
(
String
)
title
=
Column
(
String
)
content
=
Column
(
String
)
timestamp
=
Column
(
BigInteger
)
class
PromptModel
(
BaseModel
):
class
PromptModel
(
BaseModel
):
...
@@ -34,6 +31,8 @@ class PromptModel(BaseModel):
...
@@ -34,6 +31,8 @@ class PromptModel(BaseModel):
content
:
str
content
:
str
timestamp
:
int
# timestamp in epoch
timestamp
:
int
# timestamp in epoch
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -48,12 +47,8 @@ class PromptForm(BaseModel):
...
@@ -48,12 +47,8 @@ class PromptForm(BaseModel):
class
PromptsTable
:
class
PromptsTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Prompt
])
def
insert_new_prompt
(
def
insert_new_prompt
(
self
,
user_id
:
str
,
form_data
:
PromptForm
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
PromptForm
)
->
Optional
[
PromptModel
]:
)
->
Optional
[
PromptModel
]:
prompt
=
PromptModel
(
prompt
=
PromptModel
(
**
{
**
{
...
@@ -66,53 +61,48 @@ class PromptsTable:
...
@@ -66,53 +61,48 @@ class PromptsTable:
)
)
try
:
try
:
result
=
Prompt
.
create
(
**
prompt
.
model_dump
())
result
=
Prompt
(
**
prompt
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
p
rompt
return
P
rompt
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
:
except
Exception
as
e
:
return
None
return
None
def
get_prompt_by_command
(
self
,
command
:
str
)
->
Optional
[
PromptModel
]:
def
get_prompt_by_command
(
self
,
db
:
Session
,
command
:
str
)
->
Optional
[
PromptModel
]:
try
:
try
:
prompt
=
Prompt
.
get
(
Prompt
.
command
==
command
)
prompt
=
db
.
query
(
Prompt
).
filter_by
(
command
=
command
)
.
first
()
return
PromptModel
(
**
model_
to_dict
(
prompt
)
)
return
PromptModel
.
model_
validate
(
prompt
)
except
:
except
:
return
None
return
None
def
get_prompts
(
self
)
->
List
[
PromptModel
]:
def
get_prompts
(
self
,
db
:
Session
)
->
List
[
PromptModel
]:
return
[
return
[
PromptModel
.
model_validate
(
prompt
)
for
prompt
in
db
.
query
(
Prompt
).
all
()]
PromptModel
(
**
model_to_dict
(
prompt
))
for
prompt
in
Prompt
.
select
()
# .limit(limit).offset(skip)
]
def
update_prompt_by_command
(
def
update_prompt_by_command
(
self
,
command
:
str
,
form_data
:
PromptForm
self
,
db
:
Session
,
command
:
str
,
form_data
:
PromptForm
)
->
Optional
[
PromptModel
]:
)
->
Optional
[
PromptModel
]:
try
:
try
:
query
=
Prompt
.
update
(
db
.
query
(
Prompt
).
filter_by
(
command
=
command
).
update
(
title
=
form_data
.
title
,
{
content
=
form_data
.
content
,
"title"
:
form_data
.
title
,
timestamp
=
int
(
time
.
time
()),
"content"
:
form_data
.
content
,
).
where
(
Prompt
.
command
==
command
)
"timestamp"
:
int
(
time
.
time
()),
}
query
.
execute
()
)
return
self
.
get_prompt_by_command
(
db
,
command
)
prompt
=
Prompt
.
get
(
Prompt
.
command
==
command
)
return
PromptModel
(
**
model_to_dict
(
prompt
))
except
:
except
:
return
None
return
None
def
delete_prompt_by_command
(
self
,
command
:
str
)
->
bool
:
def
delete_prompt_by_command
(
self
,
db
:
Session
,
command
:
str
)
->
bool
:
try
:
try
:
query
=
Prompt
.
delete
().
where
((
Prompt
.
command
==
command
))
db
.
query
(
Prompt
).
filter_by
(
command
=
command
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Prompts
=
PromptsTable
(
DB
)
Prompts
=
PromptsTable
()
backend/apps/webui/models/tags.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Optional
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
import
json
import
json
import
uuid
import
uuid
import
time
import
time
import
logging
import
logging
from
apps.webui.internal.db
import
DB
from
sqlalchemy
import
String
,
Column
,
BigInteger
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
Base
from
config
import
SRC_LOG_LEVELS
from
config
import
SRC_LOG_LEVELS
...
@@ -20,25 +21,23 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -20,25 +21,23 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
Tag
(
Model
):
class
Tag
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"tag"
name
=
CharField
()
user_id
=
CharField
()
data
=
TextField
(
null
=
True
)
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
name
=
Column
(
String
)
user_id
=
Column
(
String
)
data
=
Column
(
String
,
nullable
=
True
)
class
ChatIdTag
(
Model
):
class
ChatIdTag
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"chatidtag"
tag_name
=
CharField
()
chat_id
=
CharField
()
user_id
=
CharField
()
timestamp
=
BigIntegerField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
tag_name
=
Column
(
String
)
chat_id
=
Column
(
String
)
user_id
=
Column
(
String
)
timestamp
=
Column
(
BigInteger
)
class
TagModel
(
BaseModel
):
class
TagModel
(
BaseModel
):
...
@@ -47,6 +46,8 @@ class TagModel(BaseModel):
...
@@ -47,6 +46,8 @@ class TagModel(BaseModel):
user_id
:
str
user_id
:
str
data
:
Optional
[
str
]
=
None
data
:
Optional
[
str
]
=
None
model_config
=
ConfigDict
(
from_attributes
=
True
)
class
ChatIdTagModel
(
BaseModel
):
class
ChatIdTagModel
(
BaseModel
):
id
:
str
id
:
str
...
@@ -55,6 +56,8 @@ class ChatIdTagModel(BaseModel):
...
@@ -55,6 +56,8 @@ class ChatIdTagModel(BaseModel):
user_id
:
str
user_id
:
str
timestamp
:
int
timestamp
:
int
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -75,37 +78,39 @@ class ChatTagsResponse(BaseModel):
...
@@ -75,37 +78,39 @@ class ChatTagsResponse(BaseModel):
class
TagTable
:
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
]:
def
insert_new_tag
(
self
,
db
:
Session
,
name
:
str
,
user_id
:
str
)
->
Optional
[
TagModel
]:
id
=
str
(
uuid
.
uuid4
())
id
=
str
(
uuid
.
uuid4
())
tag
=
TagModel
(
**
{
"id"
:
id
,
"user_id"
:
user_id
,
"name"
:
name
})
tag
=
TagModel
(
**
{
"id"
:
id
,
"user_id"
:
user_id
,
"name"
:
name
})
try
:
try
:
result
=
Tag
.
create
(
**
tag
.
model_dump
())
result
=
Tag
(
**
tag
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
t
ag
return
T
ag
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
Exception
as
e
:
except
Exception
as
e
:
return
None
return
None
def
get_tag_by_name_and_user_id
(
def
get_tag_by_name_and_user_id
(
self
,
name
:
str
,
user_id
:
str
self
,
db
:
Session
,
name
:
str
,
user_id
:
str
)
->
Optional
[
TagModel
]:
)
->
Optional
[
TagModel
]:
try
:
try
:
tag
=
Tag
.
get
(
Tag
.
name
==
name
,
Tag
.
user_id
==
user_id
)
tag
=
db
.
query
(
Tag
).
filter
(
name
=
name
,
user_id
=
user_id
)
.
first
()
return
TagModel
(
**
model_
to_dict
(
tag
)
)
return
TagModel
.
model_
validate
(
tag
)
except
Exception
as
e
:
except
Exception
as
e
:
return
None
return
None
def
add_tag_to_chat
(
def
add_tag_to_chat
(
self
,
user_id
:
str
,
form_data
:
ChatIdTagForm
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
ChatIdTagForm
)
->
Optional
[
ChatIdTagModel
]:
)
->
Optional
[
ChatIdTagModel
]:
tag
=
self
.
get_tag_by_name_and_user_id
(
form_data
.
tag_name
,
user_id
)
tag
=
self
.
get_tag_by_name_and_user_id
(
db
,
form_data
.
tag_name
,
user_id
)
if
tag
==
None
:
if
tag
==
None
:
tag
=
self
.
insert_new_tag
(
form_data
.
tag_name
,
user_id
)
tag
=
self
.
insert_new_tag
(
db
,
form_data
.
tag_name
,
user_id
)
id
=
str
(
uuid
.
uuid4
())
id
=
str
(
uuid
.
uuid4
())
chatIdTag
=
ChatIdTagModel
(
chatIdTag
=
ChatIdTagModel
(
...
@@ -118,120 +123,135 @@ class TagTable:
...
@@ -118,120 +123,135 @@ class TagTable:
}
}
)
)
try
:
try
:
result
=
ChatIdTag
.
create
(
**
chatIdTag
.
model_dump
())
result
=
ChatIdTag
(
**
chatIdTag
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
c
hatIdTag
return
C
hatIdTag
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
:
except
:
return
None
return
None
def
get_tags_by_user_id
(
self
,
user_id
:
str
)
->
List
[
TagModel
]:
def
get_tags_by_user_id
(
self
,
db
:
Session
,
user_id
:
str
)
->
List
[
TagModel
]:
tag_names
=
[
tag_names
=
[
ChatIdTagModel
(
**
model_to_dict
(
chat_id_tag
)).
tag_name
chat_id_tag
.
tag_name
for
chat_id_tag
in
ChatIdTag
.
select
()
for
chat_id_tag
in
(
.
where
(
ChatIdTag
.
user_id
==
user_id
)
db
.
query
(
ChatIdTag
)
.
filter_by
(
user_id
=
user_id
)
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
all
()
)
]
]
return
[
return
[
TagModel
(
**
model_to_dict
(
tag
))
TagModel
.
model_validate
(
tag
)
for
tag
in
Tag
.
select
()
for
tag
in
(
.
where
(
Tag
.
user_id
==
user_id
)
db
.
query
(
Tag
)
.
where
(
Tag
.
name
.
in_
(
tag_names
))
.
filter_by
(
user_id
=
user_id
)
.
filter
(
Tag
.
name
.
in_
(
tag_names
))
.
all
()
)
]
]
def
get_tags_by_chat_id_and_user_id
(
def
get_tags_by_chat_id_and_user_id
(
self
,
chat_id
:
str
,
user_id
:
str
self
,
db
:
Session
,
chat_id
:
str
,
user_id
:
str
)
->
List
[
TagModel
]:
)
->
List
[
TagModel
]:
tag_names
=
[
tag_names
=
[
ChatIdTagModel
(
**
model_to_dict
(
chat_id_tag
)).
tag_name
chat_id_tag
.
tag_name
for
chat_id_tag
in
ChatIdTag
.
select
()
for
chat_id_tag
in
(
.
where
((
ChatIdTag
.
user_id
==
user_id
)
&
(
ChatIdTag
.
chat_id
==
chat_id
))
db
.
query
(
ChatIdTag
)
.
filter_by
(
user_id
=
user_id
,
chat_id
=
chat_id
)
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
all
()
)
]
]
return
[
return
[
TagModel
(
**
model_to_dict
(
tag
))
TagModel
.
model_validate
(
tag
)
for
tag
in
Tag
.
select
()
for
tag
in
(
.
where
(
Tag
.
user_id
==
user_id
)
db
.
query
(
Tag
)
.
where
(
Tag
.
name
.
in_
(
tag_names
))
.
filter_by
(
user_id
=
user_id
)
.
filter
(
Tag
.
name
.
in_
(
tag_names
))
.
all
()
)
]
]
def
get_chat_ids_by_tag_name_and_user_id
(
def
get_chat_ids_by_tag_name_and_user_id
(
self
,
tag_name
:
str
,
user_id
:
str
self
,
db
:
Session
,
tag_name
:
str
,
user_id
:
str
)
->
Optional
[
ChatIdTagModel
]:
)
->
List
[
ChatIdTagModel
]:
return
[
return
[
ChatIdTagModel
(
**
model_to_dict
(
chat_id_tag
))
ChatIdTagModel
.
model_validate
(
chat_id_tag
)
for
chat_id_tag
in
ChatIdTag
.
select
()
for
chat_id_tag
in
(
.
where
((
ChatIdTag
.
user_id
==
user_id
)
&
(
ChatIdTag
.
tag_name
==
tag_name
))
db
.
query
(
ChatIdTag
)
.
filter_by
(
user_id
=
user_id
,
tag_name
=
tag_name
)
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
order_by
(
ChatIdTag
.
timestamp
.
desc
())
.
all
()
)
]
]
def
count_chat_ids_by_tag_name_and_user_id
(
def
count_chat_ids_by_tag_name_and_user_id
(
self
,
tag_name
:
str
,
user_id
:
str
self
,
db
:
Session
,
tag_name
:
str
,
user_id
:
str
)
->
int
:
)
->
int
:
return
(
return
db
.
query
(
ChatIdTag
).
filter_by
(
tag_name
=
tag_name
,
user_id
=
user_id
).
count
()
ChatIdTag
.
select
()
.
where
((
ChatIdTag
.
tag_name
==
tag_name
)
&
(
ChatIdTag
.
user_id
==
user_id
))
.
count
()
)
def
delete_tag_by_tag_name_and_user_id
(
self
,
tag_name
:
str
,
user_id
:
str
)
->
bool
:
def
delete_tag_by_tag_name_and_user_id
(
self
,
db
:
Session
,
tag_name
:
str
,
user_id
:
str
)
->
bool
:
try
:
try
:
query
=
ChatIdTag
.
delete
().
where
(
res
=
(
(
ChatIdTag
.
tag_name
==
tag_name
)
&
(
ChatIdTag
.
user_id
==
user_id
)
db
.
query
(
ChatIdTag
)
.
filter_by
(
tag_name
=
tag_name
,
user_id
=
user_id
)
.
delete
()
)
)
res
=
query
.
execute
()
# Remove the rows, return number of rows removed.
log
.
debug
(
f
"res:
{
res
}
"
)
log
.
debug
(
f
"res:
{
res
}
"
)
tag_count
=
self
.
count_chat_ids_by_tag_name_and_user_id
(
tag_name
,
user_id
)
tag_count
=
self
.
count_chat_ids_by_tag_name_and_user_id
(
db
,
tag_name
,
user_id
)
if
tag_count
==
0
:
if
tag_count
==
0
:
# Remove tag item from Tag col as well
# Remove tag item from Tag col as well
query
=
Tag
.
delete
().
where
(
db
.
query
(
Tag
).
filter_by
(
name
=
tag_name
,
user_id
=
user_id
).
delete
()
(
Tag
.
name
==
tag_name
)
&
(
Tag
.
user_id
==
user_id
)
)
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
Exception
as
e
:
except
Exception
as
e
:
log
.
error
(
f
"delete_tag:
{
e
}
"
)
log
.
error
(
f
"delete_tag:
{
e
}
"
)
return
False
return
False
def
delete_tag_by_tag_name_and_chat_id_and_user_id
(
def
delete_tag_by_tag_name_and_chat_id_and_user_id
(
self
,
tag_name
:
str
,
chat_id
:
str
,
user_id
:
str
self
,
db
:
Session
,
tag_name
:
str
,
chat_id
:
str
,
user_id
:
str
)
->
bool
:
)
->
bool
:
try
:
try
:
query
=
ChatIdTag
.
delete
().
where
(
res
=
(
(
ChatIdTag
.
tag_name
==
tag_name
)
db
.
query
(
ChatIdTag
)
&
(
ChatIdTag
.
chat_id
==
chat
_id
)
.
filter_by
(
tag_name
=
tag_name
,
chat_id
=
chat_id
,
user_id
=
user
_id
)
&
(
ChatIdTag
.
user_id
==
user_id
)
.
delete
(
)
)
)
res
=
query
.
execute
()
# Remove the rows, return number of rows removed.
log
.
debug
(
f
"res:
{
res
}
"
)
log
.
debug
(
f
"res:
{
res
}
"
)
tag_count
=
self
.
count_chat_ids_by_tag_name_and_user_id
(
tag_name
,
user_id
)
tag_count
=
self
.
count_chat_ids_by_tag_name_and_user_id
(
db
,
tag_name
,
user_id
)
if
tag_count
==
0
:
if
tag_count
==
0
:
# Remove tag item from Tag col as well
# Remove tag item from Tag col as well
query
=
Tag
.
delete
().
where
(
db
.
query
(
Tag
).
filter_by
(
name
=
tag_name
,
user_id
=
user_id
).
delete
()
(
Tag
.
name
==
tag_name
)
&
(
Tag
.
user_id
==
user_id
)
)
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
Exception
as
e
:
except
Exception
as
e
:
log
.
error
(
f
"delete_tag:
{
e
}
"
)
log
.
error
(
f
"delete_tag:
{
e
}
"
)
return
False
return
False
def
delete_tags_by_chat_id_and_user_id
(
self
,
chat_id
:
str
,
user_id
:
str
)
->
bool
:
def
delete_tags_by_chat_id_and_user_id
(
tags
=
self
.
get_tags_by_chat_id_and_user_id
(
chat_id
,
user_id
)
self
,
db
:
Session
,
chat_id
:
str
,
user_id
:
str
)
->
bool
:
tags
=
self
.
get_tags_by_chat_id_and_user_id
(
db
,
chat_id
,
user_id
)
for
tag
in
tags
:
for
tag
in
tags
:
self
.
delete_tag_by_tag_name_and_chat_id_and_user_id
(
self
.
delete_tag_by_tag_name_and_chat_id_and_user_id
(
tag
.
tag_name
,
chat_id
,
user_id
db
,
tag
.
tag_name
,
chat_id
,
user_id
)
)
return
True
return
True
Tags
=
TagTable
(
DB
)
Tags
=
TagTable
()
backend/apps/webui/models/tools.py
View file @
df09d083
from
pydantic
import
BaseModel
from
pydantic
import
BaseModel
,
ConfigDict
from
peewee
import
*
from
typing
import
List
,
Optional
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
import
time
import
time
import
logging
import
logging
from
apps.webui.internal.db
import
DB
,
JSONField
from
sqlalchemy
import
String
,
Column
,
BigInteger
from
sqlalchemy.orm
import
Session
from
apps.webui.internal.db
import
Base
,
JSONField
from
apps.webui.models.users
import
Users
from
apps.webui.models.users
import
Users
import
json
import
json
...
@@ -21,19 +22,18 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
...
@@ -21,19 +22,18 @@ log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
####################
class
Tool
(
Model
):
class
Tool
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"tool"
user_id
=
CharField
()
name
=
TextField
()
content
=
TextField
()
specs
=
JSONField
()
meta
=
JSONField
()
valves
=
JSONField
()
updated_at
=
BigIntegerField
()
created_at
=
BigIntegerField
()
class
Meta
:
id
=
Column
(
String
,
primary_key
=
True
)
database
=
DB
user_id
=
Column
(
String
)
name
=
Column
(
String
)
content
=
Column
(
String
)
specs
=
Column
(
JSONField
)
meta
=
Column
(
JSONField
)
valves
=
Column
(
JSONField
)
updated_at
=
Column
(
BigInteger
)
created_at
=
Column
(
BigInteger
)
class
ToolMeta
(
BaseModel
):
class
ToolMeta
(
BaseModel
):
...
@@ -51,6 +51,8 @@ class ToolModel(BaseModel):
...
@@ -51,6 +51,8 @@ class ToolModel(BaseModel):
updated_at
:
int
# timestamp in epoch
updated_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
created_at
:
int
# timestamp in epoch
model_config
=
ConfigDict
(
from_attributes
=
True
)
####################
####################
# Forms
# Forms
...
@@ -78,12 +80,9 @@ class ToolValves(BaseModel):
...
@@ -78,12 +80,9 @@ class ToolValves(BaseModel):
class
ToolsTable
:
class
ToolsTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
Tool
])
def
insert_new_tool
(
def
insert_new_tool
(
self
,
user_id
:
str
,
form_data
:
ToolForm
,
specs
:
List
[
dict
]
self
,
db
:
Session
,
user_id
:
str
,
form_data
:
ToolForm
,
specs
:
List
[
dict
]
)
->
Optional
[
ToolModel
]:
)
->
Optional
[
ToolModel
]:
tool
=
ToolModel
(
tool
=
ToolModel
(
**
{
**
{
...
@@ -96,24 +95,27 @@ class ToolsTable:
...
@@ -96,24 +95,27 @@ class ToolsTable:
)
)
try
:
try
:
result
=
Tool
.
create
(
**
tool
.
model_dump
())
result
=
Tool
(
**
tool
.
dict
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
t
ool
return
T
ool
Model
.
model_validate
(
result
)
else
:
else
:
return
None
return
None
except
Exception
as
e
:
except
Exception
as
e
:
print
(
f
"Error creating tool:
{
e
}
"
)
print
(
f
"Error creating tool:
{
e
}
"
)
return
None
return
None
def
get_tool_by_id
(
self
,
id
:
str
)
->
Optional
[
ToolModel
]:
def
get_tool_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
ToolModel
]:
try
:
try
:
tool
=
Tool
.
get
(
Tool
.
id
==
id
)
tool
=
db
.
get
(
Tool
,
id
)
return
ToolModel
(
**
model_
to_dict
(
tool
)
)
return
ToolModel
.
model_
validate
(
tool
)
except
:
except
:
return
None
return
None
def
get_tools
(
self
)
->
List
[
ToolModel
]:
def
get_tools
(
self
,
db
:
Session
)
->
List
[
ToolModel
]:
return
[
ToolModel
(
**
model_
to_dict
(
tool
)
)
for
tool
in
Tool
.
select
()]
return
[
ToolModel
.
model_
validate
(
tool
)
for
tool
in
db
.
query
(
Tool
).
all
()]
def
get_tool_valves_by_id
(
self
,
id
:
str
)
->
Optional
[
dict
]:
def
get_tool_valves_by_id
(
self
,
id
:
str
)
->
Optional
[
dict
]:
try
:
try
:
...
@@ -180,25 +182,19 @@ class ToolsTable:
...
@@ -180,25 +182,19 @@ class ToolsTable:
def
update_tool_by_id
(
self
,
id
:
str
,
updated
:
dict
)
->
Optional
[
ToolModel
]:
def
update_tool_by_id
(
self
,
id
:
str
,
updated
:
dict
)
->
Optional
[
ToolModel
]:
try
:
try
:
query
=
Tool
.
update
(
db
.
query
(
Tool
).
filter_by
(
id
=
id
).
update
(
**
updated
,
{
**
updated
,
"updated_at"
:
int
(
time
.
time
())}
updated_at
=
int
(
time
.
time
()),
)
).
where
(
Tool
.
id
==
id
)
return
self
.
get_tool_by_id
(
db
,
id
)
query
.
execute
()
tool
=
Tool
.
get
(
Tool
.
id
==
id
)
return
ToolModel
(
**
model_to_dict
(
tool
))
except
:
except
:
return
None
return
None
def
delete_tool_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_tool_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
query
=
Tool
.
delete
().
where
((
Tool
.
id
==
id
))
db
.
query
(
Tool
).
filter_by
(
id
=
id
).
delete
()
query
.
execute
()
# Remove the rows, return number of rows removed.
return
True
return
True
except
:
except
:
return
False
return
False
Tools
=
ToolsTable
(
DB
)
Tools
=
ToolsTable
()
backend/apps/webui/models/users.py
View file @
df09d083
from
pydantic
import
BaseModel
,
ConfigDict
from
pydantic
import
BaseModel
,
ConfigDict
,
parse_obj_as
from
peewee
import
*
from
playhouse.shortcuts
import
model_to_dict
from
typing
import
List
,
Union
,
Optional
from
typing
import
List
,
Union
,
Optional
import
time
import
time
from
sqlalchemy
import
String
,
Column
,
BigInteger
,
Text
from
sqlalchemy.orm
import
Session
from
utils.misc
import
get_gravatar_url
from
utils.misc
import
get_gravatar_url
from
apps.webui.internal.db
import
D
B
,
JSONField
from
apps.webui.internal.db
import
B
ase
,
JSONField
from
apps.webui.models.chats
import
Chats
from
apps.webui.models.chats
import
Chats
####################
####################
...
@@ -13,25 +15,24 @@ from apps.webui.models.chats import Chats
...
@@ -13,25 +15,24 @@ from apps.webui.models.chats import Chats
####################
####################
class
User
(
Model
):
class
User
(
Base
):
id
=
CharField
(
unique
=
True
)
__tablename__
=
"user"
name
=
CharField
()
email
=
CharField
()
role
=
CharField
()
profile_image_url
=
TextField
()
last_active_at
=
BigIntegerField
()
id
=
Column
(
String
,
primary_key
=
True
)
updated_at
=
BigIntegerField
()
name
=
Column
(
String
)
created_at
=
BigIntegerField
()
email
=
Column
(
String
)
role
=
Column
(
String
)
profile_image_url
=
Column
(
String
)
api_key
=
CharField
(
null
=
True
,
unique
=
True
)
last_active_at
=
Column
(
BigInteger
)
settings
=
JSONField
(
null
=
True
)
updated_at
=
Column
(
BigInteger
)
info
=
JSONField
(
null
=
True
)
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
:
oauth_sub
=
Column
(
Text
,
unique
=
True
)
database
=
DB
class
UserSettings
(
BaseModel
):
class
UserSettings
(
BaseModel
):
...
@@ -41,6 +42,8 @@ class UserSettings(BaseModel):
...
@@ -41,6 +42,8 @@ class UserSettings(BaseModel):
class
UserModel
(
BaseModel
):
class
UserModel
(
BaseModel
):
model_config
=
ConfigDict
(
from_attributes
=
True
)
id
:
str
id
:
str
name
:
str
name
:
str
email
:
str
email
:
str
...
@@ -76,12 +79,10 @@ class UserUpdateForm(BaseModel):
...
@@ -76,12 +79,10 @@ class UserUpdateForm(BaseModel):
class
UsersTable
:
class
UsersTable
:
def
__init__
(
self
,
db
):
self
.
db
=
db
self
.
db
.
create_tables
([
User
])
def
insert_new_user
(
def
insert_new_user
(
self
,
self
,
db
:
Session
,
id
:
str
,
id
:
str
,
name
:
str
,
name
:
str
,
email
:
str
,
email
:
str
,
...
@@ -102,30 +103,33 @@ class UsersTable:
...
@@ -102,30 +103,33 @@ class UsersTable:
"oauth_sub"
:
oauth_sub
,
"oauth_sub"
:
oauth_sub
,
}
}
)
)
result
=
User
.
create
(
**
user
.
model_dump
())
result
=
User
(
**
user
.
model_dump
())
db
.
add
(
result
)
db
.
commit
()
db
.
refresh
(
result
)
if
result
:
if
result
:
return
user
return
user
else
:
else
:
return
None
return
None
def
get_user_by_id
(
self
,
id
:
str
)
->
Optional
[
UserModel
]:
def
get_user_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
UserModel
]:
try
:
try
:
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
Exception
as
e
:
return
None
return
None
def
get_user_by_api_key
(
self
,
api_key
:
str
)
->
Optional
[
UserModel
]:
def
get_user_by_api_key
(
self
,
db
:
Session
,
api_key
:
str
)
->
Optional
[
UserModel
]:
try
:
try
:
user
=
User
.
get
(
User
.
api_key
==
api_key
)
user
=
db
.
query
(
User
).
filter_by
(
api_key
=
api_key
)
.
first
()
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
get_user_by_email
(
self
,
email
:
str
)
->
Optional
[
UserModel
]:
def
get_user_by_email
(
self
,
db
:
Session
,
email
:
str
)
->
Optional
[
UserModel
]:
try
:
try
:
user
=
User
.
get
(
User
.
email
==
email
)
user
=
db
.
query
(
User
).
filter_by
(
email
=
email
)
.
first
()
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
...
@@ -136,88 +140,94 @@ class UsersTable:
...
@@ -136,88 +140,94 @@ class UsersTable:
except
:
except
:
return
None
return
None
def
get_users
(
self
,
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
UserModel
]:
def
get_users
(
self
,
db
:
Session
,
skip
:
int
=
0
,
limit
:
int
=
50
)
->
List
[
UserModel
]:
return
[
users
=
(
UserModel
(
**
model_to_dict
(
user
))
db
.
query
(
User
)
for
user
in
User
.
select
()
# .offset(skip).limit(limit)
# .limit(limit).offset(skip)
.
all
()
]
)
return
[
UserModel
.
model_validate
(
user
)
for
user
in
users
]
def
get_num_users
(
self
)
->
Optional
[
int
]:
def
get_num_users
(
self
,
db
:
Session
)
->
Optional
[
int
]:
return
User
.
select
(
).
count
()
return
db
.
query
(
User
).
count
()
def
get_first_user
(
self
)
->
UserModel
:
def
get_first_user
(
self
,
db
:
Session
)
->
UserModel
:
try
:
try
:
user
=
User
.
select
(
).
order_by
(
User
.
created_at
).
first
()
user
=
db
.
query
(
User
).
order_by
(
User
.
created_at
).
first
()
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
update_user_role_by_id
(
self
,
id
:
str
,
role
:
str
)
->
Optional
[
UserModel
]:
def
update_user_role_by_id
(
self
,
db
:
Session
,
id
:
str
,
role
:
str
)
->
Optional
[
UserModel
]:
try
:
try
:
query
=
User
.
update
(
role
=
role
).
where
(
User
.
id
==
id
)
db
.
query
(
User
).
filter_by
(
id
=
id
)
.
update
(
{
"
role
"
:
role
}
)
query
.
execute
()
db
.
commit
()
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
update_user_profile_image_url_by_id
(
def
update_user_profile_image_url_by_id
(
self
,
id
:
str
,
profile_image_url
:
str
self
,
db
:
Session
,
id
:
str
,
profile_image_url
:
str
)
->
Optional
[
UserModel
]:
)
->
Optional
[
UserModel
]:
try
:
try
:
query
=
User
.
update
(
profile_image_url
=
profile_image_url
).
wher
e
(
db
.
query
(
User
).
filter_by
(
id
=
id
).
updat
e
(
User
.
id
==
id
{
"profile_image_url"
:
profile_image_url
}
)
)
query
.
execute
()
db
.
commit
()
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
update_user_last_active_by_id
(
self
,
id
:
str
)
->
Optional
[
UserModel
]:
def
update_user_last_active_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
UserModel
]:
try
:
try
:
query
=
User
.
update
(
last_active_at
=
int
(
time
.
time
())).
where
(
User
.
id
==
id
)
db
.
query
(
User
).
filter_by
(
id
=
id
).
update
({
"last_active_at"
:
int
(
time
.
time
())})
query
.
execute
()
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
update_user_oauth_sub_by_id
(
def
update_user_oauth_sub_by_id
(
self
,
id
:
str
,
oauth_sub
:
str
self
,
db
:
Session
,
id
:
str
,
oauth_sub
:
str
)
->
Optional
[
UserModel
]:
)
->
Optional
[
UserModel
]:
try
:
try
:
query
=
User
.
update
(
oauth_sub
=
oauth_sub
).
where
(
User
.
id
==
id
)
db
.
query
(
User
).
filter_by
(
id
=
id
).
update
({
"oauth_sub"
:
oauth_sub
})
query
.
execute
()
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
UserModel
(
**
model_
to_dict
(
user
)
)
return
UserModel
.
model_
validate
(
user
)
except
:
except
:
return
None
return
None
def
update_user_by_id
(
self
,
id
:
str
,
updated
:
dict
)
->
Optional
[
UserModel
]:
def
update_user_by_id
(
self
,
db
:
Session
,
id
:
str
,
updated
:
dict
)
->
Optional
[
UserModel
]:
try
:
try
:
query
=
User
.
update
(
**
updated
)
.
where
(
User
.
id
==
id
)
db
.
query
(
User
).
filter_by
(
id
=
id
)
.
update
(
updated
)
query
.
execute
()
db
.
commit
()
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
()
return
UserModel
(
**
model_to_dict
(
user
))
return
UserModel
.
model_validate
(
user
)
except
:
# return UserModel(**user.dict())
except
Exception
as
e
:
return
None
return
None
def
delete_user_by_id
(
self
,
id
:
str
)
->
bool
:
def
delete_user_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
bool
:
try
:
try
:
# Delete User Chats
# Delete User Chats
result
=
Chats
.
delete_chats_by_user_id
(
id
)
result
=
Chats
.
delete_chats_by_user_id
(
db
,
id
)
if
result
:
if
result
:
# Delete User
# Delete User
query
=
User
.
delete
().
where
(
User
.
id
==
id
)
db
.
query
(
User
).
filter_by
(
id
=
id
).
delete
(
)
query
.
execute
()
# Remove the rows, return number of rows removed.
db
.
commit
()
return
True
return
True
else
:
else
:
...
@@ -225,21 +235,20 @@ class UsersTable:
...
@@ -225,21 +235,20 @@ class UsersTable:
except
:
except
:
return
False
return
False
def
update_user_api_key_by_id
(
self
,
id
:
str
,
api_key
:
str
)
->
str
:
def
update_user_api_key_by_id
(
self
,
db
:
Session
,
id
:
str
,
api_key
:
str
)
->
str
:
try
:
try
:
query
=
User
.
update
(
api_key
=
api_key
).
where
(
User
.
id
==
id
)
result
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
update
({
"api_key"
:
api_key
})
result
=
query
.
execute
()
db
.
commit
()
return
True
if
result
==
1
else
False
return
True
if
result
==
1
else
False
except
:
except
:
return
False
return
False
def
get_user_api_key_by_id
(
self
,
id
:
str
)
->
Optional
[
str
]:
def
get_user_api_key_by_id
(
self
,
db
:
Session
,
id
:
str
)
->
Optional
[
str
]:
try
:
try
:
user
=
User
.
get
(
User
.
id
==
id
)
user
=
db
.
query
(
User
).
filter_by
(
id
=
id
).
first
(
)
return
user
.
api_key
return
user
.
api_key
except
:
except
Exception
as
e
:
return
None
return
None
Users
=
UsersTable
(
DB
)
Users
=
UsersTable
()
backend/apps/webui/routers/auths.py
View file @
df09d083
...
@@ -10,6 +10,7 @@ import re
...
@@ -10,6 +10,7 @@ import re
import
uuid
import
uuid
import
csv
import
csv
from
apps.webui.internal.db
import
get_db
from
apps.webui.models.auths
import
(
from
apps.webui.models.auths
import
(
SigninForm
,
SigninForm
,
SignupForm
,
SignupForm
,
...
@@ -78,10 +79,13 @@ async def get_session_user(
...
@@ -78,10 +79,13 @@ async def get_session_user(
@
router
.
post
(
"/update/profile"
,
response_model
=
UserResponse
)
@
router
.
post
(
"/update/profile"
,
response_model
=
UserResponse
)
async
def
update_profile
(
async
def
update_profile
(
form_data
:
UpdateProfileForm
,
session_user
=
Depends
(
get_current_user
)
form_data
:
UpdateProfileForm
,
session_user
=
Depends
(
get_current_user
),
db
=
Depends
(
get_db
),
):
):
if
session_user
:
if
session_user
:
user
=
Users
.
update_user_by_id
(
user
=
Users
.
update_user_by_id
(
db
,
session_user
.
id
,
session_user
.
id
,
{
"profile_image_url"
:
form_data
.
profile_image_url
,
"name"
:
form_data
.
name
},
{
"profile_image_url"
:
form_data
.
profile_image_url
,
"name"
:
form_data
.
name
},
)
)
...
@@ -100,16 +104,18 @@ async def update_profile(
...
@@ -100,16 +104,18 @@ async def update_profile(
@
router
.
post
(
"/update/password"
,
response_model
=
bool
)
@
router
.
post
(
"/update/password"
,
response_model
=
bool
)
async
def
update_password
(
async
def
update_password
(
form_data
:
UpdatePasswordForm
,
session_user
=
Depends
(
get_current_user
)
form_data
:
UpdatePasswordForm
,
session_user
=
Depends
(
get_current_user
),
db
=
Depends
(
get_db
),
):
):
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
:
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
:
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
ACTION_PROHIBITED
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
ACTION_PROHIBITED
)
if
session_user
:
if
session_user
:
user
=
Auths
.
authenticate_user
(
session_user
.
email
,
form_data
.
password
)
user
=
Auths
.
authenticate_user
(
db
,
session_user
.
email
,
form_data
.
password
)
if
user
:
if
user
:
hashed
=
get_password_hash
(
form_data
.
new_password
)
hashed
=
get_password_hash
(
form_data
.
new_password
)
return
Auths
.
update_user_password_by_id
(
user
.
id
,
hashed
)
return
Auths
.
update_user_password_by_id
(
db
,
user
.
id
,
hashed
)
else
:
else
:
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
INVALID_PASSWORD
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
INVALID_PASSWORD
)
else
:
else
:
...
@@ -122,7 +128,7 @@ async def update_password(
...
@@ -122,7 +128,7 @@ async def update_password(
@
router
.
post
(
"/signin"
,
response_model
=
SigninResponse
)
@
router
.
post
(
"/signin"
,
response_model
=
SigninResponse
)
async
def
signin
(
request
:
Request
,
response
:
Response
,
form_data
:
SigninForm
):
async
def
signin
(
request
:
Request
,
response
:
Response
,
form_data
:
SigninForm
,
db
=
Depends
(
get_db
)
):
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
:
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
:
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
not
in
request
.
headers
:
if
WEBUI_AUTH_TRUSTED_EMAIL_HEADER
not
in
request
.
headers
:
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
INVALID_TRUSTED_HEADER
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
INVALID_TRUSTED_HEADER
)
...
@@ -133,32 +139,34 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
...
@@ -133,32 +139,34 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
trusted_name
=
request
.
headers
.
get
(
trusted_name
=
request
.
headers
.
get
(
WEBUI_AUTH_TRUSTED_NAME_HEADER
,
trusted_email
WEBUI_AUTH_TRUSTED_NAME_HEADER
,
trusted_email
)
)
if
not
Users
.
get_user_by_email
(
trusted_email
.
lower
()):
if
not
Users
.
get_user_by_email
(
db
,
trusted_email
.
lower
()):
await
signup
(
await
signup
(
request
,
request
,
SignupForm
(
SignupForm
(
email
=
trusted_email
,
password
=
str
(
uuid
.
uuid4
()),
name
=
trusted_name
email
=
trusted_email
,
password
=
str
(
uuid
.
uuid4
()),
name
=
trusted_name
),
),
db
,
)
)
user
=
Auths
.
authenticate_user_by_trusted_header
(
trusted_email
)
user
=
Auths
.
authenticate_user_by_trusted_header
(
db
,
trusted_email
)
elif
WEBUI_AUTH
==
False
:
elif
WEBUI_AUTH
==
False
:
admin_email
=
"admin@localhost"
admin_email
=
"admin@localhost"
admin_password
=
"admin"
admin_password
=
"admin"
if
Users
.
get_user_by_email
(
admin_email
.
lower
()):
if
Users
.
get_user_by_email
(
db
,
admin_email
.
lower
()):
user
=
Auths
.
authenticate_user
(
admin_email
.
lower
(),
admin_password
)
user
=
Auths
.
authenticate_user
(
db
,
admin_email
.
lower
(),
admin_password
)
else
:
else
:
if
Users
.
get_num_users
()
!=
0
:
if
Users
.
get_num_users
(
db
)
!=
0
:
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EXISTING_USERS
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EXISTING_USERS
)
await
signup
(
await
signup
(
request
,
request
,
SignupForm
(
email
=
admin_email
,
password
=
admin_password
,
name
=
"User"
),
SignupForm
(
email
=
admin_email
,
password
=
admin_password
,
name
=
"User"
),
db
,
)
)
user
=
Auths
.
authenticate_user
(
admin_email
.
lower
(),
admin_password
)
user
=
Auths
.
authenticate_user
(
db
,
admin_email
.
lower
(),
admin_password
)
else
:
else
:
user
=
Auths
.
authenticate_user
(
form_data
.
email
.
lower
(),
form_data
.
password
)
user
=
Auths
.
authenticate_user
(
db
,
form_data
.
email
.
lower
(),
form_data
.
password
)
if
user
:
if
user
:
token
=
create_token
(
token
=
create_token
(
...
@@ -192,7 +200,7 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
...
@@ -192,7 +200,7 @@ async def signin(request: Request, response: Response, form_data: SigninForm):
@
router
.
post
(
"/signup"
,
response_model
=
SigninResponse
)
@
router
.
post
(
"/signup"
,
response_model
=
SigninResponse
)
async
def
signup
(
request
:
Request
,
response
:
Response
,
form_data
:
SignupForm
):
async
def
signup
(
request
:
Request
,
response
:
Response
,
form_data
:
SignupForm
,
db
=
Depends
(
get_db
)
):
if
not
request
.
app
.
state
.
config
.
ENABLE_SIGNUP
and
WEBUI_AUTH
:
if
not
request
.
app
.
state
.
config
.
ENABLE_SIGNUP
and
WEBUI_AUTH
:
raise
HTTPException
(
raise
HTTPException
(
status
.
HTTP_403_FORBIDDEN
,
detail
=
ERROR_MESSAGES
.
ACCESS_PROHIBITED
status
.
HTTP_403_FORBIDDEN
,
detail
=
ERROR_MESSAGES
.
ACCESS_PROHIBITED
...
@@ -203,17 +211,18 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
...
@@ -203,17 +211,18 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
status
.
HTTP_400_BAD_REQUEST
,
detail
=
ERROR_MESSAGES
.
INVALID_EMAIL_FORMAT
status
.
HTTP_400_BAD_REQUEST
,
detail
=
ERROR_MESSAGES
.
INVALID_EMAIL_FORMAT
)
)
if
Users
.
get_user_by_email
(
form_data
.
email
.
lower
()):
if
Users
.
get_user_by_email
(
db
,
form_data
.
email
.
lower
()):
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EMAIL_TAKEN
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EMAIL_TAKEN
)
try
:
try
:
role
=
(
role
=
(
"admin"
"admin"
if
Users
.
get_num_users
()
==
0
if
Users
.
get_num_users
(
db
)
==
0
else
request
.
app
.
state
.
config
.
DEFAULT_USER_ROLE
else
request
.
app
.
state
.
config
.
DEFAULT_USER_ROLE
)
)
hashed
=
get_password_hash
(
form_data
.
password
)
hashed
=
get_password_hash
(
form_data
.
password
)
user
=
Auths
.
insert_new_auth
(
user
=
Auths
.
insert_new_auth
(
db
,
form_data
.
email
.
lower
(),
form_data
.
email
.
lower
(),
hashed
,
hashed
,
form_data
.
name
,
form_data
.
name
,
...
@@ -267,14 +276,16 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
...
@@ -267,14 +276,16 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
@
router
.
post
(
"/add"
,
response_model
=
SigninResponse
)
@
router
.
post
(
"/add"
,
response_model
=
SigninResponse
)
async
def
add_user
(
form_data
:
AddUserForm
,
user
=
Depends
(
get_admin_user
)):
async
def
add_user
(
form_data
:
AddUserForm
,
user
=
Depends
(
get_admin_user
),
db
=
Depends
(
get_db
)
):
if
not
validate_email_format
(
form_data
.
email
.
lower
()):
if
not
validate_email_format
(
form_data
.
email
.
lower
()):
raise
HTTPException
(
raise
HTTPException
(
status
.
HTTP_400_BAD_REQUEST
,
detail
=
ERROR_MESSAGES
.
INVALID_EMAIL_FORMAT
status
.
HTTP_400_BAD_REQUEST
,
detail
=
ERROR_MESSAGES
.
INVALID_EMAIL_FORMAT
)
)
if
Users
.
get_user_by_email
(
form_data
.
email
.
lower
()):
if
Users
.
get_user_by_email
(
db
,
form_data
.
email
.
lower
()):
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EMAIL_TAKEN
)
raise
HTTPException
(
400
,
detail
=
ERROR_MESSAGES
.
EMAIL_TAKEN
)
try
:
try
:
...
@@ -282,6 +293,7 @@ async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
...
@@ -282,6 +293,7 @@ async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
print
(
form_data
)
print
(
form_data
)
hashed
=
get_password_hash
(
form_data
.
password
)
hashed
=
get_password_hash
(
form_data
.
password
)
user
=
Auths
.
insert_new_auth
(
user
=
Auths
.
insert_new_auth
(
db
,
form_data
.
email
.
lower
(),
form_data
.
email
.
lower
(),
hashed
,
hashed
,
form_data
.
name
,
form_data
.
name
,
...
@@ -312,7 +324,9 @@ async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
...
@@ -312,7 +324,9 @@ async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
@
router
.
get
(
"/admin/details"
)
@
router
.
get
(
"/admin/details"
)
async
def
get_admin_details
(
request
:
Request
,
user
=
Depends
(
get_current_user
)):
async
def
get_admin_details
(
request
:
Request
,
user
=
Depends
(
get_current_user
),
db
=
Depends
(
get_db
)
):
if
request
.
app
.
state
.
config
.
SHOW_ADMIN_DETAILS
:
if
request
.
app
.
state
.
config
.
SHOW_ADMIN_DETAILS
:
admin_email
=
request
.
app
.
state
.
config
.
ADMIN_EMAIL
admin_email
=
request
.
app
.
state
.
config
.
ADMIN_EMAIL
admin_name
=
None
admin_name
=
None
...
@@ -320,11 +334,11 @@ async def get_admin_details(request: Request, user=Depends(get_current_user)):
...
@@ -320,11 +334,11 @@ async def get_admin_details(request: Request, user=Depends(get_current_user)):
print
(
admin_email
,
admin_name
)
print
(
admin_email
,
admin_name
)
if
admin_email
:
if
admin_email
:
admin
=
Users
.
get_user_by_email
(
admin_email
)
admin
=
Users
.
get_user_by_email
(
db
,
admin_email
)
if
admin
:
if
admin
:
admin_name
=
admin
.
name
admin_name
=
admin
.
name
else
:
else
:
admin
=
Users
.
get_first_user
()
admin
=
Users
.
get_first_user
(
db
)
if
admin
:
if
admin
:
admin_email
=
admin
.
email
admin_email
=
admin
.
email
admin_name
=
admin
.
name
admin_name
=
admin
.
name
...
@@ -397,9 +411,9 @@ async def update_admin_config(
...
@@ -397,9 +411,9 @@ async def update_admin_config(
# create api key
# create api key
@
router
.
post
(
"/api_key"
,
response_model
=
ApiKey
)
@
router
.
post
(
"/api_key"
,
response_model
=
ApiKey
)
async
def
create_api_key_
(
user
=
Depends
(
get_current_user
)):
async
def
create_api_key_
(
user
=
Depends
(
get_current_user
)
,
db
=
Depends
(
get_db
)
):
api_key
=
create_api_key
()
api_key
=
create_api_key
()
success
=
Users
.
update_user_api_key_by_id
(
user
.
id
,
api_key
)
success
=
Users
.
update_user_api_key_by_id
(
db
,
user
.
id
,
api_key
)
if
success
:
if
success
:
return
{
return
{
"api_key"
:
api_key
,
"api_key"
:
api_key
,
...
@@ -410,15 +424,15 @@ async def create_api_key_(user=Depends(get_current_user)):
...
@@ -410,15 +424,15 @@ async def create_api_key_(user=Depends(get_current_user)):
# delete api key
# delete api key
@
router
.
delete
(
"/api_key"
,
response_model
=
bool
)
@
router
.
delete
(
"/api_key"
,
response_model
=
bool
)
async
def
delete_api_key
(
user
=
Depends
(
get_current_user
)):
async
def
delete_api_key
(
user
=
Depends
(
get_current_user
)
,
db
=
Depends
(
get_db
)
):
success
=
Users
.
update_user_api_key_by_id
(
user
.
id
,
None
)
success
=
Users
.
update_user_api_key_by_id
(
db
,
user
.
id
,
None
)
return
success
return
success
# get api key
# get api key
@
router
.
get
(
"/api_key"
,
response_model
=
ApiKey
)
@
router
.
get
(
"/api_key"
,
response_model
=
ApiKey
)
async
def
get_api_key
(
user
=
Depends
(
get_current_user
)):
async
def
get_api_key
(
user
=
Depends
(
get_current_user
)
,
db
=
Depends
(
get_db
)
):
api_key
=
Users
.
get_user_api_key_by_id
(
user
.
id
)
api_key
=
Users
.
get_user_api_key_by_id
(
db
,
user
.
id
)
if
api_key
:
if
api_key
:
return
{
return
{
"api_key"
:
api_key
,
"api_key"
:
api_key
,
...
...
Prev
1
2
3
Next
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment