Allow different providers for the chat and the completer (#105)

* The provider registry set the completer or chat provider separately, each having its settings

* Make the settings panel allowing multiple providers

* Fix default settings

* Handle the settings storage/retrieval (settings registry and local storage) from the parent component

* Apply suggestions from PR comments

Co-authored-by: Jeremy Tuloup <591645+jtpio@users.noreply.github.com>

* Rename type to 'ModelRole'

---------

Co-authored-by: Jeremy Tuloup <591645+jtpio@users.noreply.github.com>

Author: brichet

schema/provider-registry.json

@@ -1,5 +1,5 @@
 {
-  "title": "AI provider",
+  "title": "AI providers",
   "description": "Provider registry settings",
   "jupyter.lab.setting-icon": "@jupyterlite/ai:jupyternaut-lite",
   "jupyter.lab.setting-icon-label": "JupyterLite AI provider",
@@ -11,10 +11,16 @@
       "description": "Whether to use or not the secrets manager. If not, secrets will be stored in the browser (local storage)",
       "default": true
     },
-    "AIprovider": {
+    "UniqueProvider": {
+      "type": "boolean",
+      "title": "Use the same provider for chat and completer",
+      "description": "Whether to use only one provider for the chat and the completer.\nThis will overwrite all the settings for the completer, and copy the ones from the chat.",
+      "default": true
+    },
+    "AIproviders": {
       "type": "object",
-      "title": "AI provider",
-      "description": "The AI provider configuration",
+      "title": "AI providers",
+      "description": "The AI providers configuration",
       "default": {},
       "additionalProperties": true
     }

src/completion-provider.ts

@@ -28,7 +28,7 @@ export class CompletionProvider implements IInlineCompletionProvider {
    * Get the current completer name.
    */
   get name(): string {
-    return this._providerRegistry.currentName;
+    return this._providerRegistry.currentName('completer');
   }
 
   /**

src/default-providers/index.ts

@@ -45,7 +45,7 @@ import { IAIProvider, IAIProviderRegistry } from '../tokens';
 const AIProviders: IAIProvider[] = [
   {
     name: 'Anthropic',
-    chatModel: ChatAnthropic,
+    chat: ChatAnthropic,
     completer: AnthropicCompleter,
     settingsSchema: AnthropicSettings,
     errorMessage: (error: any) => error.error.error.message
@@ -54,29 +54,29 @@ const AIProviders: IAIProvider[] = [
     name: 'ChromeAI',
     // TODO: fix
     // @ts-expect-error: missing properties
-    chatModel: ChromeAI,
+    chat: ChromeAI,
     completer: ChromeCompleter,
     instructions: ChromeAIInstructions,
     settingsSchema: ChromeAISettings,
     compatibilityCheck: chromeAICompatibilityCheck
   },
   {
     name: 'MistralAI',
-    chatModel: ChatMistralAI,
+    chat: ChatMistralAI,
     completer: CodestralCompleter,
     instructions: MistralAIInstructions,
     settingsSchema: MistralAISettings
   },
   {
     name: 'Ollama',
-    chatModel: ChatOllama,
+    chat: ChatOllama,
     completer: OllamaCompleter,
     instructions: OllamaInstructions,
     settingsSchema: OllamaAISettings
   },
   {
     name: 'OpenAI',
-    chatModel: ChatOpenAI,
+    chat: ChatOpenAI,
     completer: OpenAICompleter,
     settingsSchema: OpenAISettings
   }
@@ -94,16 +94,16 @@ const webLLMProviderPlugin: JupyterFrontEndPlugin<void> = {
   activate: (app: JupyterFrontEnd, registry: IAIProviderRegistry) => {
     registry.add({
       name: 'WebLLM',
-      chatModel: ChatWebLLM,
+      chat: ChatWebLLM,
       completer: WebLLMCompleter,
       settingsSchema: WebLLMSettings,
       instructions: WebLLMInstructions,
       compatibilityCheck: webLLMCompatibilityCheck,
       exposeChatModel: true
     });
 
-    registry.providerChanged.connect(async (sender, args) => {
-      const { currentName, currentChatModel, chatError } = registry;
+    registry.providerChanged.connect(async (sender, role) => {
+      const { currentChatModel, chatError } = registry;
       if (currentChatModel === null) {
         Notification.emit(chatError, 'error', {
           autoClose: 2000
@@ -113,7 +113,7 @@ const webLLMProviderPlugin: JupyterFrontEndPlugin<void> = {
 
       // TODO: implement a proper way to handle models that may need to be initialized before being used.
       // Mostly applies to WebLLM and ChromeAI as they may need to download the model in the browser first.
-      if (currentName === 'WebLLM') {
+      if (registry.currentName(role) === 'WebLLM') {
         const compatibilityError = await webLLMCompatibilityCheck();
 
         if (compatibilityError) {

src/index.ts

@@ -188,7 +188,7 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
       });
 
       editorRegistry.addRenderer(
-        `${PLUGIN_IDS.providerRegistry}.AIprovider`,
+        `${PLUGIN_IDS.providerRegistry}.AIproviders`,
         aiSettingsRenderer({
           providerRegistry,
           secretsToken: token,
@@ -204,14 +204,27 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
             delete settings.schema.properties?.['UseSecretsManager'];
           }
           const updateProvider = () => {
-            // Update the settings to the AI providers.
-            const providerSettings = (settings.get('AIprovider').composite ?? {
-              provider: 'None'
-            }) as ReadonlyPartialJSONObject;
-            providerRegistry.setProvider({
-              name: providerSettings.provider as string,
-              settings: providerSettings
-            });
+            // Get the Ai provider settings.
+            const providerSettings = settings.get('AIproviders')
+              .composite as ReadonlyPartialJSONObject;
+
+            // Update completer provider.
+            if (Object.keys(providerSettings).includes('completer')) {
+              providerRegistry.setCompleterProvider(
+                providerSettings['completer'] as ReadonlyPartialJSONObject
+              );
+            } else {
+              providerRegistry.setCompleterProvider({});
+            }
+
+            // Update chat provider.
+            if (Object.keys(providerSettings).includes('chat')) {
+              providerRegistry.setChatProvider(
+                providerSettings['chat'] as ReadonlyPartialJSONObject
+              );
+            } else {
+              providerRegistry.setChatProvider({});
+            }
           };
 
           settings.changed.connect(() => updateProvider());