Skip to content

Commit 1c4ffa5

Browse files
committed
improvement(knowledge): maintain keyword projections only for search indexes
1 parent a0c93d6 commit 1c4ffa5

10 files changed

Lines changed: 486 additions & 23 deletions

‎apps/sim/lib/knowledge/__integration__/embedding-insert-batches.integration.ts‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -151,11 +151,9 @@ describe('bounded embedding insert transactions', () => {
151151
.from(embedding)
152152
.where(eq(embedding.documentId, file.documentId))
153153
).toEqual([{ id: previousId }])
154-
for (const table of ['embedding_search', 'embedding_keyword_search']) {
155-
expect(
156-
await db.$client.unsafe(`SELECT id FROM ${table} WHERE document_id = $1`, [file.documentId])
157-
).toEqual([{ id: previousId }])
158-
}
154+
expect(
155+
await db.$client`SELECT id FROM embedding_search WHERE document_id = ${file.documentId}`
156+
).toEqual([{ id: previousId }])
159157
expect(
160158
await db
161159
.select()
@@ -175,13 +173,18 @@ describe('bounded embedding insert transactions', () => {
175173
expect(await db.select().from(document).where(eq(document.id, file.documentId))).toMatchObject([
176174
{ processingStatus: 'completed', chunkCount: 205, processingError: null },
177175
])
178-
for (const table of ['embedding', 'embedding_search', 'embedding_keyword_search']) {
176+
/** A workspace knowledge base keeps no keyword projection rows: its keywords rank on `embedding`. */
177+
for (const [table, count] of [
178+
['embedding', 205],
179+
['embedding_search', 205],
180+
['embedding_keyword_search', 0],
181+
] as const) {
179182
expect(
180183
await db.$client.unsafe(
181184
`SELECT count(*)::int AS count FROM ${table} WHERE document_id = $1`,
182185
[file.documentId]
183186
)
184-
).toEqual([{ count: 205 }])
187+
).toEqual([{ count }])
185188
}
186189
expect(
187190
await db.$client`SELECT count(*)::int AS count FROM embedding_secret_provenance p

‎apps/sim/lib/knowledge/__integration__/knowledge-projection.integration.ts‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
credentialGroupEnrollment,
2525
document,
2626
embedding,
27+
embeddingKeywordSearch,
2728
embeddingKeywordTin,
2829
embeddingSearch,
2930
knowledgeBase,
@@ -383,7 +384,6 @@ beforeAll(async () => {
383384
sql.raw(`CREATE SCHEMA tin;
384385
CREATE FUNCTION tin.full_score(tid) RETURNS double precision LANGUAGE sql IMMUTABLE AS 'SELECT 1.0::float8';
385386
CREATE FUNCTION knowledge_tin_base_token(text) RETURNS text LANGUAGE sql IMMUTABLE AS $$SELECT 'kb'$$;
386-
CREATE FUNCTION knowledge_tin_membership_key(text) RETURNS bigint LANGUAGE sql IMMUTABLE AS $$SELECT hashtextextended('embedding_keyword_tin:' || $1, 0)$$;
387387
CREATE FUNCTION knowledge_tin_stream(vector tsvector) RETURNS text LANGUAGE sql IMMUTABLE AS $$
388388
SELECT coalesce(string_agg(entry.lexeme, ' ' ORDER BY position), '')
389389
FROM unnest(vector) AS entry(lexeme, positions, weights), unnest(entry.positions) AS position
@@ -400,7 +400,6 @@ afterAll(async () => {
400400
sql.raw(`DROP OPERATOR IF EXISTS ==> (text, text);
401401
DROP FUNCTION IF EXISTS tin_fixture_match(text, text);
402402
DROP FUNCTION IF EXISTS knowledge_tin_base_token(text);
403-
DROP FUNCTION IF EXISTS knowledge_tin_membership_key(text);
404403
DROP FUNCTION IF EXISTS knowledge_tin_stream(tsvector);
405404
DROP SCHEMA IF EXISTS tin CASCADE;`)
406405
)
@@ -730,6 +729,46 @@ describe('the projector', () => {
730729
expect(await markOf()).toBeUndefined()
731730
})
732731

732+
it('writes keyword rows for search-index knowledge bases only', async () => {
733+
const workspaceBaseId = generateId()
734+
const workspaceDocument = generateId()
735+
const workspaceChunk = generateId()
736+
await db.insert(knowledgeBase).values({
737+
id: workspaceBaseId,
738+
userId: ids.aliceId,
739+
workspaceId: ids.workspaceId,
740+
name: 'Workspace keyword fixture',
741+
chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 20 },
742+
})
743+
try {
744+
await db.insert(document).values({
745+
id: workspaceDocument,
746+
knowledgeBaseId: workspaceBaseId,
747+
filename: 'keyword.md',
748+
fileUrl: 'https://fixture.test/keyword',
749+
fileSize: 12,
750+
mimeType: 'text/plain',
751+
processingStatus: 'completed',
752+
})
753+
await write('async', (tx) =>
754+
tx.insert(embedding).values({
755+
...chunkRow(workspaceChunk, 0),
756+
documentId: workspaceDocument,
757+
knowledgeBaseId: workspaceBaseId,
758+
})
759+
)
760+
await project()
761+
const keywordRows = await db
762+
.select({ id: embeddingKeywordSearch.id })
763+
.from(embeddingKeywordSearch)
764+
.where(inArray(embeddingKeywordSearch.id, [chunkId, workspaceChunk]))
765+
expect(keywordRows).toEqual([{ id: chunkId }])
766+
expect(await rowAcl(embeddingSearch, workspaceChunk)).toBeDefined()
767+
} finally {
768+
await db.delete(knowledgeBase).where(eq(knowledgeBase.id, workspaceBaseId))
769+
}
770+
})
771+
733772
it.each(['sync', 'async'] as const)(
734773
'writes %s projection rows from a chunk commit only when the writer did not defer them',
735774
async (mode) => {

‎apps/sim/lib/sim-search/indexed/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ While the gate is off:
1414
- Indexed-only surfaces refuse with `SearchIndexDormantError` (a `409`): the Stats report and connecting a source that crawls into a search index. The indexed document page is not found.
1515
- A knowledge search that names a search-index knowledge base (the Knowledge block, v1, v2, Sim's knowledge tool) still answers from the documents it already holds, decided on each document exactly as a workspace knowledge base is.
1616
- Nothing crawls into search indexes: content syncs, member syncs, and processing recovery skip them (`lib/knowledge/connectors/indexing-policy.ts`).
17-
- The projector owes search-index documents nothing: their marks are released with the rest, and it writes no Tin keyword rows.
17+
- The projector owes search-index documents nothing: their marks are released with the rest, and it writes no Tin keyword rows. The GIN keyword projection follows `is_search_index` alone, so it keeps its search-index rows either way.
1818

1919
## Layout
2020

@@ -28,7 +28,7 @@ The dormant UI sits in `indexed/` folders next to the component that picks it fr
2828
## Re-enabling
2929

3030
1. Set `SIM_SEARCH_LIVE=false` in both the app and the Trigger.dev environment, and deploy. The container entrypoint (`apps/sim/bootstrap.ts`) mirrors it to `NEXT_PUBLIC_SIM_SEARCH_LIVE` for the client; crawling, processing, and projection read it in whichever process runs them.
31-
2. Confirm the Tin objects exist (`0019_tin_keyword_projection`, `0024_knowledge_projection_async`), backfill `embedding_keyword_tin` for every search-index knowledge base, and build its index.
31+
2. Confirm the keyword projection objects exist (`0019_tin_keyword_projection`, `0024_knowledge_projection_async`, `0025_scope_keyword_projections`). Both keyword projections, `embedding_keyword_search` and `embedding_keyword_tin`, hold only search-index rows, written by the chunk triggers and by the trigger on `knowledge_base.is_search_index`. Backfill both for every search-index knowledge base whose rows were removed while dormant, and build the Tin index.
3232
3. Resume and fully resync the connectors of search-index knowledge bases, so content that went stale while dormant is indexed again.
3333

3434
Projection rows written before projections carried their document's source and ACL are decided on their document until they are rewritten.

‎packages/db/knowledge-projection.ts‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,10 @@ export type SourceAclProjection = (typeof SOURCE_ACL_PROJECTIONS)[number]
106106
const mirrorsSourceAcl = (projection: KnowledgeProjection): projection is SourceAclProjection =>
107107
(SOURCE_ACL_PROJECTIONS as readonly string[]).includes(projection)
108108

109+
/** The keyword projections, which hold only the rows of search-index knowledge bases. */
110+
const holdsSearchIndexesOnly = (projection: KnowledgeProjection) =>
111+
projection === 'embedding_keyword_search' || projection === 'embedding_keyword_tin'
112+
109113
/**
110114
* The chunks one page covers: the next {@link PROJECTION_ROW_BATCH_SIZE} of the document in
111115
* chunk order, read off `emb_doc_chunk_idx`. Every projection pages the same way, so a page is
@@ -163,10 +167,16 @@ function contentPageStatement(projection: KnowledgeProjection): string {
163167
}
164168
if (projection === 'embedding_keyword_search') {
165169
const compared = ['knowledge_base_id', 'document_id', 'enabled', 'content_tsv']
166-
return `WITH ${PAGE}, written AS (
167-
INSERT INTO embedding_keyword_search AS s (${['id', ...compared].join(', ')})
168-
SELECT e.id, e.knowledge_base_id, e.document_id, e.enabled, e.content_tsv
170+
return `WITH ${PAGE}, source AS MATERIALIZED (
171+
SELECT e.id, ${compared.map((column) => `e.${column}`).join(', ')}, k.is_search_index
169172
FROM page p JOIN embedding e ON e.id = p.id
173+
JOIN knowledge_base k ON k.id = e.knowledge_base_id
174+
), removed AS (
175+
DELETE FROM embedding_keyword_search s USING source
176+
WHERE s.id = source.id AND NOT source.is_search_index
177+
), written AS (
178+
INSERT INTO embedding_keyword_search AS s (${['id', ...compared].join(', ')})
179+
SELECT id, ${compared.join(', ')} FROM source WHERE is_search_index
170180
ON CONFLICT (id) DO UPDATE SET
171181
${compared.map((column) => `${column} = EXCLUDED.${column}`).join(', ')}
172182
WHERE (${compared.map((column) => `s.${column}`).join(', ')})
@@ -249,8 +259,9 @@ export interface KnowledgeProjectionOptions {
249259
/**
250260
* Whether search-index rows are read, that is whether indexed organization search is on (see
251261
* {@link MarkScope}). Only then is the Tin keyword projection written, since it holds only
252-
* search-index rows; the other projections are written either way, and Tin is still skipped
253-
* where it is not installed.
262+
* search-index rows; the other projections are written either way, the GIN keyword projection
263+
* for search-index bases alone whatever this says, and Tin is still skipped where it is not
264+
* installed.
254265
*/
255266
searchIndexes: boolean
256267
/** Called after each page commits, for tests that interleave writes with a run. */
@@ -375,8 +386,8 @@ async function projectDocumentRows(
375386
if (Date.now() >= deadline) return { pages, written, finished: false }
376387
const page = await sql.begin(async (tx) => {
377388
await enterProjectorTransaction(tx, PROJECTION_PAGE_LOCK_TIMEOUT_MS)
378-
/** Shares the base's Tin membership lock, as the embedding trigger does, so a flip of its marker waits. */
379-
if (projection === 'embedding_keyword_tin')
389+
/** Shares the base's membership lock, as the embedding triggers do, so a flip of its marker waits. */
390+
if (holdsSearchIndexesOnly(projection))
380391
await tx`SELECT pg_advisory_xact_lock_shared(knowledge_tin_membership_key(${mark.knowledgeBaseId}))`
381392
const [row] = await tx.unsafe<
382393
Array<{ scanned: number; written: number; last_chunk: number | null }>

‎packages/db/script-migrations/0016_backfill_search_vectors.integration.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,7 @@ describe('search projection upgrade in PostgreSQL', () => {
395395
{ name: '0022_projection_source_acl_backfill' },
396396
{ name: '0023_projection_acl_skip_unfilled' },
397397
{ name: '0024_knowledge_projection_async' },
398+
{ name: '0025_scope_keyword_projections' },
398399
])
399400
const [{ complete }] = await sql`SELECT count(*)::int AS complete FROM embedding e
400401
JOIN embedding_search s ON s.id = e.id JOIN embedding_keyword_search k ON k.id = e.id

‎packages/db/script-migrations/0019_tin_keyword_projection.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { EMBEDDING_KEYWORD_TIN_INDEX } from '@sim/db/schema'
22
import { resolveMigrationDatabaseUrl } from '@sim/db/script-migrations/database-url'
33
import { type ScriptMigration, ScriptMigrationDeferred } from '@sim/db/script-migrations/types'
44
import { createLogger } from '@sim/logger'
5-
import postgres, { type Sql } from 'postgres'
5+
import postgres, { type Sql, type TransactionSql } from 'postgres'
66

77
const logger = createLogger('TinKeywordProjection')
88
const BATCH_SIZE = 500
@@ -39,6 +39,17 @@ async function tinAvailable(sql: Sql): Promise<boolean> {
3939
return rows.length > 0
4040
}
4141

42+
/**
43+
* The advisory lock key that serializes a knowledge base's keyword projection writers with a change
44+
* of its search-index marker (see {@link installProjection}).
45+
*/
46+
export async function installMembershipKey(tx: Sql | TransactionSql): Promise<void> {
47+
await tx.unsafe(`CREATE OR REPLACE FUNCTION knowledge_tin_membership_key(knowledge_base_id text)
48+
RETURNS bigint LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
49+
SELECT hashtextextended('embedding_keyword_tin:' || knowledge_base_id, 0)
50+
$$`)
51+
}
52+
4253
/**
4354
* Installs the stream and base-token functions shared by the triggers and the query path, and the
4455
* embedding and knowledge base triggers, atomically with respect to embedding writers.
@@ -75,10 +86,7 @@ export async function installProjection(sql: Sql): Promise<void> {
7586
RETURNS text LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
7687
SELECT 'zkb' || md5(knowledge_base_id)
7788
$$`)
78-
await tx.unsafe(`CREATE OR REPLACE FUNCTION knowledge_tin_membership_key(knowledge_base_id text)
79-
RETURNS bigint LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
80-
SELECT hashtextextended('embedding_keyword_tin:' || knowledge_base_id, 0)
81-
$$`)
89+
await installMembershipKey(tx)
8290
await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_embedding_keyword_tin()
8391
RETURNS trigger LANGUAGE plpgsql AS $$
8492
BEGIN

0 commit comments

Comments
 (0)