feat: 文档处理调用大模型提取关键词

api/core/indexing_runner.py

@@ -44,7 +44,7 @@ class IndexingRunner:
         self.storage = storage
         self.model_manager = ModelManager()
 
-    def run(self, dataset_documents: list[DatasetDocument]):
+    def run(self, tenant_id, dataset_documents: list[DatasetDocument]):
         """Run the indexing process."""
         for dataset_document in dataset_documents:
             try:
@@ -74,6 +74,7 @@ class IndexingRunner:
 
                 # load
                 self._load(
+                    tenant_id=tenant_id,
                     index_processor=index_processor,
                     dataset=dataset,
                     dataset_document=dataset_document,
@@ -611,6 +612,7 @@ class IndexingRunner:
 
     def _load(
         self,
+        tenant_id :str,
         index_processor: BaseIndexProcessor,
         dataset: Dataset,
         dataset_document: DatasetDocument,
@@ -637,7 +639,7 @@ class IndexingRunner:
         # create keyword index
         create_keyword_thread = threading.Thread(
             target=self._process_keyword_index,
-            args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents),
+            args=(current_app._get_current_object(), tenant_id, dataset.id, dataset_document.id, documents),
         )
         create_keyword_thread.start()
         if dataset.indexing_technique == "high_quality":
@@ -676,13 +678,13 @@ class IndexingRunner:
         )
 
     @staticmethod
-    def _process_keyword_index(flask_app, dataset_id, document_id, documents):
+    def _process_keyword_index(flask_app, tenant_id, dataset_id, document_id, documents):
         with flask_app.app_context():
             dataset = Dataset.query.filter_by(id=dataset_id).first()
             if not dataset:
                 raise ValueError("no dataset found")
             keyword = Keyword(dataset)
