Skip to content

Commit af953ff

Browse files
authored
Paginate Query History table (#3592)
* Add pagination for query history table * Fix method name * Fix mypy
1 parent 6fc52c8 commit af953ff

File tree

14 files changed

+517
-341
lines changed

14 files changed

+517
-341
lines changed

backend/ee/onyx/db/query_history.py

Lines changed: 114 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,135 @@
1-
import datetime
2-
from typing import Literal
1+
from collections.abc import Sequence
2+
from datetime import datetime
33

44
from sqlalchemy import asc
55
from sqlalchemy import BinaryExpression
66
from sqlalchemy import ColumnElement
77
from sqlalchemy import desc
8+
from sqlalchemy import distinct
89
from sqlalchemy.orm import contains_eager
910
from sqlalchemy.orm import joinedload
1011
from sqlalchemy.orm import Session
12+
from sqlalchemy.sql import case
13+
from sqlalchemy.sql import func
14+
from sqlalchemy.sql import select
15+
from sqlalchemy.sql.expression import literal
1116
from sqlalchemy.sql.expression import UnaryExpression
1217

18+
from onyx.configs.constants import QAFeedbackType
1319
from onyx.db.models import ChatMessage
20+
from onyx.db.models import ChatMessageFeedback
1421
from onyx.db.models import ChatSession
1522

16-
SortByOptions = Literal["time_sent"]
23+
24+
def _build_filter_conditions(
25+
start_time: datetime | None,
26+
end_time: datetime | None,
27+
feedback_filter: QAFeedbackType | None,
28+
) -> list[ColumnElement]:
29+
"""
30+
Helper function to build all filter conditions for chat sessions.
31+
Filters by start and end time, feedback type, and any sessions without messages.
32+
start_time: Date from which to filter
33+
end_time: Date to which to filter
34+
feedback_filter: Feedback type to filter by
35+
Returns: List of filter conditions
36+
"""
37+
conditions = []
38+
39+
if start_time is not None:
40+
conditions.append(ChatSession.time_created >= start_time)
41+
if end_time is not None:
42+
conditions.append(ChatSession.time_created <= end_time)
43+
44+
if feedback_filter is not None:
45+
feedback_subq = (
46+
select(ChatMessage.chat_session_id)
47+
.join(ChatMessageFeedback)
48+
.group_by(ChatMessage.chat_session_id)
49+
.having(
50+
case(
51+
(
52+
case(
53+
{literal(feedback_filter == QAFeedbackType.LIKE): True},
54+
else_=False,
55+
),
56+
func.bool_and(ChatMessageFeedback.is_positive),
57+
),
58+
(
59+
case(
60+
{literal(feedback_filter == QAFeedbackType.DISLIKE): True},
61+
else_=False,
62+
),
63+
func.bool_and(func.not_(ChatMessageFeedback.is_positive)),
64+
),
65+
else_=func.bool_or(ChatMessageFeedback.is_positive)
66+
& func.bool_or(func.not_(ChatMessageFeedback.is_positive)),
67+
)
68+
)
69+
)
70+
conditions.append(ChatSession.id.in_(feedback_subq))
71+
72+
return conditions
73+
74+
75+
def get_total_filtered_chat_sessions_count(
76+
db_session: Session,
77+
start_time: datetime | None,
78+
end_time: datetime | None,
79+
feedback_filter: QAFeedbackType | None,
80+
) -> int:
81+
conditions = _build_filter_conditions(start_time, end_time, feedback_filter)
82+
stmt = (
83+
select(func.count(distinct(ChatSession.id)))
84+
.select_from(ChatSession)
85+
.filter(*conditions)
86+
)
87+
return db_session.scalar(stmt) or 0
88+
89+
90+
def get_page_of_chat_sessions(
91+
start_time: datetime | None,
92+
end_time: datetime | None,
93+
db_session: Session,
94+
page_num: int,
95+
page_size: int,
96+
feedback_filter: QAFeedbackType | None = None,
97+
) -> Sequence[ChatSession]:
98+
conditions = _build_filter_conditions(start_time, end_time, feedback_filter)
99+
100+
subquery = (
101+
select(ChatSession.id, ChatSession.time_created)
102+
.filter(*conditions)
103+
.order_by(ChatSession.id, desc(ChatSession.time_created))
104+
.distinct(ChatSession.id)
105+
.limit(page_size)
106+
.offset(page_num * page_size)
107+
.subquery()
108+
)
109+
110+
stmt = (
111+
select(ChatSession)
112+
.join(subquery, ChatSession.id == subquery.c.id)
113+
.outerjoin(ChatMessage, ChatSession.id == ChatMessage.chat_session_id)
114+
.options(
115+
joinedload(ChatSession.user),
116+
joinedload(ChatSession.persona),
117+
contains_eager(ChatSession.messages).joinedload(
118+
ChatMessage.chat_message_feedbacks
119+
),
120+
)
121+
.order_by(desc(ChatSession.time_created), asc(ChatMessage.id))
122+
)
123+
124+
return db_session.scalars(stmt).unique().all()
17125

18126

19127
def fetch_chat_sessions_eagerly_by_time(
20-
start: datetime.datetime,
21-
end: datetime.datetime,
128+
start: datetime,
129+
end: datetime,
22130
db_session: Session,
23131
limit: int | None = 500,
24-
initial_time: datetime.datetime | None = None,
132+
initial_time: datetime | None = None,
25133
) -> list[ChatSession]:
26134
time_order: UnaryExpression = desc(ChatSession.time_created)
27135
message_order: UnaryExpression = asc(ChatMessage.id)

0 commit comments

Comments
 (0)