Start improving externalIngest (#2728)

* Start improving externalIngest externalIngest

Starting to hook up better basic db support. Still not very optimized yet

* Fix tests for ci

Author: mjbvz

src/platform/workspaceChunkSearch/node/codeSearch/codeSearchChunkSearch.ts

@@ -44,6 +44,7 @@ import { EmbeddingsChunkSearch } from '../embeddingsChunkSearch';
 import { TfIdfWithSemanticChunkSearch } from '../tfidfWithSemanticChunkSearch';
 import { IWorkspaceFileIndex } from '../workspaceFileIndex';
 import { AdoCodeSearchRepo, BuildIndexTriggerReason, CodeSearchRepo, CodeSearchRepoStatus, GithubCodeSearchRepo, TriggerIndexingError, TriggerRemoteIndexingError } from './codeSearchRepo';
+import { ExternalIngestClient } from './externalIngestClient';
 import { ExternalIngestIndex } from './externalIngestIndex';
 import { CodeSearchRepoTracker, RepoInfo, TrackedRepoStatus } from './repoTracker';
 import { CodeSearchDiff, CodeSearchWorkspaceDiffTracker } from './workspaceDiff';
@@ -162,7 +163,10 @@ export class CodeSearchChunkSearch extends Disposable implements IWorkspaceChunk
 		this._tfIdfChunkSearch = tfIdfChunkSearch;
 
 		this._repoTracker = this._register(instantiationService.createInstance(CodeSearchRepoTracker));
-		this._externalIngestIndex = new Lazy(() => this._register(instantiationService.createInstance(ExternalIngestIndex)));
+		this._externalIngestIndex = new Lazy(() => {
+			const client = instantiationService.createInstance(ExternalIngestClient);
+			return this._register(instantiationService.createInstance(ExternalIngestIndex, client));
+		});
 
 		this._register(this._repoTracker.onDidAddOrUpdateRepo(info => {
 			if (info.status === TrackedRepoStatus.Resolved && info.resolvedRemoteInfo) {

src/platform/workspaceChunkSearch/node/codeSearch/externalIngestClient.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { DocumentContents, GeoFilter, IngestFilter, canIngestDocument, createCodedSymbols, setupPanicHooks } from '@github/blackbird-external-ingest-utils';
+import { canIngestDocument, canIngestPathAndSize, createCodedSymbols, DocumentContents, GeoFilter, IngestFilter, setupPanicHooks } from '@github/blackbird-external-ingest-utils';
 import crypto from 'crypto';
 import fs from 'fs';
 import { posix } from 'node:path';
@@ -12,26 +12,71 @@ import { coalesce } from '../../../../util/vs/base/common/arrays';
 import { timeout } from '../../../../util/vs/base/common/async';
 import { URI } from '../../../../util/vs/base/common/uri';
 import { Range } from '../../../../util/vs/editor/common/core/range';
+import { IAuthenticationService } from '../../../authentication/common/authentication';
 import { FileChunkAndScore } from '../../../chunking/common/chunk';
 import { EmbeddingType } from '../../../embeddings/common/embeddingsComputer';
 import { ILogService } from '../../../log/common/logService';
 import { CodeSearchResult } from '../../../remoteCodeSearch/common/remoteCodeSearch';
 import { ApiClient } from './externalIngestApi';
 
+export interface ExternalIngestFile {
+	readonly uri: URI;
+	readonly docSha: Uint8Array;
+
+	read(): Promise<Uint8Array>;
+}
+
+/**
+ * Interface for the external ingest client that handles indexing and searching files.
+ */
+export interface IExternalIngestClient {
+	doInitialIndex(filesetName: string, root: URI, allFiles: AsyncIterable<ExternalIngestFile>, token: CancellationToken): Promise<void>;
+
+	listFilesets(token: CancellationToken): Promise<string[]>;
+	deleteFileset(filesetName: string, token: CancellationToken): Promise<void>;
+
+	searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, token: CancellationToken): Promise<CodeSearchResult>;
+
+	/**
+	 * Quickly checks if a file can be ingested based on its path and size.
+	 */
+	canIngestPathAndSize(filePath: string, size: number): boolean;
+
+	/**
+	 * Checks if a file can be ingested based on its path and file contents.
+	 */
+	canIngestDocument(filePath: string, data: Uint8Array): boolean;
+}
 
 // Create a shared API client with throttling (target quota usage of 80)
 // You can change this to `null` to ignore the throttle
 
-export class ExternalIngestClient {
+export class ExternalIngestClient implements IExternalIngestClient {
 	private static apiClient = new ApiClient(80);
 
 	private static readonly PROMISE_POOL_SIZE = 32;
 	private static baseUrl = 'https://api.github.com';
 
+	private readonly _ingestFilter = new IngestFilter();
+
 	constructor(
+		@IAuthenticationService private readonly _authenticationService: IAuthenticationService,
 		@ILogService private readonly logService: ILogService,
 	) { }
 
+	public async getAuthToken(): Promise<string | undefined> {
+		return (await this._authenticationService.getGitHubSession('permissive', { silent: true }))?.accessToken
+			?? (await this._authenticationService.getGitHubSession('any', { silent: true }))?.accessToken;
+	}
+
+	public canIngestPathAndSize(filePath: string, size: number): boolean {
+		return canIngestPathAndSize(this._ingestFilter, filePath, size);
+	}
+
+	public canIngestDocument(filePath: string, data: Uint8Array): boolean {
+		return canIngestDocument(this._ingestFilter, filePath, new DocumentContents(data));
+	}
+
 	private getHeaders(authToken: string): Record<string, string> {
 		const headers: Record<string, string> = {
 			'Content-Type': 'application/json',
@@ -51,56 +96,36 @@ export class ExternalIngestClient {
 		return ExternalIngestClient.apiClient.makeRequest(url, this.getHeaders(authToken), 'POST', body, token);
 	}
 
-	async doInitialIndex(authToken: string, filesetName: string, root: URI, allFiles: AsyncIterable<{ readonly uri: URI; readonly docSha: Uint8Array }>, token: CancellationToken): Promise<void> {
+	async doInitialIndex(filesetName: string, root: URI, allFiles: AsyncIterable<ExternalIngestFile>, token: CancellationToken): Promise<void> {
+		const authToken = await this.getAuthToken();
+		if (!authToken) {
+			this.logService.warn('ExternalIngestClient::doInitialIndex(): No auth token available');
+			return;
+		}
+
 		setupPanicHooks();
 
 		// Initial setup
-		const ingestFilter = new IngestFilter();
 		const mappings = new Map<string, { full: string; relative: string }>();
 		const geoFilter = new GeoFilter();
 
 
 		this.logService.info(`ExternalIngestClient::doInitialIndex(). Creating ingest for fileset: ${filesetName}`);
-		const allDocShas: Uint8Array[] = [];
 
 		this.logService.trace(`ExternalIngestClient::doInitialIndex(). Checking for ingestable files...`);
 		const ingestableCheckStart = performance.now();
-		// Figure out which documents are uploadable and insert them into the geoFilter
-		// and DocSha to path map and DocSha array.
-		const checking = new Set<Promise<void>>();
+
+		const allDocShas: Uint8Array[] = [];
 		for await (const file of allFiles) {
 			const relativePath = posix.relative(root.path, file.uri.path);
 			const full = file.uri.fsPath;
 
-			const p = (async () => {
-				this.logService.debug(`ExternalIngestClient::doInitialIndex(). Checking if file can be ingested: ${relativePath}`);
-				const fileBytes = await fs.promises.readFile(full);
-				const content = new DocumentContents(fileBytes);
-				if (canIngestDocument(ingestFilter, relativePath, content)) { // Can we do this lazily?
-					try {
-						const docSha = file.docSha; //getDocSha(relativePath, content);
-						geoFilter.push(docSha);
-						allDocShas.push(docSha);
-						// Clients of the external ingest process are required to store a mapping of docSha to
-						// document path. In this example ingestion code it is handled in memory but you might want
-						// to persist somewhere. Note that our example converts the Uin8Arrays to base64 strings
-						// since Uint8Array doesn't work as a Map key because equality is checked by reference.
-						const docShaBase64 = Buffer.from(docSha).toString('base64');
-						mappings.set(docShaBase64, { full, relative: relativePath });
-					} catch (err) {
-						throw new Error('Exception during ingest file', err);
-					}
-				}
-			})();
-			p.finally(() => {
-				checking.delete(p);
-			});
-			checking.add(p);
-			if (checking.size >= ExternalIngestClient.PROMISE_POOL_SIZE) {
-				await Promise.race(checking);
-			}
+			geoFilter.push(file.docSha);
+			allDocShas.push(file.docSha);
+
+			const docShaBase64 = Buffer.from(file.docSha).toString('base64');
+			mappings.set(docShaBase64, { full, relative: relativePath });
 		}
-		await Promise.all(checking);
 
 		this.logService.debug(`ExternalIngestClient::doInitialIndex(). Found ${mappings.size} ingestable files in ${Math.round(performance.now() - ingestableCheckStart)}ms`,);
 
@@ -280,7 +305,13 @@ export class ExternalIngestClient {
 		this.logService.debug(`requestId: '${requestId}', body: ${body}`);
 	}
 
-	async listFilesets(authToken: string, token: CancellationToken): Promise<string[]> {
+	async listFilesets(token: CancellationToken): Promise<string[]> {
+		const authToken = await this.getAuthToken();
+		if (!authToken) {
+			this.logService.warn('ExternalIngestClient::listFilesets(): No auth token available');
+			return [];
+		}
+
 		const resp = await ExternalIngestClient.apiClient.makeRequest(
 			`${ExternalIngestClient.baseUrl}/external/code/ingest`,
 			this.getHeaders(authToken),
@@ -293,7 +324,13 @@ export class ExternalIngestClient {
 		return coalesce((body.filesets ?? []).map(x => x.name));
 	}
 
-	async deleteFileset(authToken: string, filesetName: string, token: CancellationToken): Promise<void> {
+	async deleteFileset(filesetName: string, token: CancellationToken): Promise<void> {
+		const authToken = await this.getAuthToken();
+		if (!authToken) {
+			this.logService.warn('ExternalIngestClient::deleteFileset(): No auth token available');
+			return;
+		}
+
 		return this.deleteFilesetByName(authToken, filesetName, token);
 	}
 
@@ -313,7 +350,13 @@ export class ExternalIngestClient {
 		this.logService.info(`ExternalIngestClient::deleteFilesetByName(): Deleted: ${fileSetName}`);
 	}
 
-	async searchFilesets(authToken: string, filesetName: string, rootUri: URI, prompt: string, limit: number, token: CancellationToken): Promise<CodeSearchResult> {
+	async searchFilesets(filesetName: string, rootUri: URI, prompt: string, limit: number, token: CancellationToken): Promise<CodeSearchResult> {
+		const authToken = await this.getAuthToken();
+		if (!authToken) {
+			this.logService.warn('ExternalIngestClient::searchFilesets(): No auth token available');
+			return { outOfSync: false, chunks: [] };
+		}
+
 		this.logService.debug(`ExternalIngestClient::searchFilesets(): Searching fileset '${filesetName}' for prompt: '${prompt}'`);
 		const embeddingType = EmbeddingType.metis_1024_I16_Binary;
 		const resp = await this.post(authToken, '/external/embeddings/code/search', {