Merge branch 'feat/milvus-support' into deploy/dev

api/controllers/console/datasets/datasets_segments.py

@@ -16,6 +16,7 @@ from models.dataset import DocumentSegment
 
 from libs.helper import TimestampField
 from services.dataset_service import DatasetService, DocumentService, SegmentService
+from tasks.test_task import test_task
 from tasks.enable_segment_to_index_task import enable_segment_to_index_task
 from tasks.remove_segment_from_index_task import remove_segment_from_index_task
 
@@ -284,6 +285,15 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         }, 200
 
 
+class DatasetDocumentTest(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(self):
+        test_task.delay()
+        return 200
+
+
 api.add_resource(DatasetDocumentSegmentListApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments')
 api.add_resource(DatasetDocumentSegmentApi,
@@ -292,3 +302,5 @@ api.add_resource(DatasetDocumentSegmentAddApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segment')
 api.add_resource(DatasetDocumentSegmentUpdateApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>')
+api.add_resource(DatasetDocumentTest,
+                 '/datasets/test')

api/core/generator/llm_generator.py

@@ -192,6 +192,18 @@ class LLMGenerator:
         prompt = GENERATOR_QA_PROMPT
 
 
+        if isinstance(llm, BaseChatModel):
+            prompt = [SystemMessage(content=prompt), HumanMessage(content=query)]
+
+        response = llm.generate([prompt])
+        answer = response.generations[0][0].text
+        return answer.strip()
+
+    @classmethod
+    def generate_qa_document_sync(cls, llm: StreamableOpenAI, query):
+        prompt = GENERATOR_QA_PROMPT
+
+
         if isinstance(llm, BaseChatModel):
             prompt = [SystemMessage(content=prompt), HumanMessage(content=query)]
 

api/core/indexing_runner.py

@@ -274,10 +274,10 @@ class IndexingRunner:
                 # qa model document
                 llm: StreamableOpenAI = LLMBuilder.to_llm(
                     tenant_id=current_user.current_tenant_id,
-                    model_name='claude-2',
-                    max_tokens=5000
+                    model_name='gpt-3.5-turbo',
+                    max_tokens=2000
                 )
-                response = LLMGenerator.generate_qa_document(llm, preview_texts[0])
+                response = LLMGenerator.generate_qa_document_sync(llm, preview_texts[0])
                 document_qa_list = self.format_split_text(response)
                 return {
                     "total_segments": total_segments * 20,
@@ -351,10 +351,10 @@ class IndexingRunner:
                 # qa model document
                 llm: StreamableOpenAI = LLMBuilder.to_llm(
                     tenant_id=current_user.current_tenant_id,
-                    model_name='claude-2',
-                    max_tokens=5000
+                    model_name='gpt-3.5-turbo',
+                    max_tokens=2000
                 )
-                response = LLMGenerator.generate_qa_document(llm, preview_texts[0])
+                response = LLMGenerator.generate_qa_document_sync(llm, preview_texts[0])
                 document_qa_list = self.format_split_text(response)
                 return {
                     "total_segments": total_segments * 20,

web/app/components/datasets/create/step-two/index.tsx

@@ -165,10 +165,11 @@ const StepTwo = ({
       setAutomaticFileIndexingEstimate(res)
   }
 
-  const confirmChangeCustomConfig = async () => {
+  const confirmChangeCustomConfig = () => {
     setCustomFileIndexingEstimate(null)
     setShowPreview()
-    await fetchFileIndexingEstimate()
+    fetchFileIndexingEstimate()
+    setPreviewSwitched(false)
   }
 
   const getIndexing_technique = () => indexingType || indexType
@@ -346,6 +347,13 @@ const StepTwo = ({
       setDocForm(DocForm.TEXT)
   }
 
+  const changeToEconomicalType = () => {
+    if (!hasSetIndexType) {
+      setIndexType(IndexingType.ECONOMICAL)
+      setDocForm(DocForm.TEXT)
+    }
+  }
+
   const previewSwitch = async () => {
     setPreviewSwitched(true)
     if (segmentationType === SegmentType.AUTO)
@@ -545,7 +553,7 @@ const StepTwo = ({
                     hasSetIndexType && s.disabled,
                     hasSetIndexType && '!w-full',
                   )}
-                  onClick={() => !hasSetIndexType && setIndexType(IndexingType.ECONOMICAL)}
+                  onClick={changeToEconomicalType}
                 >
                   <span className={cn(s.typeIcon, s.economical)} />
                   {!hasSetIndexType && <span className={cn(s.radio)} />}
@@ -564,7 +572,7 @@ const StepTwo = ({
                 <Link className='text-[#155EEF]' href={`/datasets/${datasetId}/settings`}>{t('datasetCreation.stepTwo.datasetSettingLink')}</Link>
               </div>
             )}
-            {(!hasSetIndexType || (hasSetIndexType && indexingType === IndexingType.QUALIFIED)) && (
+            {indexType === IndexingType.QUALIFIED && (
               <div className='flex justify-between items-center mt-3 px-5 py-4 rounded-xl bg-gray-50 border border-gray-100'>
                 <div className='flex justify-center items-center w-8 h-8 rounded-lg bg-indigo-50'>
                   <MessageChatSquare className='w-4 h-4' />
@@ -691,7 +699,12 @@ const StepTwo = ({
                   ))}
                 </>
               )}
-              {!fileIndexingEstimate?.preview && !fileIndexingEstimate?.qa_preview && (
+              {previewSwitched && docForm === DocForm.QA && !fileIndexingEstimate?.qa_preview && (
+                <div className='flex items-center justify-center h-[200px]'>
+                  <Loading type='area' />
+                </div>
+              )}
+              {!previewSwitched && !fileIndexingEstimate?.preview && (
                 <div className='flex items-center justify-center h-[200px]'>
                   <Loading type='area' />
                 </div>