fix: Merge models of the same category infiniflow#2479

cike8899 · Sep 23, 2024 · f02aebc · f02aebc
1 parent 139268d
commit f02aebc
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 16 deletions.
diff --git a/web/src/components/llm-select/index.tsx b/web/src/components/llm-select/index.tsx
@@ -1,5 +1,5 @@
 import { LlmModelType } from '@/constants/knowledge';
-import { useSelectLlmOptionsByModelType } from '@/hooks/llm-hooks';
+import { useComposeLlmOptionsByModelTypes } from '@/hooks/llm-hooks';
 import { Popover, Select } from 'antd';
 import LlmSettingItems from '../llm-setting-items';
 
@@ -10,7 +10,10 @@ interface IProps {
 }
 
 const LLMSelect = ({ id, value, onChange }: IProps) => {
-  const modelOptions = useSelectLlmOptionsByModelType();
+  const modelOptions = useComposeLlmOptionsByModelTypes([
+    LlmModelType.Chat,
+    LlmModelType.Image2text,
+  ]);
 
   const content = (
     <div style={{ width: 400 }}>
@@ -29,10 +32,7 @@ const LLMSelect = ({ id, value, onChange }: IProps) => {
       destroyTooltipOnHide
     >
       <Select
-        options={[
-          ...modelOptions[LlmModelType.Chat],
-          ...modelOptions[LlmModelType.Image2text],
-        ]}
+        options={modelOptions}
         style={{ width: '100%' }}
         dropdownStyle={{ display: 'none' }}
         id={id}

diff --git a/web/src/components/llm-setting-items/index.tsx b/web/src/components/llm-setting-items/index.tsx
@@ -7,7 +7,7 @@ import { Divider, Flex, Form, InputNumber, Select, Slider, Switch } from 'antd';
 import camelCase from 'lodash/camelCase';
 
 import { useTranslate } from '@/hooks/common-hooks';
-import { useSelectLlmOptionsByModelType } from '@/hooks/llm-hooks';
+import { useComposeLlmOptionsByModelTypes } from '@/hooks/llm-hooks';
 import { useCallback, useMemo } from 'react';
 import styles from './index.less';
 
@@ -39,7 +39,10 @@ const LlmSettingItems = ({ prefix, formItemLayout = {} }: IProps) => {
 
   const memorizedPrefix = useMemo(() => (prefix ? [prefix] : []), [prefix]);
 
-  const modelOptions = useSelectLlmOptionsByModelType();
+  const modelOptions = useComposeLlmOptionsByModelTypes([
+    LlmModelType.Chat,
+    LlmModelType.Image2text,
+  ]);
 
   return (
     <>
@@ -50,13 +53,7 @@ const LlmSettingItems = ({ prefix, formItemLayout = {} }: IProps) => {
         {...formItemLayout}
         rules={[{ required: true, message: t('modelMessage') }]}
       >
-        <Select
-          options={[
-            ...modelOptions[LlmModelType.Chat],
-            ...modelOptions[LlmModelType.Image2text],
-          ]}
-          showSearch
-        />
+        <Select options={modelOptions} showSearch />
       </Form.Item>
       <Divider></Divider>
       <Form.Item

diff --git a/web/src/hooks/llm-hooks.ts b/web/src/hooks/llm-hooks.ts
@@ -14,6 +14,7 @@ import userService from '@/services/user-service';
 import { sortLLmFactoryListBySpecifiedOrder } from '@/utils/common-util';
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
 import { message } from 'antd';
+import { DefaultOptionType } from 'antd/es/select';
 import { useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 
@@ -42,7 +43,7 @@ export const useSelectLlmOptions = () => {
         label: key,
         options: value.map((x) => ({
           label: x.llm_name,
-          value: x.llm_name,
+          value: `${x.llm_name}@${x.fid}`,
           disabled: !x.available,
         })),
       };
@@ -91,6 +92,26 @@ export const useSelectLlmOptionsByModelType = () => {
   };
 };
 
+export const useComposeLlmOptionsByModelTypes = (
+  modelTypes: LlmModelType[],
+) => {
+  const allOptions = useSelectLlmOptionsByModelType();
+
+  return modelTypes.reduce<DefaultOptionType[]>((pre, cur) => {
+    const options = allOptions[cur];
+    options.forEach((x) => {
+      const item = pre.find((y) => y.label === x.label);
+      if (item) {
+        item.options.push(...x.options);
+      } else {
+        pre.push(x);
+      }
+    });
+
+    return pre;
+  }, []);
+};
+
 export const useFetchLlmFactoryList = (): ResponseGetType<IFactory[]> => {
   const { data, isFetching: loading } = useQuery({
     queryKey: ['factoryList'],