tjdify/api/core/indexing_runner.py

import datetime
import json
import re
import tempfile
import time
from pathlib import Path
from typing import Optional, List
from langchain.text_splitter import RecursiveCharacterTextSplitter

from llama_index import SimpleDirectoryReader
from llama_index.data_structs import Node
from llama_index.data_structs.node_v2 import DocumentRelationship
from llama_index.node_parser import SimpleNodeParser, NodeParser
from llama_index.readers.file.base import DEFAULT_FILE_EXTRACTOR
from core.docstore.dataset_docstore import DatesetDocumentStore
from core.index.keyword_table_index import KeywordTableIndex
from core.index.readers.html_parser import HTMLParser
from core.index.readers.markdown_parser import MarkdownParser
from core.index.readers.pdf_parser import PDFParser
from core.index.spiltter.fixed_text_splitter import FixedRecursiveCharacterTextSplitter
from core.index.vector_index import VectorIndex
from core.llm.token_calculator import TokenCalculator
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from extensions.ext_storage import storage
from models.dataset import Document, Dataset, DocumentSegment, DatasetProcessRule
from models.model import UploadFile


class IndexingRunner:

    def __init__(self, embedding_model_name: str = "text-embedding-ada-002"):
        self.storage = storage
        self.embedding_model_name = embedding_model_name

    def run(self, document: Document):
        """Run the indexing process."""
        # get dataset
        dataset = Dataset.query.filter_by(
            id=document.dataset_id
        ).first()

        if not dataset:
            raise ValueError("no dataset found")

        # load file
        text_docs = self._load_data(document)

        # get the process rule
        processing_rule = db.session.query(DatasetProcessRule). \
            filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
            first()

        # get node parser for splitting
        node_parser = self._get_node_parser(processing_rule)

        # split to nodes
        nodes = self._step_split(
            text_docs=text_docs,
            node_parser=node_parser,
            dataset=dataset,
            document=document,
            processing_rule=processing_rule
        )

        # build index
        self._build_index(
            dataset=dataset,
            document=document,
            nodes=nodes
        )

    def run_in_splitting_status(self, document: Document):
        """Run the indexing process when the index_status is splitting."""
        # get dataset
        dataset = Dataset.query.filter_by(
            id=document.dataset_id
        ).first()

        if not dataset:
            raise ValueError("no dataset found")

        # get exist document_segment list and delete
        document_segments = DocumentSegment.query.filter_by(
            dataset_id=dataset.id,
            document_id=document.id
        ).all()
        db.session.delete(document_segments)
        db.session.commit()
        # load file
        text_docs = self._load_data(document)

        # get the process rule
        processing_rule = db.session.query(DatasetProcessRule). \
            filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
            first()

        # get node parser for splitting
        node_parser = self._get_node_parser(processing_rule)

        # split to nodes
        nodes = self._step_split(
            text_docs=text_docs,
            node_parser=node_parser,
            dataset=dataset,
            document=document,
            processing_rule=processing_rule
        )

        # build index
        self._build_index(
            dataset=dataset,
            document=document,
            nodes=nodes
        )

    def run_in_indexing_status(self, document: Document):
        """Run the indexing process when the index_status is indexing."""
        # get dataset
        dataset = Dataset.query.filter_by(
            id=document.dataset_id
        ).first()

        if not dataset:
            raise ValueError("no dataset found")

        # get exist document_segment list and delete
        document_segments = DocumentSegment.query.filter_by(
            dataset_id=dataset.id,
            document_id=document.id
        ).all()
        nodes = []
        if document_segments:
            for document_segment in document_segments:
                # transform segment to node
                if document_segment.status != "completed":
                    relationships = {
                        DocumentRelationship.SOURCE: document_segment.document_id,
                    }

                    previous_segment = document_segment.previous_segment
                    if previous_segment:
                        relationships[DocumentRelationship.PREVIOUS] = previous_segment.index_node_id

                    next_segment = document_segment.next_segment
                    if next_segment:
                        relationships[DocumentRelationship.NEXT] = next_segment.index_node_id
                    node = Node(
                        doc_id=document_segment.index_node_id,
                        doc_hash=document_segment.index_node_hash,
                        text=document_segment.content,
                        extra_info=None,
                        node_info=None,
                        relationships=relationships
                    )
                    nodes.append(node)

        # build index
        self._build_index(
            dataset=dataset,
            document=document,
            nodes=nodes
        )

    def indexing_estimate(self, file_detail: UploadFile, tmp_processing_rule: dict) -> dict:
        """
        Estimate the indexing for the document.
        """
        # load data from file
        text_docs = self._load_data_from_file(file_detail)

        processing_rule = DatasetProcessRule(
            mode=tmp_processing_rule["mode"],
            rules=json.dumps(tmp_processing_rule["rules"])
        )

        # get node parser for splitting
        node_parser = self._get_node_parser(processing_rule)

        # split to nodes
        nodes = self._split_to_nodes(
            text_docs=text_docs,
            node_parser=node_parser,
            processing_rule=processing_rule
        )

        tokens = 0
        preview_texts = []
        for node in nodes:
            if len(preview_texts) < 5:
                preview_texts.append(node.get_text())

            tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())

        return {
            "total_segments": len(nodes),
            "tokens": tokens,
            "total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
            "currency": TokenCalculator.get_currency(self.embedding_model_name),
            "preview": preview_texts
        }

    def _load_data(self, document: Document) -> List[Document]:
        # load file
        if document.data_source_type != "upload_file":
            return []

        data_source_info = document.data_source_info_dict
        if not data_source_info or 'upload_file_id' not in data_source_info:
            raise ValueError("no upload file found")

        file_detail = db.session.query(UploadFile). \
            filter(UploadFile.id == data_source_info['upload_file_id']). \
            one_or_none()

        text_docs = self._load_data_from_file(file_detail)

        # update document status to splitting
        self._update_document_index_status(
            document_id=document.id,
            after_indexing_status="splitting",
            extra_update_params={
                Document.file_id: file_detail.id,
                Document.word_count: sum([len(text_doc.text) for text_doc in text_docs]),
                Document.parsing_completed_at: datetime.datetime.utcnow()
            }
        )

        # replace doc id to document model id
        for text_doc in text_docs:
            # remove invalid symbol
            text_doc.text = self.filter_string(text_doc.get_text())
            text_doc.doc_id = document.id

        return text_docs

    def filter_string(self, text):
        pattern = re.compile('[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\x80-\xFF]')
        return pattern.sub('', text)

    def _load_data_from_file(self, upload_file: UploadFile) -> List[Document]:
        with tempfile.TemporaryDirectory() as temp_dir:
            suffix = Path(upload_file.key).suffix
            filepath = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
            self.storage.download(upload_file.key, filepath)

            file_extractor = DEFAULT_FILE_EXTRACTOR.copy()
            file_extractor[".markdown"] = MarkdownParser()
            file_extractor[".md"] = MarkdownParser()
            file_extractor[".html"] = HTMLParser()
            file_extractor[".htm"] = HTMLParser()
            file_extractor[".pdf"] = PDFParser({'upload_file': upload_file})

            loader = SimpleDirectoryReader(input_files=[filepath], file_extractor=file_extractor)
            text_docs = loader.load_data()

            return text_docs

    def _get_node_parser(self, processing_rule: DatasetProcessRule) -> NodeParser:
        """
        Get the NodeParser object according to the processing rule.
        """
        if processing_rule.mode == "custom":
            # The user-defined segmentation rule
            rules = json.loads(processing_rule.rules)
            segmentation = rules["segmentation"]
            if segmentation["max_tokens"] < 50 or segmentation["max_tokens"] > 1000:
                raise ValueError("Custom segment length should be between 50 and 1000.")

            separator = segmentation["separator"]
            if separator:
                separator = separator.replace('\\n', '\n')

            character_splitter = FixedRecursiveCharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=segmentation["max_tokens"],
                chunk_overlap=0,
                fixed_separator=separator,
                separators=["\n\n", "。", ".", " ", ""]
            )
        else:
            # Automatic segmentation
            character_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=DatasetProcessRule.AUTOMATIC_RULES['segmentation']['max_tokens'],
                chunk_overlap=0,
                separators=["\n\n", "。", ".", " ", ""]
            )

        return SimpleNodeParser(text_splitter=character_splitter, include_extra_info=True)

    def _step_split(self, text_docs: List[Document], node_parser: NodeParser,
                    dataset: Dataset, document: Document, processing_rule: DatasetProcessRule) -> List[Node]:
        """
        Split the text documents into nodes and save them to the document segment.
        """
        nodes = self._split_to_nodes(
            text_docs=text_docs,
            node_parser=node_parser,
            processing_rule=processing_rule
        )

        # save node to document segment
        doc_store = DatesetDocumentStore(
            dataset=dataset,
            user_id=document.created_by,
            embedding_model_name=self.embedding_model_name,
            document_id=document.id
        )

        doc_store.add_documents(nodes)

        # update document status to indexing
        cur_time = datetime.datetime.utcnow()
        self._update_document_index_status(
            document_id=document.id,
            after_indexing_status="indexing",
            extra_update_params={
                Document.cleaning_completed_at: cur_time,
                Document.splitting_completed_at: cur_time,
            }
        )

        # update segment status to indexing
        self._update_segments_by_document(
            document_id=document.id,
            update_params={
                DocumentSegment.status: "indexing",
                DocumentSegment.indexing_at: datetime.datetime.utcnow()
            }
        )

        return nodes

    def _split_to_nodes(self, text_docs: List[Document], node_parser: NodeParser,
                        processing_rule: DatasetProcessRule) -> List[Node]:
        """
        Split the text documents into nodes.
        """
        all_nodes = []
        for text_doc in text_docs:
            # document clean
            document_text = self._document_clean(text_doc.get_text(), processing_rule)
            text_doc.text = document_text

            # parse document to nodes
            nodes = node_parser.get_nodes_from_documents([text_doc])
            nodes = [node for node in nodes if node.text is not None and node.text.strip()]
            all_nodes.extend(nodes)

        return all_nodes

    def _document_clean(self, text: str, processing_rule: DatasetProcessRule) -> str:
        """
        Clean the document text according to the processing rules.
        """
        if processing_rule.mode == "automatic":
            rules = DatasetProcessRule.AUTOMATIC_RULES
        else:
            rules = json.loads(processing_rule.rules) if processing_rule.rules else {}

        if 'pre_processing_rules' in rules:
            pre_processing_rules = rules["pre_processing_rules"]
            for pre_processing_rule in pre_processing_rules:
                if pre_processing_rule["id"] == "remove_extra_spaces" and pre_processing_rule["enabled"] is True:
                    # Remove extra spaces
                    pattern = r'\n{3,}'
                    text = re.sub(pattern, '\n\n', text)
                    pattern = r'[\t\f\r\x20\u00a0\u1680\u180e\u2000-\u200a\u202f\u205f\u3000]{2,}'
                    text = re.sub(pattern, ' ', text)
                elif pre_processing_rule["id"] == "remove_urls_emails" and pre_processing_rule["enabled"] is True:
                    # Remove email
                    pattern = r'([a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)'
                    text = re.sub(pattern, '', text)

                    # Remove URL
                    pattern = r'https?://[^\s]+'
                    text = re.sub(pattern, '', text)

        return text

    def _build_index(self, dataset: Dataset, document: Document, nodes: List[Node]) -> None:
        """
        Build the index for the document.
        """
        vector_index = VectorIndex(dataset=dataset)
        keyword_table_index = KeywordTableIndex(dataset=dataset)

        # chunk nodes by chunk size
        indexing_start_at = time.perf_counter()
        tokens = 0
        chunk_size = 100
        for i in range(0, len(nodes), chunk_size):
            # check document is paused
            self._check_document_paused_status(document.id)
            chunk_nodes = nodes[i:i + chunk_size]

            tokens += sum(
                TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text()) for node in chunk_nodes
            )

            # save vector index
            if dataset.indexing_technique == "high_quality":
                vector_index.add_nodes(chunk_nodes)

            # save keyword index
            keyword_table_index.add_nodes(chunk_nodes)

            node_ids = [node.doc_id for node in chunk_nodes]
            db.session.query(DocumentSegment).filter(
                DocumentSegment.document_id == document.id,
                DocumentSegment.index_node_id.in_(node_ids),
                DocumentSegment.status == "indexing"
            ).update({
                DocumentSegment.status: "completed",
                DocumentSegment.completed_at: datetime.datetime.utcnow()
            })

            db.session.commit()

        indexing_end_at = time.perf_counter()

        # update document status to completed
        self._update_document_index_status(
            document_id=document.id,
            after_indexing_status="completed",
            extra_update_params={
                Document.tokens: tokens,
                Document.completed_at: datetime.datetime.utcnow(),
                Document.indexing_latency: indexing_end_at - indexing_start_at,
            }
        )

    def _check_document_paused_status(self, document_id: str):
        indexing_cache_key = 'document_{}_is_paused'.format(document_id)
        result = redis_client.get(indexing_cache_key)
        if result:
            raise DocumentIsPausedException()

    def _update_document_index_status(self, document_id: str, after_indexing_status: str,
                                      extra_update_params: Optional[dict] = None) -> None:
        """
        Update the document indexing status.
        """
        count = Document.query.filter_by(id=document_id, is_paused=True).count()
        if count > 0:
            raise DocumentIsPausedException()

        update_params = {
            Document.indexing_status: after_indexing_status
        }

        if extra_update_params:
            update_params.update(extra_update_params)

        Document.query.filter_by(id=document_id).update(update_params)
        db.session.commit()

    def _update_segments_by_document(self, document_id: str, update_params: dict) -> None:
        """
        Update the document segment by document id.
        """
        DocumentSegment.query.filter_by(document_id=document_id).update(update_params)
        db.session.commit()


class DocumentIsPausedException(Exception):
    pass