feat(codewhisperer): generate code based on language of notebook cell #3346

Author: andrewyuq

src/codewhisperer/util/editorContext.ts

@@ -35,7 +35,7 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew
     )
 
     return {
-        filename: getFileName(editor),
+        filename: getFileNameForRequest(editor),
         programmingLanguage: {
             languageName:
                 runtimeLanguageContext.mapVscLanguageToCodeWhispererLanguage(editor.document.languageId) ??
@@ -51,6 +51,21 @@ export function getFileName(editor: vscode.TextEditor): string {
     return fileName.substring(0, CodeWhispererConstants.filenameCharsLimit)
 }
 
+export function getFileNameForRequest(editor: vscode.TextEditor): string {
+    const fileName = path.basename(editor.document.fileName)
+
+    // For notebook files, we want to use the programming language for each cell for the code suggestions, so change
+    // the filename sent in the request to reflect that language
+    if (fileName.endsWith('.ipynb')) {
+        const fileExtension = runtimeLanguageContext.getLanguageExtensionForNotebook(editor.document.languageId)
+        if (fileExtension !== undefined) {
+            const filenameWithNewExtension = fileName.substring(0, fileName.length - 5) + fileExtension
+            return filenameWithNewExtension.substring(0, CodeWhispererConstants.filenameCharsLimit)
+        }
+    }
+    return fileName.substring(0, CodeWhispererConstants.filenameCharsLimit)
+}
+
 export function buildListRecommendationRequest(
     editor: vscode.TextEditor,
     nextToken: string,

src/codewhisperer/util/runtimeLanguageContext.ts

@@ -16,6 +16,13 @@ export class RuntimeLanguageContext {
      */
     private supportedLanguageMap: ConstantMap<CodeWhispererConstants.SupportedLanguage, CodewhispererLanguage>
 
+    /**
+     * A map storing CodeWhisperer supported programming language with key: vscLanguageId and value: language extension
+     * Key: vscLanguageId
+     * Value: language extension
+     */
+    private supportedLanguageExtensionMap: ConstantMap<CodeWhispererConstants.SupportedLanguage, string>
+
     // A set contains vscode languageId and CodeWhispererLanguage
     private supportedLanguageSet = new Set<string>()
 
@@ -39,6 +46,25 @@ export class RuntimeLanguageContext {
             shellscript: 'shell',
             sql: 'sql',
         })
+        this.supportedLanguageExtensionMap = createConstantMap<CodeWhispererConstants.SupportedLanguage, string>({
+            java: 'java',
+            python: 'py',
+            javascriptreact: 'jsx',
+            javascript: 'js',
+            typescript: 'ts',
+            typescriptreact: 'tsx',
+            csharp: 'cs',
+            c: 'c',
+            cpp: 'cpp',
+            go: 'go',
+            kotlin: 'kt',
+            php: 'php',
+            ruby: 'rb',
+            rust: 'rs',
+            scala: 'scala',
+            shellscript: 'sh',
+            sql: 'sql',
+        })
 
         const values = Array.from(this.supportedLanguageMap.values())
         const keys = Array.from(this.supportedLanguageMap.keys())
@@ -55,6 +81,15 @@ export class RuntimeLanguageContext {
         return this.supportedLanguageMap.get(vscLanguageId) ?? undefined
     }
 
+    /**
+     * This is for notebook files map to a new filename with the corresponding language extension
+     * @param vscLanguageId : official vscode languageId
+     * @returns corresponding language extension if any, otherwise undefined
+     */
+    public getLanguageExtensionForNotebook(vscLanguageId?: string): string | undefined {
+        return this.supportedLanguageExtensionMap.get(vscLanguageId) ?? undefined
+    }
+
     /**
      * @param vscLanguageId : official vscode languageId
      * @returns An object with a field language: CodewhispererLanguage, if no corresponding CodewhispererLanguage ID, plaintext is returned

src/test/codewhisperer/util/editorContext.test.ts

@@ -64,6 +64,25 @@ describe('editorContext', function () {
         })
     })
 
+    describe('getfileNameForRequest', function () {
+        it('Should return a new filename with correct extension given a .ipynb file', function () {
+            const languageToExtension = new Map<string, string>([
+                ['python', 'py'],
+                ['rust', 'rs'],
+                ['javascript', 'js'],
+                ['typescript', 'ts'],
+                ['c', 'c'],
+            ])
+
+            languageToExtension.forEach((extension, language) => {
+                const editor = createMockTextEditor('', 'test.ipynb', language, 1, 17)
+                const actual = EditorContext.getFileNameForRequest(editor)
+                const expected = 'test.' + extension
+                assert.strictEqual(actual, expected)
+            })
+        })
+    })
+
     describe('validateRequest', function () {
         it('Should return false if request filename.length is invalid', function () {
             const req = createMockClientRequest()