mirror of
http://112.124.100.131/huang.ze/ebiz-dify-ai.git
synced 2025-12-09 19:06:51 +08:00
@@ -1,44 +1,18 @@
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from langchain.schema import Document
|
||||
|
||||
from core.index.index import IndexBuilder
|
||||
from core.rag.datasource.keyword.keyword_factory import Keyword
|
||||
from core.rag.datasource.vdb.vector_factory import Vector
|
||||
from core.rag.models.document import Document
|
||||
from models.dataset import Dataset, DocumentSegment
|
||||
|
||||
|
||||
class VectorService:
|
||||
|
||||
@classmethod
|
||||
def create_segment_vector(cls, keywords: Optional[list[str]], segment: DocumentSegment, dataset: Dataset):
|
||||
document = Document(
|
||||
page_content=segment.content,
|
||||
metadata={
|
||||
"doc_id": segment.index_node_id,
|
||||
"doc_hash": segment.index_node_hash,
|
||||
"document_id": segment.document_id,
|
||||
"dataset_id": segment.dataset_id,
|
||||
}
|
||||
)
|
||||
|
||||
# save vector index
|
||||
index = IndexBuilder.get_index(dataset, 'high_quality')
|
||||
if index:
|
||||
index.add_texts([document], duplicate_check=True)
|
||||
|
||||
# save keyword index
|
||||
index = IndexBuilder.get_index(dataset, 'economy')
|
||||
if index:
|
||||
if keywords and len(keywords) > 0:
|
||||
index.create_segment_keywords(segment.index_node_id, keywords)
|
||||
else:
|
||||
index.add_texts([document])
|
||||
|
||||
@classmethod
|
||||
def multi_create_segment_vector(cls, pre_segment_data_list: list, dataset: Dataset):
|
||||
def create_segments_vector(cls, keywords_list: Optional[list[list[str]]],
|
||||
segments: list[DocumentSegment], dataset: Dataset):
|
||||
documents = []
|
||||
for pre_segment_data in pre_segment_data_list:
|
||||
segment = pre_segment_data['segment']
|
||||
for segment in segments:
|
||||
document = Document(
|
||||
page_content=segment.content,
|
||||
metadata={
|
||||
@@ -49,30 +23,26 @@ class VectorService:
|
||||
}
|
||||
)
|
||||
documents.append(document)
|
||||
|
||||
# save vector index
|
||||
index = IndexBuilder.get_index(dataset, 'high_quality')
|
||||
if index:
|
||||
index.add_texts(documents, duplicate_check=True)
|
||||
if dataset.indexing_technique == 'high_quality':
|
||||
# save vector index
|
||||
vector = Vector(
|
||||
dataset=dataset
|
||||
)
|
||||
vector.add_texts(documents, duplicate_check=True)
|
||||
|
||||
# save keyword index
|
||||
keyword_index = IndexBuilder.get_index(dataset, 'economy')
|
||||
if keyword_index:
|
||||
keyword_index.multi_create_segment_keywords(pre_segment_data_list)
|
||||
keyword = Keyword(dataset)
|
||||
|
||||
if keywords_list and len(keywords_list) > 0:
|
||||
keyword.add_texts(documents, keyword_list=keywords_list)
|
||||
else:
|
||||
keyword.add_texts(documents)
|
||||
|
||||
@classmethod
|
||||
def update_segment_vector(cls, keywords: Optional[list[str]], segment: DocumentSegment, dataset: Dataset):
|
||||
# update segment index task
|
||||
vector_index = IndexBuilder.get_index(dataset, 'high_quality')
|
||||
kw_index = IndexBuilder.get_index(dataset, 'economy')
|
||||
# delete from vector index
|
||||
if vector_index:
|
||||
vector_index.delete_by_ids([segment.index_node_id])
|
||||
|
||||
# delete from keyword index
|
||||
kw_index.delete_by_ids([segment.index_node_id])
|
||||
|
||||
# add new index
|
||||
# format new index
|
||||
document = Document(
|
||||
page_content=segment.content,
|
||||
metadata={
|
||||
@@ -82,13 +52,20 @@ class VectorService:
|
||||
"dataset_id": segment.dataset_id,
|
||||
}
|
||||
)
|
||||
if dataset.indexing_technique == 'high_quality':
|
||||
# update vector index
|
||||
vector = Vector(
|
||||
dataset=dataset
|
||||
)
|
||||
vector.delete_by_ids([segment.index_node_id])
|
||||
vector.add_texts([document], duplicate_check=True)
|
||||
|
||||
# save vector index
|
||||
if vector_index:
|
||||
vector_index.add_texts([document], duplicate_check=True)
|
||||
# update keyword index
|
||||
keyword = Keyword(dataset)
|
||||
keyword.delete_by_ids([segment.index_node_id])
|
||||
|
||||
# save keyword index
|
||||
if keywords and len(keywords) > 0:
|
||||
kw_index.create_segment_keywords(segment.index_node_id, keywords)
|
||||
keyword.add_texts([document], keywords_list=[keywords])
|
||||
else:
|
||||
kw_index.add_texts([document])
|
||||
keyword.add_texts([document])
|
||||
|
||||
Reference in New Issue
Block a user