update crossfile context to use entire project but not limited to ope…

…n tabs

packages/core/src/amazonq/lsp/lspClient.ts

@@ -17,7 +17,19 @@ import * as jose from 'jose'
 import { Disposable, ExtensionContext } from 'vscode'
 
 import { LanguageClient, LanguageClientOptions, ServerOptions, TransportKind } from 'vscode-languageclient'
-import { GetUsageRequestType, IndexRequestType, QueryRequestType, UpdateIndexRequestType, Usage } from './types'
+import {
+    BuildIndexRequestPayload,
+    BuildIndexRequestType,
+    GetUsageRequestType,
+    IndexConfig,
+    IndexRequestType,
+    QueryInlineProjectContextRequestType,
+    QueryRequestType,
+    UpdateIndexRequestType,
+    UpdateIndexV2RequestPayload,
+    UpdateIndexV2RequestType,
+    Usage,
+} from './types'
 import { Writable } from 'stream'
 import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
 import { fs, getLogger } from '../../shared'
@@ -78,6 +90,24 @@ export class LspClient {
         }
     }
 
+    // v2
+    async indexFilesV2(paths: string[], rootPath: string, config: IndexConfig) {
+        const payload: BuildIndexRequestPayload = {
+            filePaths: paths,
+            projectRoot: rootPath,
+            config: config,
+            language: '',
+        }
+        try {
+            const encryptedRequest = await this.encrypt(JSON.stringify(payload))
+            const resp = await this.client?.sendRequest(BuildIndexRequestType, encryptedRequest)
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: indexFilesV2 error: ${e}`)
+            return undefined
+        }
+    }
+
     async query(request: string) {
         try {
             const encryptedRequest = await this.encrypt(
@@ -93,6 +123,23 @@ export class LspClient {
         }
     }
 
+    async queryBM25(query: string, path: string) {
+        try {
+            const request = JSON.stringify({
+                query: query,
+                filePath: path,
+            })
+
+            const encrpted = await this.encrypt(request)
+
+            let resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrpted)
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: query error: ${e}`)
+            return []
+        }
+    }
+
     async getLspServerUsage(): Promise<Usage | undefined> {
         if (this.client) {
             return (await this.client.sendRequest(GetUsageRequestType, '')) as Usage
@@ -113,6 +160,23 @@ export class LspClient {
             return undefined
         }
     }
+
+    // not yet account for file move
+    // v2
+    async updateIndexV2(filePath: string[], mode: 'update' | 'remove' | 'add') {
+        const payload: UpdateIndexV2RequestPayload = {
+            filePaths: filePath,
+            updateMode: mode,
+        }
+        try {
+            const encryptedRequest = await this.encrypt(JSON.stringify(payload))
+            const resp = await this.client?.sendRequest(UpdateIndexV2RequestType, encryptedRequest)
+            return resp
+        } catch (e) {
+            getLogger().error(`LspClient: updateIndexV2 error: ${e}`)
+            return undefined
+        }
+    }
 }
 /**
  * Activates the language server, this will start LSP server running over IPC protocol.
@@ -197,15 +261,30 @@ export async function activate(extensionContext: ExtensionContext) {
                 return
             }
             savedDocument = document.uri
-        })
-    )
-    toDispose.push(
+            void LspClient.instance.updateIndexV2([document.uri.fsPath], 'update')
+        }),
         vscode.window.onDidChangeActiveTextEditor((editor) => {
             if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
-                void LspClient.instance.updateIndex(savedDocument.fsPath)
+                // void LspClient.instance.updateIndexV2([editor.document.uri.fsPath], 'update')
             }
+        }),
+        vscode.workspace.onDidCreateFiles((e) => {
+            void LspClient.instance.updateIndexV2(
+                e.files.map((f) => f.fsPath),
+                'add'
+            )
+        }),
+        vscode.workspace.onDidDeleteFiles((e) => {
+            void LspClient.instance.updateIndexV2(
+                e.files.map((f) => f.fsPath),
+                'remove'
+            )
+        }),
+        vscode.workspace.onDidRenameFiles((e) => {
+            // void LspClient.instance.updateIndexV2(e.files.map((f) => f.newUri.fsPath), 'rename')
         })
     )
+
     return LspClient.instance.client.onReady().then(() => {
         const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
         toDispose.push(disposableFunc)

packages/core/src/amazonq/lsp/lspController.ts

@@ -304,7 +304,7 @@ export class LspController {
     }
 
     async buildIndex() {
-        getLogger().info(`LspController: Starting to build vector index of project`)
+        getLogger().info(`LspController: Starting LSP`)
         const start = performance.now()
         const projPaths = getProjectPaths()
         projPaths.sort()
@@ -325,13 +325,11 @@ export class LspController {
                 0
             )
             getLogger().info(`LspController: Found ${files.length} files in current project ${getProjectPaths()}`)
-            const resp = await LspClient.instance.indexFiles(
-                files.map((f) => f.fileUri.fsPath),
-                projRoot,
-                false
-            )
+            const config = CodeWhispererSettings.instance.isLocalIndexEnabled() ? 'all' : 'default'
+            const r = files.map((f) => f.fileUri.fsPath)
+            const resp = await LspClient.instance.indexFilesV2(r, projRoot, config)
             if (resp) {
-                getLogger().debug(`LspController: Finish building vector index of project`)
+                getLogger().debug(`LspController: Finish building index of project`)
                 const usage = await LspClient.instance.getLspServerUsage()
                 telemetry.amazonq_indexWorkspace.emit({
                     duration: performance.now() - start,
@@ -343,7 +341,7 @@ export class LspController {
                     credentialStartUrl: AuthUtil.instance.startUrl,
                 })
             } else {
-                getLogger().error(`LspController: Failed to build vector index of project`)
+                getLogger().error(`LspController: Failed to build index of project`)
                 telemetry.amazonq_indexWorkspace.emit({
                     duration: performance.now() - start,
                     result: 'Failed',
@@ -352,7 +350,7 @@ export class LspController {
                 })
             }
         } catch (e) {
-            getLogger().error(`LspController: Failed to build vector index of project`)
+            getLogger().error(`LspController: Failed to build index of project ${e}`)
             telemetry.amazonq_indexWorkspace.emit({
                 duration: performance.now() - start,
                 result: 'Failed',
@@ -371,12 +369,6 @@ export class LspController {
             return
         }
         setImmediate(async () => {
-            if (!CodeWhispererSettings.instance.isLocalIndexEnabled()) {
-                // only download LSP for users who did not turn on this feature
-                // do not start LSP server
-                await LspController.instance.tryInstallLsp(context)
-                return
-            }
             const ok = await LspController.instance.tryInstallLsp(context)
             if (!ok) {
                 return

packages/core/src/amazonq/lsp/types.ts

@@ -35,3 +35,40 @@ export interface Usage {
     memoryUsage: number
     cpuUsage: number
 }
+
+export type BuildIndexRequestPayload = {
+    filePaths: string[]
+    projectRoot: string
+    config: string
+    language: string
+}
+
+export type BuildIndexRequest = string
+
+export const BuildIndexRequestType: RequestType<BuildIndexRequest, any, any> = new RequestType('lsp/buildIndex')
+
+export type UpdateIndexV2Request = string
+
+export type UpdateIndexV2RequestPayload = { filePaths: string[]; updateMode: string }
+
+export const UpdateIndexV2RequestType: RequestType<UpdateIndexV2Request, any, any> = new RequestType(
+    'lsp/updateIndexV2'
+)
+
+export type QueryInlineProjectContextRequest = string
+export type QueryInlineProjectContextRequestPayload = {
+    query: string
+    filePath: string
+}
+export const QueryInlineProjectContextRequestType: RequestType<QueryInlineProjectContextRequest, any, any> =
+    new RequestType('lsp/queryInlineProjectContext')
+
+export type QueryVectorIndexRequestPayload = { query: string }
+
+export type QueryVectorIndexRequest = string
+
+export const QueryVectorIndexRequestType: RequestType<QueryVectorIndexRequest, any, any> = new RequestType(
+    'lsp/queryVectorIndex'
+)
+
+export type IndexConfig = 'all' | 'default'

packages/core/src/codewhisperer/models/constants.ts

@@ -690,6 +690,7 @@ export const crossFileContextConfig = {
     numberOfChunkToFetch: 60,
     topK: 3,
     numberOfLinesEachChunk: 10,
+    codemapMark: 'q-inline',
 }
 
 export const utgConfig = {

packages/core/src/codewhisperer/models/model.ts

@@ -48,7 +48,7 @@ export type UtgStrategy = 'ByName' | 'ByContent'
 
 export type CrossFileStrategy = 'OpenTabs_BM25'
 
-export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
+export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty' | 'LSP'
 
 export interface CodeWhispererSupplementalContext {
     isUtg: boolean

packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -16,6 +16,7 @@ import { getFileDistance } from '../../../shared/filesystemUtilities'
 import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities'
 import { getLogger } from '../../../shared/logger/logger'
 import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem } from '../../models/model'
+import { LspClient } from '../../../amazonq'
 
 type CrossFileSupportedLanguage =
     | 'java'
@@ -66,6 +67,51 @@ export async function fetchSupplementalContextForSrc(
               }
     }
 
+    // TODO:
+    if (false) {
+        return fetchSupplementalContextForSrcV1(editor, cancellationToken)
+    } else {
+        return fetchSupplementalContextForSrcV2(editor)
+    }
+}
+
+export async function fetchSupplementalContextForSrcV2(
+    editor: vscode.TextEditor
+): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
+    const inputChunkContent = getInputChunk(editor)
+
+    const bm25Response: { content: string; score: number; filePath: string }[] = await LspClient.instance.queryBM25(
+        inputChunkContent.content,
+        editor.document.uri.fsPath
+    )
+    getLogger().info(JSON.stringify(bm25Response))
+    console.log(bm25Response)
+
+    const supContextItems: CodeWhispererSupplementalContextItem[] = bm25Response
+    return {
+        supplementalContextItems: [...supContextItems],
+        strategy: 'LSP',
+    }
+}
+
+export async function fetchSupplementalContextForSrcV1(
+    editor: vscode.TextEditor,
+    cancellationToken: vscode.CancellationToken
+): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
+    const shouldProceed = shouldFetchCrossFileContext(
+        editor.document.languageId,
+        CodeWhispererUserGroupSettings.instance.userGroup
+    )
+
+    if (!shouldProceed) {
+        return shouldProceed === undefined
+            ? undefined
+            : {
+                  supplementalContextItems: [],
+                  strategy: 'Empty',
+              }
+    }
+
     const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
 
     // Step 1: Get relevant cross files to refer
@@ -91,7 +137,7 @@ export async function fetchSupplementalContextForSrc(
 
     // Step 3: Generate Input chunk (10 lines left of cursor position)
     // and Find Best K chunks w.r.t input chunk using BM25
-    const inputChunk: Chunk = getInputChunk(editor, crossFileContextConfig.numberOfLinesEachChunk)
+    const inputChunk: Chunk = getInputChunk(editor)
     const bestChunks: Chunk[] = findBestKChunkMatches(inputChunk, chunkList, crossFileContextConfig.topK)
     throwIfCancelled(cancellationToken)
 
@@ -137,7 +183,8 @@ function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: n
 /* This extract 10 lines to the left of the cursor from trigger file.
  * This will be the inputquery to bm25 matching against list of cross-file chunks
  */
-function getInputChunk(editor: vscode.TextEditor, chunkSize: number) {
+function getInputChunk(editor: vscode.TextEditor) {
+    const chunkSize = crossFileContextConfig.numberOfLinesEachChunk
     const cursorPosition = editor.selection.active
     const startLine = Math.max(cursorPosition.line - chunkSize, 0)
     const endLine = Math.max(cursorPosition.line - 1, 0)