mirror of
https://github.com/langgenius/dify.git
synced 2026-05-06 02:18:08 +08:00
fix: fix summary index bug.
This commit is contained in:
FFXN
@ -2,10 +2,11 @@
|
|||||||
|
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
|
from collections import defaultdict
|
||||||
|
|
||||||
import click
|
import click
|
||||||
from celery import shared_task
|
from celery import shared_task
|
||||||
from sqlalchemy import select
|
from sqlalchemy import and_, or_, select
|
||||||
|
|
||||||
from extensions.ext_database import db
|
from extensions.ext_database import db
|
||||||
from models.dataset import Dataset, DocumentSegment, DocumentSegmentSummary
|
from models.dataset import Dataset, DocumentSegment, DocumentSegmentSummary
|
||||||
@ -62,158 +63,254 @@ def regenerate_summary_index_task(
|
|||||||
db.session.close()
|
db.session.close()
|
||||||
return
|
return
|
||||||
|
|
||||||
# Check if summary index is enabled
|
# Check if summary index is enabled (only for summary_model change)
|
||||||
|
# For embedding_model change, we still re-vectorize existing summaries even if setting is disabled
|
||||||
summary_index_setting = dataset.summary_index_setting
|
summary_index_setting = dataset.summary_index_setting
|
||||||
if not summary_index_setting or not summary_index_setting.get("enable"):
|
if not regenerate_vectors_only:
|
||||||
logger.info(
|
# For summary_model change, require summary_index_setting to be enabled
|
||||||
click.style(
|
if not summary_index_setting or not summary_index_setting.get("enable"):
|
||||||
f"Summary index is disabled for dataset {dataset_id}",
|
logger.info(
|
||||||
fg="cyan",
|
click.style(
|
||||||
|
f"Summary index is disabled for dataset {dataset_id}",
|
||||||
|
fg="cyan",
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
db.session.close()
|
||||||
db.session.close()
|
return
|
||||||
return
|
|
||||||
|
|
||||||
# Get all documents with completed indexing status
|
|
||||||
dataset_documents = db.session.scalars(
|
|
||||||
select(DatasetDocument).where(
|
|
||||||
DatasetDocument.dataset_id == dataset_id,
|
|
||||||
DatasetDocument.indexing_status == "completed",
|
|
||||||
DatasetDocument.enabled == True,
|
|
||||||
DatasetDocument.archived == False,
|
|
||||||
)
|
|
||||||
).all()
|
|
||||||
|
|
||||||
if not dataset_documents:
|
|
||||||
logger.info(
|
|
||||||
click.style(
|
|
||||||
f"No documents found for summary regeneration in dataset {dataset_id}",
|
|
||||||
fg="cyan",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
db.session.close()
|
|
||||||
return
|
|
||||||
|
|
||||||
logger.info(
|
|
||||||
"Found %s documents for summary regeneration in dataset %s",
|
|
||||||
len(dataset_documents),
|
|
||||||
dataset_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
total_segments_processed = 0
|
total_segments_processed = 0
|
||||||
total_segments_failed = 0
|
total_segments_failed = 0
|
||||||
|
|
||||||
for dataset_document in dataset_documents:
|
if regenerate_vectors_only:
|
||||||
# Skip qa_model documents
|
# For embedding_model change: directly query all segments with existing summaries
|
||||||
if dataset_document.doc_form == "qa_model":
|
# Don't require document indexing_status == "completed"
|
||||||
continue
|
# Include summaries with status "completed" or "error" (if they have content)
|
||||||
|
segments_with_summaries = (
|
||||||
try:
|
db.session.query(DocumentSegment, DocumentSegmentSummary)
|
||||||
# Get all segments with existing summaries
|
.join(
|
||||||
segments = (
|
DocumentSegmentSummary,
|
||||||
db.session.query(DocumentSegment)
|
DocumentSegment.id == DocumentSegmentSummary.chunk_id,
|
||||||
.join(
|
|
||||||
DocumentSegmentSummary,
|
|
||||||
DocumentSegment.id == DocumentSegmentSummary.chunk_id,
|
|
||||||
)
|
|
||||||
.where(
|
|
||||||
DocumentSegment.document_id == dataset_document.id,
|
|
||||||
DocumentSegment.dataset_id == dataset_id,
|
|
||||||
DocumentSegment.status == "completed",
|
|
||||||
DocumentSegment.enabled == True,
|
|
||||||
DocumentSegmentSummary.dataset_id == dataset_id,
|
|
||||||
)
|
|
||||||
.order_by(DocumentSegment.position.asc())
|
|
||||||
.all()
|
|
||||||
)
|
)
|
||||||
|
.join(
|
||||||
|
DatasetDocument,
|
||||||
|
DocumentSegment.document_id == DatasetDocument.id,
|
||||||
|
)
|
||||||
|
.where(
|
||||||
|
DocumentSegment.dataset_id == dataset_id,
|
||||||
|
DocumentSegment.status == "completed", # Segment must be completed
|
||||||
|
DocumentSegment.enabled == True,
|
||||||
|
DocumentSegmentSummary.dataset_id == dataset_id,
|
||||||
|
DocumentSegmentSummary.summary_content.isnot(None), # Must have summary content
|
||||||
|
# Include completed summaries or error summaries (with content)
|
||||||
|
or_(
|
||||||
|
DocumentSegmentSummary.status == "completed",
|
||||||
|
DocumentSegmentSummary.status == "error",
|
||||||
|
),
|
||||||
|
DatasetDocument.enabled == True, # Document must be enabled
|
||||||
|
DatasetDocument.archived == False, # Document must not be archived
|
||||||
|
DatasetDocument.doc_form != "qa_model", # Skip qa_model documents
|
||||||
|
)
|
||||||
|
.order_by(DocumentSegment.document_id.asc(), DocumentSegment.position.asc())
|
||||||
|
.all()
|
||||||
|
)
|
||||||
|
|
||||||
if not segments:
|
if not segments_with_summaries:
|
||||||
continue
|
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"Regenerating summaries for %s segments in document %s",
|
click.style(
|
||||||
len(segments),
|
f"No segments with summaries found for re-vectorization in dataset {dataset_id}",
|
||||||
dataset_document.id,
|
fg="cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
db.session.close()
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Found %s segments with summaries for re-vectorization in dataset %s",
|
||||||
|
len(segments_with_summaries),
|
||||||
|
dataset_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Group by document for logging
|
||||||
|
segments_by_document = defaultdict(list)
|
||||||
|
for segment, summary_record in segments_with_summaries:
|
||||||
|
segments_by_document[segment.document_id].append((segment, summary_record))
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Segments grouped into %s documents for re-vectorization",
|
||||||
|
len(segments_by_document),
|
||||||
|
)
|
||||||
|
|
||||||
|
for document_id, segment_summary_pairs in segments_by_document.items():
|
||||||
|
logger.info(
|
||||||
|
"Re-vectorizing summaries for %s segments in document %s",
|
||||||
|
len(segment_summary_pairs),
|
||||||
|
document_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
for segment in segments:
|
for segment, summary_record in segment_summary_pairs:
|
||||||
try:
|
try:
|
||||||
# Get existing summary record
|
# Delete old vector
|
||||||
summary_record = (
|
if summary_record.summary_index_node_id:
|
||||||
db.session.query(DocumentSegmentSummary)
|
try:
|
||||||
.filter_by(
|
from core.rag.datasource.vdb.vector_factory import Vector
|
||||||
chunk_id=segment.id,
|
|
||||||
dataset_id=dataset_id,
|
|
||||||
)
|
|
||||||
.first()
|
|
||||||
)
|
|
||||||
|
|
||||||
if not summary_record:
|
vector = Vector(dataset)
|
||||||
logger.warning("Summary record not found for segment %s, skipping", segment.id)
|
vector.delete_by_ids([summary_record.summary_index_node_id])
|
||||||
continue
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
if regenerate_vectors_only:
|
"Failed to delete old summary vector for segment %s: %s",
|
||||||
# Only regenerate vectors (for embedding_model change)
|
segment.id,
|
||||||
# Delete old vector
|
str(e),
|
||||||
if summary_record.summary_index_node_id:
|
)
|
||||||
try:
|
|
||||||
from core.rag.datasource.vdb.vector_factory import Vector
|
|
||||||
|
|
||||||
vector = Vector(dataset)
|
|
||||||
vector.delete_by_ids([summary_record.summary_index_node_id])
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning(
|
|
||||||
"Failed to delete old summary vector for segment %s: %s",
|
|
||||||
segment.id,
|
|
||||||
str(e),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Re-vectorize with new embedding model
|
|
||||||
SummaryIndexService.vectorize_summary(summary_record, segment, dataset)
|
|
||||||
db.session.commit()
|
|
||||||
else:
|
|
||||||
# Regenerate both summary content and vectors (for summary_model change)
|
|
||||||
SummaryIndexService.generate_and_vectorize_summary(segment, dataset, summary_index_setting)
|
|
||||||
db.session.commit()
|
|
||||||
|
|
||||||
|
# Re-vectorize with new embedding model
|
||||||
|
SummaryIndexService.vectorize_summary(summary_record, segment, dataset)
|
||||||
|
db.session.commit()
|
||||||
total_segments_processed += 1
|
total_segments_processed += 1
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(
|
logger.error(
|
||||||
"Failed to regenerate summary for segment %s: %s",
|
"Failed to re-vectorize summary for segment %s: %s",
|
||||||
segment.id,
|
segment.id,
|
||||||
str(e),
|
str(e),
|
||||||
exc_info=True,
|
exc_info=True,
|
||||||
)
|
)
|
||||||
total_segments_failed += 1
|
total_segments_failed += 1
|
||||||
# Update summary record with error status
|
# Update summary record with error status
|
||||||
if summary_record:
|
summary_record.status = "error"
|
||||||
summary_record.status = "error"
|
summary_record.error = f"Re-vectorization failed: {str(e)}"
|
||||||
summary_record.error = f"Regeneration failed: {str(e)}"
|
db.session.add(summary_record)
|
||||||
db.session.add(summary_record)
|
db.session.commit()
|
||||||
db.session.commit()
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
except Exception as e:
|
else:
|
||||||
logger.error(
|
# For summary_model change: require document indexing_status == "completed"
|
||||||
"Failed to process document %s for summary regeneration: %s",
|
# Get all documents with completed indexing status
|
||||||
dataset_document.id,
|
dataset_documents = db.session.scalars(
|
||||||
str(e),
|
select(DatasetDocument).where(
|
||||||
exc_info=True,
|
DatasetDocument.dataset_id == dataset_id,
|
||||||
|
DatasetDocument.indexing_status == "completed",
|
||||||
|
DatasetDocument.enabled == True,
|
||||||
|
DatasetDocument.archived == False,
|
||||||
)
|
)
|
||||||
continue
|
).all()
|
||||||
|
|
||||||
|
if not dataset_documents:
|
||||||
|
logger.info(
|
||||||
|
click.style(
|
||||||
|
f"No documents found for summary regeneration in dataset {dataset_id}",
|
||||||
|
fg="cyan",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
db.session.close()
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Found %s documents for summary regeneration in dataset %s",
|
||||||
|
len(dataset_documents),
|
||||||
|
dataset_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
for dataset_document in dataset_documents:
|
||||||
|
# Skip qa_model documents
|
||||||
|
if dataset_document.doc_form == "qa_model":
|
||||||
|
continue
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Get all segments with existing summaries
|
||||||
|
segments = (
|
||||||
|
db.session.query(DocumentSegment)
|
||||||
|
.join(
|
||||||
|
DocumentSegmentSummary,
|
||||||
|
DocumentSegment.id == DocumentSegmentSummary.chunk_id,
|
||||||
|
)
|
||||||
|
.where(
|
||||||
|
DocumentSegment.document_id == dataset_document.id,
|
||||||
|
DocumentSegment.dataset_id == dataset_id,
|
||||||
|
DocumentSegment.status == "completed",
|
||||||
|
DocumentSegment.enabled == True,
|
||||||
|
DocumentSegmentSummary.dataset_id == dataset_id,
|
||||||
|
)
|
||||||
|
.order_by(DocumentSegment.position.asc())
|
||||||
|
.all()
|
||||||
|
)
|
||||||
|
|
||||||
|
if not segments:
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Regenerating summaries for %s segments in document %s",
|
||||||
|
len(segments),
|
||||||
|
dataset_document.id,
|
||||||
|
)
|
||||||
|
|
||||||
|
for segment in segments:
|
||||||
|
try:
|
||||||
|
# Get existing summary record
|
||||||
|
summary_record = (
|
||||||
|
db.session.query(DocumentSegmentSummary)
|
||||||
|
.filter_by(
|
||||||
|
chunk_id=segment.id,
|
||||||
|
dataset_id=dataset_id,
|
||||||
|
)
|
||||||
|
.first()
|
||||||
|
)
|
||||||
|
|
||||||
|
if not summary_record:
|
||||||
|
logger.warning("Summary record not found for segment %s, skipping", segment.id)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Regenerate both summary content and vectors (for summary_model change)
|
||||||
|
SummaryIndexService.generate_and_vectorize_summary(segment, dataset, summary_index_setting)
|
||||||
|
db.session.commit()
|
||||||
|
total_segments_processed += 1
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to regenerate summary for segment %s: %s",
|
||||||
|
segment.id,
|
||||||
|
str(e),
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
total_segments_failed += 1
|
||||||
|
# Update summary record with error status
|
||||||
|
if summary_record:
|
||||||
|
summary_record.status = "error"
|
||||||
|
summary_record.error = f"Regeneration failed: {str(e)}"
|
||||||
|
db.session.add(summary_record)
|
||||||
|
db.session.commit()
|
||||||
|
continue
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
"Failed to process document %s for summary regeneration: %s",
|
||||||
|
dataset_document.id,
|
||||||
|
str(e),
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
|
||||||
end_at = time.perf_counter()
|
end_at = time.perf_counter()
|
||||||
logger.info(
|
if regenerate_vectors_only:
|
||||||
click.style(
|
logger.info(
|
||||||
f"Summary index regeneration completed for dataset {dataset_id}: "
|
click.style(
|
||||||
f"{total_segments_processed} segments processed successfully, "
|
f"Summary re-vectorization completed for dataset {dataset_id}: "
|
||||||
f"{total_segments_failed} segments failed, "
|
f"{total_segments_processed} segments processed successfully, "
|
||||||
f"total documents: {len(dataset_documents)}, "
|
f"{total_segments_failed} segments failed, "
|
||||||
f"latency: {end_at - start_at:.2f}s",
|
f"latency: {end_at - start_at:.2f}s",
|
||||||
fg="green",
|
fg="green",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.info(
|
||||||
|
click.style(
|
||||||
|
f"Summary index regeneration completed for dataset {dataset_id}: "
|
||||||
|
f"{total_segments_processed} segments processed successfully, "
|
||||||
|
f"{total_segments_failed} segments failed, "
|
||||||
|
f"latency: {end_at - start_at:.2f}s",
|
||||||
|
fg="green",
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
|
||||||
|
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Regenerate summary index failed for dataset %s", dataset_id)
|
logger.exception("Regenerate summary index failed for dataset %s", dataset_id)
|
||||||
|
|||||||
Reference in New Issue
Block a user