refac/enh: add performance indexes

Co-Authored-By: decent-engineer-decent-datascientist <77806775+decent-engineer-decent-datascientist@users.noreply.github.com>
This commit is contained in:
Timothy Jaeryang Baek 2025-08-19 03:24:10 +04:00
parent 96643f5b6d
commit b3a95f40fc
5 changed files with 71 additions and 3 deletions

View file

@ -0,0 +1,46 @@
"""Add indexes
Revision ID: 018012973d35
Revises: d31026856c01
Create Date: 2025-08-13 03:00:00.000000
"""
from alembic import op
import sqlalchemy as sa
revision = "018012973d35"
down_revision = "d31026856c01"
branch_labels = None
depends_on = None
def upgrade():
# Chat table indexes
op.create_index("folder_id_idx", "chat", ["folder_id"])
op.create_index("user_id_pinned_idx", "chat", ["user_id", "pinned"])
op.create_index("user_id_archived_idx", "chat", ["user_id", "archived"])
op.create_index("updated_at_user_id_idx", "chat", ["updated_at", "user_id"])
op.create_index("folder_id_user_id_idx", "chat", ["folder_id", "user_id"])
# Tag table index
op.create_index("user_id_idx", "tag", ["user_id"])
# Function table index
op.create_index("is_global_idx", "function", ["is_global"])
def downgrade():
# Chat table indexes
op.drop_index("folder_id_idx", table_name="chat")
op.drop_index("user_id_pinned_idx", table_name="chat")
op.drop_index("user_id_archived_idx", table_name="chat")
op.drop_index("updated_at_user_id_idx", table_name="chat")
op.drop_index("folder_id_user_id_idx", table_name="chat")
# Tag table index
op.drop_index("user_id_idx", table_name="tag")
# Function table index
op.drop_index("is_global_idx", table_name="function")

View file

@ -10,7 +10,7 @@ from open_webui.models.folders import Folders
from open_webui.env import SRC_LOG_LEVELS
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON, Index
from sqlalchemy import or_, func, select, and_, text
from sqlalchemy.sql import exists
from sqlalchemy.sql.expression import bindparam
@ -41,6 +41,20 @@ class Chat(Base):
meta = Column(JSON, server_default="{}")
folder_id = Column(Text, nullable=True)
__table_args__ = (
# Performance indexes for common queries
# WHERE folder_id = ...
Index("folder_id_idx", "folder_id"),
# WHERE user_id = ... AND pinned = ...
Index("user_id_pinned_idx", "user_id", "pinned"),
# WHERE user_id = ... AND archived = ...
Index("user_id_archived_idx", "user_id", "archived"),
# WHERE user_id = ... ORDER BY updated_at DESC
Index("updated_at_user_id_idx", "updated_at", "user_id"),
# WHERE folder_id = ... AND user_id = ...
Index("folder_id_user_id_idx", "folder_id", "user_id"),
)
class ChatModel(BaseModel):
model_config = ConfigDict(from_attributes=True)

View file

@ -6,7 +6,7 @@ from open_webui.internal.db import Base, JSONField, get_db
from open_webui.models.users import Users
from open_webui.env import SRC_LOG_LEVELS
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Boolean, Column, String, Text
from sqlalchemy import BigInteger, Boolean, Column, String, Text, Index
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])
@ -31,6 +31,8 @@ class Function(Base):
updated_at = Column(BigInteger)
created_at = Column(BigInteger)
__table_args__ = (Index("is_global_idx", "is_global"),)
class FunctionMeta(BaseModel):
description: Optional[str] = None

View file

@ -8,7 +8,7 @@ from open_webui.internal.db import Base, get_db
from open_webui.env import SRC_LOG_LEVELS
from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Column, String, JSON, PrimaryKeyConstraint
from sqlalchemy import BigInteger, Column, String, JSON, PrimaryKeyConstraint, Index
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])
@ -24,6 +24,11 @@ class Tag(Base):
user_id = Column(String)
meta = Column(JSON, nullable=True)
__table_args__ = (
PrimaryKeyConstraint("id", "user_id", name="pk_id_user_id"),
Index("user_id_idx", "user_id"),
)
# Unique constraint ensuring (id, user_id) is unique, not just the `id` column
__table_args__ = (PrimaryKeyConstraint("id", "user_id", name="pk_id_user_id"),)

View file

@ -76,6 +76,7 @@ def get_async_tool_function_and_apply_extra_params(
# https://github.com/googleapis/python-genai/issues/907
async def new_function(*args, **kwargs):
return await partial_func(*args, **kwargs)
else:
# Make it a coroutine function when it is not already
async def new_function(*args, **kwargs):