fix rendering of model IDs with a colon in their name (#704)

jupyterlab · Mar 28, 2024 · 144cc9b · 144cc9b
1 parent 35202d0
commit 144cc9b
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 40 deletions.
diff --git a/packages/jupyter-ai/src/components/chat-settings.tsx b/packages/jupyter-ai/src/components/chat-settings.tsx
@@ -23,6 +23,7 @@ import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { minifyUpdate } from './settings/minify';
 import { useStackingAlert } from './mui-extras/stacking-alert';
 import { RendermimeMarkdown } from './rendermime-markdown';
+import { getProviderId, getModelLocalId } from '../utils';
 
 type ChatSettingsProps = {
   rmRegistry: IRenderMimeRegistry;
@@ -391,23 +392,6 @@ export function ChatSettings(props: ChatSettingsProps): JSX.Element {
   );
 }
 
-function getProviderId(globalModelId: string) {
-  if (!globalModelId) {
-    return null;
-  }
-
-  return globalModelId.split(':')[0];
-}
-
-function getModelLocalId(globalModelId: string) {
-  if (!globalModelId) {
-    return null;
-  }
-
-  const components = globalModelId.split(':').slice(1);
-  return components.join(':');
-}
-
 function getProvider(
   globalModelId: string,
   providers: AiService.ListProvidersResponse

diff --git a/packages/jupyter-ai/src/components/settings/use-server-info.ts b/packages/jupyter-ai/src/components/settings/use-server-info.ts
@@ -1,5 +1,6 @@
 import { useState, useEffect, useMemo, useCallback } from 'react';
 import { AiService } from '../../handler';
+import { getProviderId, getModelLocalId } from '../../utils';
 
 type ServerInfoProperties = {
   config: AiService.DescribeConfigResponse;
@@ -63,7 +64,7 @@ export function useServerInfo(): ServerInfo {
         lmGid === null ? null : getProvider(lmGid, lmProviders);
       const emProvider =
         emGid === null ? null : getProvider(emGid, emProviders);
-      const lmLocalId = lmGid === null ? '' : getLocalId(lmGid);
+      const lmLocalId = (lmGid && getModelLocalId(lmGid)) ?? '';
       setServerInfoProps({
         config,
         lmProviders,
@@ -135,25 +136,3 @@ function getProvider(
   const provider = providers.providers.find(p => p.id === providerId);
   return provider ?? null;
 }
-
-function getProviderId(gid: string) {
-  if (!gid) {
-    return null;
-  }
-
-  const components = gid.split(':');
-  if (components.length < 2) {
-    return null;
-  }
-
-  return components[0];
-}
-
-function getLocalId(gid: string) {
-  const components = gid.split(':');
-  if (components.length < 2) {
-    return '';
-  }
-
-  return components[1];
-}
diff --git a/packages/jupyter-ai/src/utils.ts b/packages/jupyter-ai/src/utils.ts
@@ -48,3 +48,26 @@ export function getCellIndex(notebook: Notebook, cellId: string): number {
   );
   return idx === undefined ? -1 : idx;
 }
+
+/**
+ * Obtain the provider ID component from a model ID.
+ */
+export function getProviderId(globalModelId: string): string | null {
+  if (!globalModelId) {
+    return null;
+  }
+
+  return globalModelId.split(':')[0];
+}
+
+/**
+ * Obtain the model name component from a model ID.
+ */
+export function getModelLocalId(globalModelId: string): string | null {
+  if (!globalModelId) {
+    return null;
+  }
+
+  const components = globalModelId.split(':').slice(1);
+  return components.join(':');
+}