Add experimental reranker service

src/extension/extension/vscode-node/services.ts

@@ -62,6 +62,7 @@ import { IWorkspaceMutationManager } from '../../../platform/testing/common/work
 import { ISetupTestsDetector, SetupTestsDetector } from '../../../platform/testing/node/setupTestDetector';
 import { ITestDepsResolver, TestDepsResolver } from '../../../platform/testing/node/testDepsResolver';
 import { ITokenizerProvider, TokenizerProvider } from '../../../platform/tokenizer/node/tokenizer';
+import { IRerankerService, RerankerService } from '../../../platform/workspaceChunkSearch/common/rerankerService';
 import { IWorkspaceChunkSearchService, WorkspaceChunkSearchService } from '../../../platform/workspaceChunkSearch/node/workspaceChunkSearchService';
 import { IWorkspaceFileIndex, WorkspaceFileIndex } from '../../../platform/workspaceChunkSearch/node/workspaceFileIndex';
 import { IInstantiationServiceBuilder } from '../../../util/common/services';
@@ -197,6 +198,7 @@ export function registerServices(builder: IInstantiationServiceBuilder, extensio
 	builder.define(ICodeSearchAuthenticationService, new SyncDescriptor(VsCodeCodeSearchAuthenticationService));
 	builder.define(ITodoListContextProvider, new SyncDescriptor(TodoListContextProvider));
 	builder.define(IGithubAvailableEmbeddingTypesService, new SyncDescriptor(GithubAvailableEmbeddingTypesService));
+	builder.define(IRerankerService, new SyncDescriptor(RerankerService));
 }
 
 function setupMSFTExperimentationService(builder: IInstantiationServiceBuilder, extensionContext: ExtensionContext) {

src/platform/workspaceChunkSearch/common/rerankerService.ts

@@ -0,0 +1,113 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { createServiceIdentifier } from '../../../util/common/services';
+import { CancellationToken } from '../../../util/vs/base/common/cancellation';
+import { raceCancellationError } from '../../../util/vs/base/common/async';
+import { FileChunkAndScore } from '../../chunking/common/chunk';
+import { ILogService } from '../../log/common/logService';
+import { IExperimentationService } from '../../telemetry/common/nullExperimentationService';
+
+export const IRerankerService = createServiceIdentifier<IRerankerService>('IRerankerService');
+
+export interface IRerankerService {
+	readonly _serviceBrand: undefined;
+	/**
+	 * Re-rank a list of file chunks for a natural language query.
+	 */
+	rerank(query: string, documents: readonly FileChunkAndScore[], token: CancellationToken): Promise<readonly FileChunkAndScore[]>;
+	/**
+	 * Whether the remote reranker endpoint is available
+	 */
+	readonly isAvailable: boolean;
+}
+
+
+interface RemoteRerankResultEntry {
+	readonly index: number;
+	readonly relevance_score?: number;
+}
+interface RemoteRerankResponse {
+	readonly results?: readonly RemoteRerankResultEntry[];
+}
+
+function buildQueryPrompt(userQuery: string): string {
+	return '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n'
+		+ '<|im_start|>user\n'
+		+ '<Instruct>: Given a web search query, retrieve relevant passages that answer the query\n'
+		+ `<Query>: ${userQuery}\n`;
+}
+
+function wrapDocument(text: string): string {
+	return `<Document>: ${text}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n`;
+}
+
+export class RerankerService implements IRerankerService {
+	declare readonly _serviceBrand: undefined;
+
+	constructor(
+		@ILogService private readonly _logService: ILogService,
+		@IExperimentationService private readonly _expService: IExperimentationService
+	) { }
+
+	private get _endpoint(): string | undefined {
+		return this._expService.getTreatmentVariable<string>('rerankEndpointUrl')?.trim();
+	}
+
+	get isAvailable(): boolean {
+		return !!this._endpoint;
+	}
+
+	async rerank(query: string, documents: readonly FileChunkAndScore[], token: CancellationToken): Promise<readonly FileChunkAndScore[]> {
+		if (!documents.length || !this.isAvailable || !this._endpoint) {
+			return documents;
+		}
+
+		const payload = {
+			query: buildQueryPrompt(query),
+			documents: documents.map(d => wrapDocument(d.chunk.text))
+		};
+
+		try {
+			const response = await raceCancellationError(fetch(this._endpoint, {
+				method: 'POST',
+				headers: { 'Content-Type': 'application/json' },
+				body: JSON.stringify(payload)
+			}), token);
+
+			if (!response.ok) {
+				this._logService.error(`RerankerService::rerank request failed. status=${response.status}`);
+				throw new Error(`Reranker request failed with status ${response.status}`);
+			}
+
+			const json = await raceCancellationError(response.json() as Promise<RemoteRerankResponse>, token);
+			const results = json.results;
+			if (!Array.isArray(results) || results.length === 0) {
+				throw new Error('Reranker returned no results');
+			}
+
+			// Sort descending by relevance (higher score = more relevant). If scores missing, treat as 0.
+			const sorted = [...results].sort((a, b) => (b.relevance_score ?? 0) - (a.relevance_score ?? 0));
+			const used = new Set<number>();
+			const reordered: FileChunkAndScore[] = [];
+			for (const entry of sorted) {
+				if (typeof entry.index === 'number' && entry.index >= 0 && entry.index < documents.length && !used.has(entry.index)) {
+					used.add(entry.index);
+					reordered.push(documents[entry.index]);
+				}
+			}
+			// Preserve any documents that were not returned by the reranker (defensive)
+			for (let i = 0; i < documents.length; i++) {
+				if (!used.has(i)) {
+					reordered.push(documents[i]);
+				}
+			}
+			return reordered;
+		} catch (e) {
+			this._logService.error(e, 'RerankerService::rerank exception');
+			throw e;
+		}
+	}
+}

src/platform/workspaceChunkSearch/common/workspaceChunkSearch.ts

@@ -71,6 +71,7 @@ export interface StrategySearchSizing {
 
 export interface WorkspaceChunkSearchOptions {
 	readonly globPatterns?: GlobIncludeOptions;
+	readonly enableRerank?: boolean;
 }
 
 export interface StrategySearchResult {

src/platform/workspaceChunkSearch/node/workspaceChunkSearchService.ts

@@ -35,6 +35,7 @@
 import { ITelemetryService } from '../../telemetry/common/telemetry';
 import { getWorkspaceFileDisplayPath, IWorkspaceService } from '../../workspace/common/workspaceService';
 import { IGithubAvailableEmbeddingTypesService } from '../common/githubAvailableEmbeddingTypes';
+import { IRerankerService } from '../common/rerankerService';
 import { IWorkspaceChunkSearchStrategy, StrategySearchResult, StrategySearchSizing, WorkspaceChunkQuery, WorkspaceChunkQueryWithEmbeddings, WorkspaceChunkSearchOptions, WorkspaceChunkSearchStrategyId, WorkspaceSearchAlert } from '../common/workspaceChunkSearch';
 import { CodeSearchChunkSearch, CodeSearchRemoteIndexState } from './codeSearchChunkSearch';
 import { EmbeddingsChunkSearch, LocalEmbeddingsIndexState, LocalEmbeddingsIndexStatus } from './embeddingsChunkSearch';
@@ -234,6 +235,7 @@
 		@IExperimentationService private readonly _experimentationService: IExperimentationService,
 		@IIgnoreService private readonly _ignoreService: IIgnoreService,
 		@ILogService private readonly _logService: ILogService,
+		@IRerankerService private readonly _rerankerService: IRerankerService,
 		@ISimulationTestContext private readonly _simulationTestContext: ISimulationTestContext,
 		@ITelemetryService private readonly _telemetryService: ITelemetryService,
 		@IVSCodeExtensionContext private readonly _extensionContext: IVSCodeExtensionContext,
@@ -417,6 +419,23 @@
 					chunks: filteredChunks
 				}
 			};
+
+			// If explicit rerank is enabled, use the remote reranker
+			if (options.enableRerank && this._rerankerService.isAvailable) {
+				try {
+					const queryString = await query.resolveQuery(token);
+					const reranked = await this._rerankerService.rerank(queryString, filteredResult.result.chunks, token);
+					return {
+						chunks: reranked.slice(0, this.getMaxChunks(sizing)),
+						isFullWorkspace: filteredResult.result.isFullWorkspace,
+						alerts: filteredResult.result.alerts,
+						strategy: filteredResult.strategy,
+					};
+				} catch (e) {
+					this._logService.error(e, 'Reranker service failed; falling back to local rerank');
+				}
+			}
+
 			return this.rerankResultIfNeeded(queryWithEmbeddings, filteredResult, this.getMaxChunks(sizing), telemetryInfo, progress, token);
 		}, (execTime, status) => {
 			/* __GDPR__