-            keyword.create(documents)
+            keyword.create(tenant_id, documents)
             if dataset.indexing_technique != "high_quality":
                 document_ids = [document.metadata["doc_id"] for document in documents]
                 db.session.query(DocumentSegment).filter(

api/core/llm_generator/llm_generator.py

@@ -11,6 +11,7 @@ from core.llm_generator.prompts import (
     JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE,
     PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE,
     WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE,
+    EXTRACTE_KEYWORDS_PROMPT,
 )
 from core.model_manager import ModelManager
 from core.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
@@ -309,3 +310,25 @@ class LLMGenerator:
 
         answer = response.message.content
         return answer.strip()
+
+    @classmethod
+    def extract_keyword(cls, tenant_id: str, query, document_language: str="Chinese"):
+        prompt = EXTRACTE_KEYWORDS_PROMPT.format(language=document_language)
+
+        model_manager = ModelManager()
+        model_instance = model_manager.get_default_model_instance(
+            tenant_id=tenant_id,
+            model_type=ModelType.LLM,
+        )
+
+        prompt_messages = [SystemPromptMessage(content=prompt), UserPromptMessage(content=query)]
+
+        response = model_instance.invoke_llm(
+            prompt_messages=prompt_messages, model_parameters={"temperature": 0.01, "max_tokens": 2000}, stream=False
+        )
+
+        answer = response.message.content
+        logging.info(f"LLM answer: {answer}")
+        # answer = ["keyword1", "keyword2", "keyword3"]
+        answer = answer.split(":")[-1].strip().split("\n")
+        return [word.split("-")[-1].strip() for word in answer if word.split("-")[-1].strip()!="" and "关键词" not in word and "Keyword" not in word]

api/core/llm_generator/prompts.py

@@ -150,6 +150,19 @@ GENERATOR_QA_PROMPT = (
     "<QA Pairs>"
 )
 
+EXTRACTE_KEYWORDS_PROMPT = (
+    "<Task> The user will send a long text. Extract keywords from the long text. Please think step by step."
+    "Step 1: Understand and summarize the main content of this text.\n"
+    "Step 2: What key information or concepts are mentioned in this text?\n"
+    "Step 3: Decompose or combine multiple pieces of information and concepts.\n"
+    "Step 4: Extract keywords based on these key information and concepts, like the time, topic, location, event.\n"
+    "<Constraints> The keywords should cover the main points of the text, similar or redundant keywords should be consolidated. Avoid irrelevant or noisy terms."
+    "You must answer in {language}, in a style that is clear and detailed in {language}."
+    " No language other than {language} should be used. \n"
+    "<Format> Use the following format: keyword1\nkeyword2\nkeyword3\n...\n"
+    "<Keywords>"
+)
+
 WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE = """
 Here is a task description for which I would like you to create a high-quality prompt template for:
 <task_description>

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -1,8 +1,11 @@
 import json
 from collections import defaultdict
 from typing import Any, Optional
+import threading
+import logging
 
 from pydantic import BaseModel
+from flask import Flask, current_app
 
 from configs import dify_config
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler
@@ -23,21 +26,51 @@ class Jieba(BaseKeyword):
         super().__init__(dataset)
         self._config = KeywordTableConfig()
 
-    def create(self, texts: list[Document], **kwargs) -> BaseKeyword:
+    def create(self, tenant_id, texts: list[Document], **kwargs) -> BaseKeyword:
         lock_name = "keyword_indexing_lock_{}".format(self.dataset.id)
         with redis_client.lock(lock_name, timeout=600):
+            
             keyword_table_handler = JiebaKeywordTableHandler()
             keyword_table = self._get_dataset_keyword_table()
-            for text in texts:
-                keywords = keyword_table_handler.extract_keywords(
-                    text.page_content, self._config.max_keywords_per_chunk
-                )
-                self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords))
-                keyword_table = self._add_text_to_keyword_table(keyword_table, text.metadata["doc_id"], list(keywords))
+            
+            # for text in texts:
+            #     keywords = keyword_table_handler.extract_keywords(
+            #         tenant_id, text.page_content, self._config.max_keywords_per_chunk
+            #     )
+            #     self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords))
+            #     keyword_table = self._add_text_to_keyword_table(keyword_table, text.metadata["doc_id"], list(keywords))
+
+            for i in range(0, len(texts), 10):
+                threads = []
+                sub_texts = texts[i : i + 10]
+                for text in sub_texts:
+                    extract_keywords_thread = threading.Thread(
+                        target=self.extract_keywords,
+                        kwargs={
+                            "flask_app": current_app._get_current_object(),
+                            "tenant_id": tenant_id,
+                            "keyword_table_handler": keyword_table_handler,
+                            "keyword_table": keyword_table,
+                            "text": text,
+                        },
+                    )
+                    threads.append(extract_keywords_thread)
+                    extract_keywords_thread.start()
+                for thread in threads:
+                    thread.join()
 
+            logging.info(f"keyword_table: {keyword_table}")
             self._save_dataset_keyword_table(keyword_table)
 
             return self
+    
+    def extract_keywords(self, flask_app: Flask, tenant_id, keyword_table_handler, keyword_table, text: list[Document]):
+        with flask_app.app_context():
+            keywords = keyword_table_handler.extract_keywords(
+                tenant_id, text.page_content, self._config.max_keywords_per_chunk
+            )
+            self._update_segment_keywords(self.dataset.id, text.metadata["doc_id"], list(keywords))
+            keyword_table = self._add_text_to_keyword_table(keyword_table, text.metadata["doc_id"], list(keywords))
 
     def add_texts(self, texts: list[Document], **kwargs):
         lock_name = "keyword_indexing_lock_{}".format(self.dataset.id)

api/core/rag/datasource/keyword/jieba/jieba_keyword_table_handler.py

@@ -3,6 +3,8 @@ from typing import Optional
 
 import jieba
 from jieba.analyse import default_tfidf
+import logging
+from core.llm_generator.llm_generator import LLMGenerator
 
 from core.rag.datasource.keyword.jieba.stopwords import STOPWORDS
 
@@ -11,14 +13,18 @@ class JiebaKeywordTableHandler:
     def __init__(self):
         default_tfidf.stop_words = STOPWORDS
 
-    def extract_keywords(self, text: str, max_keywords_per_chunk: Optional[int] = 10) -> set[str]:
+    def extract_keywords(self, tenant_id, text: str, max_keywords_per_chunk: Optional[int] = 10) -> set[str]:
         """Extract keywords with JIEBA tfidf."""
-        keywords = jieba.analyse.extract_tags(
-            sentence=text,
-            topK=max_keywords_per_chunk,
-        )
-
-        return set(self._expand_tokens_with_subtokens(keywords))
+        # keywords = jieba.analyse.extract_tags(
+        #     sentence=text,
+        #     topK=max_keywords_per_chunk,
+        # )
+        keywords = LLMGenerator.extract_keyword(tenant_id, text)
+        logging.log(logging.INFO, f"tenant_id: {tenant_id}")
+        logging.log(logging.INFO, f"Extracted keywords: {keywords}")
+
+        # return set(self._expand_tokens_with_subtokens(keywords))
+        return set(keywords)
 
     def _expand_tokens_with_subtokens(self, tokens: set[str]) -> set[str]:
         """Get subtokens from a list of tokens., filtering for stopwords."""

api/core/rag/datasource/keyword/keyword_factory.py

@@ -27,8 +27,8 @@ class Keyword:
             case _:
                 raise ValueError(f"Keyword store {keyword_type} is not supported.")
 
-    def create(self, texts: list[Document], **kwargs):
-        self._keyword_processor.create(texts, **kwargs)
+    def create(self, tenant_id, texts: list[Document], **kwargs):
+        self._keyword_processor.create(tenant_id, texts, **kwargs)
 
     def add_texts(self, texts: list[Document], **kwargs):
         self._keyword_processor.add_texts(texts, **kwargs)

api/services/dataset_service.py

@@ -672,6 +672,7 @@ class DocumentService:
 
     @staticmethod
     def save_document_with_dataset_id(
+        tenant_id: str, 
         dataset: Dataset,
         document_data: dict,
         account: Account | Any,
@@ -917,7 +918,7 @@ class DocumentService:
 
                 # trigger async task
                 if document_ids:
-                    document_indexing_task.delay(dataset.id, document_ids)
+                    document_indexing_task.delay(tenant_id, dataset.id, document_ids)
                 if duplicate_document_ids:
                     duplicate_document_indexing_task.delay(dataset.id, duplicate_document_ids)
 
@@ -1141,7 +1142,7 @@ class DocumentService:
         db.session.add(dataset)
         db.session.flush()
 
-        documents, batch = DocumentService.save_document_with_dataset_id(dataset, document_data, account)
+        documents, batch = DocumentService.save_document_with_dataset_id(tenant_id, dataset, document_data, account)
 
         cut_length = 18
         cut_name = documents[0].name[:cut_length]

api/tasks/document_indexing_task.py

@@ -13,7 +13,7 @@ from services.feature_service import FeatureService
 
 
 @shared_task(queue="dataset")
-def document_indexing_task(dataset_id: str, document_ids: list):
+def document_indexing_task(tenant_id: str, dataset_id: str, document_ids: list):
     """
     Async process document
     :param dataset_id:
@@ -71,7 +71,7 @@ def document_indexing_task(dataset_id: str, document_ids: list):
 
     try:
         indexing_runner = IndexingRunner()
-        indexing_runner.run(documents)
+        indexing_runner.run(tenant_id, documents)
         end_at = time.perf_counter()
         logging.info(click.style("Processed dataset: {} latency: {}".format(dataset_id, end_at - start_at), fg="green"))
     except DocumentIsPausedError as ex:

docker/docker-compose.yaml

@@ -283,7 +283,7 @@ x-shared-env: &shared-api-worker-env
 services:
   # API service
   api:
-    image: langgenius/dify-api:0.11.1
+    image: seec-dify-api:0.1.1
     restart: always
     environment:
       # Use the shared environment variables.
@@ -303,7 +303,7 @@ services:
   # worker service
   # The Celery worker for processing the queue.
   worker:
-    image: langgenius/dify-api:0.11.1
+    image: seec-dify-api:0.1.1
     restart: always
     environment:
       # Use the shared environment variables.