feat: collection model

2023-04-28 13:10:12 +08:00
parent 08ae4073bd
commit 5b9185159d
15 changed files with 286 additions and 86 deletions
--- a/src/pages/api/model/share/collection.ts
+++ b/src/pages/api/model/share/collection.ts
@@ -0,0 +1,44 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, Collection, Model } from '@/service/mongo';
+import { authToken } from '@/service/utils/tools';
+
+/* 模型收藏切换 */
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    const { modelId } = req.query as { modelId: string };
+
+    if (!modelId) {
+      throw new Error('缺少参数');
+    }
+    // 凭证校验
+    const userId = await authToken(req.headers.authorization);
+
+    await connectToDatabase();
+
+    const collectionRecord = await Collection.findOne({
+      userId,
+      modelId
+    });
+
+    if (collectionRecord) {
+      await Collection.findByIdAndRemove(collectionRecord._id);
+    } else {
+      await Collection.create({
+        userId,
+        modelId
+      });
+    }
+
+    await Model.findByIdAndUpdate(modelId, {
+      'share.collection': await Collection.countDocuments({ modelId })
+    });
+
+    jsonRes(res);
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}
--- a/src/pages/api/model/share/getCollection.ts
+++ b/src/pages/api/model/share/getCollection.ts
@@ -0,0 +1,38 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { connectToDatabase, Collection } from '@/service/mongo';
+import { authToken } from '@/service/utils/tools';
+import type { ShareModelItem } from '@/types/model';
+
+/* 获取模型列表 */
+export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    // 凭证校验
+    const userId = await authToken(req.headers.authorization);
+
+    await connectToDatabase();
+
+    // get my collections
+    const collections = await Collection.find({
+      userId
+    }).populate('modelId', '_id avatar name userId share');
+
+    jsonRes<ShareModelItem[]>(res, {
+      data: collections
+        .map((item: any) => ({
+          _id: item.modelId?._id,
+          avatar: item.modelId?.avatar || '',
+          name: item.modelId?.name || '',
+          userId: item.modelId?.userId || '',
+          share: item.modelId?.share || {},
+          isCollection: true
+        }))
+        .filter((item) => item.share.isShare)
+    });
+  } catch (err) {
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+}
--- a/src/pages/api/model/share/getModels.ts
+++ b/src/pages/api/model/share/getModels.ts
@@ -1,8 +1,7 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { jsonRes } from '@/service/response';
-import { connectToDatabase } from '@/service/mongo';
+import { connectToDatabase, Collection, Model } from '@/service/mongo';
 import { authToken } from '@/service/utils/tools';
-import { Model } from '@/service/models/model';
 import type { PagingData } from '@/types';
 import type { ShareModelItem } from '@/types/model';

@@ -30,19 +29,29 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
    };

    // 根据分享的模型
