feat: Add cross-encoder re-ranking after hybrid search (CF-1317)
Add rerank() function calling LiteLLM /v1/rerank endpoint (Cohere-compatible). Plugged into all 3 search functions (sessions, session-docs, archives) after RRF merge. Disabled by default via RERANK_ENABLED env var. Graceful fallback to RRF-only ranking on API failure. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
// Project archives operations for database-backed archival
|
||||
|
||||
import { query, queryOne, execute } from '../db.js';
|
||||
import { getEmbedding, formatEmbedding, generateContentHash, rrfMerge } from '../embeddings.js';
|
||||
import { getEmbedding, formatEmbedding, generateContentHash, rrfMerge, rerank } from '../embeddings.js';
|
||||
|
||||
type ArchiveType = 'session' | 'research' | 'audit' | 'investigation' | 'completed' | 'migration';
|
||||
|
||||
@@ -193,10 +193,30 @@ export async function archiveSearch(args: ArchiveSearchArgs): Promise<string> {
|
||||
let finalIds: number[];
|
||||
let searchLabel: string;
|
||||
|
||||
let rerankScores: Map<number, number> | null = null;
|
||||
|
||||
if (search_mode === 'hybrid' && vectorIds.length > 0 && keywordIds.length > 0) {
|
||||
const merged = rrfMerge(vectorIds, keywordIds);
|
||||
finalIds = merged.slice(0, limit).map(m => m.id as number);
|
||||
finalIds = merged.map(m => m.id as number);
|
||||
searchLabel = 'hybrid';
|
||||
|
||||
// Cross-encoder re-ranking (CF-1317)
|
||||
const docs = finalIds.map(id => {
|
||||
const r = vectorRows.get(id) || keywordRows.get(id);
|
||||
return (r as any)?.title || '';
|
||||
});
|
||||
const reranked = await rerank(searchQuery, docs, limit);
|
||||
if (reranked) {
|
||||
rerankScores = new Map();
|
||||
const reorderedIds = reranked.map(r => {
|
||||
rerankScores!.set(finalIds[r.index], r.relevance_score);
|
||||
return finalIds[r.index];
|
||||
});
|
||||
finalIds = reorderedIds;
|
||||
searchLabel = 'hybrid+rerank';
|
||||
} else {
|
||||
finalIds = finalIds.slice(0, limit);
|
||||
}
|
||||
} else if (vectorIds.length > 0) {
|
||||
finalIds = vectorIds;
|
||||
searchLabel = 'vector';
|
||||
@@ -212,9 +232,12 @@ export async function archiveSearch(args: ArchiveSearchArgs): Promise<string> {
|
||||
for (const id of finalIds) {
|
||||
const a = vectorRows.get(id) || keywordRows.get(id);
|
||||
if (!a) continue;
|
||||
const sim = vectorRows.has(id) ? ` (${Math.round((vectorRows.get(id)!).similarity * 100)}% match)` : '';
|
||||
const simParts: string[] = [];
|
||||
if (vectorRows.has(id)) simParts.push(`${Math.round((vectorRows.get(id)!).similarity * 100)}% match`);
|
||||
if (rerankScores?.has(id)) simParts.push(`rerank: ${rerankScores.get(id)!.toFixed(2)}`);
|
||||
const scores = simParts.length > 0 ? ` (${simParts.join(', ')})` : '';
|
||||
const sizeStr = a.file_size ? ` (${Math.round(a.file_size / 1024)}KB)` : '';
|
||||
lines.push(`**[${a.archive_type}]** ${a.title}${sim}`);
|
||||
lines.push(`**[${a.archive_type}]** ${a.title}${scores}`);
|
||||
lines.push(` Archived: ${a.archived_at}${sizeStr}`);
|
||||
if (a.original_path) {
|
||||
lines.push(` Path: ${a.original_path}`);
|
||||
|
||||
Reference in New Issue
Block a user