Quellcode durchsuchen

refac/enh: add performance indexes

Co-Authored-By: decent-engineer-decent-datascientist <77806775+decent-engineer-decent-datascientist@users.noreply.github.com>
Timothy Jaeryang Baek vor 1 Monat
Ursprung
Commit
b3a95f40fc

+ 46 - 0
backend/open_webui/migrations/versions/018012973d35_add_indexes.py

@@ -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")

+ 15 - 1
backend/open_webui/models/chats.py

@@ -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)

+ 3 - 1
backend/open_webui/models/functions.py

@@ -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

+ 6 - 1
backend/open_webui/models/tags.py

@@ -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"),)
 

+ 1 - 0
backend/open_webui/utils/tools.py

@@ -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):