first-update

easy-dataset-main/hooks/useGenerateDataset.js

@@ -0,0 +1,135 @@
+import { useCallback } from 'react';
+import { toast } from 'sonner';
+import i18n from '@/lib/i18n';
+import axios from 'axios';
+import { useAtomValue } from 'jotai/index';
+import { selectedModelInfoAtom } from '@/lib/store';
+import { useTranslation } from 'react-i18next';
+
+export function useGenerateDataset() {
+  const model = useAtomValue(selectedModelInfoAtom);
+  const { t } = useTranslation();
+
+  const generateSingleDataset = useCallback(
+    async ({ projectId, questionId, questionInfo, imageId, imageName }) => {
+      // 获取模型参数
+      if (!model) {
+        toast.error(t('models.configNotFound'));
+        return null;
+      }
+
+      // 判断是否为图片问题
+      const isImageQuestion = !!imageId;
+
+      // 调用API生成数据集
+      const currentLanguage = i18n.language === 'zh-CN' ? '中文' : 'en';
+
+      if (isImageQuestion) {
+        // 图片问题：调用图片数据集生成接口
+        toast.promise(
+          axios.post(`/api/projects/${projectId}/images/datasets`, {
+            imageName,
+            question: { question: questionInfo, id: questionId },
+            model,
+            language: currentLanguage
+          }),
+          {
+            loading: t('datasets.generating'),
+            description: `图片：【${imageName}】\n问题：【${questionInfo}】`,
+            position: 'top-right',
+            success: data => {
+              return '生成数据集成功';
+            },
+            error: error => {
+              return t('datasets.generateFailed', { error: error.response?.data?.error });
+            }
+          }
+        );
+      } else {
+        // 文本问题：调用普通数据集生成接口
+        toast.promise(
+          axios.post(`/api/projects/${projectId}/datasets`, {
+            questionId,
+            model,
+            language: currentLanguage
+          }),
+          {
+            loading: t('datasets.generating'),
+            description: `问题：【${questionInfo}】`,
+            position: 'top-right',
+            success: data => {
+              return '生成数据集成功';
+            },
+            error: error => {
+              return t('datasets.generateFailed', { error: error.response?.data?.error });
+            }
+          }
+        );
+      }
+    },
+    [model, t]
+  );
+
+  const generateMultipleDataset = useCallback(
+    async (projectId, questions) => {
+      let completed = 0;
+      const total = questions.length;
+      // 显示带进度的Loading
+      const loadingToastId = toast.loading(`正在处理请求 (${completed}/${total})...`, { position: 'top-right' });
+
+      // 处理每个请求
+      const processRequest = async question => {
+        try {
+          const isImageQuestion = !!question.imageId;
+          let response;
+
+          if (isImageQuestion) {
+            // 图片问题
+            response = await axios.post(`/api/projects/${projectId}/images/datasets`, {
+              imageName: question.imageName,
+              question,
+              model,
+              language: i18n.language === 'zh-CN' ? '中文' : 'en'
+            });
+          } else {
+            // 文本问题
+            response = await axios.post(`/api/projects/${projectId}/datasets`, {
+              questionId: question.id,
+              model,
+              language: i18n.language === 'zh-CN' ? '中文' : 'en'
+            });
+          }
+
+          const data = response.data;
+          completed++;
+          toast.success(`${question.question} 完成`, { position: 'top-right' });
+          toast.loading(`正在处理请求 (${completed}/${total})...`, { id: loadingToastId });
+          return data;
+        } catch (error) {
+          completed++;
+          toast.error(`${question.question} 失败`, {
+            description: error.message,
+            position: 'top-right'
+          });
+          toast.loading(`正在处理请求 (${completed}/${total})...`, { id: loadingToastId });
+          throw error;
+        }
+      };
+
+      try {
+        const results = await Promise.allSettled(questions.map(req => processRequest(req)));
+        // 全部完成后更新Loading为完成状态
+        toast.success(`全部请求处理完成 (成功: ${results.filter(r => r.status === 'fulfilled').length}/${total})`, {
+          id: loadingToastId,
+          position: 'top-right'
+        });
+        return results;
+      } catch {
+        // Promise.allSettled不会进入catch，这里只是保险
+      }
+    },
+    [model, t]
+  );
+
+  return { generateSingleDataset, generateMultipleDataset };
+}