70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import asc, desc, nulls_last, select
|
|
|
|
from letta.orm.run import Run as RunModel
|
|
from letta.services.helpers.agent_manager_helper import _cursor_filter
|
|
from letta.settings import DatabaseChoice, settings
|
|
|
|
|
|
async def _apply_pagination_async(
|
|
query,
|
|
before: Optional[str],
|
|
after: Optional[str],
|
|
session,
|
|
ascending: bool = True,
|
|
sort_by: str = "created_at",
|
|
) -> any:
|
|
# Determine the sort column
|
|
if sort_by == "last_run_completion":
|
|
sort_column = RunModel.last_run_completion
|
|
sort_nulls_last = True # TODO: handle this as a query param eventually
|
|
else:
|
|
sort_column = RunModel.created_at
|
|
sort_nulls_last = False
|
|
|
|
if after:
|
|
result = (await session.execute(select(sort_column, RunModel.id).where(RunModel.id == after))).first()
|
|
if result:
|
|
after_sort_value, after_id = result
|
|
# SQLite does not support as granular timestamping, so we need to round the timestamp
|
|
if settings.database_engine is DatabaseChoice.SQLITE and isinstance(after_sort_value, datetime):
|
|
after_sort_value = after_sort_value.strftime("%Y-%m-%d %H:%M:%S")
|
|
query = query.where(
|
|
_cursor_filter(
|
|
sort_column,
|
|
RunModel.id,
|
|
after_sort_value,
|
|
after_id,
|
|
forward=not ascending,
|
|
nulls_last=sort_nulls_last,
|
|
)
|
|
)
|
|
|
|
if before:
|
|
result = (await session.execute(select(sort_column, RunModel.id).where(RunModel.id == before))).first()
|
|
if result:
|
|
before_sort_value, before_id = result
|
|
# SQLite does not support as granular timestamping, so we need to round the timestamp
|
|
if settings.database_engine is DatabaseChoice.SQLITE and isinstance(before_sort_value, datetime):
|
|
before_sort_value = before_sort_value.strftime("%Y-%m-%d %H:%M:%S")
|
|
query = query.where(
|
|
_cursor_filter(
|
|
sort_column,
|
|
RunModel.id,
|
|
before_sort_value,
|
|
before_id,
|
|
forward=ascending,
|
|
nulls_last=sort_nulls_last,
|
|
)
|
|
)
|
|
|
|
# Apply ordering
|
|
order_fn = asc if ascending else desc
|
|
query = query.order_by(
|
|
nulls_last(order_fn(sort_column)) if sort_nulls_last else order_fn(sort_column),
|
|
order_fn(RunModel.id),
|
|
)
|
|
return query
|