lablup · graphite-app · Apr 24, 2025 · Apr 24, 2025
diff --git a/react/src/components/Chat/ChatCard.tsx b/react/src/components/Chat/ChatCard.tsx
@@ -11,20 +11,23 @@ import {
   ChatType,
   Model,
 } from './ChatModel';
-import { CustomModelAlert, CustomModelForm } from './CustomModelForm';
-import {
-  ChatCard_endpoint$data,
-  ChatCard_endpoint$key,
-} from './__generated__/ChatCard_endpoint.graphql';
+import { CustomModelForm } from './CustomModelForm';
+import { ChatCard_endpoint$key } from './__generated__/ChatCard_endpoint.graphql';
 import { createOpenAI } from '@ai-sdk/openai';
 import { useChat } from '@ai-sdk/react';
 import { extractReasoningMiddleware, streamText, wrapLanguageModel } from 'ai';
-import { Alert, Card, CardProps, FormInstance } from 'antd';
+import { Alert, App, Card, CardProps } from 'antd';
 import { createStyles } from 'antd-style';
 import graphql from 'babel-plugin-relay/macro';
-import { isEmpty } from 'lodash';
 import _ from 'lodash';
-import React, { useEffect, useMemo, useRef, useState } from 'react';
+import React, {
+  startTransition,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+  useTransition,
+} from 'react';
 import { useFragment } from 'react-relay';
 
 interface ChatCardProps extends CardProps, ChatLifecycleEventType {
@@ -80,38 +83,76 @@ function useEndpoint(selectedEndpoint?: ChatCard_endpoint$key | null) {
   return { endpoint, setEndpoint } as const;
 }
 
+function createModelsURL(baseURL: string) {
+  const { origin, port, pathname: path } = new URL(baseURL.trim());
+  const host = port.length > 0 ? `${origin}:${port}` : origin;
+  const normalizedPath = path === '/' ? '/models' : `${path}/models`;
+
+  return new URL(normalizedPath, host).toString();
+}
+
 function useModels(
   provider: ChatProviderType,
   fetchKey: string,
-  endpoint?: ChatCard_endpoint$data | null,
+  baseURL?: string,
+  token?: string,
 ) {
+  const { t } = useTranslation();
+  const getModelsErrorMessage = (status?: number) => {
+    switch (status) {
+      case 401:
+        return t('error.UnauthorizedToken');
+      case 404:
+        return t('error.NotFoundBasePath');
+      case 500:
+        return t('error.InternalServerError');
+      case 503:
+        return t('error.ServiceUnavailable');
+      default:
+        return t('error.UnknownError');
+    }
+  };
+
   const { data: modelsResult } = useSuspenseTanQuery<{
     data: Array<Model>;
   }>({
-    queryKey: ['models', fetchKey, endpoint?.endpoint_id],
-    queryFn: () => {
-      return endpoint?.url
-        ? fetch(
-            new URL(
-              provider.basePath + '/models',
-              endpoint?.url ?? undefined,
-            ).toString(),
-          )
-            .then((res) => res.json())
-            .catch((e) => ({ data: [] }))
-        : Promise.resolve({ data: [] });
+    queryKey: ['models', fetchKey, baseURL, token],
+    queryFn: async () => {
+      try {
+        if (baseURL) {
+          const url = createModelsURL(baseURL);
+          const authToken = token || provider.apiKey;
+          const res = await fetch(url, {
+            headers: {
+              Authorization: authToken ? `Bearer ${authToken}` : '',
+            },
+          });
+
+      const url = createModelsURL(baseURL);
+      const authToken = provider.apiKey;
+      const res = await fetch(url, {
+        headers: {
+          Authorization: authToken ? `Bearer ${authToken}` : '',
+        },
+      });
+
+      if (!res.ok) {
+        return { data: [], error: res.status };
+      }
+
+      return await res.json();
     },
   });
 
-  const models = _.map(modelsResult?.data, (m) => ({
+  const models = _.map(modelsResult?.data || [], (m) => ({
     id: m.id,
     name: m.id,
   })) as BAIModel[];
 
   const selectedModelId = useMemo(
     () =>
       provider.modelId &&
-      _.includes(_.map(modelsResult?.data, 'id'), provider.modelId)
+      _.includes(_.map(modelsResult?.data || [], 'id'), provider.modelId)
         ? provider.modelId
         : (modelsResult?.data?.[0]?.id ?? 'custom'),
     [modelsResult?.data, provider.modelId],
@@ -153,6 +194,10 @@ const ChatHeader = React.memo(PureChatHeader, (prev, next) => {
 
 const ChatInput = React.memo(PureChatInput);
 
+function createBaseURL(basePath: string, endpointUrl?: string | null) {
+  return endpointUrl ? new URL(basePath, endpointUrl).toString() : undefined;
+}
+
 const ChatCard: React.FC<ChatCardProps> = ({
   chat,
   selectedEndpoint,
@@ -164,37 +209,26 @@ const ChatCard: React.FC<ChatCardProps> = ({
   const {
     styles: { chatCard: chatCardStyle, alert: alertStyle, ...chatCardStyles },
   } = useStyles();
-  const formRef = useRef<FormInstance>(null);
+  const [isPendingUpdate, startUpdateTransition] = useTransition();
+
   const dropContainerRef = useRef<HTMLDivElement>(null);
   const [fetchKey, updateFetchKey] = useUpdatableState('first');
   const [startTime, setStartTime] = useState<number | null>(null);
 
   const { endpoint, setEndpoint } = useEndpoint(selectedEndpoint);
+  const [baseURL, setBaseURL] = useState<string | undefined>(
+    createBaseURL(chat.provider.basePath, endpoint?.url),
+  );
+  const [token, setToken] = useState<string | undefined>();
   const { models, modelId, setModelId } = useModels(
     chat.provider,
     fetchKey,
-    endpoint,
+    baseURL,
+    token,
   );
   const { agents, agent, setAgent } = useAgents(chat.provider);
   const [sync, setSync] = useState(chat.sync);
 
-  const baseURL = endpoint?.url
-    ? new URL(chat.provider.basePath, endpoint?.url ?? undefined).toString()
-    : undefined;
-
-  const allowCustomModel = isEmpty(models);
-  const providerSettings = {
-    baseURL: allowCustomModel
-      ? formRef.current?.getFieldValue('baseURL')
-      : baseURL,
-    modelId: allowCustomModel
-      ? formRef.current?.getFieldValue('modelId')
-      : modelId,
-    apiKey: allowCustomModel
-      ? formRef.current?.getFieldValue('token')
-      : chat.provider.apiKey,
-  };
-
   const {
     error,
     messages,
@@ -216,13 +250,13 @@ const ChatCard: React.FC<ChatCardProps> = ({
       if (fetchOnClient || modelId === 'custom') {
         const body = JSON.parse(init?.body as string);
         const provider = createOpenAI({
-          baseURL: providerSettings.baseURL,
-          apiKey: providerSettings.apiKey || 'dummy',
+          baseURL: baseURL,
+          apiKey: token || chat.provider.apiKey || 'dummy',
         });
         const result = streamText({
           abortSignal: init?.signal || undefined,
           model: wrapLanguageModel({
-            model: provider(providerSettings.modelId),
+            model: provider(modelId),
             middleware: extractReasoningMiddleware({ tagName: 'think' }),
           }),
           messages: body?.messages,
@@ -234,13 +268,22 @@ const ChatCard: React.FC<ChatCardProps> = ({
         return result.toDataStreamResponse({
           sendReasoning: true,
         });
-      } else {
-        return fetch(input, init);
       }
+
+      return fetch(input, init);
     },
   });
 
+  useEffect(() => {
+    startTransition(() => {
+      setBaseURL(createBaseURL(chat.provider.basePath, endpoint?.url));
+      setToken(undefined);
+      updateFetchKey('first');
+    });
+  }, [endpoint?.url, chat.provider.basePath, updateFetchKey]);
+
   const isStreaming = status === 'streaming' || status === 'submitted';
+
   return (
     <Card
       variant="outlined"
@@ -249,7 +292,6 @@ const ChatCard: React.FC<ChatCardProps> = ({
       title={
         <ChatHeader
           chat={chat}
-          allowCustomModel={allowCustomModel}
           models={models}
           modelId={modelId}
           setModelId={setModelId}
@@ -269,19 +311,21 @@ const ChatCard: React.FC<ChatCardProps> = ({
       }
       ref={dropContainerRef}
     >
-      {allowCustomModel ? (
+      {_.isEmpty(models) && (
         <CustomModelForm
-          modelId={modelId}
           baseURL={baseURL}
-          formRef={formRef}
-          allowCustomModel={allowCustomModel}
-          alert={
-            formRef && (
-              <CustomModelAlert onClick={() => updateFetchKey(baseURL)} />
-            )
-          }
+          token={token}
+          endpointId={endpoint?.endpoint_id}
+          loading={isPendingUpdate}
+          onSubmit={(data) => {
+            startUpdateTransition(() => {
+              updateFetchKey();
+              setBaseURL(data.baseURL);
+              setToken(data.token);
+            });
+          }}
         />
-      ) : null}
+      )}
       {!_.isEmpty(error?.message) ? (
         <Alert
           message={error?.message}

diff --git a/react/src/components/Chat/ChatHeader.tsx b/react/src/components/Chat/ChatHeader.tsx
@@ -14,6 +14,7 @@ import {
 import { Message } from '@ai-sdk/react';
 import { CloseOutlined, MoreOutlined, PlusOutlined } from '@ant-design/icons';
 import { Dropdown, Button, theme, MenuProps, Typography, Switch } from 'antd';
+import { isEmpty } from 'lodash';
 import { Scale as ScaleIcon, Eraser as EraserIcon } from 'lucide-react';
 import React, { useState, startTransition } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -40,7 +41,6 @@ const SyncSwitch: React.FC<SyncSwitchProps> = ({ sync, onClick }) => {
 interface ChatHeaderProps extends ChatLifecycleEventType {
   chat: ChatType;
   showCompareMenuItem?: boolean;
-  allowCustomModel?: boolean;
   closable?: boolean;
   models: BAIModel[];
   modelId: string;
@@ -59,7 +59,6 @@ interface ChatHeaderProps extends ChatLifecycleEventType {
 const ChatHeader: React.FC<ChatHeaderProps> = ({
   chat,
   showCompareMenuItem,
-  allowCustomModel,
   closable,
   models,
   modelId,
@@ -169,16 +168,17 @@ const ChatHeader: React.FC<ChatHeaderProps> = ({
           value={endpoint?.endpoint_id}
           popupMatchSelectWidth={false}
         />
-        <ModelSelect
-          models={models}
-          value={modelId}
-          onChange={(modelId) => {
-            startTransition(() => {
-              setModelId(modelId);
-            });
-          }}
-          allowCustomModel={allowCustomModel}
-        />
+        {!isEmpty(models) && (
+          <ModelSelect
+            models={models}
+            value={modelId}
+            onChange={(modelId) => {
+              startTransition(() => {
+                setModelId(modelId);
+              });
+            }}
+          />
+        )}
       </Flex>
       <Flex gap={'xs'}>
         {closable && (

diff --git a/react/src/components/Chat/ChatModel.tsx b/react/src/components/Chat/ChatModel.tsx
@@ -3,6 +3,7 @@ import {
   ChatCard_endpoint$data,
   ChatCard_endpoint$key,
 } from './__generated__/ChatCard_endpoint.graphql';
+import i18n from 'i18next';
 
 export type ChatProviderType = {
   baseURL?: string;
@@ -77,3 +78,32 @@ export type ChatOptions = {
   agents: AIAgent[];
   agentId?: string;
 };
+
+export class ChatModelError extends Error {
+  status: number;
+
+  constructor(status: number) {
+    super(ChatModelError.errorMessage(status));
+    this.name = 'ModelsError';
+    this.status = status;
+  }
+
+  static errorMessage(status: number) {
+    const messageKey = (() => {
+      switch (status) {
+        case 401:
+          return 'error.UnauthorizedToken';
+        case 404:
+          return 'error.NotFoundBasePath';
+        case 500:
+          return 'error.InternalServerError';
+        case 503:
+          return 'error.ServiceUnavailable';
+        default:
+          return 'error.UnknownError';
+      }
+    })();
+
+    return i18n.t(messageKey, { status });
+  }
+}