Underground-Digital/Workflow-Engine

1from typing import Optional2 3from core.rag.datasource.keyword.keyword_factory import Keyword4from core.rag.datasource.vdb.vector_factory import Vector5from core.rag.models.document import Document6from models.dataset import Dataset, DocumentSegment7 8 9class VectorService:10    @classmethod11    def create_segments_vector(12        cls, keywords_list: Optional[list[list[str]]], segments: list[DocumentSegment], dataset: Dataset13    ):14        documents = []15        for segment in segments:16            document = Document(17                page_content=segment.content,18                metadata={19                    "doc_id": segment.index_node_id,20                    "doc_hash": segment.index_node_hash,21                    "document_id": segment.document_id,22                    "dataset_id": segment.dataset_id,23                },24            )25            documents.append(document)26        if dataset.indexing_technique == "high_quality":27            # save vector index28            vector = Vector(dataset=dataset)29            vector.add_texts(documents, duplicate_check=True)30 31        # save keyword index32        keyword = Keyword(dataset)33 34        if keywords_list and len(keywords_list) > 0:35            keyword.add_texts(documents, keywords_list=keywords_list)36        else:37            keyword.add_texts(documents)38 39    @classmethod40    def update_segment_vector(cls, keywords: Optional[list[str]], segment: DocumentSegment, dataset: Dataset):41        # update segment index task42 43        # format new index44        document = Document(45            page_content=segment.content,46            metadata={47                "doc_id": segment.index_node_id,48                "doc_hash": segment.index_node_hash,49                "document_id": segment.document_id,50                "dataset_id": segment.dataset_id,51            },52        )53        if dataset.indexing_technique == "high_quality":54            # update vector index55            vector = Vector(dataset=dataset)56            vector.delete_by_ids([segment.index_node_id])57            vector.add_texts([document], duplicate_check=True)58 59        # update keyword index60        keyword = Keyword(dataset)61        keyword.delete_by_ids([segment.index_node_id])62 63        # save keyword index64        if keywords and len(keywords) > 0:65            keyword.add_texts([document], keywords_list=[keywords])66        else:67            keyword.add_texts([document])68