admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: refactor conversation pagination to use SQLAlchemy session manag… (#11956)

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN-
2024-12-22 10:39:29 +08:00
committed by GitHub
parent 366857cd26
commit 3d07a94bd7
6 changed files with 78 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
api/services/conversation_service.py
View File

@@ -1,8 +1,9 @@
from collections.abc import Callable
from collections.abc import Callable, Sequence
from datetime import UTC, datetime
from typing import Optional, Union
from sqlalchemy import asc, desc, or_
from sqlalchemy import asc, desc, func, or_, select
from sqlalchemy.orm import Session
from core.app.entities.app_invoke_entities import InvokeFrom
from core.llm_generator.llm_generator import LLMGenerator
@@ -18,19 +19,21 @@ class ConversationService:
@classmethod
def pagination_by_last_id(
cls,
*,
session: Session,
app_model: App,
user: Optional[Union[Account, EndUser]],
last_id: Optional[str],
limit: int,
invoke_from: InvokeFrom,
include_ids: Optional[list] = None,
exclude_ids: Optional[list] = None,
include_ids: Optional[Sequence[str]] = None,
exclude_ids: Optional[Sequence[str]] = None,
sort_by: str = "-updated_at",
) -> InfiniteScrollPagination:
if not user:
return InfiniteScrollPagination(data=[], limit=limit, has_more=False)
base_query = db.session.query(Conversation).filter(
stmt = select(Conversation).where(
Conversation.is_deleted == False,
Conversation.app_id == app_model.id,
Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
@@ -38,37 +41,40 @@ class ConversationService:
Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
or_(Conversation.invoke_from.is_(None), Conversation.invoke_from == invoke_from.value),
)
if include_ids is not None:
base_query = base_query.filter(Conversation.id.in_(include_ids))
stmt = stmt.where(Conversation.id.in_(include_ids))
if exclude_ids is not None:
base_query = base_query.filter(~Conversation.id.in_(exclude_ids))
stmt = stmt.where(~Conversation.id.in_(exclude_ids))
# define sort fields and directions
sort_field, sort_direction = cls._get_sort_params(sort_by)
if last_id:
last_conversation = base_query.filter(Conversation.id == last_id).first()
last_conversation = session.scalar(stmt.where(Conversation.id == last_id))
if not last_conversation:
raise LastConversationNotExistsError()
# build filters based on sorting
filter_condition = cls._build_filter_condition(sort_field, sort_direction, last_conversation)
base_query = base_query.filter(filter_condition)
base_query = base_query.order_by(sort_direction(getattr(Conversation, sort_field)))
conversations = base_query.limit(limit).all()
filter_condition = cls._build_filter_condition(
sort_field=sort_field,
sort_direction=sort_direction,
reference_conversation=last_conversation,
)
stmt = stmt.where(filter_condition)
query_stmt = stmt.order_by(sort_direction(getattr(Conversation, sort_field))).limit(limit)
conversations = session.scalars(query_stmt).all()
has_more = False
if len(conversations) == limit:
current_page_last_conversation = conversations[-1]
rest_filter_condition = cls._build_filter_condition(
sort_field, sort_direction, current_page_last_conversation, is_next_page=True
sort_field=sort_field,
sort_direction=sort_direction,
reference_conversation=current_page_last_conversation,
)
rest_count = base_query.filter(rest_filter_condition).count()
count_stmt = stmt.where(rest_filter_condition)
count_stmt = select(func.count()).select_from(count_stmt.subquery())
rest_count = session.scalar(count_stmt) or 0
if rest_count > 0:
has_more = True
@@ -81,11 +87,9 @@ class ConversationService:
return sort_by, asc
@classmethod
def _build_filter_condition(
cls, sort_field: str, sort_direction: Callable, reference_conversation: Conversation, is_next_page: bool = False
):
def _build_filter_condition(cls, sort_field: str, sort_direction: Callable, reference_conversation: Conversation):
field_value = getattr(reference_conversation, sort_field)
if (sort_direction == desc and not is_next_page) or (sort_direction == asc and is_next_page):
if sort_direction == desc:
return getattr(Conversation, sort_field) < field_value
else:
return getattr(Conversation, sort_field) > field_value

18
api/services/web_conversation_service.py
View File

@@ -1,5 +1,8 @@
from typing import Optional, Union
from sqlalchemy import select
from sqlalchemy.orm import Session
from core.app.entities.app_invoke_entities import InvokeFrom
from extensions.ext_database import db
from libs.infinite_scroll_pagination import InfiniteScrollPagination
@@ -13,6 +16,8 @@ class WebConversationService:
@classmethod
def pagination_by_last_id(
cls,
*,
session: Session,
app_model: App,
user: Optional[Union[Account, EndUser]],
last_id: Optional[str],
@@ -23,24 +28,25 @@ class WebConversationService:
) -> InfiniteScrollPagination:
include_ids = None
exclude_ids = None
if pinned is not None:
pinned_conversations = (
db.session.query(PinnedConversation)
.filter(
if pinned is not None and user:
stmt = (
select(PinnedConversation.conversation_id)
.where(
PinnedConversation.app_id == app_model.id,
PinnedConversation.created_by_role == ("account" if isinstance(user, Account) else "end_user"),
PinnedConversation.created_by == user.id,
)
.order_by(PinnedConversation.created_at.desc())
.all()
)
pinned_conversation_ids = [pc.conversation_id for pc in pinned_conversations]
pinned_conversation_ids = session.scalars(stmt).all()
if pinned:
include_ids = pinned_conversation_ids
else:
exclude_ids = pinned_conversation_ids
return ConversationService.pagination_by_last_id(
session=session,
app_model=app_model,
user=user,
last_id=last_id,