mutil thread

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,