feat(search) Фаза B PR-1: семантический слой — TEI e5-small + гибрид вектор+RRF (#530) #567

Merged
vvzvlad merged 1 commits from feat/530-search-semantic into develop 2026-07-12 20:39:44 +03:00
16 changed files with 1355 additions and 47 deletions
+36
View File
@@ -198,6 +198,42 @@ MCP_DOCMOST_PASSWORD=
# A slow/hung embeddings endpoint fails after this and the batch continues.
# AI_EMBEDDING_TIMEOUT_MS=120000
# ─── #530 Semantic search (Phase B) ──────────────────────────────────────────
# The GLOBAL embedding provider used by search (query embed) AND the indexer
# (document embed) when a workspace has NO embedding provider of its own. It is
# an OpenAI-compatible endpoint — the docker-compose `embeddings` (TEI) sidecar.
# A workspace that configures its own embedding provider OVERRIDES all of this.
# When neither resolves, search runs lexical-only (semantic.reason=no-provider).
EMBEDDING_ENDPOINT=http://embeddings:80/v1
EMBEDDING_MODEL=intfloat/multilingual-e5-small
# Dummy — the self-hosted TEI sidecar is keyless.
EMBEDDING_API_KEY=unused
EMBEDDING_DIMENSIONS=384
# PLACEHOLDER — set to a PINNED intfloat/multilingual-e5-small commit sha (used
# both as the TEI --revision and as part of the embedding fingerprint). Keep this
# in lockstep with docker-compose; a bump changes the fingerprint (PR-2 handles
# the generational swap/GC — PR-1 uses a single active fingerprint).
EMBEDDING_REVISION=REPLACE_WITH_PINNED_E5_SMALL_COMMIT_SHA
# e5 models require these input prefixes. Empty them for a non-e5 model.
EMBEDDING_QUERY_PREFIX="query: "
EMBEDDING_DOC_PREFIX="passage: "
# Per-request timeout (ms) for the interactive SEARCH query embed — much shorter
# than the batch indexer timeout: a slow/hung sidecar degrades search fast to the
# lexical-only path instead of blocking the request. Default 800.
# SEARCH_EMBED_TIMEOUT_MS=800
# RRF weight of the vector leg relative to each lexical leg (1.0 = equal). Default 1.0.
# SEARCH_VECTOR_WEIGHT=1.0
# Vector top-N pulled into the fused candidate union per request. Default 50.
# SEARCH_VECTOR_CANDIDATES=50
# Per-statement timeout (ms) bounding ONLY the fused vector query (the brute-force
# KNN seq scan — there is no ANN index). If the vector scan exceeds this it is
# cancelled and search degrades to lexical-only (never hangs the request). Scoped
# per-query (SET LOCAL), so it never affects other queries. Default 2000.
# SEARCH_VECTOR_STATEMENT_TIMEOUT_MS=2000
# Kill-switch: set to `off` to disable the semantic layer entirely (search then
# runs the byte-identical Phase-A lexical path). Default on.
# SEARCH_SEMANTIC=on
# Silence timeout (ms) for streaming chat/agent AI calls AND external-MCP traffic.
# Bounds time-to-first-byte and the gap BETWEEN chunks (NOT the total turn length),
# so an arbitrarily long turn that keeps streaming is never cut. Finite so a hung
@@ -32,6 +32,14 @@ describe('EmbeddingIndexerService.reindexWorkspace fail-fast', () => {
const pageEmbeddingRepo = {};
const aiService = {
getEmbeddingModel: jest.fn().mockResolvedValue('some-model'),
// #530: reindexWorkspace's pre-check now resolves the provider (workspace
// or global). Resolve it so the batch control flow under test proceeds.
resolveEmbeddingProvider: jest.fn().mockResolvedValue({
model: 'some-model',
queryPrefix: '',
docPrefix: '',
fingerprint: 'fp-test',
}),
};
// Progress is a best-effort cosmetic store; mock its async methods so the
// batch control flow can be tested without Redis.
@@ -108,6 +116,14 @@ describe('EmbeddingIndexerService.reindexWorkspace progress', () => {
const pageEmbeddingRepo = {};
const aiService = {
getEmbeddingModel: jest.fn().mockResolvedValue('some-model'),
// #530: reindexWorkspace's pre-check now resolves the provider (workspace
// or global). Resolve it so the batch control flow under test proceeds.
resolveEmbeddingProvider: jest.fn().mockResolvedValue({
model: 'some-model',
queryPrefix: '',
docPrefix: '',
fingerprint: 'fp-test',
}),
};
const reindexProgress = {
start: jest.fn().mockResolvedValue(undefined),
@@ -174,7 +190,7 @@ describe('EmbeddingIndexerService.reindexWorkspace progress', () => {
const { service, aiService, reindexProgress } = makeService();
// Embeddings not configured: reindexWorkspace returns early WITHOUT starting
// a fresh record, but the finally must still clear the enqueue-time seed.
aiService.getEmbeddingModel = jest
aiService.resolveEmbeddingProvider = jest
.fn()
.mockRejectedValue(new AiEmbeddingNotConfiguredException());
@@ -187,3 +203,87 @@ describe('EmbeddingIndexerService.reindexWorkspace progress', () => {
expect(reindexProgress.clear).toHaveBeenCalledWith(WORKSPACE_ID);
});
});
/**
* #530 PR-1: reindexPage must (a) prepend the provider's DOC prefix to each chunk
* BEFORE embedding (so stored vectors live in the same prefixed space as a
* prefixed query), and (b) stamp the ACTIVE fingerprint on every inserted row (so
* search only fuses same-generation vectors). Uses lightweight mocks; the tx is
* stubbed to run its callback inline.
*/
describe('EmbeddingIndexerService.reindexPage doc-prefix + fingerprint (#530)', () => {
const WORKSPACE_ID = 'ws-1';
const SPACE_ID = 'space-1';
const PAGE_ID = 'page-1';
function makeService(docPrefix: string) {
const pageRepo = {
findById: jest.fn().mockResolvedValue({
id: PAGE_ID,
workspaceId: WORKSPACE_ID,
spaceId: SPACE_ID,
title: 'Заголовок',
// No ProseMirror content -> the plain-text fallback path (single chunk).
content: null,
textContent: 'простой текст страницы',
deletedAt: null,
}),
};
const insertChunks = jest.fn().mockResolvedValue(undefined);
const pageEmbeddingRepo = {
deleteByPage: jest.fn().mockResolvedValue(undefined),
insertChunks,
};
const embedWithModel = jest.fn().mockResolvedValue([[0.1, 0.2, 0.3]]);
const aiService = {
resolveEmbeddingProvider: jest.fn().mockResolvedValue({
model: { modelId: 'e5-small' },
queryPrefix: 'query: ',
docPrefix,
fingerprint: 'fp-gen-1',
}),
embedWithModel,
};
const reindexProgress = {};
// Stub the tx so executeTx runs its callback inline against a fake trx.
const db = {
transaction: () => ({ execute: (cb: any) => cb({}) }),
};
const service = new EmbeddingIndexerService(
pageRepo as unknown as PageRepo,
pageEmbeddingRepo as unknown as PageEmbeddingRepo,
aiService as unknown as AiService,
reindexProgress as unknown as EmbeddingReindexProgressService,
db as unknown as KyselyDB,
);
return { service, embedWithModel, insertChunks };
}
it('prepends the doc prefix to each chunk and stamps the fingerprint on rows', async () => {
const { service, embedWithModel, insertChunks } = makeService('passage: ');
await service.reindexPage(PAGE_ID);
// Embedded values are DOC-prefixed; the model is the resolved provider model.
expect(embedWithModel).toHaveBeenCalledTimes(1);
const [modelArg, wsArg, valuesArg] = embedWithModel.mock.calls[0];
expect(modelArg).toEqual({ modelId: 'e5-small' });
expect(wsArg).toBe(WORKSPACE_ID);
expect(valuesArg).toEqual(['passage: простой текст страницы']);
// Inserted rows carry the active fingerprint and the ORIGINAL (un-prefixed)
// content (the prefix is an embedding-space artifact, not stored text).
expect(insertChunks).toHaveBeenCalledTimes(1);
const rows = insertChunks.mock.calls[0][0];
expect(rows).toHaveLength(1);
expect(rows[0].fingerprint).toBe('fp-gen-1');
expect(rows[0].content).toBe('простой текст страницы');
expect(rows[0].modelName).toBe('e5-small');
});
it('does not prefix when the provider has an empty doc prefix', async () => {
const { service, embedWithModel } = makeService('');
await service.reindexPage(PAGE_ID);
const [, , valuesArg] = embedWithModel.mock.calls[0];
expect(valuesArg).toEqual(['простой текст страницы']);
});
});
@@ -108,15 +108,24 @@ export class EmbeddingIndexerService {
return;
}
// Resolve embeddings config WITHOUT crashing the queue when unconfigured.
// Resolve the embeddings provider WITHOUT crashing the queue when
// unconfigured. #530: resolveEmbeddingProvider prefers the workspace provider
// and falls back to the GLOBAL env provider (TEI sidecar), and yields the
// doc-prefix + the active fingerprint stored per row so search filters by the
// active generation.
let modelName = 'unknown';
let provider: Awaited<
ReturnType<AiService['resolveEmbeddingProvider']>
>;
try {
const model = await this.aiService.getEmbeddingModel(workspaceId);
provider = await this.aiService.resolveEmbeddingProvider(workspaceId);
// Record the model id per row so a future migration can detect + re-index
// rows produced by a different model (see the migration header). The SDK
// type is `string | EmbeddingModel{V2,V3}`; model objects carry `modelId`.
modelName =
typeof model === 'string' ? model : (model.modelId ?? 'unknown');
typeof provider.model === 'string'
? provider.model
: (provider.model.modelId ?? 'unknown');
} catch (err) {
if (err instanceof AiEmbeddingNotConfiguredException) {
// No embeddings provider for this workspace: NO-OP (§6.7). The page can
@@ -145,8 +154,18 @@ export class EmbeddingIndexerService {
return;
}
// Embed all chunks in one batch.
const vectors = await this.aiService.embedTexts(workspaceId, chunks);
// #530: prepend the provider's DOC prefix to each chunk (e5-style
// "passage: "; empty for a non-e5 provider) so stored vectors live in the
// same prefixed space as a prefixed query, then embed with the RESOLVED
// provider model (which may be the global TEI sidecar, not a workspace one).
const prefixedChunks = provider.docPrefix
? chunks.map((c) => provider.docPrefix + c)
: chunks;
const vectors = await this.aiService.embedWithModel(
provider.model,
workspaceId,
prefixedChunks,
);
// The column is dimension-agnostic, so ANY model dimension is stored as-is.
// Defensive sanity check only: all chunks of ONE page come from the SAME
@@ -170,6 +189,7 @@ export class EmbeddingIndexerService {
vectors,
{ pageId, workspaceId, spaceId },
modelName,
provider.fingerprint,
);
// HARD replace in one transaction: delete then insert so search never
@@ -216,7 +236,9 @@ export class EmbeddingIndexerService {
// (seeded at enqueue time); the finally cleans that too.
try {
try {
await this.aiService.getEmbeddingModel(workspaceId);
// #530: resolve via the same path reindexPage uses (workspace provider,
// else the global TEI sidecar) so a global-only deployment is NOT skipped.
await this.aiService.resolveEmbeddingProvider(workspaceId);
} catch (err) {
if (err instanceof AiEmbeddingNotConfiguredException) {
this.logger.log(
@@ -348,6 +370,7 @@ export class EmbeddingIndexerService {
vectors: number[][],
ids: { pageId: string; workspaceId: string; spaceId: string },
modelName: string,
fingerprint: string | null,
): PageEmbeddingChunkRow[] {
const rows: PageEmbeddingChunkRow[] = [];
let cursor = 0;
@@ -370,6 +393,8 @@ export class EmbeddingIndexerService {
// Provenance for a future re-index sweep on model change.
modelName,
modelDimensions: embedding.length,
// #530: the active generation fingerprint this row belongs to.
fingerprint,
embedding,
});
}
@@ -32,10 +32,27 @@ export class SearchResultDto {
matchedTerms: string[];
}
// The paginated envelope (A5). `total` is the EXACT permission-filtered count of
// pages matching the positive lexical query (fail-closed). `hasMore` is true when
// more results exist WITHIN the fusion window; `truncatedAtCap` signals the match
// set exceeded CANDIDATE_CAP and the tail is unreachable by pagination.
// #530 Phase B — the semantic (vector) layer's per-request status.
// - `state`: 'full' when a provider resolved and the vector arm ran this
// request; 'off' otherwise (kill-switch, no provider, or a degrade).
// - `available`: whether the vector arm actually contributed this request.
// - `reason`: why the arm did NOT run — 'no-provider' (no embedding provider
// resolved) or 'degraded' (sidecar down / query embed timed out).
// - `indexed`/`total`: reserved for PR-2 coverage reporting (unused in PR-1).
export class SearchSemanticDto {
state: 'full' | 'off';
available: boolean;
reason?: 'no-provider' | 'degraded';
indexed?: number;
total?: number;
}
// The paginated envelope (A5). `total` is the permission-filtered count of the
// candidate UNION (lexical ∪ vector top-N, fail-closed) — a deliberate, #530
// documented change from Phase A's exact-lexical count (on a semantic degrade it
// falls back to exactly that lexical count). `hasMore` is true when more results
// exist WITHIN the fusion window; `truncatedAtCap` signals the match set exceeded
// CANDIDATE_CAP and the tail is unreachable by pagination.
export class SearchResponseDto {
items: SearchResultDto[];
total: number;
@@ -53,4 +70,7 @@ export class SearchResponseDto {
mode: 'or' | 'and';
match: string;
};
// Absent on the early-exit paths (empty/garbage query); present once search
// actually runs. Optional so those short-circuit responses stay valid.
semantic?: SearchSemanticDto;
}
@@ -1,8 +1,15 @@
import { Module } from '@nestjs/common';
import { SearchController } from './search.controller';
import { SearchService } from './search.service';
import { AiModule } from '../../integrations/ai/ai.module';
/**
* #530: AiModule supplies AiService (embedQuery / resolveEmbeddingProvider).
* PageEmbeddingRepo is provided by the @Global DatabaseModule, so it is injected
* into SearchService without an explicit import here.
*/
@Module({
imports: [AiModule],
controllers: [SearchController],
providers: [SearchService],
exports: [SearchService],
@@ -48,6 +48,10 @@ describe('SearchService.searchPage — scope-security early returns', () => {
shareRepo as any,
spaceMemberRepo as any,
pagePermissionRepo as any,
// #530 semantic path off for these query-mode unit tests: a stubbed
// embedQuery that throws makes the service degrade to the lexical path.
{ embedQuery: jest.fn().mockRejectedValue(new Error('no embed')) } as any,
{ vectorCandidateArm: jest.fn() } as any,
);
return { service, pageRepo, shareRepo, spaceMemberRepo, pagePermissionRepo };
}
@@ -11,6 +11,8 @@ describe('SearchService', () => {
{} as any, // shareRepo
{} as any, // spaceMemberRepo
{} as any, // pagePermissionRepo
{} as any, // aiService
{} as any, // pageEmbeddingRepo
);
expect(service).toBeDefined();
});
@@ -61,6 +63,8 @@ describe('SearchService.searchSuggestions — onlyTemplates filter', () => {
shareRepo as any,
spaceMemberRepo as any,
pagePermissionRepo as any,
{} as any, // aiService (searchSuggestions never embeds)
{} as any, // pageEmbeddingRepo
);
return { service, db, pageBuilder };
+247 -28
View File
@@ -1,8 +1,9 @@
import { Injectable } from '@nestjs/common';
import { Injectable, Logger } from '@nestjs/common';
import { SearchDTO, SearchSuggestionDTO } from './dto/search.dto';
import {
SearchResponseDto,
SearchResultDto,
SearchSemanticDto,
} from './dto/search-response.dto';
import { InjectKysely } from 'nestjs-kysely';
import { KyselyDB } from '@docmost/db/types/kysely.types';
@@ -11,6 +12,10 @@ import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { SpaceMemberRepo } from '@docmost/db/repos/space/space-member.repo';
import { ShareRepo } from '@docmost/db/repos/share/share.repo';
import { PagePermissionRepo } from '@docmost/db/repos/page/page-permission.repo';
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
import { AiService } from '../../integrations/ai/ai.service';
import { AiEmbeddingNotConfiguredException } from '../../integrations/ai/ai-embedding-not-configured.exception';
import { isStatementTimeout } from '@docmost/db/utils';
import {
ParsedQuery,
ParsedTerm,
@@ -59,14 +64,42 @@ function defaultBooleanMode(): SearchBooleanMode {
return process.env.SEARCH_MODE === 'and' ? 'and' : 'or';
}
// #530 semantic env knobs. SEARCH_SEMANTIC=off is the kill-switch (hybrid is
// otherwise transparent — no per-request mode flag). SEARCH_VECTOR_WEIGHT tunes
// the vector RRF leg's contribution (default 1.0, i.e. equal to each lexical
// leg). SEARCH_VECTOR_CANDIDATES caps the vector top-N pulled into the union.
function semanticKillSwitchOff(): boolean {
return process.env.SEARCH_SEMANTIC === 'off';
}
function getVectorWeight(): number {
const raw = Number(process.env.SEARCH_VECTOR_WEIGHT);
return Number.isFinite(raw) && raw >= 0 ? raw : 1.0;
}
function getVectorCandidates(): number {
const raw = Number(process.env.SEARCH_VECTOR_CANDIDATES);
return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : 50;
}
// Safety net for the brute-force vector scan (no ANN index — see
// vectorCandidateArm): a per-statement timeout bounding ONLY the fused vector
// query, so a pathological seq scan is cancelled (SQLSTATE 57014) and search
// degrades to lexical instead of hanging the request. Default 2000ms.
function getVectorStatementTimeoutMs(): number {
const raw = Number(process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS);
return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : 2000;
}
@Injectable()
export class SearchService {
private readonly logger = new Logger(SearchService.name);
constructor(
@InjectKysely() private readonly db: KyselyDB,
private pageRepo: PageRepo,
private shareRepo: ShareRepo,
private spaceMemberRepo: SpaceMemberRepo,
private pagePermissionRepo: PagePermissionRepo,
private aiService: AiService,
private pageEmbeddingRepo: PageEmbeddingRepo,
) {}
// === #529 SQL fragment builders (parameterized AST, never string-concat) =====
@@ -292,33 +325,102 @@ export class SearchService {
const scopeSql = sql`(${sql.join(scopePreds, sql` AND `)})`;
const candidateSql = sql`(${scopeSql} AND ${this.candidatePredicate(parsed, titleOnly)})`;
// --- Semantic (vector) arm: embed the query, degrade gracefully. (#530) ----
// The try/catch wraps ONLY embedQuery + vector-arm construction. On ANY throw
// (TEI down / 800ms timeout / no provider) we omit the vector arm entirely and
// run the byte-identical Phase-A lexical path — never a 500 from the sidecar.
// The permission filter below stays OUTSIDE this try (a permission error must
// 500, never fail-open). Hybrid is transparent; SEARCH_SEMANTIC=off disables it.
let vectorArm: RawBuilder<unknown> | null = null;
let semanticAvailable = false;
let semanticReason: 'no-provider' | 'degraded' | undefined;
if (!semanticKillSwitchOff()) {
try {
const { vector, fingerprint } = await this.aiService.embedQuery(
opts.workspaceId,
rawQuery,
);
vectorArm = this.pageEmbeddingRepo.vectorCandidateArm({
queryEmbedding: vector,
dimensions: vector.length,
fingerprint,
scope: scopeSql,
limit: getVectorCandidates(),
// Filter page_embeddings by (immutable) workspace_id directly, so the
// composite index bites on its leading column and candidates are
// workspace-scoped at the embedding level (#530 review WARNING 2). Space
// scoping stays on the pages join only — page_embeddings.space_id can be
// STALE after a cross-space move, so filtering it here would drop a moved
// page's vector hit (re-review regression fix).
workspaceId: scope.workspaceId,
});
semanticAvailable = true;
} catch (err) {
if (err instanceof AiEmbeddingNotConfiguredException) {
// No embedding provider is the DEFAULT state of any deployment without
// the TEI sidecar — normal, not a problem. Log at DEBUG so it never
// floods WARN once per search request.
semanticReason = 'no-provider';
this.logger.debug(
`search.semantic.no-provider workspace=${opts.workspaceId}`,
);
} else {
// A provider IS configured but the embed call failed/timed out — a real
// sidecar problem worth a WARN. Structured, greppable event (fixed code)
// — no query text/secrets.
semanticReason = 'degraded';
this.logger.warn(
`search.semantic.degraded reason=degraded workspace=${opts.workspaceId}`,
);
}
}
}
// --- Ranked candidate ids (ALL matches, RRF order). -----------------------
// Two independent branches ranked SEPARATELY (FTS by ts_rank_cd, substring by
// tier) then fused with RRF: score = Σ 1/(k + rank_branch). Deterministic
// ORDER BY rrf DESC, id so pagination never dupes/skips (A4, acceptance #8).
const rankedRows = await sql<{ id: string }>`
WITH candidates AS (
SELECT pages.id AS id,
${this.ftsScoreExpr(parsed)} AS fts_score,
${this.subTierExpr(parsed, titleOnly)} AS sub_tier
FROM pages
WHERE ${candidateSql}
),
ranked AS (
SELECT id, fts_score, sub_tier,
row_number() OVER (ORDER BY fts_score DESC NULLS LAST, id) AS rn_fts,
row_number() OVER (ORDER BY sub_tier DESC NULLS LAST, id) AS rn_sub
FROM candidates
)
SELECT id
FROM ranked
ORDER BY
(CASE WHEN fts_score IS NOT NULL THEN 1.0/(${RRF_K} + rn_fts) ELSE 0 END)
+ (CASE WHEN sub_tier IS NOT NULL THEN 1.0/(${RRF_K} + rn_sub) ELSE 0 END) DESC,
id ASC
`.execute(this.db);
// Lexical branches ranked SEPARATELY (FTS by ts_rank_cd, substring by tier);
// when a query vector is available a third VECTOR branch is fused over the
// UNION of lexical + vector candidates. RRF: score = Σ w/(k + rank_branch).
// Deterministic ORDER BY rrf DESC, id so pagination never dupes/skips (A4,
// acceptance #8). `total` is thus the union size (lexical ∪ vector top-N),
// post-permission-filter — a documented change from Phase A's exact lexical
// count (#530). On degrade the query is byte-identical to Phase A.
//
// The fused query is bounded by a per-statement timeout (see runRankedQuery):
// if the brute-force vector scan is cancelled (SQLSTATE 57014) we degrade to
// the byte-identical lexical path — search returns lexical results, never
// hangs, never 500s. Only the timeout is caught here; any OTHER SQL error
// propagates. The permission filter below stays outside, unchanged.
let orderedIds: string[];
try {
orderedIds = await this.runRankedQuery(
candidateSql,
parsed,
titleOnly,
vectorArm,
);
} catch (err) {
if (vectorArm && isStatementTimeout(err)) {
semanticAvailable = false;
semanticReason = 'degraded';
this.logger.warn(
`search.semantic.degraded reason=degraded workspace=${opts.workspaceId} (vector statement timeout)`,
);
// Re-run the exact Phase-A lexical path (no vector arm, no timeout).
orderedIds = await this.runRankedQuery(
candidateSql,
parsed,
titleOnly,
null,
);
} else {
throw err;
}
}
let orderedIds = rankedRows.rows.map((r) => r.id);
const semantic: SearchSemanticDto = {
state: semanticAvailable ? 'full' : 'off',
available: semanticAvailable,
...(semanticReason ? { reason: semanticReason } : {}),
};
// --- Permission filter (fail-closed, exact total). ------------------------
// filterAccessiblePageIds runs the #348 hasRestricted pre-check internally:
@@ -355,14 +457,131 @@ export class SearchService {
const hasMore = offset + pageIds.length < window.length;
if (pageIds.length === 0) {
return { items: [], total, hasMore, truncatedAtCap, offset, query: queryMeta };
return {
items: [],
total,
hasMore,
truncatedAtCap,
offset,
query: queryMeta,
semantic,
};
}
// --- Detail fetch for the page slice only (ts_headline/snippet are costly, so
// compute them ONLY for the returned rows), preserving RRF order. -------
const items = await this.fetchDetails(pageIds, parsed, titleOnly);
return { items, total, hasMore, truncatedAtCap, offset, query: queryMeta };
return {
items,
total,
hasMore,
truncatedAtCap,
offset,
query: queryMeta,
semantic,
};
}
/**
* Run the ranked-candidate-ids query and return the ids in RRF order (#530).
*
* When `vectorArm` is null (degrade / semantic off / no query vector) this runs
* the BYTE-IDENTICAL Phase-A 2-branch lexical query — the semantic layer must
* be a pure superset that never changes lexical behaviour on degrade.
*
* When a vector arm is supplied it is UNION ALL'd with the lexical arm into one
* `candidates` set; `agg` collapses each page to its best per-branch signal
* (MAX fts_score / MAX sub_tier / MIN vec_distance); `ranked` assigns a
* per-branch row_number; the final ORDER BY fuses the three legs with RRF,
* adding the vector leg ONLY for rows that have a vec_distance (a lexical-only
* hit contributes 0 to the vector leg, and vice-versa). W_VEC weights the
* vector leg.
*/
private async runRankedQuery(
candidateSql: RawBuilder<unknown>,
parsed: ParsedQuery,
titleOnly: boolean,
vectorArm: RawBuilder<unknown> | null,
): Promise<string[]> {
if (!vectorArm) {
// Phase-A lexical path — DO NOT change (byte-identical on degrade, #529).
const rankedRows = await sql<{ id: string }>`
WITH candidates AS (
SELECT pages.id AS id,
${this.ftsScoreExpr(parsed)} AS fts_score,
${this.subTierExpr(parsed, titleOnly)} AS sub_tier
FROM pages
WHERE ${candidateSql}
),
ranked AS (
SELECT id, fts_score, sub_tier,
row_number() OVER (ORDER BY fts_score DESC NULLS LAST, id) AS rn_fts,
row_number() OVER (ORDER BY sub_tier DESC NULLS LAST, id) AS rn_sub
FROM candidates
)
SELECT id
FROM ranked
ORDER BY
(CASE WHEN fts_score IS NOT NULL THEN 1.0/(${RRF_K} + rn_fts) ELSE 0 END)
+ (CASE WHEN sub_tier IS NOT NULL THEN 1.0/(${RRF_K} + rn_sub) ELSE 0 END) DESC,
id ASC
`.execute(this.db);
return rankedRows.rows.map((r) => r.id);
}
// #530 fused 3-branch RRF over lexical ∪ vector candidates. The lexical arm
// adds NULL::float vec_distance to stay UNION-compatible with the vector arm.
const wVec = getVectorWeight();
// Bound the brute-force vector scan with a per-statement timeout, scoped to
// THIS query only. `set_config(..., is_local := true)` = SET LOCAL semantics:
// it applies within this transaction and auto-resets at COMMIT/ROLLBACK, so
// it can NEVER leak onto the next query that reuses this pooled connection.
// On a cancellation the driver raises SQLSTATE 57014, which searchPage catches
// to degrade to lexical-only. Running inside a transaction also guarantees the
// SET and the ranked query share one connection.
const timeoutMs = getVectorStatementTimeoutMs();
return this.db.transaction().execute(async (trx) => {
await sql`SELECT set_config('statement_timeout', ${String(timeoutMs)}, true)`.execute(
trx,
);
const rankedRows = await sql<{ id: string }>`
WITH candidates AS (
(SELECT pages.id AS id,
${this.ftsScoreExpr(parsed)} AS fts_score,
${this.subTierExpr(parsed, titleOnly)} AS sub_tier,
NULL::float AS vec_distance
FROM pages
WHERE ${candidateSql})
UNION ALL
(${vectorArm})
),
agg AS (
SELECT id,
MAX(fts_score) AS fts_score,
MAX(sub_tier) AS sub_tier,
MIN(vec_distance) AS vec_distance
FROM candidates
GROUP BY id
),
ranked AS (
SELECT id, fts_score, sub_tier, vec_distance,
row_number() OVER (ORDER BY fts_score DESC NULLS LAST, id) AS rn_fts,
row_number() OVER (ORDER BY sub_tier DESC NULLS LAST, id) AS rn_sub,
row_number() OVER (ORDER BY vec_distance ASC NULLS LAST, id) AS rn_vec
FROM agg
)
SELECT id
FROM ranked
ORDER BY
(CASE WHEN fts_score IS NOT NULL THEN 1.0/(${RRF_K} + rn_fts) ELSE 0 END)
+ (CASE WHEN sub_tier IS NOT NULL THEN 1.0/(${RRF_K} + rn_sub) ELSE 0 END)
+ (CASE WHEN vec_distance IS NOT NULL THEN ${wVec}::float/(${RRF_K} + rn_vec) ELSE 0 END) DESC,
id ASC
`.execute(trx);
return rankedRows.rows.map((r) => r.id);
});
}
// Resolve the search scope: explicit space, the authed user's member spaces, or
@@ -0,0 +1,49 @@
import { type Kysely, sql } from 'kysely';
/**
* #530 Search Phase B (PR-1): add the embedding FINGERPRINT column.
*
* The fingerprint is a deterministic id of the embedding configuration that
* produced a row — model id + revision + query/doc prefix scheme + dimensions
* (see AiService.computeEmbeddingFingerprint). Search filters vector candidates
* by the workspace's ACTIVE fingerprint so a revision bump or a prefix-scheme
* change never fuses incompatible vectors into results.
*
* PR-1 scope: this migration only ADDS the column (nullable — existing rows stay
* NULL = legacy) and the composite index the vector-candidate scan uses. The full
* generational swap / GC lifecycle (target-fingerprint reindex, atomic flip,
* old-generation GC) is deliberately deferred to PR-2.
*
* Independent migration (per the #363 crash-loop net): it creates ONLY its own
* objects and never touches another migration's tables/indexes, so a partial
* failure cannot leave a shared object half-built.
*/
export async function up(db: Kysely<any>): Promise<void> {
// Nullable text column. Existing rows keep NULL (legacy generation) and are
// simply not matched by the active-fingerprint filter until re-indexed.
await sql`
ALTER TABLE page_embeddings
ADD COLUMN IF NOT EXISTS fingerprint text
`.execute(db);
// Composite btree supporting the scoped, dimension- + fingerprint-filtered
// vector-candidate scan (workspace_id + space_id + fingerprint + model_dimensions).
await db.schema
.createIndex('idx_page_embeddings_ws_space_fp_dim')
.ifNotExists()
.on('page_embeddings')
.columns(['workspace_id', 'space_id', 'fingerprint', 'model_dimensions'])
.execute();
}
export async function down(db: Kysely<any>): Promise<void> {
await db.schema
.dropIndex('idx_page_embeddings_ws_space_fp_dim')
.ifExists()
.execute();
await sql`
ALTER TABLE page_embeddings
DROP COLUMN IF EXISTS fingerprint
`.execute(db);
}
@@ -1,6 +1,6 @@
import { Injectable } from '@nestjs/common';
import { InjectKysely } from 'nestjs-kysely';
import { sql } from 'kysely';
import { RawBuilder, sql } from 'kysely';
import * as pgvector from 'pgvector';
import { KyselyDB, KyselyTransaction } from '../../types/kysely.types';
import { dbOrTx } from '../../utils';
@@ -35,6 +35,9 @@ export interface PageEmbeddingChunkRow {
content: string;
modelName: string;
modelDimensions: number;
// #530 PR-1: the active embedding fingerprint (see computeEmbeddingFingerprint).
// null only when no provider resolves (the indexer no-ops in that case).
fingerprint: string | null;
embedding: number[];
}
@@ -120,6 +123,7 @@ export class PageEmbeddingRepo {
content: row.content,
modelName: row.modelName,
modelDimensions: row.modelDimensions,
fingerprint: row.fingerprint,
// pgvector.toSql -> '[1,2,3]'; cast the bound literal to vector.
embedding: sql`${pgvector.toSql(row.embedding)}::vector`,
})),
@@ -127,6 +131,82 @@ export class PageEmbeddingRepo {
.execute();
}
/**
* #530: build the VECTOR candidate arm for SearchService's fused RRF union — a
* page-level nearest-neighbour sub-select. Returns a `sql` fragment (not an
* executed query) so the caller can UNION ALL it with the lexical arm inside a
* single ranked-ids query.
*
* SCALING BOUNDARY: the column is dimension-agnostic, so it carries NO ANN
* index — this is a brute-force O(N) KNN seq scan with `<=>`. Accepted for the
* small-tenant / homelab fork target (ANN + a pinned-dimension column are
* deferred). It is NOT unbounded, though: the caller (SearchService) runs this
* fused query under a per-statement timeout (SEARCH_VECTOR_STATEMENT_TIMEOUT_MS)
* and DEGRADES to lexical-only on a 57014 cancellation, so a pathological scan
* can never hang the interactive search request.
*
* The arm collapses a page's chunks to its best (MIN) cosine distance:
* SELECT pages.id, NULL fts_score, NULL sub_tier,
* MIN(pe.embedding <=> $qvec) AS vec_distance
* FROM page_embeddings pe JOIN pages ON pages.id = pe.page_id
* WHERE <scope> AND pe.model_dimensions = $dim AND pe.fingerprint = $fp
* GROUP BY pages.id ORDER BY vec_distance LIMIT $limit
*
* `scope` is SearchService's shared scope predicate (workspace + space/id set +
* creator + descendants + deleted_at), referencing the `pages` table — it must
* be spliced verbatim so the vector candidate set mirrors the lexical scope
* EXACTLY, minus the lexical text predicate (vector candidates need not match
* text). The vector is bound via pgvector's `toSql(...)::vector`, and both the
* dimension and the ACTIVE fingerprint are filtered so `<=>` only ever compares
* compatible, same-generation vectors (pgvector errors on a dimension
* mismatch; a fingerprint mismatch would fuse incomparable vectors).
*
* `workspaceId` is ALSO filtered directly on `page_embeddings` — not only via
* the pages join. It is IMMUTABLE (there are no cross-workspace page moves), so
* an embedding row's workspace_id always matches its page's, and this drops NO
* legitimate hit; it lets the composite index idx_page_embeddings_ws_space_fp_dim
* (workspace_id, space_id, fingerprint, model_dimensions) bite from its LEADING
* column, and workspace-scopes candidates at the embedding level (defense in
* depth).
*
* SPACE is deliberately NOT filtered on page_embeddings: `page_embeddings.space_id`
* is stamped at index time and is NOT updated when a page is MOVED between spaces
* (movePageToSpace does not reindex, and PAGE_MOVED_TO_SPACE has no reindex
* consumer), so a moved page's rows carry the OLD space until the next reindex. A
* `page_embeddings.space_id = ANY(scope)` predicate would then wrongly drop a
* legitimate vector hit for a page moved INTO the searched space. Space scoping
* is therefore enforced ONLY by the join to `pages` (pages.space_id ∈ scope,
* always current) — the same way the lexical arm scopes, so no space leak.
*
* The NULL fts_score / sub_tier columns keep the arm UNION-compatible with the
* lexical arm's column list.
*/
vectorCandidateArm(params: {
queryEmbedding: number[];
dimensions: number;
fingerprint: string;
scope: RawBuilder<unknown>;
limit: number;
workspaceId: string;
}): RawBuilder<unknown> {
const qvec = sql`${pgvector.toSql(params.queryEmbedding)}::vector`;
return sql`
SELECT pages.id AS id,
NULL::float AS fts_score,
NULL::int AS sub_tier,
MIN(page_embeddings.embedding <=> ${qvec}) AS vec_distance
FROM page_embeddings
JOIN pages ON pages.id = page_embeddings.page_id
WHERE ${params.scope}
AND page_embeddings.workspace_id = ${params.workspaceId}
AND page_embeddings.model_dimensions = ${params.dimensions}
AND page_embeddings.fingerprint = ${params.fingerprint}
GROUP BY pages.id
ORDER BY vec_distance ASC
LIMIT ${params.limit}
`;
}
/**
* Cosine search over the embeddings, scoped to a workspace AND a set of
* spaces the caller may read (see semanticSearch access-scoping). Orders by
@@ -7,6 +7,11 @@ export interface PageEmbeddings {
spaceId: string;
modelName: string;
modelDimensions: number;
// #530 PR-1: the active embedding fingerprint (model id + revision + prefix
// scheme + dimensions). NULL on legacy rows written before this column existed.
// Search filters candidates by the ACTIVE fingerprint so a revision/prefix
// change never mixes incompatible vectors.
fingerprint: string | null;
workspaceId: string;
// Nullable: page-body embeddings have no attachment (only attachment chunks set it).
attachmentId: string | null;
+17
View File
@@ -108,6 +108,23 @@ export function isUniqueViolation(err: unknown): boolean {
return (err as { code?: unknown } | null | undefined)?.code === PG_UNIQUE_VIOLATION;
}
/** Postgres `query_canceled` SQLSTATE — raised when statement_timeout fires. */
const PG_QUERY_CANCELED = '57014';
/**
* Whether `err` is a Postgres statement-timeout cancellation (SQLSTATE `57014`).
* Used by #530 semantic search to degrade a slow vector scan to lexical-only
* instead of hanging/500-ing. Matches on the SQLSTATE (the driver surfaces it as
* `.code`), with a message fallback for any wrapper that drops the code.
*/
export function isStatementTimeout(err: unknown): boolean {
if ((err as { code?: unknown } | null | undefined)?.code === PG_QUERY_CANCELED) {
return true;
}
const msg = err instanceof Error ? err.message : '';
return /canceling statement due to statement timeout/i.test(msg);
}
/**
* The name of the UNIQUE index/constraint a `23505` error violated, or
* undefined. The `kysely-postgres-js` / `postgres@3.x` driver surfaces it as
@@ -0,0 +1,143 @@
import { AiService, computeEmbeddingFingerprint } from './ai.service';
/**
* #530 Search Phase B (PR-1) unit coverage for the embedding layer:
* - computeEmbeddingFingerprint: a revision bump / prefix toggle MUST change the
* fingerprint even when the bare model name + dimension are unchanged (it is
* deliberately NOT the bare model name);
* - embedQuery: applies the QUERY prefix and embeds under the short SEARCH embed
* timeout (not the long batch timeout);
* - embedWithModel: a hung model + short timeout degrades with a clear timeout
* error (the search caller then falls back to the lexical path).
*/
describe('computeEmbeddingFingerprint (#530)', () => {
const base = {
modelId: 'intfloat/multilingual-e5-small',
revision: 'sha-aaa',
queryPrefix: 'query: ',
docPrefix: 'passage: ',
dimensions: 384,
};
it('is deterministic for identical inputs', () => {
expect(computeEmbeddingFingerprint(base)).toBe(
computeEmbeddingFingerprint({ ...base }),
);
});
it('CHANGES on a revision bump (same model name + dimension)', () => {
const a = computeEmbeddingFingerprint(base);
const b = computeEmbeddingFingerprint({ ...base, revision: 'sha-bbb' });
expect(b).not.toBe(a);
});
it('CHANGES on a query-prefix toggle (same model name + dimension + revision)', () => {
const a = computeEmbeddingFingerprint(base);
const b = computeEmbeddingFingerprint({ ...base, queryPrefix: '' });
expect(b).not.toBe(a);
});
it('CHANGES on a doc-prefix toggle', () => {
const a = computeEmbeddingFingerprint(base);
const b = computeEmbeddingFingerprint({ ...base, docPrefix: '' });
expect(b).not.toBe(a);
});
it('CHANGES on a dimension change', () => {
const a = computeEmbeddingFingerprint(base);
const b = computeEmbeddingFingerprint({ ...base, dimensions: 768 });
expect(b).not.toBe(a);
});
it('is NOT the bare model name (guards the mutation)', () => {
// If the fingerprint were the bare model name, a revision bump would not
// change it — the tests above would red. Assert it directly too.
expect(computeEmbeddingFingerprint(base)).not.toBe(base.modelId);
});
it('distinguishes swapped prefixes — ("q","") != ("","q")', () => {
const a = computeEmbeddingFingerprint({
...base,
queryPrefix: 'q',
docPrefix: '',
});
const b = computeEmbeddingFingerprint({
...base,
queryPrefix: '',
docPrefix: 'q',
});
expect(a).not.toBe(b);
});
});
describe('AiService.embedQuery (#530)', () => {
function makeService(): AiService {
// The constructor body only stores its deps; embedQuery is fully driven by
// the spied resolveEmbeddingProvider + embedWithModel below.
return new AiService({} as any, {} as any, {} as any);
}
it('prepends the query prefix and returns the provider fingerprint', async () => {
const svc = makeService();
jest.spyOn(svc, 'resolveEmbeddingProvider').mockResolvedValue({
model: { modelId: 'm' } as any,
queryPrefix: 'query: ',
docPrefix: 'passage: ',
fingerprint: 'fp-active',
});
const embedSpy = jest
.spyOn(svc, 'embedWithModel')
.mockResolvedValue([[0.1, 0.2, 0.3]]);
const res = await svc.embedQuery('ws-1', 'кофе');
// The single embedded value carries the query prefix.
expect(embedSpy).toHaveBeenCalledTimes(1);
const [, wsArg, valuesArg, timeoutArg] = embedSpy.mock.calls[0];
expect(wsArg).toBe('ws-1');
expect(valuesArg).toEqual(['query: кофе']);
// Bound by the SHORT search-embed timeout (default 800ms), not the 120000ms
// batch timeout.
expect(timeoutArg).toBe(800);
expect(res).toEqual({ vector: [0.1, 0.2, 0.3], fingerprint: 'fp-active' });
});
it('propagates a no-provider error (drives the no-provider degrade)', async () => {
const svc = makeService();
const { AiEmbeddingNotConfiguredException } = await import(
'./ai-embedding-not-configured.exception'
);
jest
.spyOn(svc, 'resolveEmbeddingProvider')
.mockRejectedValue(new AiEmbeddingNotConfiguredException());
await expect(svc.embedQuery('ws-1', 'x')).rejects.toBeInstanceOf(
AiEmbeddingNotConfiguredException,
);
});
});
describe('AiService.embedWithModel timeout (#530)', () => {
it('degrades a hung model with a timeout error under the short window', async () => {
const svc = new AiService({} as any, {} as any, {} as any);
// A model whose doEmbed never resolves but honours the abort signal, so the
// AbortSignal.timeout fires and embedWithModel raises a clear timeout error.
const hangingModel: any = {
specificationVersion: 'v2',
provider: 'fake',
modelId: 'fake',
maxEmbeddingsPerCall: 100,
supportsParallelCalls: true,
doEmbed: ({ abortSignal }: { abortSignal?: AbortSignal }) =>
new Promise((_resolve, reject) => {
abortSignal?.addEventListener('abort', () => {
const err = new Error('The operation was aborted');
err.name = 'AbortError';
reject(err);
});
}),
};
await expect(
svc.embedWithModel(hangingModel, 'ws-1', ['x'], 30),
).rejects.toThrow(/timed out/i);
});
});
+162 -7
View File
@@ -23,6 +23,50 @@ import {
import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
import { SecretBoxService } from '../crypto/secret-box';
import { AiDriver } from './ai.types';
import { createHash } from 'node:crypto';
/**
* A resolved embedding provider for #530 semantic search. `model` is the AI SDK
* embedding model; `queryPrefix`/`docPrefix` are prepended to a query / a stored
* chunk respectively (e5-style `"query: "` / `"passage: "`, empty for a non-e5
* provider); `fingerprint` is the deterministic id of the whole configuration.
*/
export interface ResolvedEmbeddingProvider {
model: EmbeddingModel;
queryPrefix: string;
docPrefix: string;
fingerprint: string;
}
/**
* Deterministic embedding FINGERPRINT (#530). Encodes model id + revision +
* prefix scheme + dimensions so that ANY of them changing (a revision bump, a
* prefix toggle) yields a DIFFERENT fingerprint even when the bare model name and
* dimension are unchanged. Deliberately NOT the bare model name: two rows from
* the same model but a different revision/prefix must not be fused together.
*
* Exported as a pure function so it can be unit-tested in isolation and reused by
* the indexer without a service instance.
*/
export function computeEmbeddingFingerprint(parts: {
modelId: string;
revision: string;
queryPrefix: string;
docPrefix: string;
dimensions: number | null;
}): string {
// The two prefixes are SEPARATE keys (not a concatenated string): JSON.stringify
// escapes each independently, so the prefix scheme ("a","") is distinct from
// ("","a") with no separator/collision hazard even if a prefix contains spaces.
const canonical = JSON.stringify({
m: parts.modelId,
r: parts.revision,
q: parts.queryPrefix,
d: parts.docPrefix,
dim: parts.dimensions ?? 0,
});
return createHash('sha256').update(canonical).digest('hex').slice(0, 32);
}
/**
* Optional chat-model override carried by an agent role (`ai_agent_roles.
@@ -362,11 +406,112 @@ export class AiService {
async embedTexts(workspaceId: string, texts: string[]): Promise<number[][]> {
if (texts.length === 0) return [];
const model = await this.getEmbeddingModel(workspaceId);
// Bound the embedding call: a slow/hung embeddings endpoint must fail loudly
// (and let the caller move on to the next page) instead of blocking forever.
// The single signal caps the WHOLE call, including the SDK's internal
// retries/backoff (embedMany defaults to maxRetries: 2).
const timeoutMs = AiService.embeddingTimeoutMs();
return this.embedWithModel(model, workspaceId, texts);
}
/**
* #530: resolve the embedding provider for a workspace. Prefers the workspace's
* own configured embedding provider; falls back to the GLOBAL env provider (a
* TEI sidecar via the OpenAI-compatible path) when the workspace has none.
* Returns the model + the query/doc prefixes + the config fingerprint. Throws
* AiEmbeddingNotConfiguredException when NEITHER resolves, so callers can drive
* a `no-provider` degrade path.
*/
async resolveEmbeddingProvider(
workspaceId: string,
): Promise<ResolvedEmbeddingProvider> {
// 1. Per-workspace provider (uses the workspace's own creds/endpoint). When
// it is not configured getEmbeddingModel throws the not-configured
// exception; we swallow ONLY that and fall through to the global provider.
try {
const model = await this.getEmbeddingModel(workspaceId);
const modelId =
typeof model === 'string' ? model : (model.modelId ?? 'unknown');
// A per-workspace (typically non-e5) provider gets no e5-style prefixes.
return {
model,
queryPrefix: '',
docPrefix: '',
fingerprint: computeEmbeddingFingerprint({
modelId,
revision: 'workspace',
queryPrefix: '',
docPrefix: '',
dimensions: null,
}),
};
} catch (err) {
if (!(err instanceof AiEmbeddingNotConfiguredException)) throw err;
}
// 2. Global env provider (TEI sidecar). TEI is OpenAI-compatible, so reuse
// the existing openai path — no new SDK. A dummy key is fine for a
// keyless self-hosted sidecar.
const endpoint = process.env.EMBEDDING_ENDPOINT?.trim();
const globalModel = process.env.EMBEDDING_MODEL?.trim();
if (endpoint && globalModel) {
const model = createOpenAI({
baseURL: endpoint,
apiKey: process.env.EMBEDDING_API_KEY || 'unused',
}).textEmbeddingModel(globalModel);
const queryPrefix = process.env.EMBEDDING_QUERY_PREFIX ?? '';
const docPrefix = process.env.EMBEDDING_DOC_PREFIX ?? '';
const dimRaw = Number(process.env.EMBEDDING_DIMENSIONS);
const dimensions = Number.isFinite(dimRaw) && dimRaw > 0 ? dimRaw : null;
return {
model,
queryPrefix,
docPrefix,
fingerprint: computeEmbeddingFingerprint({
modelId: globalModel,
revision: process.env.EMBEDDING_REVISION ?? '',
queryPrefix,
docPrefix,
dimensions,
}),
};
}
// Neither resolved -> drives semantic.reason=no-provider.
throw new AiEmbeddingNotConfiguredException();
}
/**
* #530: embed a SEARCH QUERY. Resolves the provider, prepends its query prefix,
* and embeds the single value under its OWN short timeout
* (SEARCH_EMBED_TIMEOUT_MS, default 800ms) — NOT the long batch-indexing
* timeout — so a slow/hung sidecar degrades search fast. Throws on
* timeout/error (the caller degrades to the lexical-only path). Returns the
* vector plus the active fingerprint used to filter candidate rows.
*/
async embedQuery(
workspaceId: string,
text: string,
): Promise<{ vector: number[]; fingerprint: string }> {
const provider = await this.resolveEmbeddingProvider(workspaceId);
const [vector] = await this.embedWithModel(
provider.model,
workspaceId,
[provider.queryPrefix + text],
AiService.searchEmbedTimeoutMs(),
);
return { vector, fingerprint: provider.fingerprint };
}
/**
* Embed values with an EXPLICIT model, bounded by `timeoutMs` (default: the
* batch-indexing timeout). Shared core of embedTexts / embedQuery: a slow/hung
* embeddings endpoint must fail loudly instead of blocking forever. The single
* signal caps the WHOLE call, including the SDK's internal retries/backoff
* (embedMany defaults to maxRetries: 2).
*/
async embedWithModel(
model: EmbeddingModel,
workspaceId: string,
texts: string[],
timeoutMs: number = AiService.embeddingTimeoutMs(),
): Promise<number[][]> {
if (texts.length === 0) return [];
const signal = AbortSignal.timeout(timeoutMs);
try {
const { embeddings } = await embedMany({
@@ -391,8 +536,8 @@ export class AiService {
if (signal.aborted && abortLike) {
throw new Error(
`Embedding request timed out after ${timeoutMs}ms ` +
`(workspace ${workspaceId}, ${texts.length} chunk(s)). ` +
`Increase AI_EMBEDDING_TIMEOUT_MS or check the embeddings endpoint.`,
`(workspace ${workspaceId}, ${texts.length} value(s)). ` +
`Increase the embedding timeout or check the embeddings endpoint.`,
);
}
throw err;
@@ -408,6 +553,16 @@ export class AiService {
return Number.isFinite(raw) && raw > 0 ? raw : 120_000;
}
/**
* #530: per-call timeout for the interactive SEARCH query embed. Much shorter
* than the batch-indexing timeout (a search request cannot wait 2 minutes on
* the sidecar). Configurable via SEARCH_EMBED_TIMEOUT_MS; default 800ms.
*/
private static searchEmbedTimeoutMs(): number {
const raw = Number(process.env.SEARCH_EMBED_TIMEOUT_MS);
return Number.isFinite(raw) && raw > 0 ? raw : 800;
}
// Build a tiny valid WAV (mono, 16-bit PCM, 16 kHz, ~1s of silence), used only
// as a connectivity probe for the STT endpoint in testConnection.
private static silentWavProbe(): Uint8Array {
@@ -1,7 +1,10 @@
import { randomUUID } from 'node:crypto';
import { Kysely, sql } from 'kysely';
import { Logger } from '@nestjs/common';
import { SearchService } from 'src/core/search/search.service';
import { PageRepo } from '@docmost/db/repos/page/page.repo';
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
import { AiEmbeddingNotConfiguredException } from 'src/integrations/ai/ai-embedding-not-configured.exception';
import {
getTestDb,
destroyTestDb,
@@ -53,10 +56,19 @@ describe('SearchService #529 lexical overhaul [integration]', () => {
// Service wired to the real DB + real PageRepo (recursive descendants) with
// stubbed space-membership + permission repos so a test controls scope and the
// permission filter explicitly. `accessibleIds` (when set) is the KEEP list.
// #530: a FAKE AiService whose embedQuery is fully controlled per test. By
// DEFAULT it throws AiEmbeddingNotConfiguredException, so the semantic arm is
// omitted and every Phase-A case runs the BYTE-IDENTICAL lexical path — the
// semantic layer must never change lexical behaviour on degrade. Semantic cases
// pass `embedVector` (a deterministic 384-dim query vector) or `embedThrows`.
function buildService(opts?: {
userSpaceIds?: string[];
accessibleIds?: string[] | null;
filterThrows?: boolean;
embedVector?: number[];
embedFingerprint?: string;
embedThrows?: 'no-provider' | 'degraded';
slowVectorArm?: boolean;
}): SearchService {
const pageRepo = new PageRepo(db as any, null as any, null as any);
const spaceMemberRepo = {
@@ -71,15 +83,86 @@ describe('SearchService #529 lexical overhaul [integration]', () => {
: pageIds;
},
};
const aiService = {
embedQuery: async () => {
if (opts?.embedVector) {
return {
vector: opts.embedVector,
fingerprint: opts.embedFingerprint ?? ACTIVE_FP,
};
}
if (opts?.embedThrows === 'degraded') {
const err = new Error('embedding request timed out after 1ms');
err.name = 'TimeoutError';
throw err;
}
// Default (and explicit 'no-provider'): no embedding provider resolved.
throw new AiEmbeddingNotConfiguredException();
},
};
// Real repo on the migrated test DB — exercises the actual vector-candidate
// SQL (pgvector <=>, fingerprint + dimension filters). When slowVectorArm is
// set, wrap the arm so it sleeps well past the (tiny, test-set) statement
// timeout, forcing a 57014 cancellation to exercise the degrade-to-lexical net.
const realRepo = new PageEmbeddingRepo(db as any);
const pageEmbeddingRepo = opts?.slowVectorArm
? {
vectorCandidateArm: (p: any) => {
const inner = realRepo.vectorCandidateArm(p);
return sql`SELECT * FROM (${inner}) AS slowarm WHERE (SELECT true FROM pg_sleep(0.5))`;
},
}
: realRepo;
return new SearchService(
db as any,
pageRepo as any,
{} as any,
spaceMemberRepo as any,
pagePermissionRepo as any,
aiService as any,
pageEmbeddingRepo as any,
);
}
// Active embedding fingerprint used by planted rows + the fake query embed.
const ACTIVE_FP = 'fp-active-530';
// A deterministic unit-ish 384-dim vector with a single 1 at `concept`. Two
// vectors of the same concept have cosine distance 0; different concepts are
// orthogonal (distance 1), so nearest-neighbour ordering is fully controlled.
function conceptVec(concept: number): number[] {
const v = new Array(384).fill(0);
v[concept % 384] = 1;
return v;
}
// Plant one chunk embedding for a page via the REAL repo (exercises the #530
// fingerprint insert path). Dimension is fixed at 384 to match conceptVec.
async function plantEmbedding(
pageId: string,
vector: number[],
fingerprint: string = ACTIVE_FP,
planSpaceId: string = spaceId,
): Promise<void> {
const repo = new PageEmbeddingRepo(db as any);
await repo.insertChunks([
{
pageId,
workspaceId,
spaceId: planSpaceId,
attachmentId: null,
chunkIndex: 0,
chunkStart: 0,
chunkLength: 1,
content: 'planted chunk',
modelName: 'test-model',
modelDimensions: 384,
fingerprint,
embedding: vector,
},
]);
}
const search = (service: SearchService, params: any) =>
service.searchPage(params, { userId: 'u-1', workspaceId }) as any;
@@ -609,4 +692,324 @@ describe('SearchService #529 lexical overhaul [integration]', () => {
expect(hit2).toBeDefined();
expect(hit2.path).toEqual([]);
});
// === #530 Phase B (PR-1): semantic fusion ==================================
// Deterministic fake embedder (no live model) + real PageEmbeddingRepo on the
// migrated DB with planted vectors + the active fingerprint, so cosine ordering
// is fully controlled. Each case runs in its OWN space so planted rows never
// leak across tests.
describe('#530 semantic fusion', () => {
// 1. Vector-only hit: nearest to a page that has NO lexical match for the
// query term — it appears via the vector arm, and a pure-lexical run does
// not return it. Mutation (a): drop the vector arm -> this reddens.
it('#530-1 a vector-only hit appears; a pure-lexical run would not return it', async () => {
const s = (await createSpace(db, workspaceId)).id;
const page = await insertPage({
title: 'нейтральный вект заголовок',
textContent: 'содержимое без искомого термина',
spaceId: s,
});
await plantEmbedding(page, conceptVec(1), ACTIVE_FP, s);
const svc = buildService({ embedVector: conceptVec(1), userSpaceIds: [s] });
const res = await search(svc, { query: 'квазонимикс', spaceId: s });
expect(res.items.map((i: any) => i.id)).toContain(page);
expect(res.semantic.available).toBe(true);
expect(res.semantic.state).toBe('full');
// Pure-lexical (default degrade): the vector-only page is absent.
const lex = buildService({ userSpaceIds: [s] });
const res2 = await search(lex, { query: 'квазонимикс', spaceId: s });
expect(res2.items.map((i: any) => i.id)).not.toContain(page);
expect(res2.semantic.available).toBe(false);
});
// 2. Sidecar down: embedQuery throws -> results = the lexical set; semantic
// unavailable; the fixed 'search.semantic.degraded' event is logged.
it('#530-2 sidecar-down degrades to lexical and logs search.semantic.degraded', async () => {
const s = (await createSpace(db, workspaceId)).id;
const lexHit = await insertPage({
title: 'семдвамаркер лексический',
textContent: 'обычный текст',
spaceId: s,
});
const warnSpy = jest
.spyOn(Logger.prototype, 'warn')
.mockImplementation(() => undefined as any);
try {
const svc = buildService({ embedThrows: 'degraded', userSpaceIds: [s] });
const res = await search(svc, { query: 'семдвамаркер', spaceId: s });
expect(res.items.map((i: any) => i.id)).toContain(lexHit);
expect(res.semantic.available).toBe(false);
expect(res.semantic.reason).toBe('degraded');
expect(warnSpy).toHaveBeenCalledWith(
expect.stringContaining('search.semantic.degraded'),
);
} finally {
warnSpy.mockRestore();
}
});
// 3. No provider: resolveEmbeddingProvider yields none -> lexical results,
// reason no-provider, no 500.
it('#530-3 no-provider yields lexical results with reason no-provider', async () => {
const s = (await createSpace(db, workspaceId)).id;
const hit = await insertPage({
title: 'семтримаркер тема',
textContent: 'текст',
spaceId: s,
});
const svc = buildService({ embedThrows: 'no-provider', userSpaceIds: [s] });
const res = await search(svc, { query: 'семтримаркер', spaceId: s });
expect(res.items.map((i: any) => i.id)).toContain(hit);
expect(res.semantic.available).toBe(false);
expect(res.semantic.reason).toBe('no-provider');
});
// 4. Permission over the union: a restricted vector-only hit is dropped from
// items AND total; with the permission filter throwing the call rejects
// (fail-closed, never fail-open). Mutation (b): move filterAccessiblePageIds
// inside the semantic try/catch -> the reject assertion reddens.
it('#530-4 permission filters a vector-only hit from items AND total (fail-closed)', async () => {
const s = (await createSpace(db, workspaceId)).id;
const visibleLex = await insertPage({
title: 'семчетмаркер видимый',
textContent: 'доступный текст',
spaceId: s,
});
const hiddenVec = await insertPage({
title: 'скрытая вект страница',
textContent: 'содержимое без искомого термина',
spaceId: s,
});
await plantEmbedding(hiddenVec, conceptVec(4), ACTIVE_FP, s);
// Query hits visibleLex lexically; hiddenVec only via the vector arm.
const svc = buildService({
embedVector: conceptVec(4),
accessibleIds: [visibleLex],
userSpaceIds: [s],
});
const res = await search(svc, { query: 'семчетмаркер', spaceId: s });
const ids = res.items.map((i: any) => i.id);
expect(ids).toContain(visibleLex);
expect(ids).not.toContain(hiddenVec);
// The hidden vector hit does not leak into total either.
expect(res.total).toBe(1);
expect(ids).toHaveLength(res.total);
// A permission-query error PROPAGATES (never a fail-open empty result).
const boom = buildService({
embedVector: conceptVec(4),
filterThrows: true,
userSpaceIds: [s],
});
await expect(
search(boom, { query: 'семчетмаркер', spaceId: s }),
).rejects.toThrow(/permission query failed/);
});
// 5. Hung sidecar under a short embed timeout: graceful degrade (never a 500).
it('#530-5 a hung sidecar under a short timeout degrades gracefully', async () => {
const s = (await createSpace(db, workspaceId)).id;
const hit = await insertPage({
title: 'семпятьмаркер тема',
textContent: 'текст',
spaceId: s,
});
process.env.SEARCH_EMBED_TIMEOUT_MS = '1';
try {
const svc = buildService({ embedThrows: 'degraded', userSpaceIds: [s] });
const res = await search(svc, { query: 'семпятьмаркер', spaceId: s });
expect(res.items.map((i: any) => i.id)).toContain(hit);
expect(res.semantic.available).toBe(false);
expect(res.semantic.reason).toBe('degraded');
} finally {
delete process.env.SEARCH_EMBED_TIMEOUT_MS;
}
});
// 6. A page matched BOTH lexically and by vector is fused (de-duped), not
// returned twice — the agg CTE collapses it to one row.
it('#530-6 a page matched lexically AND by vector appears once', async () => {
const s = (await createSpace(db, workspaceId)).id;
const both = await insertPage({
title: 'семшестьмаркер общий',
textContent: 'и лексика и вектор',
spaceId: s,
});
await plantEmbedding(both, conceptVec(6), ACTIVE_FP, s);
const svc = buildService({ embedVector: conceptVec(6), userSpaceIds: [s] });
const res = await search(svc, { query: 'семшестьмаркер', spaceId: s });
const ids = res.items.map((i: any) => i.id);
expect(ids.filter((id: string) => id === both)).toHaveLength(1);
expect(res.total).toBe(1);
});
// 7. Fingerprint isolation: a planted row under a DIFFERENT fingerprint is not
// a vector candidate for the active-fingerprint query (guards mixing
// generations). It only appears if it also matches lexically (it does not).
it('#530-7 a stale-fingerprint vector row is not fused into results', async () => {
const s = (await createSpace(db, workspaceId)).id;
const stale = await insertPage({
title: 'нейтральный семь заголовок',
textContent: 'без искомого термина совсем',
spaceId: s,
});
// Same concept vector as the query, but an OLD generation fingerprint.
await plantEmbedding(stale, conceptVec(7), 'fp-stale-old', s);
const svc = buildService({ embedVector: conceptVec(7), userSpaceIds: [s] });
const res = await search(svc, { query: 'квазонимиксseven', spaceId: s });
expect(res.items.map((i: any) => i.id)).not.toContain(stale);
});
// 8. Stale embedding space_id (moved page): a page currently in space B whose
// embedding rows still carry space A (movePageToSpace does NOT reindex) is
// STILL a vector hit when searching B — space scoping is enforced only by
// the pages join (pages.space_id, always current), NOT by a stale
// page_embeddings.space_id. Guards against re-adding that predicate (which
// would wrongly drop the moved page's vector-only hit).
it('#530-8 a vector hit with a STALE embedding space_id (moved page) is still returned', async () => {
const newSpace = (await createSpace(db, workspaceId)).id;
const oldSpace = (await createSpace(db, workspaceId)).id;
// The page currently lives in newSpace (as after a move into newSpace).
const page = await insertPage({
title: 'перемещённая вект страница',
textContent: 'содержимое без искомого термина',
spaceId: newSpace,
});
// Its embedding was stamped with the OLD space and never reindexed.
await plantEmbedding(page, conceptVec(8), ACTIVE_FP, oldSpace);
const svc = buildService({
embedVector: conceptVec(8),
userSpaceIds: [newSpace],
});
const res = await search(svc, {
query: 'квазонимиксeight',
spaceId: newSpace,
});
// Survives despite page_embeddings.space_id (oldSpace) != pages.space_id.
expect(res.items.map((i: any) => i.id)).toContain(page);
expect(res.semantic.available).toBe(true);
});
// 9. Statement-timeout safety net: a pathologically slow vector scan is
// cancelled (SQLSTATE 57014) and search degrades to lexical-only — the
// lexical hit survives, the vector-only hit vanishes, no 500/hang, and the
// degraded event is logged. Mutation: make the fallback re-throw instead of
// degrading -> this test reds (the call rejects with 57014).
it('#530-9 a vector-query timeout degrades to lexical (57014), never 500', async () => {
const s = (await createSpace(db, workspaceId)).id;
const lexHit = await insertPage({
title: 'семдевятьмаркер лексический',
textContent: 'обычный текст',
spaceId: s,
});
const vecOnly = await insertPage({
title: 'нейтральный девять заголовок',
textContent: 'без искомого термина',
spaceId: s,
});
await plantEmbedding(vecOnly, conceptVec(9), ACTIVE_FP, s);
const warnSpy = jest
.spyOn(Logger.prototype, 'warn')
.mockImplementation(() => undefined as any);
// Tiny per-statement timeout so the 0.5s sleeping arm is cancelled.
process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS = '50';
try {
const svc = buildService({
embedVector: conceptVec(9),
slowVectorArm: true,
userSpaceIds: [s],
});
const res = await search(svc, { query: 'семдевятьмаркер', spaceId: s });
// Lexical result survives; the vector-only page does NOT (arm cancelled).
const ids = res.items.map((i: any) => i.id);
expect(ids).toContain(lexHit);
expect(ids).not.toContain(vecOnly);
expect(res.semantic.available).toBe(false);
expect(res.semantic.reason).toBe('degraded');
expect(warnSpy).toHaveBeenCalledWith(
expect.stringContaining('search.semantic.degraded'),
);
} finally {
delete process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS;
warnSpy.mockRestore();
}
});
// 10. SET LOCAL scoping: the search's statement timeout must NOT leak onto a
// later query on the same pooled connection. After a search that set a
// tiny timeout, a deliberately-slow standalone query still completes.
it('#530-10 the per-query statement timeout does not leak to later queries', async () => {
const s = (await createSpace(db, workspaceId)).id;
await insertPage({
title: 'семдесятьмаркер тема',
textContent: 'текст',
spaceId: s,
});
process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS = '50';
try {
const svc = buildService({
embedVector: conceptVec(10),
slowVectorArm: true,
userSpaceIds: [s],
});
// This search trips + resets the local timeout (via SET LOCAL semantics).
await search(svc, { query: 'семдесятьмаркер', spaceId: s });
// A subsequent query that sleeps 200ms must NOT be cancelled — proving the
// 50ms search timeout did not persist on the connection (it would raise
// 57014 if it had leaked). SET LOCAL auto-resets at the search tx end.
await expect(sql`SELECT pg_sleep(0.2)`.execute(db)).resolves.toBeDefined();
} finally {
delete process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS;
}
});
// 11. COMMIT-path no-leak (locks is_local=true). #530-10 only covers the
// ROLLBACK path (a timed-out search rolls back — a plain session SET would
// ALSO be undone on rollback, so it would NOT catch is_local true->false).
// Here a NORMAL, FAST semantic search COMMITs its transaction; a leaked
// session-level statement_timeout would then persist on the pooled
// connection and cancel later queries. We prove it does NOT: after the
// committing search (bounded at a small 100ms), a 300ms query succeeds.
// Because postgres.js pools, we run several sequential slow queries so at
// least one reuses the connection the committed search ran on (which is
// where a leaked timeout would live); ALL must succeed. Mutation: flip
// is_local true->false -> the committed 100ms timeout leaks and the
// follow-up pg_sleep is cancelled (57014), reddening this test.
it('#530-11 a COMMITTED search does not leak its statement timeout (is_local)', async () => {
const s = (await createSpace(db, workspaceId)).id;
const page = await insertPage({
title: 'семодиннадцать вект страница',
textContent: 'содержимое без искомого термина',
spaceId: s,
});
await plantEmbedding(page, conceptVec(11), ACTIVE_FP, s);
// Small bound: the fast vector arm still completes (<100ms) so the search
// COMMITs, but small enough that a LEAKED timeout would cancel the 300ms
// follow-ups below.
process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS = '100';
try {
const svc = buildService({
embedVector: conceptVec(11),
userSpaceIds: [s],
});
const res = await search(svc, { query: 'семодиннадцать', spaceId: s });
// The (committing) vector path actually ran.
expect(res.semantic.available).toBe(true);
// Each 300ms query must SUCCEED — a leaked 100ms timeout would cancel the
// one that reuses the committed search's connection. 8 sweeps cover the
// whole pool (max 5 conns), so a leak on any connection is caught.
for (let i = 0; i < 8; i++) {
await expect(
sql`SELECT pg_sleep(0.3)`.execute(db),
).resolves.toBeDefined();
}
} finally {
delete process.env.SEARCH_VECTOR_STATEMENT_TIMEOUT_MS;
}
});
});
});
+41
View File
@@ -4,11 +4,26 @@ services:
depends_on:
- db
- redis
- embeddings
environment:
APP_URL: 'http://localhost:3000'
APP_SECRET: 'REPLACE_WITH_LONG_SECRET'
DATABASE_URL: 'postgresql://docmost:STRONG_DB_PASSWORD@db:5432/docmost'
REDIS_URL: 'redis://redis:6379'
# #530 semantic search: the GLOBAL embedding provider (the `embeddings` TEI
# sidecar below). A workspace that configures its own embedding provider
# overrides this. TEI is OpenAI-compatible, so the endpoint is /v1.
EMBEDDING_ENDPOINT: 'http://embeddings:80/v1'
EMBEDDING_MODEL: 'intfloat/multilingual-e5-small'
EMBEDDING_API_KEY: 'unused'
EMBEDDING_DIMENSIONS: '384'
# MUST match the --revision the sidecar pins (see EMBEDDING_REVISION in
# .env). A revision bump changes the embedding fingerprint (PR-2 swaps
# generations); keep the two in lockstep.
EMBEDDING_REVISION: '${EMBEDDING_REVISION}'
# e5 models require these input prefixes; empty them for a non-e5 model.
EMBEDDING_QUERY_PREFIX: 'query: '
EMBEDDING_DOC_PREFIX: 'passage: '
ports:
- "3000:3000"
restart: unless-stopped
@@ -39,7 +54,33 @@ services:
volumes:
- redis_data:/data
# #530 Text Embeddings Inference (TEI) sidecar — the GLOBAL embedding provider
# for semantic search. OpenAI-compatible, reached only on the internal network
# (no published port). The model + revision are PINNED so the embedding
# fingerprint is stable; set EMBEDDING_REVISION to a real commit sha in .env.
embeddings:
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9
command:
- "--model-id"
- "intfloat/multilingual-e5-small"
- "--revision"
- "${EMBEDDING_REVISION}"
restart: unless-stopped
# Cache downloaded model weights so a restart does not re-download them.
volumes:
- tei-models:/data
# GET /health is TEI's readiness probe. (Drop this block if your TEI image
# variant ships without curl.)
healthcheck:
test: ["CMD-SHELL", "curl -fsS http://localhost:80/health || exit 1"]
interval: 30s
timeout: 5s
retries: 5
# Model download on first boot can be slow; don't flap as unhealthy meanwhile.
start_period: 120s
volumes:
docmost:
db_data:
redis_data:
tei-models: