336 lines
10 KiB
TypeScript
336 lines
10 KiB
TypeScript
/**
|
||
* HybridSearchService — combines FTS5 keyword search with vector similarity
|
||
* search using Reciprocal Rank Fusion (RRF) to produce a hybrid ranking.
|
||
*
|
||
* When no embedding provider is configured (or alpha = 0), the service
|
||
* transparently falls back to FTS5-only mode with zero overhead.
|
||
*
|
||
* Configuration model:
|
||
* alpha = 0.0 → FTS5 only
|
||
* alpha = 0.5 → balanced hybrid (default)
|
||
* alpha = 1.0 → vector only
|
||
*/
|
||
|
||
import type Database from 'better-sqlite3';
|
||
import type { EmbeddingProvider } from '../embeddings/provider.js';
|
||
import { SnippetSearchResult, SnippetRepositoryRef } from '$lib/server/models/search-result.js';
|
||
import { SnippetEntity } from '$lib/server/models/snippet.js';
|
||
import { SearchResultMapper } from '$lib/server/mappers/search-result.mapper.js';
|
||
import { SearchService } from './search.service.js';
|
||
import { VectorSearch } from './vector.search.js';
|
||
import { reciprocalRankFusion } from './rrf.js';
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public interfaces
|
||
// ---------------------------------------------------------------------------
|
||
|
||
export interface HybridSearchOptions {
|
||
repositoryId: string;
|
||
versionId?: string;
|
||
type?: 'code' | 'info';
|
||
/** Maximum number of results to return. Default: 20. */
|
||
limit?: number;
|
||
/**
|
||
* Blend weight between FTS5 and vector search.
|
||
* 0.0 = FTS5 only, 1.0 = vector only, 0.5 = balanced.
|
||
* Default: 0.5.
|
||
*/
|
||
alpha?: number;
|
||
/**
|
||
* Search mode: 'auto' (default), 'keyword', 'semantic', or 'hybrid'.
|
||
* Overrides alpha when set to 'keyword' (forces 0) or 'semantic' (forces 1).
|
||
*/
|
||
searchMode?: 'auto' | 'keyword' | 'semantic' | 'hybrid';
|
||
/**
|
||
* Embedding profile ID for vector search.
|
||
* Default: 'local-default'.
|
||
*/
|
||
profileId?: string;
|
||
}
|
||
|
||
/**
|
||
* Global search configuration stored in the settings table under
|
||
* `search_config`.
|
||
*/
|
||
export interface SearchConfig {
|
||
/** Blend weight (0.0–1.0). Default: 0.5. */
|
||
alpha: number;
|
||
/** Maximum results per search. Default: 20. */
|
||
maxResults: number;
|
||
/** True when an embedding provider is configured. */
|
||
enableHybrid: boolean;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Raw DB row used when re-fetching snippets by ID
|
||
// ---------------------------------------------------------------------------
|
||
|
||
interface RawSnippetById extends SnippetEntity {
|
||
repo_id: string;
|
||
repo_title: string;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// HybridSearchService
|
||
// ---------------------------------------------------------------------------
|
||
|
||
export class HybridSearchService {
|
||
private readonly vectorSearch: VectorSearch;
|
||
|
||
constructor(
|
||
private readonly db: Database.Database,
|
||
private readonly searchService: SearchService,
|
||
private readonly embeddingProvider: EmbeddingProvider | null
|
||
) {
|
||
this.vectorSearch = new VectorSearch(db);
|
||
}
|
||
|
||
/**
|
||
* Execute a hybrid search combining FTS5 and (optionally) vector search.
|
||
*
|
||
* Search modes:
|
||
* - 'keyword' : FTS5-only (alpha = 0)
|
||
* - 'semantic' : Vector-only (alpha = 1), skips FTS entirely
|
||
* - 'hybrid' : Balanced RRF fusion (alpha = 0.5 by default)
|
||
* - 'auto' : Auto-selects: semantic if embedding provider available and FTS
|
||
* yields no results on the preprocessed query. Falls back to FTS
|
||
* for punctuation-heavy queries.
|
||
*
|
||
* When embeddingProvider is null or alpha is 0, the method returns FTS5 results
|
||
* directly without embedding the query.
|
||
*
|
||
* @param query - Raw search string (preprocessing handled by SearchService).
|
||
* @param options - Search parameters including repositoryId and alpha blend.
|
||
* @returns Object with ranked results array and the search mode actually used.
|
||
*/
|
||
async search(
|
||
query: string,
|
||
options: HybridSearchOptions
|
||
): Promise<{ results: SnippetSearchResult[]; searchModeUsed: string }> {
|
||
const limit = options.limit ?? 20;
|
||
const mode = options.searchMode ?? 'auto';
|
||
|
||
// Resolve alpha from searchMode
|
||
let alpha: number;
|
||
switch (mode) {
|
||
case 'keyword':
|
||
alpha = 0;
|
||
break;
|
||
case 'semantic':
|
||
alpha = 1;
|
||
break;
|
||
case 'hybrid':
|
||
alpha = options.alpha ?? 0.5;
|
||
break;
|
||
default:
|
||
// 'auto'
|
||
alpha = options.alpha ?? 0.5;
|
||
}
|
||
|
||
// Semantic mode: skip FTS entirely and use vector search only.
|
||
if (mode === 'semantic') {
|
||
if (!this.embeddingProvider || !query.trim()) {
|
||
return { results: [], searchModeUsed: 'semantic' };
|
||
}
|
||
|
||
const embeddings = await this.embeddingProvider.embed([query]);
|
||
if (embeddings.length === 0) {
|
||
return { results: [], searchModeUsed: 'semantic' };
|
||
}
|
||
|
||
const queryEmbedding = embeddings[0].values;
|
||
const vectorResults = this.vectorSearch.vectorSearch(queryEmbedding, {
|
||
repositoryId: options.repositoryId,
|
||
versionId: options.versionId,
|
||
profileId: options.profileId,
|
||
limit
|
||
});
|
||
|
||
const topIds = vectorResults.slice(0, limit).map((r) => r.snippetId);
|
||
return {
|
||
results: this.fetchSnippetsByIds(
|
||
topIds,
|
||
options.repositoryId,
|
||
options.versionId,
|
||
options.type
|
||
),
|
||
searchModeUsed: 'semantic'
|
||
};
|
||
}
|
||
|
||
// FTS5 mode (keyword) or hybrid/auto modes: try FTS first.
|
||
const ftsResults = this.searchService.searchSnippets(query, {
|
||
repositoryId: options.repositoryId,
|
||
versionId: options.versionId,
|
||
type: options.type,
|
||
limit: limit * 3 // wider candidate pool for fusion
|
||
});
|
||
|
||
// Degenerate cases: no provider or pure FTS5 mode.
|
||
if (!this.embeddingProvider || alpha === 0) {
|
||
return { results: ftsResults.slice(0, limit), searchModeUsed: 'keyword' };
|
||
}
|
||
|
||
// For auto/hybrid modes: if FTS yielded results, use them; otherwise try vector.
|
||
// This handles punctuation-heavy queries that normalize to empty after preprocessing.
|
||
const hasFtsResults = ftsResults.length > 0;
|
||
|
||
if (!hasFtsResults) {
|
||
// No FTS results: try vector search as a fallback in auto/hybrid modes.
|
||
if (!query.trim()) {
|
||
// Query is empty; no point embedding it.
|
||
return { results: [], searchModeUsed: 'keyword_fallback' };
|
||
}
|
||
|
||
const embeddings = await this.embeddingProvider.embed([query]);
|
||
|
||
// If provider fails (Noop returns empty array), we're done.
|
||
if (embeddings.length === 0) {
|
||
return { results: [], searchModeUsed: 'keyword_fallback' };
|
||
}
|
||
|
||
const queryEmbedding = embeddings[0].values;
|
||
const vectorResults = this.vectorSearch.vectorSearch(queryEmbedding, {
|
||
repositoryId: options.repositoryId,
|
||
versionId: options.versionId,
|
||
profileId: options.profileId,
|
||
limit
|
||
});
|
||
|
||
const topIds = vectorResults.slice(0, limit).map((r) => r.snippetId);
|
||
return {
|
||
results: this.fetchSnippetsByIds(
|
||
topIds,
|
||
options.repositoryId,
|
||
options.versionId,
|
||
options.type
|
||
),
|
||
searchModeUsed: 'keyword_fallback'
|
||
};
|
||
}
|
||
|
||
// FTS has results: use RRF to blend with vector search (if alpha < 1).
|
||
const embeddings = await this.embeddingProvider.embed([query]);
|
||
|
||
// Provider may be a Noop (returns empty array) — fall back to FTS gracefully.
|
||
if (embeddings.length === 0) {
|
||
return { results: ftsResults.slice(0, limit), searchModeUsed: 'keyword' };
|
||
}
|
||
|
||
const queryEmbedding = embeddings[0].values;
|
||
|
||
const vectorResults = this.vectorSearch.vectorSearch(queryEmbedding, {
|
||
repositoryId: options.repositoryId,
|
||
versionId: options.versionId,
|
||
profileId: options.profileId,
|
||
limit: limit * 3
|
||
});
|
||
|
||
// Pure vector mode: skip RRF and return vector results directly.
|
||
if (alpha === 1) {
|
||
const topIds = vectorResults.slice(0, limit).map((r) => r.snippetId);
|
||
return {
|
||
results: this.fetchSnippetsByIds(
|
||
topIds,
|
||
options.repositoryId,
|
||
options.versionId,
|
||
options.type
|
||
),
|
||
searchModeUsed: 'semantic'
|
||
};
|
||
}
|
||
|
||
// Build ranked lists for RRF. Score field is unused by RRF — only
|
||
// the array index (rank) matters.
|
||
const ftsRanked = ftsResults.map((r, i) => ({ id: r.snippet.id, score: i }));
|
||
const vecRanked = vectorResults.map((r, i) => ({ id: r.snippetId, score: i }));
|
||
|
||
const fused = reciprocalRankFusion(ftsRanked, vecRanked);
|
||
|
||
const topIds = fused.slice(0, limit).map((r) => r.id);
|
||
return {
|
||
results: this.fetchSnippetsByIds(
|
||
topIds,
|
||
options.repositoryId,
|
||
options.versionId,
|
||
options.type
|
||
),
|
||
searchModeUsed: 'hybrid'
|
||
};
|
||
}
|
||
|
||
// -------------------------------------------------------------------------
|
||
// Private helpers
|
||
// -------------------------------------------------------------------------
|
||
|
||
/**
|
||
* Load full snippet + repository data for the given ordered snippet IDs.
|
||
*
|
||
* Results are returned in the same order as `ids` so callers receive the
|
||
* RRF-ranked list intact. Snippets not found in the database (or filtered
|
||
* out by optional type constraint) are silently omitted.
|
||
*/
|
||
private fetchSnippetsByIds(
|
||
ids: string[],
|
||
repositoryId: string,
|
||
versionId?: string,
|
||
type?: 'code' | 'info'
|
||
): SnippetSearchResult[] {
|
||
if (ids.length === 0) return [];
|
||
|
||
const placeholders = ids.map(() => '?').join(', ');
|
||
const params: unknown[] = [...ids, repositoryId];
|
||
let versionClause = '';
|
||
let typeClause = '';
|
||
if (versionId !== undefined) {
|
||
versionClause = ' AND s.version_id = ?';
|
||
params.push(versionId);
|
||
}
|
||
if (type !== undefined) {
|
||
typeClause = ' AND s.type = ?';
|
||
params.push(type);
|
||
}
|
||
|
||
const rows = this.db
|
||
.prepare<unknown[], RawSnippetById>(
|
||
`SELECT
|
||
s.id, s.document_id, s.repository_id, s.version_id, s.type,
|
||
s.title, s.content, s.language, s.breadcrumb, s.token_count,
|
||
s.created_at,
|
||
r.id AS repo_id,
|
||
r.title AS repo_title
|
||
FROM snippets s
|
||
JOIN repositories r ON r.id = s.repository_id
|
||
WHERE s.id IN (${placeholders})
|
||
AND s.repository_id = ?${versionClause}${typeClause}`
|
||
)
|
||
.all(...params) as RawSnippetById[];
|
||
|
||
// Build a map for O(1) lookup, then reconstruct in rank order.
|
||
const byId = new Map<string, RawSnippetById>();
|
||
for (const row of rows) {
|
||
byId.set(row.id, row);
|
||
}
|
||
|
||
const results: SnippetSearchResult[] = [];
|
||
for (const id of ids) {
|
||
const row = byId.get(id);
|
||
if (!row) continue;
|
||
|
||
results.push(
|
||
new SnippetSearchResult({
|
||
snippet: SearchResultMapper.snippetFromEntity(
|
||
new SnippetEntity(row),
|
||
new SnippetRepositoryRef({ id: row.repo_id, title: row.repo_title }),
|
||
0
|
||
).snippet,
|
||
score: 0,
|
||
repository: new SnippetRepositoryRef({ id: row.repo_id, title: row.repo_title })
|
||
})
|
||
);
|
||
}
|
||
|
||
return results;
|
||
}
|
||
}
|