profile: improve group history batching

src/main/kotlin/profile/ProfileHistoryReader.kt

@@ -9,8 +9,17 @@ import java.io.File
 import java.security.MessageDigest
 import java.sql.Connection
 import java.sql.DriverManager
+import java.sql.PreparedStatement
 import java.sql.ResultSet
 
+internal object ProfileConversationWindowDefaults {
+    const val IDLE_GAP_SECONDS = 20 * 60
+    const val TARGET_CONTENT_CHARS = 30_000
+    const val MAX_MESSAGES = 800
+    const val MAX_CONTENT_CHARS = 70_000
+    const val MAX_PACKED_SPAN_SECONDS = 24 * 60 * 60
+}
+
 class ProfileHistoryReader(private val databaseFile: File) {
     data class TimeBounds(val startTime: Int, val endTime: Int)
     data class GroupTimeBounds(
@@ -77,31 +86,96 @@ class ProfileHistoryReader(private val databaseFile: File) {
     }
 
     fun listGroupTimeBounds(): List<GroupTimeBounds> = openReadConnection().use { connection ->
-        connection.prepareStatement(
+        val pairs = connection.prepareStatement(
             """
-            SELECT bot_id, target_id, MIN(time) AS min_time, MAX(time) AS max_time
+            SELECT bot_id, target_id
             FROM message_record
-            WHERE kind = ? AND recalled = 0 AND target_id > 0
+            WHERE kind = ? AND target_id > 0
             GROUP BY bot_id, target_id
-            ORDER BY max_time DESC, target_id ASC, bot_id ASC
             """.trimIndent()
         ).use { statement ->
             statement.setInt(1, MessageSourceKind.GROUP.ordinal)
             statement.executeQuery().use { results ->
-                val seenGroupIds = hashSetOf<Long>()
                 buildList {
                     while (results.next()) {
-                        val groupId = results.getLong("target_id")
+                        add(results.getLong("bot_id") to results.getLong("target_id"))
-                        if (!seenGroupIds.add(groupId)) continue
+                    }
-                        add(
+                }
-                            GroupTimeBounds(
+            }
-                                botId = results.getLong("bot_id"),
+        }
+
+        val selectedByGroup = hashMapOf<Long, GroupTimeBounds>()
+        connection.prepareStatement(GROUP_ENDPOINT_TIME_SQL.format("ASC", "ASC")).use { earliest ->
+            connection.prepareStatement(GROUP_ENDPOINT_TIME_SQL.format("DESC", "DESC")).use { latest ->
+                pairs.forEach { (botId, groupId) ->
+                    val startTime = queryGroupEndpointTime(earliest, botId, groupId) ?: return@forEach
+                    val endTime = queryGroupEndpointTime(latest, botId, groupId)?.safeNextSecond()
+                        ?: return@forEach
+                    val candidate = GroupTimeBounds(
+                        botId = botId,
                         groupId = groupId,
-                                startTime = results.getInt("min_time"),
+                        startTime = startTime,
-                                endTime = results.getInt("max_time").safeNextSecond(),
+                        endTime = endTime,
                     )
+                    val current = selectedByGroup[groupId]
+                    if (current == null ||
+                        candidate.endTime > current.endTime ||
+                        (candidate.endTime == current.endTime && candidate.botId < current.botId)
+                    ) {
+                        selectedByGroup[groupId] = candidate
+                    }
+                }
+            }
+        }
+        selectedByGroup.values.sortedWith(
+            compareByDescending<GroupTimeBounds> { it.endTime }
+                .thenBy { it.groupId }
+                .thenBy { it.botId }
         )
     }
+
+    private fun queryGroupEndpointTime(
+        statement: PreparedStatement,
+        botId: Long,
+        groupId: Long,
+    ): Int? {
+        statement.setLong(1, botId)
+        statement.setInt(2, MessageSourceKind.GROUP.ordinal)
+        statement.setLong(3, groupId)
+        return statement.executeQuery().use { results ->
+            if (results.next()) results.getInt("time") else null
+        }
+    }
+
+    fun filterGroupRangesByMinimumMessageCount(
+        ranges: List<GroupTimeBounds>,
+        minimumMessages: Int,
+    ): List<GroupTimeBounds> {
+        if (ranges.isEmpty()) return emptyList()
+        val requiredMessages = minimumMessages.coerceAtLeast(1)
+        return openReadConnection().use { connection ->
+            connection.prepareStatement(
+                """
+                SELECT 1
+                FROM message_record
+                WHERE bot_id = ? AND target_id = ? AND kind = ? AND recalled = 0
+                  AND time >= ? AND time < ?
+                ORDER BY time ASC, id ASC
+                LIMIT ?
+                """.trimIndent()
+            ).use { statement ->
+                ranges.filter { range ->
+                    statement.setLong(1, range.botId)
+                    statement.setLong(2, range.groupId)
+                    statement.setInt(3, MessageSourceKind.GROUP.ordinal)
+                    statement.setInt(4, range.startTime)
+                    statement.setInt(5, range.endTime)
+                    statement.setInt(6, requiredMessages)
+                    statement.executeQuery().use { results ->
+                        var messages = 0
+                        while (messages < requiredMessages && results.next()) messages++
+                        messages >= requiredMessages
+                    }
                 }
             }
         }
@@ -230,9 +304,18 @@ class ProfileHistoryReader(private val databaseFile: File) {
         snapshotEndTime: Int,
         messageLimit: Int,
         maxMessageChars: Int,
+        idleGapSeconds: Int = ProfileConversationWindowDefaults.IDLE_GAP_SECONDS,
+        targetContentChars: Int = ProfileConversationWindowDefaults.TARGET_CONTENT_CHARS,
+        maxMessages: Int = ProfileConversationWindowDefaults.MAX_MESSAGES,
+        maxContentChars: Int = ProfileConversationWindowDefaults.MAX_CONTENT_CHARS,
+        maxPackedSpanSeconds: Int = ProfileConversationWindowDefaults.MAX_PACKED_SPAN_SECONDS,
     ): ConversationProfileBatch? {
         require(startTime <= snapshotEndTime) { "startTime must not be after snapshotEndTime" }
         if (startTime == snapshotEndTime) return null
+        val targetMessages = messageLimit.coerceAtLeast(1)
+        val adaptiveGap = idleGapSeconds.coerceAtLeast(0)
+        val hardMessageLimit = maxMessages.coerceAtLeast(targetMessages)
+        val adaptiveWindow = adaptiveGap > 0
         return openReadConnection().use { connection ->
             val firstPage = queryOldestConversationMessages(
                 connection = connection,
@@ -240,11 +323,24 @@ class ProfileHistoryReader(private val databaseFile: File) {
                 groupId = groupId,
                 startTime = startTime,
                 endTime = snapshotEndTime,
-                limit = messageLimit.coerceAtLeast(1),
+                limit = if (adaptiveWindow) hardMessageLimit.safeIncrement() else targetMessages,
             )
             if (firstPage.isEmpty()) return@use null
 
-            val endTime = firstPage.last().time.safeNextSecond()
+            val endTime = if (adaptiveWindow) {
+                selectConversationWindowEndTime(
+                    records = firstPage,
+                    idleGapSeconds = adaptiveGap,
+                    targetMessages = targetMessages,
+                    targetContentChars = targetContentChars.coerceAtLeast(1),
+                    maxMessages = hardMessageLimit,
+                    maxContentChars = maxContentChars.coerceAtLeast(targetContentChars.coerceAtLeast(1)),
+                    maxPackedSpanSeconds = maxPackedSpanSeconds.coerceAtLeast(1),
+                    maxMessageChars = maxMessageChars,
+                )
+            } else {
+                firstPage.last().time.safeNextSecond()
+            }
             val records = queryOldestConversationMessages(
                 connection = connection,
                 botId = botId,
@@ -253,10 +349,58 @@ class ProfileHistoryReader(private val databaseFile: File) {
                 endTime = endTime,
                 limit = Int.MAX_VALUE,
             )
-            createConversationBatch(botId, groupId, startTime, endTime, records, maxMessageChars)
+            createConversationBatch(
+                botId = botId,
+                groupId = groupId,
+                startTime = startTime,
+                endTime = endTime,
+                records = records,
+                maxMessageChars = maxMessageChars,
+                episodeGapSeconds = adaptiveGap.takeIf { adaptiveWindow },
+            )
         }
     }
 
+    private fun selectConversationWindowEndTime(
+        records: List<ChatMessageRecord>,
+        idleGapSeconds: Int,
+        targetMessages: Int,
+        targetContentChars: Int,
+        maxMessages: Int,
+        maxContentChars: Int,
+        maxPackedSpanSeconds: Int,
+        maxMessageChars: Int,
+    ): Int {
+        var selectedMessages = 0
+        var selectedChars = 0L
+        val targetContentLimit = targetContentChars.toLong()
+        val hardContentLimit = maxContentChars.toLong().coerceAtLeast(targetContentLimit)
+        val firstTime = records.first().time
+        var previousTime = firstTime
+
+        records.forEachIndexed { index, record ->
+            val timeBoundary = index > 0 && record.time != previousTime
+            if (timeBoundary && selectedMessages > 0) {
+                val startsNewEpisode = record.time - previousTime > idleGapSeconds
+                val nextChars = estimatePromptChars(record, maxMessageChars)
+                val exceedsHardSize = selectedMessages >= maxMessages || selectedChars + nextChars > hardContentLimit
+                val exceedsPackedSpan = startsNewEpisode && record.time - firstTime > maxPackedSpanSeconds
+                val reachedTargetAtEpisodeBoundary = startsNewEpisode &&
+                        (selectedMessages >= targetMessages || selectedChars >= targetContentLimit)
+                if (exceedsHardSize || exceedsPackedSpan || reachedTargetAtEpisodeBoundary) {
+                    return records[index - 1].time.safeNextSecond()
+                }
+            }
+            selectedMessages++
+            selectedChars += estimatePromptChars(record, maxMessageChars)
+            previousTime = record.time
+        }
+        return records.last().time.safeNextSecond()
+    }
+
+    private fun estimatePromptChars(record: ChatMessageRecord, maxMessageChars: Int): Int =
+        record.code.length.coerceAtMost(maxMessageChars.coerceAtLeast(80))
+
     private fun queryTargetMessages(
         connection: Connection,
         userId: Long,
@@ -352,6 +496,7 @@ class ProfileHistoryReader(private val databaseFile: File) {
         endTime: Int,
         records: List<ChatMessageRecord>,
         maxMessageChars: Int,
+        episodeGapSeconds: Int? = null,
     ): ConversationProfileBatch {
         val participantIds = buildSet {
             add(botId)
@@ -367,12 +512,19 @@ class ProfileHistoryReader(private val databaseFile: File) {
                 .sorted()
                 .forEachIndexed { index, participantId -> put(participantId, "U${index + 1}") }
         }
+        var episodeIndex = 1
+        var previousTime: Int? = null
         val promptMessages = records.mapIndexed { index, record ->
+            val lastTime = previousTime
+            if (episodeGapSeconds != null && lastTime != null && record.time - lastTime > episodeGapSeconds) {
+                episodeIndex++
+            }
+            previousTime = record.time
             ProfilePromptMessage(
                 record = record,
                 text = ProfileMessageRenderer.render(record, aliases, maxMessageChars.coerceAtLeast(80)),
                 evidenceRef = index + 1,
-                episodeIndex = 1,
+                episodeIndex = episodeIndex,
             )
         }
         return ConversationProfileBatch(
@@ -655,4 +807,16 @@ class ProfileHistoryReader(private val databaseFile: File) {
         .joinToString("") { byte -> "%02x".format(byte) }
 
     private fun Int.safeNextSecond(): Int = if (this == Int.MAX_VALUE) this else this + 1
+
+    private fun Int.safeIncrement(): Int = if (this == Int.MAX_VALUE) this else this + 1
+
+    private companion object {
+        const val GROUP_ENDPOINT_TIME_SQL = """
+            SELECT time
+            FROM message_record
+            WHERE bot_id = ? AND kind = ? AND target_id = ? AND recalled = 0
+            ORDER BY time %s, id %s
+            LIMIT 1
+        """
+    }
 }

src/test/kotlin/profile/ProfileHistoryReaderTest.kt

@@ -10,6 +10,81 @@ import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
 
 class ProfileHistoryReaderTest {
+    @Test
+    fun filtersPendingGroupRangesByValidMessageCount() {
+        val directory = Files.createTempDirectory("jchatgpt-profile-pending-groups-test-")
+        val database = directory.resolve("history.sqlite")
+        try {
+            DriverManager.getConnection("jdbc:sqlite:${database.absolutePathString()}").use { connection ->
+                connection.createStatement().use { statement ->
+                    statement.executeUpdate(
+                        """
+                        CREATE TABLE message_record(
+                            id INTEGER PRIMARY KEY AUTOINCREMENT,
+                            bot_id INTEGER NOT NULL,
+                            from_id INTEGER NOT NULL,
+                            target_id INTEGER NOT NULL,
+                            ids TEXT,
+                            internal_ids TEXT,
+                            time INTEGER NOT NULL,
+                            kind INTEGER NOT NULL,
+                            code TEXT NOT NULL,
+                            recalled INTEGER NOT NULL DEFAULT 0
+                        )
+                        """.trimIndent()
+                    )
+                }
+                connection.prepareStatement(
+                    "INSERT INTO message_record(" +
+                            "bot_id, from_id, target_id, time, kind, code, recalled" +
+                            ") VALUES (?, 100, ?, ?, ?, '[]', ?)"
+                ).use { statement ->
+                    fun insert(
+                        botId: Long,
+                        groupId: Long,
+                        time: Int,
+                        kind: MessageSourceKind = MessageSourceKind.GROUP,
+                        recalled: Int = 0,
+                    ) {
+                        statement.setLong(1, botId)
+                        statement.setLong(2, groupId)
+                        statement.setInt(3, time)
+                        statement.setInt(4, kind.ordinal)
+                        statement.setInt(5, recalled)
+                        statement.executeUpdate()
+                    }
+
+                    insert(1, 10, 100)
+                    insert(1, 10, 110, recalled = 1)
+                    insert(1, 10, 120)
+                    insert(1, 10, 130)
+                    insert(1, 20, 100)
+                    insert(1, 20, 110, kind = MessageSourceKind.FRIEND)
+                    insert(2, 30, 100)
+                    insert(2, 30, 110)
+                }
+            }
+
+            val reader = ProfileHistoryReader(database.toFile())
+            val ranges = listOf(
+                ProfileHistoryReader.GroupTimeBounds(1, 10, 110, 131),
+                ProfileHistoryReader.GroupTimeBounds(1, 20, 90, 120),
+                ProfileHistoryReader.GroupTimeBounds(1, 30, 90, 120),
+            )
+
+            assertEquals(
+                listOf(10L),
+                reader.filterGroupRangesByMinimumMessageCount(ranges, 2).map { it.groupId },
+            )
+            assertEquals(
+                listOf(10L, 20L),
+                reader.filterGroupRangesByMinimumMessageCount(ranges, 1).map { it.groupId },
+            )
+        } finally {
+            directory.toFile().deleteRecursively()
+        }
+    }
+
     @Test
     fun keepsEqualTimestampTargetMessagesTogetherAndLoadsGroupContext() {
         val directory = Files.createTempDirectory("jchatgpt-profile-history-test-")
@@ -136,6 +211,7 @@ class ProfileHistoryReaderTest {
                     snapshotEndTime = groupBounds.endTime,
                     messageLimit = 2,
                     maxMessageChars = 200,
+                    idleGapSeconds = 0,
                 )
             )
             assertEquals(111, firstGroupBatch.endTime)
@@ -148,10 +224,80 @@ class ProfileHistoryReaderTest {
                     snapshotEndTime = groupBounds.endTime,
                     messageLimit = 2,
                     maxMessageChars = 200,
+                    idleGapSeconds = 0,
                 )
             )
             assertEquals(201, secondGroupBatch.endTime)
             assertEquals(listOf(130, 200), secondGroupBatch.messages.map { it.record.time })
+
+            val defaultAdaptiveGroupBatch = assertNotNull(
+                reader.loadNextConversationBatch(
+                    botId = groupBounds.botId,
+                    groupId = 10,
+                    startTime = groupBounds.startTime,
+                    snapshotEndTime = groupBounds.endTime,
+                    messageLimit = 2,
+                    maxMessageChars = 200,
+                )
+            )
+            assertEquals(201, defaultAdaptiveGroupBatch.endTime)
+            assertEquals(listOf(100, 110, 110, 130, 200), defaultAdaptiveGroupBatch.messages.map { it.record.time })
+            assertEquals(listOf(1, 1, 1, 1, 1), defaultAdaptiveGroupBatch.messages.map { it.episodeIndex })
+
+            val adaptiveGroupBatch = assertNotNull(
+                reader.loadNextConversationBatch(
+                    botId = groupBounds.botId,
+                    groupId = 10,
+                    startTime = groupBounds.startTime,
+                    snapshotEndTime = groupBounds.endTime,
+                    messageLimit = 2,
+                    maxMessageChars = 200,
+                    idleGapSeconds = 30,
+                    targetContentChars = 10_000,
+                    maxMessages = 20,
+                    maxContentChars = 100_000,
+                    maxPackedSpanSeconds = 1_000,
+                )
+            )
+            assertEquals(131, adaptiveGroupBatch.endTime)
+            assertEquals(listOf(100, 110, 110, 130), adaptiveGroupBatch.messages.map { it.record.time })
+            assertEquals(listOf(1, 1, 1, 1), adaptiveGroupBatch.messages.map { it.episodeIndex })
+
+            val packedSmallEpisodes = assertNotNull(
+                reader.loadNextConversationBatch(
+                    botId = groupBounds.botId,
+                    groupId = 10,
+                    startTime = groupBounds.startTime,
+                    snapshotEndTime = groupBounds.endTime,
+                    messageLimit = 10,
+                    maxMessageChars = 200,
+                    idleGapSeconds = 30,
+                    targetContentChars = 10_000,
+                    maxMessages = 20,
+                    maxContentChars = 100_000,
+                    maxPackedSpanSeconds = 1_000,
+                )
+            )
+            assertEquals(201, packedSmallEpisodes.endTime)
+            assertEquals(listOf(1, 1, 1, 1, 2), packedSmallEpisodes.messages.map { it.episodeIndex })
+
+            val hardLimitedConversation = assertNotNull(
+                reader.loadNextConversationBatch(
+                    botId = groupBounds.botId,
+                    groupId = 10,
+                    startTime = groupBounds.startTime,
+                    snapshotEndTime = groupBounds.endTime,
+                    messageLimit = 2,
+                    maxMessageChars = 200,
+                    idleGapSeconds = 30,
+                    targetContentChars = 10_000,
+                    maxMessages = 3,
+                    maxContentChars = 100_000,
+                    maxPackedSpanSeconds = 1_000,
+                )
+            )
+            assertEquals(111, hardLimitedConversation.endTime)
+            assertEquals(listOf(100, 110, 110), hardLimitedConversation.messages.map { it.record.time })
         } finally {
             directory.toFile().deleteRecursively()
         }