From 6899775e4e3b845b7f126255c1c790ff7ee0f5f2 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Thu, 11 Jul 2024 17:41:51 +0800
Subject: [PATCH] feat: add model fetch and search

---
 .../settings/components/ModelListPopup.tsx    | 89 ++++++++++++++++---
 src/renderer/src/services/api.ts              | 10 +++
 2 files changed, 85 insertions(+), 14 deletions(-)

diff --git a/src/renderer/src/pages/settings/components/ModelListPopup.tsx b/src/renderer/src/pages/settings/components/ModelListPopup.tsx
index 6b0ea8da..d2d2118c 100644
--- a/src/renderer/src/pages/settings/components/ModelListPopup.tsx
+++ b/src/renderer/src/pages/settings/components/ModelListPopup.tsx
@@ -1,13 +1,16 @@
-import { Avatar, Button, Empty, Modal } from 'antd'
-import { useState } from 'react'
-import { TopView } from '../../../components/TopView'
-import { Model, Provider } from '@renderer/types'
-import { groupBy, isEmpty, uniqBy } from 'lodash'
-import styled from 'styled-components'
-import { MinusOutlined, PlusOutlined } from '@ant-design/icons'
-import { useProvider } from '@renderer/hooks/useProvider'
+import { LoadingOutlined, MinusOutlined, PlusOutlined } from '@ant-design/icons'
 import { SYSTEM_MODELS } from '@renderer/config/models'
+import { useProvider } from '@renderer/hooks/useProvider'
+import { fetchModels } from '@renderer/services/api'
 import { getModelLogo } from '@renderer/services/provider'
+import { Model, Provider } from '@renderer/types'
+import { getDefaultGroupName, runAsyncFunction } from '@renderer/utils'
+import { Avatar, Button, Empty, Flex, Modal } from 'antd'
+import Search from 'antd/es/input/Search'
+import { groupBy, isEmpty, uniqBy } from 'lodash'
+import { useEffect, useState } from 'react'
+import styled from 'styled-components'
+import { TopView } from '../../../components/TopView'
 
 interface ShowParams {
   provider: Provider
@@ -20,10 +23,18 @@ interface Props extends ShowParams {
 const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
   const [open, setOpen] = useState(true)
   const { provider, models, addModel, removeModel } = useProvider(_provider.id)
+  const [listModels, setListModels] = useState<Model[]>([])
+  const [loading, setLoading] = useState(false)
+  const [searchText, setSearchText] = useState('')
 
   const systemModels = SYSTEM_MODELS[_provider.id] || []
-  const allModels = uniqBy([...systemModels, ...models], 'id')
-  const systemModelGroups = groupBy(allModels, 'group')
+  const allModels = uniqBy([...systemModels, ...listModels, ...models], 'id')
+
+  const list = searchText
+    ? allModels.filter((model) => model.id.toLocaleLowerCase().includes(searchText.toLocaleLowerCase()))
+    : allModels
+
+  const modelGroups = groupBy(list, 'group')
 
   const onOk = () => {
     setOpen(false)
@@ -45,9 +56,40 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
     removeModel(model)
   }
 
+  useEffect(() => {
+    runAsyncFunction(async () => {
+      try {
+        setLoading(true)
+        const models = await fetchModels(_provider)
+        setListModels(
+          models.map((model) => ({
+            id: model.id,
+            // @ts-ignore name
+            name: model.name || model.id,
+            provider: _provider.id,
+            group: getDefaultGroupName(model.id)
+          }))
+        )
+        setLoading(false)
+      } catch (error) {
+        setLoading(false)
+      }
+    })
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  const ModalHeader = () => {
+    return (
+      <Flex>
+        <ModelHeaderTitle>{provider.name} Models</ModelHeaderTitle>
+        {loading && <LoadingOutlined size={20} />}
+      </Flex>
+    )
+  }
+
   return (
     <Modal
-      title={String(provider.name + ' Models').toUpperCase()}
+      title={<ModalHeader />}
       open={open}
       onOk={onOk}
       onCancel={onCancel}
@@ -58,11 +100,14 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
         content: { padding: 0 },
         header: { padding: 22, paddingBottom: 15 }
       }}>
+      <SearchContainer>
+        <Search placeholder="Search model id or name" allowClear onSearch={setSearchText} />
+      </SearchContainer>
       <ListContainer>
-        {Object.keys(systemModelGroups).map((group) => (
+        {Object.keys(modelGroups).map((group) => (
           <div key={group}>
             <ListHeader key={group}>{group}</ListHeader>
-            {systemModelGroups[group].map((model) => {
+            {modelGroups[group].map((model) => {
               const hasModel = provider.models.find((m) => m.id === model.id)
               return (
                 <ListItem key={model.id}>
@@ -82,12 +127,21 @@ const PopupContainer: React.FC<Props> = ({ provider: _provider, resolve }) => {
             })}
           </div>
         ))}
-        {isEmpty(allModels) && <Empty image={Empty.PRESENTED_IMAGE_SIMPLE} description="No models" />}
+        {isEmpty(list) && <Empty image={Empty.PRESENTED_IMAGE_SIMPLE} description="No models" />}
       </ListContainer>
     </Modal>
   )
 }
 
+const SearchContainer = styled.div`
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+  justify-content: space-between;
+  padding: 0 22px;
+  padding-bottom: 20px;
+`
+
 const ListContainer = styled.div`
   max-height: 70vh;
   overflow-y: scroll;
@@ -128,6 +182,13 @@ const ListItemName = styled.div`
   margin-left: 6px;
 `
 
+const ModelHeaderTitle = styled.div`
+  color: #fff;
+  font-size: 18px;
+  font-weight: 600;
+  margin-right: 10px;
+`
+
 export default class ModelListPopup {
   static topviewId = 0
   static hide() {
diff --git a/src/renderer/src/services/api.ts b/src/renderer/src/services/api.ts
index e5d784ee..d51dd2f7 100644
--- a/src/renderer/src/services/api.ts
+++ b/src/renderer/src/services/api.ts
@@ -100,3 +100,13 @@ export async function fetchMessagesSummary({ messages, assistant }: FetchMessage
 
   return response.choices[0].message?.content
 }
+
+export async function fetchModels(provider: Provider) {
+  try {
+    const openaiProvider = getOpenAiProvider(provider)
+    const response = await openaiProvider.models.list()
+    return response.data
+  } catch (error) {
+    return []
+  }
+}