diff --git a/assistant/src/memory/conversation-crud.ts b/assistant/src/memory/conversation-crud.ts index 3a22596fb9..b119e91524 100644 --- a/assistant/src/memory/conversation-crud.ts +++ b/assistant/src/memory/conversation-crud.ts @@ -10,6 +10,7 @@ import { gte, inArray, isNull, + lt, lte, sql, } from "drizzle-orm"; @@ -1112,6 +1113,77 @@ export function getMessages(conversationId: string): MessageRow[] { .map(parseMessage); } +export interface PaginatedMessagesResult { + messages: MessageRow[]; + hasMore: boolean; +} + +export function getMessagesPaginated( + conversationId: string, + limit: number | undefined, + beforeTimestamp?: number, +): PaginatedMessagesResult { + const db = getDb(); + + if (limit === undefined) { + const conditions = [eq(messages.conversationId, conversationId)]; + if (beforeTimestamp !== undefined) { + conditions.push(lt(messages.createdAt, beforeTimestamp)); + } + const rows = db + .select() + .from(messages) + .where(and(...conditions)) + .orderBy(asc(messages.createdAt)) + .all() + .map(parseMessage); + return { messages: rows, hasMore: false }; + } + + const conditions = [eq(messages.conversationId, conversationId)]; + if (beforeTimestamp !== undefined) { + conditions.push(lt(messages.createdAt, beforeTimestamp)); + } + + const rows = db + .select() + .from(messages) + .where(and(...conditions)) + .orderBy(desc(messages.createdAt)) + .limit(limit + 1) + .all() + .map(parseMessage); + + const hasMore = rows.length > limit; + if (hasMore) { + rows.splice(limit); + } + rows.reverse(); + + return { messages: rows, hasMore }; +} + +export function getLastAssistantTimestampBefore( + conversationId: string, + beforeTimestamp: number, +): number { + const db = getDb(); + const row = db + .select({ createdAt: messages.createdAt }) + .from(messages) + .where( + and( + eq(messages.conversationId, conversationId), + eq(messages.role, "assistant"), + lt(messages.createdAt, beforeTimestamp), + ), + ) + .orderBy(desc(messages.createdAt)) + .limit(1) + .get(); + return row?.createdAt ?? 0; +} + /** Fetch a single message by ID, optionally scoped to a specific conversation. */ export function getMessageById( messageId: string, diff --git a/assistant/src/memory/db-init.ts b/assistant/src/memory/db-init.ts index 36a058d5be..ad4b98efe4 100644 --- a/assistant/src/memory/db-init.ts +++ b/assistant/src/memory/db-init.ts @@ -89,6 +89,7 @@ import { migrateLlmRequestLogMessageId, migrateLlmRequestLogProvider, migrateMemoryItemSupersession, + migrateMessagesConversationCreatedAtIndex, migrateMessagesFtsBackfill, migrateNormalizePhoneIdentities, migrateNotificationDeliveryThreadDecision, @@ -520,6 +521,9 @@ export function initializeDb(): void { // 93. Strip `integration:` prefix from provider_key across OAuth tables migrateStripIntegrationPrefixFromProviderKeys(database); + // 94. Composite index on messages(conversation_id, created_at) for paginated history queries + migrateMessagesConversationCreatedAtIndex(database); + validateMigrationState(database); if (process.env.BUN_TEST === "1") { diff --git a/assistant/src/memory/migrations/196-messages-conversation-created-at-index.ts b/assistant/src/memory/migrations/196-messages-conversation-created-at-index.ts new file mode 100644 index 0000000000..ce91449e22 --- /dev/null +++ b/assistant/src/memory/migrations/196-messages-conversation-created-at-index.ts @@ -0,0 +1,9 @@ +import type { DrizzleDb } from "../db-connection.js"; + +export function migrateMessagesConversationCreatedAtIndex( + database: DrizzleDb, +): void { + database.run( + /*sql*/ `CREATE INDEX IF NOT EXISTS idx_messages_conversation_created_at ON messages(conversation_id, created_at)`, + ); +} diff --git a/assistant/src/memory/migrations/index.ts b/assistant/src/memory/migrations/index.ts index 637462ee28..d5182f3e83 100644 --- a/assistant/src/memory/migrations/index.ts +++ b/assistant/src/memory/migrations/index.ts @@ -134,6 +134,7 @@ export { migrateContactsUserFileColumn } from "./192-contacts-user-file-column.j export { migrateAddSourceTypeColumns } from "./193-add-source-type-columns.js"; export { migrateCreateMemoryRecallLogs } from "./194-memory-recall-logs.js"; export { migrateOAuthProvidersPingConfig } from "./195-oauth-providers-ping-config.js"; +export { migrateMessagesConversationCreatedAtIndex } from "./196-messages-conversation-created-at-index.js"; export { migrateStripIntegrationPrefixFromProviderKeys } from "./196-strip-integration-prefix-from-provider-keys.js"; export { MIGRATION_REGISTRY,