-    const models = await Model.find(where, '_id avatar name userId share')
-      .sort({
-        'share.collection': -1
-      })
-      .limit(pageSize)
-      .skip((pageNum - 1) * pageSize);
+    const [models, total] = await Promise.all([
+      Model.find(where, '_id avatar name userId share')
+        .sort({
+          'share.collection': -1
+        })
+        .limit(pageSize)
+        .skip((pageNum - 1) * pageSize),
+      Model.countDocuments(where)
+    ]);

    jsonRes<PagingData<ShareModelItem>>(res, {
      data: {
        pageNum,
        pageSize,
-        data: models,
-        total: await Model.countDocuments(where)
+        data: models.map((item) => ({
+          _id: item._id,
+          avatar: item.avatar,
+          name: item.name,
+          userId: item.userId,
+          share: item.share,
+          isCollection: false
+        })),
+        total
      }
    });
  } catch (err) {
--- a/src/pages/chat/components/SlideBar.tsx
+++ b/src/pages/chat/components/SlideBar.tsx
@@ -1,4 +1,4 @@
-import React, { useRef, useEffect } from 'react';
+import React, { useRef, useEffect, useMemo } from 'react';
 import { AddIcon, ChatIcon, DeleteIcon, MoonIcon, SunIcon } from '@chakra-ui/icons';
 import {
  Box,
@@ -22,6 +22,7 @@ import { getToken } from '@/utils/user';
 import MyIcon from '@/components/Icon';
 import WxConcat from '@/components/WxConcat';
 import { getChatHistory, delChatHistoryById } from '@/api/chat';
+import { getCollectionModels } from '@/api/model';
 import { modelList } from '@/constants/model';

 const SlideBar = ({
@@ -45,10 +46,30 @@ const SlideBar = ({
    cacheTime: 5 * 60 * 1000
  });

+  const { data: collectionModels = [] } = useQuery([getCollectionModels], getCollectionModels);
+
+  const models = useMemo(() => {
+    const myModelList = myModels.map((item) => ({
+      id: item._id,
+      name: item.name,
+      icon: modelList.find((model) => model.model === item?.service?.modelName)?.icon || 'model'
+    }));
+    const collectionList = collectionModels
+      .map((item) => ({
+        id: item._id,
+        name: item.name,
+        icon: 'collectionSolid' as any
+      }))
+      .filter((model) => !myModelList.find((item) => item.id === model.id));
+
+    return myModelList.concat(collectionList);
+  }, [collectionModels, myModels]);
+
  const { data: chatHistory = [], mutate: loadChatHistory } = useMutation({
    mutationFn: getChatHistory
  });

+  // update history
  useEffect(() => {
    if (chatId && preChatId.current === '') {
      loadChatHistory();
@@ -56,8 +77,11 @@ const SlideBar = ({
    preChatId.current = chatId;
  }, [chatId, loadChatHistory]);

+  // init history
  useEffect(() => {
-    loadChatHistory();
+    setTimeout(() => {
+      loadChatHistory();
+    }, 1000);
  }, [loadChatHistory]);

  const RenderHistory = () => (
@@ -165,9 +189,9 @@ const SlideBar = ({
        {isSuccess && (
          <>
            <Box>
-              {myModels.map((item) => (
+              {models.map((item) => (
                <Flex
-                  key={item._id}
+                  key={item.id}
                  alignItems={'center'}
                  p={3}
                  borderRadius={'md'}
@@ -178,28 +202,19 @@ const SlideBar = ({
                  }}
                  fontSize={'xs'}
                  border={'1px solid transparent'}
-                  {...(item._id === modelId
+                  {...(item.id === modelId
                    ? {
                        borderColor: 'rgba(255,255,255,0.5)',
                        backgroundColor: 'rgba(255,255,255,0.1)'
                      }
                    : {})}
                  onClick={async () => {
-                    if (item._id === modelId) return;
-                    resetChat(item._id);
+                    if (item.id === modelId) return;
+                    resetChat(item.id);
                    onClose();
                  }}
                >
-                  <MyIcon
-                    name={
-                      modelList.find((model) => model.model === item.service.modelName)?.icon ||
-                      'model'
-                    }
-                    mr={2}
-                    color={'white'}
-                    w={'16px'}
-                    h={'16px'}
-                  />
+                  <MyIcon name={item.icon} mr={2} color={'white'} w={'16px'} h={'16px'} />
                  <Box className={'textEllipsis'} flex={'1 0 0'} w={0}>
                    {item.name}
                  </Box>
--- a/src/pages/chat/index.tsx
+++ b/src/pages/chat/index.tsx
@@ -467,7 +467,7 @@ const Chat = ({ modelId, chatId }: { modelId: string; chatId: string }) => {
              borderBottom={'1px solid rgba(0,0,0,0.1)'}
            >
              <Flex maxW={'750px'} m={'auto'} alignItems={'flex-start'}>
-                <Menu>
+                <Menu autoSelect={false}>
                  <MenuButton as={Box} mr={media(4, 1)} cursor={'pointer'}>
                    <Image
                      src={item.obj === 'Human' ? '/icon/human.png' : chatData.avatar}
--- a/src/pages/model/detail/components/ModelEditForm.tsx
+++ b/src/pages/model/detail/components/ModelEditForm.tsx
@@ -80,9 +80,9 @@ const ModelEditForm = ({
          <Image
            src={getValues('avatar')}
            alt={'avatar'}
-            w={'36px'}
-            h={'36px'}
-            objectFit={'contain'}
+            w={['28px', '36px']}
+            h={['28px', '36px']}
+            objectFit={'cover'}
            cursor={'pointer'}
            title={'点击切换头像'}
            onClick={onOpenSelectFile}
--- a/src/pages/model/share/components/list.tsx
+++ b/src/pages/model/share/components/list.tsx
@@ -1,35 +1,56 @@
-import React, { Dispatch, useCallback } from 'react';
+import React from 'react';
 import { Card, Box, Flex, Image, Button } from '@chakra-ui/react';
 import type { ShareModelItem } from '@/types/model';
 import { useRouter } from 'next/router';
 import MyIcon from '@/components/Icon';
 import styles from '../index.module.scss';

-const ShareModelList = ({ models }: { models: ShareModelItem[] }) => {
+const ShareModelList = ({
+  models = [],
+  onclickCollection
+}: {
+  models: ShareModelItem[];
+  onclickCollection: (modelId: string) => void;
+}) => {
  const router = useRouter();

  return (
    <>
      {models.map((model) => (
-        <Card key={model._id} p={4}>
+        <Box
+          key={model._id}
+          p={4}
+          border={'1px solid'}
+          borderColor={'gray.200'}
+          borderRadius={'md'}
+        >
          <Flex alignItems={'center'}>
            <Image
              src={model.avatar}
              alt={'avatar'}
-              width={'36px'}
-              height={'36px'}
-              objectFit={'contain'}
+              w={['28px', '36px']}
+              h={['28px', '36px']}
+              objectFit={'cover'}
            />
            <Box fontWeight={'bold'} fontSize={'lg'} ml={5}>
              {model.name}
            </Box>
          </Flex>
-          <Box className={styles.intro} my={4} fontSize={'sm'}>
+          <Box className={styles.intro} my={4} fontSize={'sm'} color={'blackAlpha.600'}>
            {model.share.intro || '这个模型没有介绍~'}
          </Box>
          <Flex justifyContent={'space-between'}>
-            <Flex alignItems={'center'} cursor={'pointer'}>
-              <MyIcon mr={1} name={'collectionLight'} w={'16px'} />
+            <Flex
+              alignItems={'center'}
+              cursor={'pointer'}
+              color={model.isCollection ? 'blue.600' : 'alphaBlack.700'}
+              onClick={() => onclickCollection(model._id)}
+            >
+              <MyIcon
+                mr={1}
+                name={model.isCollection ? 'collectionSolid' : 'collectionLight'}
+                w={'16px'}
+              />
              {model.share.collection}
            </Flex>
            <Box>
@@ -53,7 +74,7 @@ const ShareModelList = ({ models }: { models: ShareModelItem[] }) => {
              )}
            </Box>
          </Flex>
-        </Card>
+        </Box>
      ))}
    </>
  );
--- a/src/pages/model/share/index.tsx
+++ b/src/pages/model/share/index.tsx
@@ -1,24 +1,20 @@
-import React, { useState, useRef } from 'react';
+import React, { useState, useRef, useCallback, useMemo } from 'react';
 import { Box, Flex, Card, Grid, Input } from '@chakra-ui/react';
 import { useLoading } from '@/hooks/useLoading';
-import { getShareModelList } from '@/api/model';
+import { getShareModelList, triggerModelCollection, getCollectionModels } from '@/api/model';
 import { usePagination } from '@/hooks/usePagination';
 import type { ShareModelItem } from '@/types/model';

 import ShareModelList from './components/list';
+import { useQuery } from '@tanstack/react-query';

 const modelList = () => {
-  const { Loading, setIsLoading } = useLoading();
+  const { Loading } = useLoading();
  const lastSearch = useRef('');
  const [searchText, setSearchText] = useState('');

  /* 加载模型 */
-  const {
-    data: models,
-    isLoading,
-    Pagination,
-    getData
-  } = usePagination<ShareModelItem>({
+  const { data, isLoading, Pagination, getData, pageNum } = usePagination<ShareModelItem>({
    api: getShareModelList,
    pageSize: 20,
    params: {
@@ -26,47 +22,90 @@ const modelList = () => {
    }
  });

+  const { data: collectionModels = [], refetch: refetchCollection } = useQuery(
+    [getCollectionModels],
+    getCollectionModels
+  );
+
+  const models = useMemo(() => {
+    if (!collectionModels) return [];
+    return data.map((model) => ({
+      ...model,
+      isCollection: !!collectionModels.find((item) => item._id === model._id)
+    }));
+  }, [collectionModels, data]);
+
+  const onclickCollection = useCallback(
+    async (modelId: string) => {
+      try {
+        await triggerModelCollection(modelId);
+        getData(pageNum);
+        refetchCollection();
+      } catch (error) {
+        console.log(error);
+      }
+    },
+    [getData, pageNum, refetchCollection]
+  );
+
  return (
-    <Box position={'relative'}>
-      {/* 头部 */}
+    <>
      <Card px={6} py={3}>
        <Flex alignItems={'center'} justifyContent={'space-between'}>
          <Box fontWeight={'bold'} fontSize={'xl'}>
-            模型共享市场
+            我收藏的模型
          </Box>
-          <Box flex={1}>(Beta)</Box>
-          <Input
-            maxW={'240px'}
-            size={'sm'}
-            value={searchText}
-            placeholder="搜索模型，回车确认"
-            onChange={(e) => setSearchText(e.target.value)}
-            onBlur={() => {
-              if (searchText === lastSearch.current) return;
-              getData(1);
-              lastSearch.current = searchText;
-            }}
-            onKeyDown={(e) => {
-              if (searchText === lastSearch.current) return;
-              if (e.key === 'Enter') {
-                getData(1);
-                lastSearch.current = searchText;
-              }
-            }}
-          />
        </Flex>
+        {collectionModels.length == 0 && (
+          <Box textAlign={'center'} pt={3}>
+            还没有收藏模型~
+          </Box>
+        )}
+        <Grid templateColumns={['1fr', '1fr 1fr', '1fr 1fr 1fr']} gridGap={4} mt={4}>
+          <ShareModelList models={collectionModels} onclickCollection={onclickCollection} />
+        </Grid>
      </Card>

-      <Grid templateColumns={['1fr', '1fr 1fr']} gridGap={4} mt={4}>
-        <ShareModelList models={models} />
-      </Grid>
-
-      <Box mt={4}>
-        <Pagination />
-      </Box>
+      <Card mt={5} px={6} py={3}>
+        <Box display={['block', 'flex']} alignItems={'center'} justifyContent={'space-between'}>
+          <Box fontWeight={'bold'} flex={1} fontSize={'xl'}>
+            模型共享市场{' '}
+            <Box as={'span'} fontWeight={'normal'} fontSize={'md'}>
+              (Beta)
+            </Box>
+          </Box>
+          <Box mt={[2, 0]} textAlign={'right'}>
+            <Input
+              maxW={'240px'}
+              size={'sm'}
+              value={searchText}
+              placeholder="搜索模型，回车确认"
+              onChange={(e) => setSearchText(e.target.value)}
+              onBlur={() => {
+                if (searchText === lastSearch.current) return;
+                getData(1);
+                lastSearch.current = searchText;
+              }}
+              onKeyDown={(e) => {
+                if (searchText === lastSearch.current) return;
+                if (e.key === 'Enter') {
+                  getData(1);
+                  lastSearch.current = searchText;
+                }
+              }}
+            />
+          </Box>
+        </Box>
+        <Grid templateColumns={['1fr', '1fr 1fr', '1fr 1fr 1fr']} gridGap={4} mt={4}>
+          <ShareModelList models={models} onclickCollection={onclickCollection} />
+        </Grid>
+        <Box mt={4}>
+          <Pagination />
+        </Box>
+      </Card>

      <Loading loading={isLoading} />
-    </Box>
+    </>
  );
 };