feat: add support for ollama #221 (#260)

### What problem does this PR solve?

add support for ollama

Issue link:#221

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

web/src/components/svg-icon.tsx

@@ -21,13 +21,16 @@ try {
 interface IProps extends IconComponentProps {
   name: string;
   width: string | number;
+  height?: string | number;
 }
 
-const SvgIcon = ({ name, width, ...restProps }: IProps) => {
+const SvgIcon = ({ name, width, height, ...restProps }: IProps) => {
   const ListItem = routeList.find((item) => item.name === name);
   return (
     <Icon
-      component={() => <img src={ListItem?.value} alt="" width={width} />}
+      component={() => (
+        <img src={ListItem?.value} alt="" width={width} height={height} />
+      )}
       {...(restProps as any)}
     />
   );

web/src/hooks/llmHooks.ts

@@ -4,6 +4,7 @@ import {
   IMyLlmValue,
   IThirdOAIModelCollection,
 } from '@/interfaces/database/llm';
+import { IAddLlmRequestBody } from '@/interfaces/request/llm';
 import { useCallback, useEffect, useMemo } from 'react';
 import { useDispatch, useSelector } from 'umi';
 
@@ -206,3 +207,19 @@ export const useSaveTenantInfo = () => {
 
   return saveTenantInfo;
 };
+
+export const useAddLlm = () => {
+  const dispatch = useDispatch();
+
+  const saveTenantInfo = useCallback(
+    (requestBody: IAddLlmRequestBody) => {
+      return dispatch<any>({
+        type: 'settingModel/add_llm',
+        payload: requestBody,
+      });
+    },
+    [dispatch],
+  );
+
+  return saveTenantInfo;
+};

web/src/interfaces/common.ts

@@ -7,3 +7,11 @@ export interface BaseState {
   pagination: Pagination;
   searchString: string;
 }
+
+export interface IModalProps<T> {
+  showModal?(): void;
+  hideModal(): void;
+  visible: boolean;
+  loading?: boolean;
+  onOk?(payload?: T): Promise<void> | void;
+}

web/src/interfaces/request/llm.ts

@@ -0,0 +1,6 @@
+export interface IAddLlmRequestBody {
+  llm_factory: string; // Ollama
+  llm_name: string;
+  model_type: string;
+  api_base?: string; // chat|embedding|speech2text|image2text
+}

web/src/locales/en.ts

@@ -390,6 +390,14 @@ export default {
         'The default ASR model all the newly created knowledgebase will use. Use this model to translate voices to corresponding text.',
       workspace: 'Workspace',
       upgrade: 'Upgrade',
+      addLlmTitle: 'Add LLM',
+      modelName: 'Model name',
+      modelNameMessage: 'Please input your model name!',
+      modelType: 'Model type',
+      modelTypeMessage: 'Please input your model type!',
+      addLlmBaseUrl: 'Base url',
+      baseUrlNameMessage: 'Please input your base url!',
+      vision: 'Does it support Vision?',
     },
     message: {
       registered: 'Registered!',

web/src/locales/zh.ts

@@ -375,6 +375,14 @@ export default {
         '所有新创建的知识库都将使用默认的 ASR 模型。 使用此模型将语音翻译为相应的文本。',
       workspace: '工作空间',
       upgrade: '升级',
+      addLlmTitle: '添加 LLM',
+      modelName: '模型名称',
+      modelType: '模型类型',
+      addLlmBaseUrl: '基础 Url',
+      vision: '是否支持 Vision',
+      modelNameMessage: '请输入模型名称！',
+      modelTypeMessage: '请输入模型类型！',
+      baseUrlNameMessage: '请输入基础 Url！',
     },
     message: {
       registered: '注册成功',

web/src/pages/user-setting/model.ts

@@ -151,6 +151,17 @@ const model: DvaModel<SettingModelState> = {
       }
       return retcode;
     },
+    *add_llm({ payload = {} }, { call, put }) {
+      const { data } = yield call(userService.add_llm, payload);
+      const { retcode } = data;
+      if (retcode === 0) {
+        message.success(i18n.t('message.modified'));
+
+        yield put({ type: 'my_llm' });
+        yield put({ type: 'factories_list' });
+      }
+      return retcode;
+    },
   },
 };
 export default model;

web/src/pages/user-setting/setting-model/api-key-modal/index.tsx

@@ -46,8 +46,10 @@ const ApiKeyModal = ({
   };
 
   useEffect(() => {
-    form.setFieldValue('api_key', initialValue);
-  }, [initialValue, form]);
+    if (visible) {
+      form.setFieldValue('api_key', initialValue);
+    }
+  }, [initialValue, form, visible]);
 
   return (
     <Modal

web/src/pages/user-setting/setting-model/hooks.ts

@@ -2,6 +2,7 @@ import { useSetModalState } from '@/hooks/commonHooks';
 import {
   IApiKeySavingParams,
   ISystemModelSettingSavingParams,
+  useAddLlm,
   useFetchLlmList,
   useSaveApiKey,
   useSaveTenantInfo,
@@ -12,6 +13,7 @@ import {
   useFetchTenantInfo,
   useSelectTenantInfo,
 } from '@/hooks/userSettingHook';
+import { IAddLlmRequestBody } from '@/interfaces/request/llm';
 import { useCallback, useEffect, useState } from 'react';
 
 type SavingParamsState = Omit<IApiKeySavingParams, 'api_key'>;
@@ -127,3 +129,31 @@ export const useSelectModelProvidersLoading = () => {
 
   return loading;
 };
+
+export const useSubmitOllama = () => {
+  const loading = useOneNamespaceEffectsLoading('settingModel', ['add_llm']);
+  const addLlm = useAddLlm();
+  const {
+    visible: llmAddingVisible,
+    hideModal: hideLlmAddingModal,
+    showModal: showLlmAddingModal,
+  } = useSetModalState();
+
+  const onLlmAddingOk = useCallback(
+    async (payload: IAddLlmRequestBody) => {
+      const ret = await addLlm(payload);
+      if (ret === 0) {
+        hideLlmAddingModal();
+      }
+    },
+    [hideLlmAddingModal, addLlm],
+  );
+
+  return {
+    llmAddingLoading: loading,
+    onLlmAddingOk,
+    llmAddingVisible,
+    hideLlmAddingModal,
+    showLlmAddingModal,
+  };
+};

web/src/pages/user-setting/setting-model/index.tsx

@@ -1,15 +1,11 @@
 import { ReactComponent as MoreModelIcon } from '@/assets/svg/more-model.svg';
+import SvgIcon from '@/components/svg-icon';
 import { useSetModalState, useTranslate } from '@/hooks/commonHooks';
 import {
   LlmItem,
   useFetchLlmFactoryListOnMount,
   useFetchMyLlmListOnMount,
 } from '@/hooks/llmHooks';
-import { ReactComponent as MoonshotIcon } from '@/icons/moonshot.svg';
-import { ReactComponent as OpenAiIcon } from '@/icons/openai.svg';
-import { ReactComponent as TongYiIcon } from '@/icons/tongyi.svg';
-import { ReactComponent as WenXinIcon } from '@/icons/wenxin.svg';
-import { ReactComponent as ZhiPuIcon } from '@/icons/zhipu.svg';
 import { SettingOutlined, UserOutlined } from '@ant-design/icons';
 import {
   Avatar,
@@ -33,24 +29,27 @@ import ApiKeyModal from './api-key-modal';
 import {
   useSelectModelProvidersLoading,
   useSubmitApiKey,
+  useSubmitOllama,
   useSubmitSystemModelSetting,
 } from './hooks';
-import SystemModelSettingModal from './system-model-setting-modal';
-
 import styles from './index.less';
+import OllamaModal from './ollama-modal';
+import SystemModelSettingModal from './system-model-setting-modal';
 
 const IconMap = {
-  'Tongyi-Qianwen': TongYiIcon,
-  Moonshot: MoonshotIcon,
-  OpenAI: OpenAiIcon,
-  'ZHIPU-AI': ZhiPuIcon,
-  文心一言: WenXinIcon,
+  'Tongyi-Qianwen': 'tongyi',
+  Moonshot: 'moonshot',
+  OpenAI: 'openai',
+  'ZHIPU-AI': 'zhipu',
+  文心一言: 'wenxin',
+  Ollama: 'ollama',
 };
 
 const LlmIcon = ({ name }: { name: string }) => {
-  const Icon = IconMap[name as keyof typeof IconMap];
-  return Icon ? (
-    <Icon width={48} height={48}></Icon>
+  const icon = IconMap[name as keyof typeof IconMap];
+
+  return icon ? (
+    <SvgIcon name={`llm/${icon}`} width={48} height={48}></SvgIcon>
   ) : (
     <Avatar shape="square" size="large" icon={<UserOutlined />} />
   );
@@ -90,7 +89,7 @@ const ModelCard = ({ item, clickApiKey }: IModelCardProps) => {
           <Col span={12} className={styles.factoryOperationWrapper}>
             <Space size={'middle'}>
               <Button onClick={handleApiKeyClick}>
-                API-Key
+                {item.name === 'Ollama' ? t('addTheModel') : 'API-Key'}
                 <SettingOutlined />
               </Button>
               <Button onClick={handleShowMoreClick}>
@@ -142,16 +141,31 @@ const UserSettingModel = () => {
     showSystemSettingModal,
   } = useSubmitSystemModelSetting();
   const { t } = useTranslate('setting');
+  const {
+    llmAddingVisible,
+    hideLlmAddingModal,
+    showLlmAddingModal,
+    onLlmAddingOk,
+    llmAddingLoading,
+  } = useSubmitOllama();
 
   const handleApiKeyClick = useCallback(
     (llmFactory: string) => {
-      showApiKeyModal({ llm_factory: llmFactory });
+      if (llmFactory === 'Ollama') {
+        showLlmAddingModal();
+      } else {
+        showApiKeyModal({ llm_factory: llmFactory });
+      }
     },
-    [showApiKeyModal],
+    [showApiKeyModal, showLlmAddingModal],
   );
 
   const handleAddModel = (llmFactory: string) => () => {
-    handleApiKeyClick(llmFactory);
+    if (llmFactory === 'Ollama') {
+      showLlmAddingModal();
+    } else {
+      handleApiKeyClick(llmFactory);
+    }
   };
 
   const items: CollapseProps['items'] = [
@@ -216,7 +230,7 @@ const UserSettingModel = () => {
             clickButton={showSystemSettingModal}
           ></SettingTitle>
           <Divider></Divider>
-          <Collapse defaultActiveKey={['1']} ghost items={items} />
+          <Collapse defaultActiveKey={['1', '2']} ghost items={items} />
         </section>
       </Spin>
       <ApiKeyModal
@@ -233,6 +247,12 @@ const UserSettingModel = () => {
         hideModal={hideSystemSettingModal}
         loading={saveSystemModelSettingLoading}
       ></SystemModelSettingModal>
+      <OllamaModal
+        visible={llmAddingVisible}
+        hideModal={hideLlmAddingModal}
+        onOk={onLlmAddingOk}
+        loading={llmAddingLoading}
+      ></OllamaModal>
     </>
   );
 };

web/src/pages/user-setting/setting-model/ollama-modal/index.tsx

@@ -0,0 +1,96 @@
+import { useTranslate } from '@/hooks/commonHooks';
+import { IModalProps } from '@/interfaces/common';
+import { IAddLlmRequestBody } from '@/interfaces/request/llm';
+import { Form, Input, Modal, Select, Switch } from 'antd';
+import omit from 'lodash/omit';
+
+type FieldType = IAddLlmRequestBody & { vision: boolean };
+
+const { Option } = Select;
+
+const OllamaModal = ({
+  visible,
+  hideModal,
+  onOk,
+  loading,
+}: IModalProps<IAddLlmRequestBody>) => {
+  const [form] = Form.useForm<FieldType>();
+
+  const { t } = useTranslate('setting');
+
+  const handleOk = async () => {
+    const values = await form.validateFields();
+    const modelType =
+      values.model_type === 'chat' && values.vision
+        ? 'image2text'
+        : values.model_type;
+
+    const data = {
+      ...omit(values, ['vision']),
+      model_type: modelType,
+      llm_factory: 'Ollama',
+    };
+    console.info(data);
+
+    onOk?.(data);
+  };
+
+  return (
+    <Modal
+      title={t('addLlmTitle')}
+      open={visible}
+      onOk={handleOk}
+      onCancel={hideModal}
+      okButtonProps={{ loading }}
+    >
+      <Form
+        name="basic"
+        style={{ maxWidth: 600 }}
+        autoComplete="off"
+        layout={'vertical'}
+        form={form}
+      >
+        <Form.Item<FieldType>
+          label={t('modelType')}
+          name="model_type"
+          initialValue={'chat'}
+          rules={[{ required: true, message: t('modelTypeMessage') }]}
+        >
+          <Select placeholder={t('modelTypeMessage')}>
+            <Option value="chat">chat</Option>
+            <Option value="embedding">embedding</Option>
+          </Select>
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={t('modelName')}
+          name="llm_name"
+          rules={[{ required: true, message: t('modelNameMessage') }]}
+        >
+          <Input placeholder={t('modelNameMessage')} />
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={t('addLlmBaseUrl')}
+          name="api_base"
+          rules={[{ required: true, message: t('baseUrlNameMessage') }]}
+        >
+          <Input placeholder={t('baseUrlNameMessage')} />
+        </Form.Item>
+        <Form.Item noStyle dependencies={['model_type']}>
+          {({ getFieldValue }) =>
+            getFieldValue('model_type') === 'chat' && (
+              <Form.Item
+                label={t('vision')}
+                valuePropName="checked"
+                name={'vision'}
+              >
+                <Switch />
+              </Form.Item>
+            )
+          }
+        </Form.Item>
+      </Form>
+    </Modal>
+  );
+};
+
+export default OllamaModal;

web/src/pages/user-setting/setting-model/system-model-setting-modal/index.tsx

@@ -30,8 +30,10 @@ const SystemModelSettingModal = ({
   };
 
   useEffect(() => {
-    form.setFieldsValue(initialValues);
-  }, [form, initialValues]);
+    if (visible) {
+      form.setFieldsValue(initialValues);
+    }
+  }, [form, initialValues, visible]);
 
   const onFormLayoutChange = () => {};
 

web/src/services/userService.ts

@@ -14,6 +14,7 @@ const {
   my_llm,
   set_api_key,
   set_tenant_info,
+  add_llm,
 } = api;
 
 const methods = {
@@ -61,6 +62,10 @@ const methods = {
     url: set_api_key,
     method: 'post',
   },
+  add_llm: {
+    url: add_llm,
+    method: 'post',
+  },
 } as const;
 
 const userService = registerServer<keyof typeof methods>(methods, request);

web/src/utils/api.ts

@@ -17,6 +17,7 @@ export default {
   llm_list: `${api_host}/llm/list`,
   my_llm: `${api_host}/llm/my_llms`,
   set_api_key: `${api_host}/llm/set_api_key`,
+  add_llm: `${api_host}/llm/add_llm`,
 
   //知识库管理
   kb_list: `${api_host}/kb/list`,