From 7fbe5e373e34dae649215bbf35d6ec1dd0e87b54 Mon Sep 17 00:00:00 2001
From: "David L. Qiu" <david@qiu.dev>
Date: Wed, 27 Mar 2024 09:14:48 -0700
Subject: [PATCH] fix rendering of model IDs with a colon in their name

---
 .../src/components/chat-settings.tsx          | 18 +------------
 .../components/settings/use-server-info.ts    | 25 ++-----------------
 packages/jupyter-ai/src/utils.ts              | 23 +++++++++++++++++
 3 files changed, 26 insertions(+), 40 deletions(-)

diff --git a/packages/jupyter-ai/src/components/chat-settings.tsx b/packages/jupyter-ai/src/components/chat-settings.tsx
index 9650334d0..889342a24 100644
--- 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
index 87c54c474..80c77fbed 100644
--- 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
index 691cc18b3..e21b47c6b 100644
--- 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(':');
+}