Fix handling of languages (#4)

jtpio · Jun 4, 2024 · 6c74893 · 6c74893
1 parent a6bf199
commit 6c74893
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 11 deletions.
diff --git a/package.json b/package.json
@@ -57,6 +57,7 @@
         "@connectrpc/connect-web": "^1.4.0",
         "@jupyterlab/application": "^4.0.0",
         "@jupyterlab/apputils": "^4.0.0",
+        "@jupyterlab/codemirror": "^4.0.0",
         "@jupyterlab/completer": "^4.0.0",
         "@jupyterlab/settingregistry": "^4.0.0",
         "@lumino/coreutils": "^2.1.2"

diff --git a/src/codeium.ts b/src/codeium.ts
@@ -5,6 +5,7 @@ import { UUID } from '@lumino/coreutils';
 
 import { createPromiseClient } from '@connectrpc/connect';
 import { LanguageServerService } from './api/proto/exa/language_server_pb/language_server_connect';
+import { Language } from './api/proto/exa/codeium_common_pb/codeium_common_pb';
 import { createConnectTransport } from '@connectrpc/connect-web';
 import {
   Document,
@@ -33,6 +34,8 @@ export async function getCodeiumCompletions({
   config: ICodeiumConfig;
   otherDocuments: PartialMessage<Document>[];
 }) {
+  const lang = config.language;
+  const language = Language[lang?.toUpperCase() as keyof typeof Language];
   return await client.getCompletions(
     {
       metadata: {
@@ -48,8 +51,8 @@ export async function getCodeiumCompletions({
       document: {
         text: text,
         cursorOffset: BigInt(cursorOffset),
-        language: config.language,
-        editorLanguage: 'python',
+        language,
+        editorLanguage: lang ?? 'python',
         lineEnding: '\n'
       },
       editorOptions: {

diff --git a/src/config.ts b/src/config.ts
@@ -1,8 +1,6 @@
 // Some of this code is based on the code from https://github.com/val-town/codemirror-codeium/blob/main/src/config.ts
 // licensed under the ISC License: https://github.com/val-town/codemirror-codeium/blob/main/LICENSE
 
-import { Language } from './api/proto/exa/codeium_common_pb/codeium_common_pb';
-
 export interface ICodeiumConfig {
   /**
    * Codeium API key
@@ -12,7 +10,7 @@ export interface ICodeiumConfig {
   /**
    * The programming language of the given document.
    */
-  language?: Language;
+  language?: string;
   /**
    * Time in millseconds after typing to fetch
    * completions from codeium

diff --git a/src/index.ts b/src/index.ts
@@ -2,22 +2,27 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
-
 import { Dialog, showDialog } from '@jupyterlab/apputils';
+import { IEditorLanguageRegistry } from '@jupyterlab/codemirror';
 import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { CodeiumProvider } from './provider';
 
 const plugin: JupyterFrontEndPlugin<void> = {
   id: 'jupyterlab-codeium:inline-provider',
   autoStart: true,
-  requires: [ICompletionProviderManager, ISettingRegistry],
+  requires: [
+    ICompletionProviderManager,
+    IEditorLanguageRegistry,
+    ISettingRegistry
+  ],
   activate: (
     app: JupyterFrontEnd,
     manager: ICompletionProviderManager,
+    editorLanguageRegistry: IEditorLanguageRegistry,
     settingRegistry: ISettingRegistry
   ): void => {
-    const provider = new CodeiumProvider();
+    const provider = new CodeiumProvider({ editorLanguageRegistry });
     manager.registerInlineProvider(provider);
 
     settingRegistry

diff --git a/src/provider.ts b/src/provider.ts
@@ -4,12 +4,18 @@ import {
   IInlineCompletionProvider
 } from '@jupyterlab/completer';
 
+import { IEditorLanguageRegistry } from '@jupyterlab/codemirror';
+
 import { getCodeiumCompletions, simplifyCompletions } from './codeium';
 
 export class CodeiumProvider implements IInlineCompletionProvider {
   readonly identifier = 'codeium';
   readonly name = 'Codeium';
 
+  constructor(options: CodeiumProvider.IOptions) {
+    this._editorLanguageRegistry = options.editorLanguageRegistry;
+  }
+
   set apiKey(apiKey: string) {
     this._apiKey = apiKey;
   }
@@ -18,12 +24,14 @@ export class CodeiumProvider implements IInlineCompletionProvider {
     request: CompletionHandler.IRequest,
     context: IInlineCompletionContext
   ) {
-    const { text, offset: cursorOffset } = request;
+    const { text, offset: cursorOffset, mimeType } = request;
+    const language = this._editorLanguageRegistry.findByMIME(mimeType ?? '');
     const results = await getCodeiumCompletions({
       text,
       cursorOffset,
       config: {
-        apiKey: this._apiKey
+        apiKey: this._apiKey,
+        language: language?.support?.language.name
       },
       otherDocuments: []
     });
@@ -40,4 +48,11 @@ export class CodeiumProvider implements IInlineCompletionProvider {
   }
 
   private _apiKey = '';
+  private _editorLanguageRegistry: IEditorLanguageRegistry;
+}
+
+export namespace CodeiumProvider {
+  export interface IOptions {
+    editorLanguageRegistry: IEditorLanguageRegistry;
+  }
 }
diff --git a/yarn.lock b/yarn.lock
@@ -732,7 +732,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyterlab/codemirror@npm:^4.2.1":
+"@jupyterlab/codemirror@npm:^4.0.0, @jupyterlab/codemirror@npm:^4.2.1":
   version: 4.2.1
   resolution: "@jupyterlab/codemirror@npm:4.2.1"
   dependencies:
@@ -3976,6 +3976,7 @@ __metadata:
     "@jupyterlab/application": ^4.0.0
     "@jupyterlab/apputils": ^4.0.0
     "@jupyterlab/builder": ^4.0.0
+    "@jupyterlab/codemirror": ^4.0.0
     "@jupyterlab/completer": ^4.0.0
     "@jupyterlab/settingregistry": ^4.0.0
     "@lumino/coreutils": ^2.1.2