diff --git a/apps/server/src/orchestration-v2/ThreadSearch.test.ts b/apps/server/src/orchestration-v2/ThreadSearch.test.ts index f12355a49e554..1a84936737baa 100644 --- a/apps/server/src/orchestration-v2/ThreadSearch.test.ts +++ b/apps/server/src/orchestration-v2/ThreadSearch.test.ts @@ -178,6 +178,95 @@ it.layer(layerTest)("ThreadSearch", (it) => { assert.deepEqual((yield* search.search({ query: "ne%le" })).matches, []); }), ); + + it.effect("keeps the winning message and thread ordering when limiting matches", () => + Effect.gen(function* () { + const projections = yield* ProjectionStore.ProjectionStoreV2; + const search = yield* ThreadSearch.ThreadSearch; + const project = ProjectId.make("project:search-ranking"); + yield* createProject(project); + const older = ThreadId.make("thread:ranking-a"); + const newerFirst = ThreadId.make("thread:ranking-b"); + const newerLast = ThreadId.make("thread:ranking-z"); + const assistant = ThreadId.make("thread:ranking-assistant"); + const events: ReadonlyArray = [ + thread(older, project), + message(older, "ranking-older", "user", "ranking-probe older thread"), + thread(newerFirst, project), + message(newerFirst, "ranking-z-tie", "user", "ranking-probe losing tie", { minute: 2 }), + message(newerFirst, "ranking-a-tie", "user", "ranking-probe winning tie", { minute: 2 }), + message(newerFirst, "ranking-old-user", "user", "ranking-probe older message"), + message(newerFirst, "ranking-nonmatch", "user", "unrelated newer message", { minute: 4 }), + message(newerFirst, "ranking-streaming", "user", "ranking-probe unfinished", { + minute: 5, + streaming: true, + }), + message(newerFirst, "ranking-long-assistant", "assistant", "ranking-probe ".repeat(1024), { + minute: 6, + }), + thread(newerLast, project), + message(newerLast, "ranking-newer-last", "user", "ranking-probe second thread", { + minute: 6, + }), + thread(assistant, project), + message(assistant, "ranking-assistant", "assistant", "ranking-probe assistant only", { + minute: 8, + }), + ]; + yield* Effect.forEach(events, projections.apply, { discard: true }); + + const result = yield* search.search({ query: "ranking-probe" }); + assert.deepEqual( + result.matches.map((match) => [match.threadId, match.source, match.snippet]), + [ + [newerFirst, "user", "ranking-probe winning tie"], + [newerLast, "user", "ranking-probe second thread"], + [older, "user", "ranking-probe older thread"], + [assistant, "assistant", "ranking-probe assistant only"], + ], + ); + assert.deepEqual( + (yield* search.search({ query: "ranking-probe", limit: 2 })).matches, + result.matches.slice(0, 2), + ); + }), + ); + + it.effect("preserves literal wildcard matching and snippets from a long winning message", () => + Effect.gen(function* () { + const projections = yield* ProjectionStore.ProjectionStoreV2; + const search = yield* ThreadSearch.ThreadSearch; + const project = ProjectId.make("project:search-literals"); + const threadId = ThreadId.make("thread:search-literals"); + const decoy = ThreadId.make("thread:search-literals-decoy"); + yield* createProject(project); + yield* Effect.forEach( + [ + thread(threadId, project), + message( + threadId, + "literal-user", + "user", + `${"prefix ".repeat(80)}literal!_% marker${" suffix".repeat(80)}`, + ), + message(threadId, "literal-assistant", "assistant", "literal!_% losing answer", { + minute: 2, + }), + thread(decoy, project), + message(decoy, "literal-decoy", "user", "literal!XY marker"), + ], + projections.apply, + { discard: true }, + ); + + const result = yield* search.search({ query: "LITERAL!_%" }); + assert.lengthOf(result.matches, 1); + assert.strictEqual(result.matches[0]?.threadId, threadId); + assert.strictEqual(result.matches[0]?.source, "user"); + assert.include(result.matches[0]?.snippet ?? "", "literal!_% marker"); + assert.isAtMost(result.matches[0]?.snippet.length ?? 0, 240); + }), + ); }); it.effect("searches on the reader worker while the writer is busy", () => diff --git a/apps/server/src/orchestration-v2/ThreadSearch.ts b/apps/server/src/orchestration-v2/ThreadSearch.ts index 0862c5160dad3..fe16b3b2159f4 100644 --- a/apps/server/src/orchestration-v2/ThreadSearch.ts +++ b/apps/server/src/orchestration-v2/ThreadSearch.ts @@ -72,19 +72,27 @@ export class ThreadSearch extends Context.Service< export const make = Effect.gen(function* () { const sql = yield* SqlClient.SqlClient; + // Carrying message bodies through ranking can spill large histories to temporary files. const searchRows = SqlSchema.findAll({ Request: SearchRequest, Result: SearchRow, execute: ({ pattern, limit }) => sql` - WITH candidate AS ( + WITH ranked AS ( SELECT threads.thread_id, threads.project_id, - messages.role, - json_extract(messages.payload_json, '$.text') AS match_text, - messages.created_at AS message_created_at, + messages.role AS source, messages.message_id, - threads.updated_at AS thread_updated_at + messages.created_at AS message_created_at, + CASE messages.role WHEN 'user' THEN 0 ELSE 1 END AS match_rank, + threads.updated_at AS thread_updated_at, + ROW_NUMBER() OVER ( + PARTITION BY threads.thread_id + ORDER BY + CASE messages.role WHEN 'user' THEN 0 ELSE 1 END ASC, + messages.created_at DESC, + messages.message_id ASC + ) AS thread_match_rank FROM orchestration_v2_projection_messages AS messages INNER JOIN orchestration_v2_projection_threads AS threads ON threads.thread_id = messages.thread_id @@ -97,34 +105,23 @@ export const make = Effect.gen(function* () { AND messages.role IN ('user', 'assistant') AND json_extract(messages.payload_json, '$.text') LIKE ${pattern} ESCAPE '!' ), - ranked AS ( - SELECT - thread_id, - project_id, - role AS source, - match_text, - message_created_at, - CASE role WHEN 'user' THEN 0 ELSE 1 END AS match_rank, - thread_updated_at, - ROW_NUMBER() OVER ( - PARTITION BY thread_id - ORDER BY - CASE role WHEN 'user' THEN 0 ELSE 1 END ASC, - message_created_at DESC, - message_id ASC - ) AS thread_match_rank - FROM candidate + matches AS ( + SELECT * + FROM ranked + WHERE thread_match_rank = 1 + ORDER BY match_rank ASC, thread_updated_at DESC, thread_id ASC + LIMIT ${limit} ) SELECT - thread_id AS "threadId", - project_id AS "projectId", - source, - match_text AS "matchText", - message_created_at AS "messageCreatedAt" - FROM ranked - WHERE thread_match_rank = 1 - ORDER BY match_rank ASC, thread_updated_at DESC, thread_id ASC - LIMIT ${limit} + matches.thread_id AS "threadId", + matches.project_id AS "projectId", + matches.source, + json_extract(messages.payload_json, '$.text') AS "matchText", + matches.message_created_at AS "messageCreatedAt" + FROM matches + INNER JOIN orchestration_v2_projection_messages AS messages + ON messages.message_id = matches.message_id + ORDER BY matches.match_rank ASC, matches.thread_updated_at DESC, matches.thread_id ASC `, });