feat: 异步同步知识库分块检索索引
This commit is contained in:
@@ -53,6 +53,10 @@
|
||||
<groupId>com.easyagents</groupId>
|
||||
<artifactId>easy-agents-spring-boot-starter</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>io.milvus</groupId>
|
||||
<artifactId>milvus-sdk-java</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.google.re2j</groupId>
|
||||
<artifactId>re2j</artifactId>
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
package tech.easyflow.ai.config;
|
||||
|
||||
import com.easyagents.store.milvus.MilvusClientManager;
|
||||
import jakarta.annotation.PreDestroy;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.context.annotation.Lazy;
|
||||
|
||||
/**
|
||||
* 应用级 Milvus 客户端池。
|
||||
*/
|
||||
@Component
|
||||
@Lazy
|
||||
public class AiMilvusClientManager extends MilvusClientManager {
|
||||
|
||||
public AiMilvusClientManager(AiMilvusConfig config) {
|
||||
super(config);
|
||||
}
|
||||
|
||||
@PreDestroy
|
||||
public void destroy() {
|
||||
close();
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,13 @@ public class AiMilvusConfig extends MilvusVectorStoreConfig {
|
||||
config.setPassword(getPassword());
|
||||
config.setAutoCreateCollection(isAutoCreateCollection());
|
||||
config.setDefaultCollectionName(collectionName);
|
||||
config.setPoolMaxTotal(getPoolMaxTotal());
|
||||
config.setPoolMaxTotalPerKey(getPoolMaxTotalPerKey());
|
||||
config.setPoolMaxIdlePerKey(getPoolMaxIdlePerKey());
|
||||
config.setPoolMinIdlePerKey(getPoolMinIdlePerKey());
|
||||
config.setPoolMaxWaitMillis(getPoolMaxWaitMillis());
|
||||
config.setPoolEvictionIntervalMillis(getPoolEvictionIntervalMillis());
|
||||
config.setPoolMinEvictableIdleMillis(getPoolMinEvictableIdleMillis());
|
||||
return config;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,16 +25,23 @@ public class RagHealthIndicator {
|
||||
public static class RagMilvusHealthIndicator extends CachedHealthIndicatorSupport implements HealthIndicator {
|
||||
|
||||
private final AiMilvusConfig aiMilvusConfig;
|
||||
private final AiMilvusClientManager milvusClientManager;
|
||||
|
||||
/**
|
||||
* 创建 Milvus 健康检查器。
|
||||
*
|
||||
* @param aiMilvusConfig Milvus 配置
|
||||
* @param milvusClientManager 应用级 Milvus 客户端池
|
||||
* @param healthProperties RAG 健康检查配置
|
||||
*/
|
||||
public RagMilvusHealthIndicator(AiMilvusConfig aiMilvusConfig, RagHealthProperties healthProperties) {
|
||||
public RagMilvusHealthIndicator(
|
||||
AiMilvusConfig aiMilvusConfig,
|
||||
AiMilvusClientManager milvusClientManager,
|
||||
RagHealthProperties healthProperties
|
||||
) {
|
||||
super(healthProperties);
|
||||
this.aiMilvusConfig = aiMilvusConfig;
|
||||
this.milvusClientManager = milvusClientManager;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -51,8 +58,10 @@ public class RagHealthIndicator {
|
||||
protected Health doHealthCheck() {
|
||||
MilvusVectorStore vectorStore = null;
|
||||
try {
|
||||
milvusClientManager.reconfigureIfNeeded(aiMilvusConfig);
|
||||
vectorStore = new MilvusVectorStore(
|
||||
aiMilvusConfig.copyForCollection("__rag_health_probe__")
|
||||
aiMilvusConfig.copyForCollection("__rag_health_probe__"),
|
||||
milvusClientManager
|
||||
);
|
||||
if (vectorStore.checkAvailable()) {
|
||||
return Health.up().withDetail("uri", aiMilvusConfig.getUri()).build();
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
/**
|
||||
* 分块检索索引同步状态常量。
|
||||
*/
|
||||
public final class DocumentChunkSyncState {
|
||||
|
||||
private static final String PARENT_LOCK_PREFIX =
|
||||
"easyflow:lock:document-chunk-parent:";
|
||||
private static final String SYNC_LOCK_PREFIX =
|
||||
"easyflow:lock:document-chunk-sync:";
|
||||
|
||||
public static final String SYNCED = "SYNCED";
|
||||
public static final String PENDING = "PENDING";
|
||||
public static final String FAILED = "FAILED";
|
||||
|
||||
public static final String TASK_RUNNING = "RUNNING";
|
||||
public static final String TASK_SUCCEEDED = "SUCCEEDED";
|
||||
public static final String TASK_SUPERSEDED = "SUPERSEDED";
|
||||
|
||||
public static final String OPERATION_UPSERT = "UPSERT";
|
||||
public static final String OPERATION_DELETE = "DELETE";
|
||||
|
||||
public static String parentLockKey(Object documentId) {
|
||||
return PARENT_LOCK_PREFIX + documentId;
|
||||
}
|
||||
|
||||
public static String syncLockKey(Object chunkId) {
|
||||
return SYNC_LOCK_PREFIX + chunkId;
|
||||
}
|
||||
|
||||
private DocumentChunkSyncState() {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,414 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import com.easyagents.core.model.embedding.EmbeddingModel;
|
||||
import com.easyagents.core.model.embedding.EmbeddingOptions;
|
||||
import com.easyagents.core.store.DocumentStore;
|
||||
import com.easyagents.core.store.StoreOptions;
|
||||
import com.easyagents.core.store.StoreResult;
|
||||
import com.easyagents.search.engine.service.DocumentSearcher;
|
||||
import com.easyagents.search.engine.service.KeywordSearchMetadataKeys;
|
||||
import com.easyagents.store.milvus.MilvusVectorStore;
|
||||
import com.easyagents.store.milvus.MilvusVectorStoreConfig;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.support.TransactionTemplate;
|
||||
import tech.easyflow.ai.config.AiMilvusClientManager;
|
||||
import tech.easyflow.ai.config.AiMilvusConfig;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentChunkSyncTask;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.entity.Model;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkSyncTaskMapper;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.ai.service.ModelService;
|
||||
import tech.easyflow.ai.support.DocumentStoreLifecycleSupport;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.time.Duration;
|
||||
import java.util.Collections;
|
||||
import java.util.Date;
|
||||
import java.util.UUID;
|
||||
|
||||
/**
|
||||
* 持久化分块索引同步任务的投递、执行和恢复。
|
||||
*/
|
||||
@Service
|
||||
public class DocumentChunkSyncTaskAppService {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(DocumentChunkSyncTaskAppService.class);
|
||||
private static final int MAX_ATTEMPTS = 5;
|
||||
private static final int DISPATCH_LIMIT = 100;
|
||||
private static final long REDISPATCH_MILLIS = 5_000L;
|
||||
private static final long LEASE_MILLIS = 300_000L;
|
||||
|
||||
private final DocumentChunkSyncTaskMapper taskMapper;
|
||||
private final DocumentChunkMapper chunkMapper;
|
||||
private final DocumentCollectionService collectionService;
|
||||
private final ModelService modelService;
|
||||
private final SearcherFactory searcherFactory;
|
||||
private final DocumentChunkSyncTaskProducer producer;
|
||||
private final PlatformTransactionManager transactionManager;
|
||||
private final RedisLockExecutor redisLockExecutor;
|
||||
private final AiMilvusConfig milvusConfig;
|
||||
private final ObjectProvider<AiMilvusClientManager> milvusClientManagerProvider;
|
||||
|
||||
public DocumentChunkSyncTaskAppService(
|
||||
DocumentChunkSyncTaskMapper taskMapper,
|
||||
DocumentChunkMapper chunkMapper,
|
||||
DocumentCollectionService collectionService,
|
||||
ModelService modelService,
|
||||
SearcherFactory searcherFactory,
|
||||
DocumentChunkSyncTaskProducer producer,
|
||||
PlatformTransactionManager transactionManager,
|
||||
RedisLockExecutor redisLockExecutor,
|
||||
AiMilvusConfig milvusConfig,
|
||||
ObjectProvider<AiMilvusClientManager> milvusClientManagerProvider
|
||||
) {
|
||||
this.taskMapper = taskMapper;
|
||||
this.chunkMapper = chunkMapper;
|
||||
this.collectionService = collectionService;
|
||||
this.modelService = modelService;
|
||||
this.searcherFactory = searcherFactory;
|
||||
this.producer = producer;
|
||||
this.transactionManager = transactionManager;
|
||||
this.redisLockExecutor = redisLockExecutor;
|
||||
this.milvusConfig = milvusConfig;
|
||||
this.milvusClientManagerProvider = milvusClientManagerProvider;
|
||||
}
|
||||
|
||||
public DocumentChunkSyncTask createTask(
|
||||
DocumentChunk chunk,
|
||||
DocumentCollection collection,
|
||||
String operation,
|
||||
long version,
|
||||
Date now
|
||||
) {
|
||||
taskMapper.supersedeOlder(chunk.getId(), version, now);
|
||||
DocumentChunkSyncTask task = new DocumentChunkSyncTask();
|
||||
task.setChunkId(chunk.getId());
|
||||
task.setDocumentId(chunk.getDocumentId());
|
||||
task.setDocumentCollectionId(chunk.getDocumentCollectionId());
|
||||
task.setVectorCollection(collection.getVectorStoreCollection());
|
||||
task.setOperation(operation);
|
||||
task.setSyncVersion(version);
|
||||
task.setStatus(DocumentChunkSyncState.PENDING);
|
||||
task.setAttemptCount(0);
|
||||
task.setNextRetryAt(now);
|
||||
task.setCreated(now);
|
||||
task.setModified(now);
|
||||
if (taskMapper.insert(task) <= 0) {
|
||||
throw new IllegalStateException("创建分块索引同步任务失败");
|
||||
}
|
||||
return task;
|
||||
}
|
||||
|
||||
public void dispatchBestEffort(BigInteger taskId) {
|
||||
try {
|
||||
producer.send(taskId);
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.warn("分块索引同步消息投递失败，等待数据库补投: taskId={}", taskId, exception);
|
||||
}
|
||||
}
|
||||
|
||||
public void dispatchPendingTasks() {
|
||||
Date now = new Date();
|
||||
taskMapper.recoverExpired(now);
|
||||
Date redispatchBefore = new Date(now.getTime() - REDISPATCH_MILLIS);
|
||||
for (DocumentChunkSyncTask task : taskMapper.selectPendingDue(
|
||||
now,
|
||||
redispatchBefore,
|
||||
DISPATCH_LIMIT
|
||||
)) {
|
||||
if (taskMapper.markDispatched(task.getId(), now, redispatchBefore) <= 0) {
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
producer.send(task.getId());
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.warn("补投分块索引同步消息失败: taskId={}", task.getId(), exception);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public void handleTask(BigInteger taskId) {
|
||||
DocumentChunkSyncTask snapshot = taskMapper.selectOneById(taskId);
|
||||
if (snapshot == null || !DocumentChunkSyncState.PENDING.equals(snapshot.getStatus())) {
|
||||
return;
|
||||
}
|
||||
redisLockExecutor.executeWithRenewingLock(
|
||||
DocumentChunkSyncState.syncLockKey(snapshot.getChunkId()),
|
||||
Duration.ofSeconds(3),
|
||||
Duration.ofMinutes(5),
|
||||
() -> {
|
||||
claimAndExecute(taskId, snapshot.getChunkId());
|
||||
return null;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private void claimAndExecute(BigInteger taskId, BigInteger chunkId) {
|
||||
Date now = new Date();
|
||||
String token = UUID.randomUUID().toString();
|
||||
if (taskMapper.claim(
|
||||
taskId,
|
||||
token,
|
||||
new Date(now.getTime() + LEASE_MILLIS),
|
||||
now
|
||||
) <= 0) {
|
||||
return;
|
||||
}
|
||||
new TransactionTemplate(transactionManager).executeWithoutResult(status ->
|
||||
executeOwnedTask(taskId, chunkId, token)
|
||||
);
|
||||
}
|
||||
|
||||
private void executeOwnedTask(
|
||||
BigInteger taskId,
|
||||
BigInteger chunkId,
|
||||
String token
|
||||
) {
|
||||
// FOR UPDATE 必须是本事务的第一次读取。MySQL REPEATABLE READ 下，
|
||||
// 若先做普通查询会建立旧快照，使删除或新版本建单已经提交后仍读取旧分块。
|
||||
taskMapper.lockChunkTasks(chunkId);
|
||||
DocumentChunkSyncTask task = taskMapper.selectOneById(taskId);
|
||||
if (task == null
|
||||
|| !DocumentChunkSyncState.TASK_RUNNING.equals(task.getStatus())
|
||||
|| !token.equals(task.getExecutionToken())) {
|
||||
return;
|
||||
}
|
||||
if (isSuperseded(task)) {
|
||||
finishTask(task, token, DocumentChunkSyncState.TASK_SUPERSEDED, null, null);
|
||||
return;
|
||||
}
|
||||
try {
|
||||
synchronizeIndexes(task);
|
||||
finishSuccess(task, token);
|
||||
} catch (IndexSyncException exception) {
|
||||
LOG.warn("分块索引同步失败: taskId={}, chunkId={}, code={}",
|
||||
task.getId(), task.getChunkId(), exception.code, exception.getCause());
|
||||
finishFailure(task, token, exception.code, exception.getMessage());
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("分块索引同步发生未分类异常: taskId={}, chunkId={}",
|
||||
task.getId(), task.getChunkId(), exception);
|
||||
finishFailure(task, token, "INDEX_SYNC_FAILED", "检索索引同步失败");
|
||||
}
|
||||
}
|
||||
|
||||
private boolean isSuperseded(DocumentChunkSyncTask task) {
|
||||
DocumentChunk chunk = chunkMapper.selectOneById(task.getChunkId());
|
||||
if (DocumentChunkSyncState.OPERATION_DELETE.equals(task.getOperation())) {
|
||||
return chunk != null;
|
||||
}
|
||||
return chunk == null
|
||||
|| chunk.getIndexSyncVersion() == null
|
||||
|| chunk.getIndexSyncVersion().longValue() != task.getSyncVersion().longValue()
|
||||
|| !DocumentChunkSyncState.PENDING.equals(chunk.getIndexSyncStatus());
|
||||
}
|
||||
|
||||
private void synchronizeIndexes(DocumentChunkSyncTask task) {
|
||||
if (DocumentChunkSyncState.OPERATION_DELETE.equals(task.getOperation())) {
|
||||
synchronizeDelete(task);
|
||||
return;
|
||||
}
|
||||
synchronizeUpsert(task);
|
||||
}
|
||||
|
||||
private void synchronizeUpsert(DocumentChunkSyncTask task) {
|
||||
DocumentChunk chunk = chunkMapper.selectOneById(task.getChunkId());
|
||||
DocumentCollection collection = collectionService.getById(task.getDocumentCollectionId());
|
||||
if (chunk == null || collection == null) {
|
||||
throw new IndexSyncException("INDEX_SOURCE_MISSING", "分块或知识库不存在", null);
|
||||
}
|
||||
StoreContext context = prepareUpsertContext(collection, task.getVectorCollection());
|
||||
try {
|
||||
com.easyagents.core.document.Document document = toSearchDocument(chunk, task.getDocumentCollectionId());
|
||||
StoreResult vectorResult = context.documentStore.update(
|
||||
Collections.singletonList(document),
|
||||
context.storeOptions
|
||||
);
|
||||
if (vectorResult == null || !vectorResult.isSuccess()) {
|
||||
throw new IndexSyncException("VECTOR_UPSERT_FAILED", "向量索引更新失败", null);
|
||||
}
|
||||
if (context.searcher != null
|
||||
&& !context.searcher.addDocuments(Collections.singletonList(document))) {
|
||||
throw new IndexSyncException("KEYWORD_UPSERT_FAILED", "关键词索引更新失败", null);
|
||||
}
|
||||
} catch (IndexSyncException exception) {
|
||||
throw exception;
|
||||
} catch (RuntimeException exception) {
|
||||
throw new IndexSyncException("INDEX_UPSERT_FAILED", "检索索引更新失败", exception);
|
||||
} finally {
|
||||
DocumentStoreLifecycleSupport.closeQuietly(context.documentStore);
|
||||
}
|
||||
}
|
||||
|
||||
private void synchronizeDelete(DocumentChunkSyncTask task) {
|
||||
MilvusVectorStoreConfig storeConfig = milvusConfig.copyForCollection(task.getVectorCollection());
|
||||
AiMilvusClientManager clientManager = milvusClientManagerProvider.getObject();
|
||||
clientManager.reconfigureIfNeeded(milvusConfig);
|
||||
DocumentStore documentStore = new MilvusVectorStore(
|
||||
storeConfig,
|
||||
clientManager
|
||||
);
|
||||
StoreOptions storeOptions = StoreOptions.ofCollectionName(task.getVectorCollection());
|
||||
DocumentSearcher searcher = searcherFactory.getSearcher();
|
||||
try {
|
||||
StoreResult result = documentStore.delete(
|
||||
Collections.singletonList(task.getChunkId().toString()),
|
||||
storeOptions
|
||||
);
|
||||
if (result == null || !result.isSuccess()) {
|
||||
throw new IndexSyncException("VECTOR_DELETE_FAILED", "向量索引删除失败", null);
|
||||
}
|
||||
if (searcher != null && !searcher.deleteDocument(task.getChunkId())) {
|
||||
throw new IndexSyncException("KEYWORD_DELETE_FAILED", "关键词索引删除失败", null);
|
||||
}
|
||||
} catch (IndexSyncException exception) {
|
||||
throw exception;
|
||||
} catch (RuntimeException exception) {
|
||||
throw new IndexSyncException("INDEX_DELETE_FAILED", "检索索引删除失败", exception);
|
||||
} finally {
|
||||
DocumentStoreLifecycleSupport.closeQuietly(documentStore);
|
||||
}
|
||||
}
|
||||
|
||||
private StoreContext prepareUpsertContext(
|
||||
DocumentCollection collection,
|
||||
String vectorCollection
|
||||
) {
|
||||
DocumentStore documentStore = collection.toDocumentStore();
|
||||
if (documentStore == null) {
|
||||
throw new IndexSyncException("VECTOR_STORE_MISSING", "知识库没有配置向量库", null);
|
||||
}
|
||||
try {
|
||||
Model model = modelService.getModelInstance(collection.getVectorEmbedModelId());
|
||||
if (model == null) {
|
||||
throw new IndexSyncException("EMBEDDING_MODEL_MISSING", "知识库没有配置向量模型", null);
|
||||
}
|
||||
EmbeddingModel embeddingModel = model.toEmbeddingModel();
|
||||
documentStore.setEmbeddingModel(embeddingModel);
|
||||
StoreOptions options = StoreOptions.ofCollectionName(vectorCollection);
|
||||
EmbeddingOptions embeddingOptions = new EmbeddingOptions();
|
||||
embeddingOptions.setModel(model.getModelName());
|
||||
embeddingOptions.setDimensions(collection.getDimensionOfVectorModel());
|
||||
options.setEmbeddingOptions(embeddingOptions);
|
||||
options.setIndexName(vectorCollection);
|
||||
return new StoreContext(documentStore, options, searcherFactory.getSearcher());
|
||||
} catch (RuntimeException exception) {
|
||||
DocumentStoreLifecycleSupport.closeQuietly(documentStore);
|
||||
throw exception;
|
||||
}
|
||||
}
|
||||
|
||||
private com.easyagents.core.document.Document toSearchDocument(
|
||||
DocumentChunk chunk,
|
||||
BigInteger knowledgeId
|
||||
) {
|
||||
com.easyagents.core.document.Document document =
|
||||
com.easyagents.core.document.Document.of(chunk.getContent());
|
||||
document.setId(chunk.getId());
|
||||
document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, knowledgeId.toString());
|
||||
return document;
|
||||
}
|
||||
|
||||
private void finishSuccess(DocumentChunkSyncTask task, String token) {
|
||||
new TransactionTemplate(transactionManager).executeWithoutResult(status -> {
|
||||
if (taskMapper.finishOwned(
|
||||
task.getId(), token, DocumentChunkSyncState.TASK_SUCCEEDED, null, null, new Date()
|
||||
) > 0 && DocumentChunkSyncState.OPERATION_UPSERT.equals(task.getOperation())) {
|
||||
chunkMapper.updateSyncState(
|
||||
task.getChunkId(), task.getSyncVersion(), DocumentChunkSyncState.SYNCED, null, null
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private void finishTask(
|
||||
DocumentChunkSyncTask task,
|
||||
String token,
|
||||
String status,
|
||||
String errorCode,
|
||||
String errorMessage
|
||||
) {
|
||||
taskMapper.finishOwned(task.getId(), token, status, errorCode, errorMessage, new Date());
|
||||
}
|
||||
|
||||
private void finishFailure(
|
||||
DocumentChunkSyncTask task,
|
||||
String token,
|
||||
String errorCode,
|
||||
String errorMessage
|
||||
) {
|
||||
int attempts = task.getAttemptCount() == null ? 1 : task.getAttemptCount();
|
||||
boolean deleteOperation = DocumentChunkSyncState.OPERATION_DELETE.equals(
|
||||
task.getOperation()
|
||||
);
|
||||
// 删除后已没有页面实体承载手动重试入口，因此清理任务必须保留为持久化 tombstone。
|
||||
boolean exhausted = !deleteOperation && attempts >= MAX_ATTEMPTS;
|
||||
String nextStatus = exhausted ? DocumentChunkSyncState.FAILED : DocumentChunkSyncState.PENDING;
|
||||
long delaySeconds = Math.min(1L << Math.min(attempts, 6), 60L);
|
||||
Date now = new Date();
|
||||
Date nextRetryAt = exhausted ? now : new Date(now.getTime() + delaySeconds * 1_000L);
|
||||
new TransactionTemplate(transactionManager).executeWithoutResult(status -> {
|
||||
if (taskMapper.failOrRetryOwned(
|
||||
task.getId(), token, nextStatus, nextRetryAt, errorCode, errorMessage, now
|
||||
) > 0 && exhausted && DocumentChunkSyncState.OPERATION_UPSERT.equals(task.getOperation())) {
|
||||
chunkMapper.updateSyncState(
|
||||
task.getChunkId(), task.getSyncVersion(), DocumentChunkSyncState.FAILED,
|
||||
errorCode, errorMessage
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
public DocumentChunk retry(BigInteger chunkId, long syncVersion) {
|
||||
Date now = new Date();
|
||||
DocumentChunk chunk = chunkMapper.selectOneById(chunkId);
|
||||
if (chunk == null || chunk.getIndexSyncVersion() == null
|
||||
|| chunk.getIndexSyncVersion().longValue() != syncVersion
|
||||
|| !DocumentChunkSyncState.FAILED.equals(chunk.getIndexSyncStatus())) {
|
||||
throw new IllegalStateException("分块同步状态已变化，请刷新后重试");
|
||||
}
|
||||
DocumentChunkSyncTask task = taskMapper.selectFailed(chunkId, syncVersion);
|
||||
if (task == null) {
|
||||
throw new IllegalStateException("未找到可重试的索引同步任务");
|
||||
}
|
||||
new TransactionTemplate(transactionManager).executeWithoutResult(status -> {
|
||||
if (taskMapper.retryFailed(task.getId(), chunkId, syncVersion, now) <= 0) {
|
||||
throw new IllegalStateException("索引同步任务状态已变化");
|
||||
}
|
||||
if (chunkMapper.updateSyncState(
|
||||
chunkId, syncVersion, DocumentChunkSyncState.PENDING, null, null
|
||||
) <= 0) {
|
||||
throw new IllegalStateException("分块同步状态已变化");
|
||||
}
|
||||
});
|
||||
dispatchBestEffort(task.getId());
|
||||
chunk.setIndexSyncStatus(DocumentChunkSyncState.PENDING);
|
||||
chunk.setIndexSyncErrorCode(null);
|
||||
chunk.setIndexSyncErrorMessage(null);
|
||||
return chunk;
|
||||
}
|
||||
|
||||
private record StoreContext(
|
||||
DocumentStore documentStore,
|
||||
StoreOptions storeOptions,
|
||||
DocumentSearcher searcher
|
||||
) {
|
||||
}
|
||||
|
||||
private static final class IndexSyncException extends RuntimeException {
|
||||
private final String code;
|
||||
|
||||
private IndexSyncException(String code, String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
this.code = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import com.alibaba.fastjson2.JSON;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.stereotype.Component;
|
||||
import tech.easyflow.common.mq.config.MQProperties;
|
||||
import tech.easyflow.common.mq.core.MQConsumerHandler;
|
||||
import tech.easyflow.common.mq.core.MQDeferException;
|
||||
import tech.easyflow.common.mq.core.MQMessage;
|
||||
import tech.easyflow.common.mq.core.MQSubscription;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 分块索引同步消息消费者。
|
||||
*/
|
||||
@Component
|
||||
public class DocumentChunkSyncTaskConsumer implements MQConsumerHandler {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(DocumentChunkSyncTaskConsumer.class);
|
||||
|
||||
private final DocumentChunkSyncTaskAppService appService;
|
||||
private final MQProperties mqProperties;
|
||||
|
||||
public DocumentChunkSyncTaskConsumer(
|
||||
DocumentChunkSyncTaskAppService appService,
|
||||
MQProperties mqProperties
|
||||
) {
|
||||
this.appService = appService;
|
||||
this.mqProperties = mqProperties;
|
||||
}
|
||||
|
||||
@Override
|
||||
public MQSubscription subscription() {
|
||||
MQSubscription subscription = new MQSubscription();
|
||||
subscription.setTopic(DocumentChunkSyncTaskProducer.TOPIC);
|
||||
subscription.setConsumerGroup(DocumentChunkSyncTaskProducer.GROUP);
|
||||
subscription.setShardCount(Math.max(
|
||||
mqProperties.getRedis().getChatPersistShardCount(),
|
||||
1
|
||||
));
|
||||
subscription.setBatchEnabled(false);
|
||||
return subscription;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void handle(List<MQMessage> messages) {
|
||||
for (MQMessage message : messages) {
|
||||
DocumentChunkSyncTaskMessage event = JSON.parseObject(
|
||||
message.getBody(),
|
||||
DocumentChunkSyncTaskMessage.class
|
||||
);
|
||||
if (event == null || event.getTaskId() == null) {
|
||||
LOG.warn("忽略非法分块索引同步消息: messageId={}", message.getMessageId());
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
appService.handleTask(event.getTaskId());
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("分块索引同步任务状态处理失败: taskId={}", event.getTaskId(), exception);
|
||||
throw new MQDeferException("分块索引同步任务暂时无法处理", exception);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import java.io.Serializable;
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
|
||||
/**
|
||||
* 分块索引同步消息。
|
||||
*/
|
||||
public class DocumentChunkSyncTaskMessage implements Serializable {
|
||||
private BigInteger taskId;
|
||||
private Date occurredAt;
|
||||
|
||||
public BigInteger getTaskId() { return taskId; }
|
||||
public void setTaskId(BigInteger taskId) { this.taskId = taskId; }
|
||||
public Date getOccurredAt() { return occurredAt; }
|
||||
public void setOccurredAt(Date occurredAt) { this.occurredAt = occurredAt; }
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.scheduling.annotation.Scheduled;
|
||||
import org.springframework.stereotype.Component;
|
||||
import tech.easyflow.common.cache.DistributedScheduledLock;
|
||||
|
||||
/**
|
||||
* 补投待同步任务并回收过期租约。
|
||||
*/
|
||||
@Component
|
||||
public class DocumentChunkSyncTaskMonitor {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(DocumentChunkSyncTaskMonitor.class);
|
||||
|
||||
private final DocumentChunkSyncTaskAppService appService;
|
||||
|
||||
public DocumentChunkSyncTaskMonitor(DocumentChunkSyncTaskAppService appService) {
|
||||
this.appService = appService;
|
||||
}
|
||||
|
||||
@Scheduled(fixedDelayString = "${easyflow.ai.document-chunk-sync.dispatch-interval:2s}",
|
||||
initialDelayString = "${easyflow.ai.document-chunk-sync.dispatch-interval:2s}")
|
||||
@DistributedScheduledLock(key = "easyflow:schedule:document-chunk-index-sync", leaseSeconds = 2L)
|
||||
public void dispatchPendingTasks() {
|
||||
try {
|
||||
appService.dispatchPendingTasks();
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("分块索引同步补投失败", exception);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import com.alibaba.fastjson2.JSON;
|
||||
import org.springframework.stereotype.Service;
|
||||
import tech.easyflow.common.mq.core.MQMessage;
|
||||
import tech.easyflow.common.mq.core.MQProducer;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
|
||||
/**
|
||||
* 分块索引同步消息生产者。
|
||||
*/
|
||||
@Service
|
||||
public class DocumentChunkSyncTaskProducer {
|
||||
|
||||
static final String TOPIC = "document-chunk-index-sync";
|
||||
static final String GROUP = "document-chunk-index-sync-group";
|
||||
|
||||
private final MQProducer mqProducer;
|
||||
|
||||
public DocumentChunkSyncTaskProducer(MQProducer mqProducer) {
|
||||
this.mqProducer = mqProducer;
|
||||
}
|
||||
|
||||
public void send(BigInteger taskId) {
|
||||
Date now = new Date();
|
||||
DocumentChunkSyncTaskMessage event = new DocumentChunkSyncTaskMessage();
|
||||
event.setTaskId(taskId);
|
||||
event.setOccurredAt(now);
|
||||
MQMessage message = new MQMessage();
|
||||
message.setMessageId("chunk-sync-" + taskId);
|
||||
message.setTopic(TOPIC);
|
||||
message.setKey(String.valueOf(taskId));
|
||||
message.setCreatedAt(now);
|
||||
message.setBody(JSON.toJSONString(event));
|
||||
mqProducer.send(message);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
|
||||
import java.math.BigInteger;
|
||||
|
||||
/**
|
||||
* 兼容 success 字段的异步分块更新结果。
|
||||
*/
|
||||
public record DocumentChunkAsyncUpdateResult(
|
||||
boolean success,
|
||||
BigInteger id,
|
||||
String indexSyncStatus,
|
||||
Long indexSyncVersion
|
||||
) {
|
||||
public static DocumentChunkAsyncUpdateResult from(DocumentChunk chunk) {
|
||||
return new DocumentChunkAsyncUpdateResult(
|
||||
true,
|
||||
chunk.getId(),
|
||||
chunk.getIndexSyncStatus(),
|
||||
chunk.getIndexSyncVersion()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import java.io.Serializable;
|
||||
import java.math.BigInteger;
|
||||
|
||||
/**
|
||||
* 分块正文更新请求，只允许修改 Markdown 内容。
|
||||
*/
|
||||
public class DocumentChunkContentUpdateRequest implements Serializable {
|
||||
|
||||
private BigInteger id;
|
||||
private String content;
|
||||
|
||||
public BigInteger getId() {
|
||||
return id;
|
||||
}
|
||||
|
||||
public void setId(BigInteger id) {
|
||||
this.id = id;
|
||||
}
|
||||
|
||||
public String getContent() {
|
||||
return content;
|
||||
}
|
||||
|
||||
public void setContent(String content) {
|
||||
this.content = content;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import java.io.Serializable;
|
||||
import java.math.BigInteger;
|
||||
|
||||
/**
|
||||
* 分块删除结果。
|
||||
*/
|
||||
public record DocumentChunkDeleteResult(
|
||||
BigInteger id,
|
||||
BigInteger documentId,
|
||||
long remainingChunkCount
|
||||
) implements Serializable {
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import java.math.BigInteger;
|
||||
|
||||
/**
|
||||
* 重试分块索引同步请求。
|
||||
*/
|
||||
public class DocumentChunkSyncRetryRequest {
|
||||
private BigInteger id;
|
||||
private Long indexSyncVersion;
|
||||
|
||||
public BigInteger getId() { return id; }
|
||||
public void setId(BigInteger id) { this.id = id; }
|
||||
public Long getIndexSyncVersion() { return indexSyncVersion; }
|
||||
public void setIndexSyncVersion(Long indexSyncVersion) { this.indexSyncVersion = indexSyncVersion; }
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import java.math.BigInteger;
|
||||
|
||||
/**
|
||||
* 分块索引同步状态。
|
||||
*/
|
||||
public record DocumentChunkSyncStatus(
|
||||
BigInteger id,
|
||||
String indexSyncStatus,
|
||||
Long indexSyncVersion,
|
||||
String indexSyncErrorCode,
|
||||
String indexSyncErrorMessage
|
||||
) {
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package tech.easyflow.ai.dto;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 批量查询分块索引同步状态。
|
||||
*/
|
||||
public class DocumentChunkSyncStatusRequest {
|
||||
private BigInteger documentId;
|
||||
private List<BigInteger> ids;
|
||||
|
||||
public BigInteger getDocumentId() { return documentId; }
|
||||
public void setDocumentId(BigInteger documentId) { this.documentId = documentId; }
|
||||
public List<BigInteger> getIds() { return ids; }
|
||||
public void setIds(List<BigInteger> ids) { this.ids = ids; }
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package tech.easyflow.ai.entity;
|
||||
|
||||
import com.mybatisflex.annotation.Column;
|
||||
import com.mybatisflex.annotation.Id;
|
||||
import com.mybatisflex.annotation.KeyType;
|
||||
import com.mybatisflex.annotation.Table;
|
||||
import tech.easyflow.common.entity.DateEntity;
|
||||
|
||||
import java.io.Serializable;
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
|
||||
/**
|
||||
* 文档分块检索索引同步任务。
|
||||
*/
|
||||
@Table("tb_document_chunk_sync_task")
|
||||
public class DocumentChunkSyncTask extends DateEntity implements Serializable {
|
||||
|
||||
@Id(keyType = KeyType.Generator, value = "snowFlakeId")
|
||||
private BigInteger id;
|
||||
private BigInteger chunkId;
|
||||
private BigInteger documentId;
|
||||
private BigInteger documentCollectionId;
|
||||
private String vectorCollection;
|
||||
private String operation;
|
||||
private Long syncVersion;
|
||||
private String status;
|
||||
private Integer attemptCount;
|
||||
private Date nextRetryAt;
|
||||
private Date lastDispatchedAt;
|
||||
private String executionToken;
|
||||
private Date leaseUntil;
|
||||
private String errorCode;
|
||||
private String errorMessage;
|
||||
private Date created;
|
||||
private BigInteger createdBy;
|
||||
private Date modified;
|
||||
private BigInteger modifiedBy;
|
||||
|
||||
public BigInteger getId() { return id; }
|
||||
public void setId(BigInteger id) { this.id = id; }
|
||||
public BigInteger getChunkId() { return chunkId; }
|
||||
public void setChunkId(BigInteger chunkId) { this.chunkId = chunkId; }
|
||||
public BigInteger getDocumentId() { return documentId; }
|
||||
public void setDocumentId(BigInteger documentId) { this.documentId = documentId; }
|
||||
public BigInteger getDocumentCollectionId() { return documentCollectionId; }
|
||||
public void setDocumentCollectionId(BigInteger documentCollectionId) { this.documentCollectionId = documentCollectionId; }
|
||||
public String getVectorCollection() { return vectorCollection; }
|
||||
public void setVectorCollection(String vectorCollection) { this.vectorCollection = vectorCollection; }
|
||||
public String getOperation() { return operation; }
|
||||
public void setOperation(String operation) { this.operation = operation; }
|
||||
public Long getSyncVersion() { return syncVersion; }
|
||||
public void setSyncVersion(Long syncVersion) { this.syncVersion = syncVersion; }
|
||||
public String getStatus() { return status; }
|
||||
public void setStatus(String status) { this.status = status; }
|
||||
public Integer getAttemptCount() { return attemptCount; }
|
||||
public void setAttemptCount(Integer attemptCount) { this.attemptCount = attemptCount; }
|
||||
public Date getNextRetryAt() { return nextRetryAt; }
|
||||
public void setNextRetryAt(Date nextRetryAt) { this.nextRetryAt = nextRetryAt; }
|
||||
public Date getLastDispatchedAt() { return lastDispatchedAt; }
|
||||
public void setLastDispatchedAt(Date lastDispatchedAt) { this.lastDispatchedAt = lastDispatchedAt; }
|
||||
public String getExecutionToken() { return executionToken; }
|
||||
public void setExecutionToken(String executionToken) { this.executionToken = executionToken; }
|
||||
public Date getLeaseUntil() { return leaseUntil; }
|
||||
public void setLeaseUntil(Date leaseUntil) { this.leaseUntil = leaseUntil; }
|
||||
public String getErrorCode() { return errorCode; }
|
||||
public void setErrorCode(String errorCode) { this.errorCode = errorCode; }
|
||||
public String getErrorMessage() { return errorMessage; }
|
||||
public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; }
|
||||
@Override public Date getCreated() { return created; }
|
||||
@Override public void setCreated(Date created) { this.created = created; }
|
||||
public BigInteger getCreatedBy() { return createdBy; }
|
||||
public void setCreatedBy(BigInteger createdBy) { this.createdBy = createdBy; }
|
||||
@Override public Date getModified() { return modified; }
|
||||
@Override public void setModified(Date modified) { this.modified = modified; }
|
||||
public BigInteger getModifiedBy() { return modifiedBy; }
|
||||
public void setModifiedBy(BigInteger modifiedBy) { this.modifiedBy = modifiedBy; }
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import com.easyagents.store.milvus.MilvusVectorStore;
|
||||
import com.easyagents.store.milvus.MilvusVectorStoreConfig;
|
||||
import com.mybatisflex.annotation.Table;
|
||||
import tech.easyflow.ai.config.AiMilvusConfig;
|
||||
import tech.easyflow.ai.config.AiMilvusClientManager;
|
||||
import tech.easyflow.ai.chattime.availability.ChatTimeToolAvailabilityContext;
|
||||
import tech.easyflow.ai.easyagents.tool.DocumentCollectionTool;
|
||||
import tech.easyflow.ai.entity.base.DocumentCollectionBase;
|
||||
@@ -103,8 +104,10 @@ public class DocumentCollection extends DocumentCollectionBase implements Visibi
|
||||
|
||||
private DocumentStore milvusStore() {
|
||||
AiMilvusConfig aiMilvusConfig = SpringContextUtil.getBean(AiMilvusConfig.class);
|
||||
AiMilvusClientManager clientManager = SpringContextUtil.getBean(AiMilvusClientManager.class);
|
||||
clientManager.reconfigureIfNeeded(aiMilvusConfig);
|
||||
MilvusVectorStoreConfig milvusVectorStoreConfig = aiMilvusConfig.copyForCollection(this.getVectorStoreCollection());
|
||||
return new MilvusVectorStore(milvusVectorStoreConfig);
|
||||
return new MilvusVectorStore(milvusVectorStoreConfig, clientManager);
|
||||
}
|
||||
|
||||
public Tool toFunction(boolean needEnglishName) {
|
||||
|
||||
@@ -46,6 +46,18 @@ public class DocumentChunkBase implements Serializable {
|
||||
@Column(typeHandler = FastjsonTypeHandler.class, comment = "扩展元信息")
|
||||
private Map<String, Object> options;
|
||||
|
||||
@Column(comment = "检索索引同步状态")
|
||||
private String indexSyncStatus;
|
||||
|
||||
@Column(comment = "检索索引同步版本")
|
||||
private Long indexSyncVersion;
|
||||
|
||||
@Column(comment = "脱敏同步错误码")
|
||||
private String indexSyncErrorCode;
|
||||
|
||||
@Column(comment = "脱敏同步错误摘要")
|
||||
private String indexSyncErrorMessage;
|
||||
|
||||
public BigInteger getId() {
|
||||
return id;
|
||||
}
|
||||
@@ -94,4 +106,36 @@ public class DocumentChunkBase implements Serializable {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
public String getIndexSyncStatus() {
|
||||
return indexSyncStatus;
|
||||
}
|
||||
|
||||
public void setIndexSyncStatus(String indexSyncStatus) {
|
||||
this.indexSyncStatus = indexSyncStatus;
|
||||
}
|
||||
|
||||
public Long getIndexSyncVersion() {
|
||||
return indexSyncVersion;
|
||||
}
|
||||
|
||||
public void setIndexSyncVersion(Long indexSyncVersion) {
|
||||
this.indexSyncVersion = indexSyncVersion;
|
||||
}
|
||||
|
||||
public String getIndexSyncErrorCode() {
|
||||
return indexSyncErrorCode;
|
||||
}
|
||||
|
||||
public void setIndexSyncErrorCode(String indexSyncErrorCode) {
|
||||
this.indexSyncErrorCode = indexSyncErrorCode;
|
||||
}
|
||||
|
||||
public String getIndexSyncErrorMessage() {
|
||||
return indexSyncErrorMessage;
|
||||
}
|
||||
|
||||
public void setIndexSyncErrorMessage(String indexSyncErrorMessage) {
|
||||
this.indexSyncErrorMessage = indexSyncErrorMessage;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -2,6 +2,12 @@ package tech.easyflow.ai.mapper;
|
||||
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import com.mybatisflex.core.BaseMapper;
|
||||
import org.apache.ibatis.annotations.Param;
|
||||
import org.apache.ibatis.annotations.Select;
|
||||
import org.apache.ibatis.annotations.Update;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 映射层。
|
||||
@@ -11,4 +17,23 @@ import com.mybatisflex.core.BaseMapper;
|
||||
*/
|
||||
public interface DocumentChunkMapper extends BaseMapper<DocumentChunk> {
|
||||
|
||||
@Update("UPDATE tb_document_chunk SET index_sync_status=#{status}, "
|
||||
+ "index_sync_error_code=#{errorCode}, index_sync_error_message=#{errorMessage} "
|
||||
+ "WHERE id=#{id} AND index_sync_version=#{version}")
|
||||
int updateSyncState(@Param("id") BigInteger id,
|
||||
@Param("version") long version,
|
||||
@Param("status") String status,
|
||||
@Param("errorCode") String errorCode,
|
||||
@Param("errorMessage") String errorMessage);
|
||||
|
||||
@Select("<script>SELECT id, document_id AS documentId, "
|
||||
+ "document_collection_id AS documentCollectionId, "
|
||||
+ "index_sync_status AS indexSyncStatus, index_sync_version AS indexSyncVersion, "
|
||||
+ "index_sync_error_code AS indexSyncErrorCode, "
|
||||
+ "index_sync_error_message AS indexSyncErrorMessage "
|
||||
+ "FROM tb_document_chunk WHERE document_id=#{documentId} AND id IN "
|
||||
+ "<foreach collection='ids' item='id' open='(' separator=',' close=')'>#{id}</foreach>"
|
||||
+ "</script>")
|
||||
List<DocumentChunk> selectSyncStates(@Param("documentId") BigInteger documentId,
|
||||
@Param("ids") List<BigInteger> ids);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
package tech.easyflow.ai.mapper;
|
||||
|
||||
import com.mybatisflex.core.BaseMapper;
|
||||
import org.apache.ibatis.annotations.Param;
|
||||
import org.apache.ibatis.annotations.Select;
|
||||
import org.apache.ibatis.annotations.Update;
|
||||
import tech.easyflow.ai.entity.DocumentChunkSyncTask;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.Date;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 文档分块索引同步任务映射层。
|
||||
*/
|
||||
public interface DocumentChunkSyncTaskMapper extends BaseMapper<DocumentChunkSyncTask> {
|
||||
|
||||
String SELECT_COLUMNS = "id, chunk_id AS chunkId, document_id AS documentId, "
|
||||
+ "document_collection_id AS documentCollectionId, "
|
||||
+ "vector_collection AS vectorCollection, operation, sync_version AS syncVersion, "
|
||||
+ "status, attempt_count AS attemptCount, next_retry_at AS nextRetryAt, "
|
||||
+ "last_dispatched_at AS lastDispatchedAt, execution_token AS executionToken, "
|
||||
+ "lease_until AS leaseUntil, error_code AS errorCode, error_message AS errorMessage, "
|
||||
+ "created, created_by AS createdBy, modified, modified_by AS modifiedBy";
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status='SUPERSEDED', "
|
||||
+ "execution_token=NULL, lease_until=NULL, modified=#{now} "
|
||||
+ "WHERE chunk_id=#{chunkId} AND sync_version < #{syncVersion} "
|
||||
+ "AND status IN ('PENDING','RUNNING','FAILED')")
|
||||
int supersedeOlder(@Param("chunkId") BigInteger chunkId,
|
||||
@Param("syncVersion") long syncVersion,
|
||||
@Param("now") Date now);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status='SUPERSEDED', "
|
||||
+ "execution_token=NULL, lease_until=NULL, modified=#{now} "
|
||||
+ "WHERE chunk_id=#{chunkId} AND status IN ('PENDING','RUNNING','FAILED')")
|
||||
int supersedeChunk(@Param("chunkId") BigInteger chunkId,
|
||||
@Param("now") Date now);
|
||||
|
||||
/**
|
||||
* 锁定同一分块的任务版本范围，保证外部索引写入与新版本建单按版本串行。
|
||||
*/
|
||||
@Select("SELECT id FROM tb_document_chunk_sync_task WHERE chunk_id=#{chunkId} "
|
||||
+ "ORDER BY sync_version, id FOR UPDATE")
|
||||
List<BigInteger> lockChunkTasks(@Param("chunkId") BigInteger chunkId);
|
||||
|
||||
@Select("SELECT " + SELECT_COLUMNS + " FROM tb_document_chunk_sync_task "
|
||||
+ "WHERE status='PENDING' AND next_retry_at <= #{now} "
|
||||
+ "AND (last_dispatched_at IS NULL OR last_dispatched_at <= #{redispatchBefore}) "
|
||||
+ "ORDER BY next_retry_at, id LIMIT #{limit}")
|
||||
List<DocumentChunkSyncTask> selectPendingDue(@Param("now") Date now,
|
||||
@Param("redispatchBefore") Date redispatchBefore,
|
||||
@Param("limit") int limit);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET last_dispatched_at=#{now}, modified=#{now} "
|
||||
+ "WHERE id=#{id} AND status='PENDING' AND next_retry_at <= #{now} "
|
||||
+ "AND (last_dispatched_at IS NULL OR last_dispatched_at <= #{redispatchBefore})")
|
||||
int markDispatched(@Param("id") BigInteger id,
|
||||
@Param("now") Date now,
|
||||
@Param("redispatchBefore") Date redispatchBefore);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status='PENDING', "
|
||||
+ "execution_token=NULL, lease_until=NULL, next_retry_at=#{now}, "
|
||||
+ "last_dispatched_at=NULL, modified=#{now} "
|
||||
+ "WHERE status='RUNNING' AND lease_until <= #{now}")
|
||||
int recoverExpired(@Param("now") Date now);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status='RUNNING', "
|
||||
+ "attempt_count=attempt_count + 1, execution_token=#{token}, "
|
||||
+ "lease_until=#{leaseUntil}, error_code=NULL, error_message=NULL, modified=#{now} "
|
||||
+ "WHERE id=#{id} AND status='PENDING' AND next_retry_at <= #{now}")
|
||||
int claim(@Param("id") BigInteger id,
|
||||
@Param("token") String token,
|
||||
@Param("leaseUntil") Date leaseUntil,
|
||||
@Param("now") Date now);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status=#{status}, "
|
||||
+ "execution_token=NULL, lease_until=NULL, error_code=#{errorCode}, "
|
||||
+ "error_message=#{errorMessage}, modified=#{now} "
|
||||
+ "WHERE id=#{id} AND status='RUNNING' AND execution_token=#{token}")
|
||||
int finishOwned(@Param("id") BigInteger id,
|
||||
@Param("token") String token,
|
||||
@Param("status") String status,
|
||||
@Param("errorCode") String errorCode,
|
||||
@Param("errorMessage") String errorMessage,
|
||||
@Param("now") Date now);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status=#{status}, "
|
||||
+ "execution_token=NULL, lease_until=NULL, next_retry_at=#{nextRetryAt}, "
|
||||
+ "last_dispatched_at=NULL, error_code=#{errorCode}, "
|
||||
+ "error_message=#{errorMessage}, modified=#{now} "
|
||||
+ "WHERE id=#{id} AND status='RUNNING' AND execution_token=#{token}")
|
||||
int failOrRetryOwned(@Param("id") BigInteger id,
|
||||
@Param("token") String token,
|
||||
@Param("status") String status,
|
||||
@Param("nextRetryAt") Date nextRetryAt,
|
||||
@Param("errorCode") String errorCode,
|
||||
@Param("errorMessage") String errorMessage,
|
||||
@Param("now") Date now);
|
||||
|
||||
@Update("UPDATE tb_document_chunk_sync_task SET status='PENDING', "
|
||||
+ "attempt_count=0, next_retry_at=#{now}, last_dispatched_at=NULL, "
|
||||
+ "execution_token=NULL, lease_until=NULL, error_code=NULL, error_message=NULL, modified=#{now} "
|
||||
+ "WHERE id=#{id} AND chunk_id=#{chunkId} AND sync_version=#{syncVersion} "
|
||||
+ "AND status='FAILED'")
|
||||
int retryFailed(@Param("id") BigInteger id,
|
||||
@Param("chunkId") BigInteger chunkId,
|
||||
@Param("syncVersion") long syncVersion,
|
||||
@Param("now") Date now);
|
||||
|
||||
@Select("SELECT " + SELECT_COLUMNS + " FROM tb_document_chunk_sync_task WHERE chunk_id=#{chunkId} "
|
||||
+ "AND sync_version=#{syncVersion} AND status='FAILED' ORDER BY id DESC LIMIT 1")
|
||||
DocumentChunkSyncTask selectFailed(@Param("chunkId") BigInteger chunkId,
|
||||
@Param("syncVersion") long syncVersion);
|
||||
}
|
||||
@@ -1,10 +1,12 @@
|
||||
package tech.easyflow.ai.service;
|
||||
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import com.mybatisflex.core.service.IService;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.dto.DocumentChunkDeleteResult;
|
||||
import tech.easyflow.ai.dto.DocumentChunkSyncStatus;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 服务层。
|
||||
@@ -14,5 +16,36 @@ import java.math.BigInteger;
|
||||
*/
|
||||
public interface DocumentChunkService extends IService<DocumentChunk> {
|
||||
|
||||
boolean removeChunk(DocumentCollection knowledge, BigInteger chunkId);
|
||||
/**
|
||||
* 更新分块的 Markdown 正文，并创建后台检索索引同步任务。
|
||||
*
|
||||
* @param knowledgeId 知识库 ID
|
||||
* @param chunkId 分块 ID
|
||||
* @param markdown Markdown 正文
|
||||
* @return 更新后的分块
|
||||
*/
|
||||
DocumentChunk updateContent(BigInteger knowledgeId, BigInteger chunkId, String markdown);
|
||||
|
||||
/**
|
||||
* 删除分块，并创建后台检索索引清理任务。
|
||||
*
|
||||
* @param knowledgeId 知识库 ID
|
||||
* @param chunkId 分块 ID
|
||||
* @return 删除结果
|
||||
*/
|
||||
DocumentChunkDeleteResult deleteChunk(BigInteger knowledgeId, BigInteger chunkId);
|
||||
|
||||
/**
|
||||
* 手动重试当前版本的检索索引同步。
|
||||
*/
|
||||
DocumentChunk retryIndexSync(BigInteger knowledgeId, BigInteger chunkId, long syncVersion);
|
||||
|
||||
/**
|
||||
* 批量读取文档内分块的索引同步状态。
|
||||
*/
|
||||
List<DocumentChunkSyncStatus> listIndexSyncStatus(
|
||||
BigInteger knowledgeId,
|
||||
BigInteger documentId,
|
||||
List<BigInteger> chunkIds
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,36 +1,365 @@
|
||||
package tech.easyflow.ai.service.impl;
|
||||
|
||||
import com.easyagents.search.engine.service.DocumentSearcher;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.service.DocumentChunkService;
|
||||
import com.easyagents.rag.core.BgeM3ChunkSafety;
|
||||
import com.easyagents.rag.core.RagDefaults;
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import com.mybatisflex.spring.service.impl.ServiceImpl;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.support.TransactionTemplate;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncState;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncTaskAppService;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
import tech.easyflow.ai.dto.DocumentChunkDeleteResult;
|
||||
import tech.easyflow.ai.dto.DocumentChunkSyncStatus;
|
||||
import tech.easyflow.ai.entity.Document;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentChunkSyncTask;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentChunkService;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.time.Duration;
|
||||
import java.util.Date;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
/**
|
||||
* 服务层实现。
|
||||
*
|
||||
* @author michael
|
||||
* @since 2024-08-23
|
||||
* 分块服务层实现。
|
||||
*/
|
||||
@Service
|
||||
public class DocumentChunkServiceImpl extends ServiceImpl<DocumentChunkMapper, DocumentChunk> implements DocumentChunkService {
|
||||
public class DocumentChunkServiceImpl
|
||||
extends ServiceImpl<DocumentChunkMapper, DocumentChunk>
|
||||
implements DocumentChunkService {
|
||||
|
||||
@Autowired
|
||||
private SearcherFactory searcherFactory;
|
||||
public static final int DOCUMENT_CHUNK_EMPTY_REQUIRES_DELETE = 42901;
|
||||
public static final int DOCUMENT_CHUNK_LOCK_UNAVAILABLE = 42905;
|
||||
|
||||
private static final Duration LOCK_WAIT_TIMEOUT = Duration.ofSeconds(5);
|
||||
private static final Duration LOCK_LEASE_TIMEOUT = Duration.ofSeconds(30);
|
||||
private static final Pattern MARKDOWN_IMAGE_PATTERN = Pattern.compile(
|
||||
"!\\[([^\\]]*)\\]\\((?:[^()\\r\\n]|\\([^()\\r\\n]*\\))*\\)"
|
||||
);
|
||||
private static final Pattern HTML_IMAGE_PATTERN = Pattern.compile(
|
||||
"<img\\b[^>]*>",
|
||||
Pattern.CASE_INSENSITIVE
|
||||
);
|
||||
private static final Pattern HTML_ALT_PATTERN = Pattern.compile(
|
||||
"\\balt\\s*=\\s*(?:\"([^\"]*)\"|'([^']*)')",
|
||||
Pattern.CASE_INSENSITIVE
|
||||
);
|
||||
|
||||
private final DocumentChunkMapper documentChunkMapper;
|
||||
private final DocumentMapper documentMapper;
|
||||
private final DocumentCollectionService documentCollectionService;
|
||||
private final DocumentChunkSyncTaskAppService syncTaskAppService;
|
||||
private final TransactionTemplate transactionTemplate;
|
||||
private final RedisLockExecutor redisLockExecutor;
|
||||
|
||||
public DocumentChunkServiceImpl(
|
||||
DocumentChunkMapper documentChunkMapper,
|
||||
DocumentMapper documentMapper,
|
||||
DocumentCollectionService documentCollectionService,
|
||||
DocumentChunkSyncTaskAppService syncTaskAppService,
|
||||
PlatformTransactionManager transactionManager,
|
||||
RedisLockExecutor redisLockExecutor
|
||||
) {
|
||||
this.documentChunkMapper = documentChunkMapper;
|
||||
this.documentMapper = documentMapper;
|
||||
this.documentCollectionService = documentCollectionService;
|
||||
this.syncTaskAppService = syncTaskAppService;
|
||||
this.transactionTemplate = new TransactionTemplate(transactionManager);
|
||||
this.redisLockExecutor = redisLockExecutor;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean removeChunk(DocumentCollection knowledge, BigInteger chunkId) {
|
||||
DocumentSearcher searcher = searcherFactory.getSearcher();
|
||||
// 删除搜索引擎中的数据
|
||||
if (searcher == null){
|
||||
return true;
|
||||
public DocumentChunk updateContent(
|
||||
BigInteger knowledgeId,
|
||||
BigInteger chunkId,
|
||||
String markdown
|
||||
) {
|
||||
if (markdown == null || markdown.trim().isEmpty()) {
|
||||
throw new BusinessException(
|
||||
422,
|
||||
DOCUMENT_CHUNK_EMPTY_REQUIRES_DELETE,
|
||||
"分块内容为空，请删除该分块"
|
||||
);
|
||||
}
|
||||
return searcher.deleteDocument(chunkId);
|
||||
DocumentChunk snapshot = requireChunk(knowledgeId, chunkId);
|
||||
UpdateOutcome outcome = withDocumentLock(
|
||||
snapshot.getDocumentId(),
|
||||
() -> updateContentLocked(knowledgeId, chunkId, markdown)
|
||||
);
|
||||
syncTaskAppService.dispatchBestEffort(outcome.taskId());
|
||||
return outcome.chunk();
|
||||
}
|
||||
|
||||
private UpdateOutcome updateContentLocked(
|
||||
BigInteger knowledgeId,
|
||||
BigInteger chunkId,
|
||||
String markdown
|
||||
) {
|
||||
DocumentChunk current = requireChunk(knowledgeId, chunkId);
|
||||
DocumentCollection collection = requireKnowledge(knowledgeId);
|
||||
String searchableContent = toSearchableContent(markdown);
|
||||
assertWithinEmbeddingLimit(chunkId, searchableContent);
|
||||
Map<String, Object> options = current.getOptions() == null
|
||||
? new HashMap<>()
|
||||
: new HashMap<>(current.getOptions());
|
||||
options.put(DocumentImportKeys.KEY_DOCUMENT_RENDER_MARKDOWN, markdown);
|
||||
long nextVersion = current.getIndexSyncVersion() == null
|
||||
? 1L
|
||||
: current.getIndexSyncVersion() + 1L;
|
||||
|
||||
UpdateOutcome outcome = transactionTemplate.execute(status -> {
|
||||
DocumentChunkSyncTask task = syncTaskAppService.createTask(
|
||||
current,
|
||||
collection,
|
||||
DocumentChunkSyncState.OPERATION_UPSERT,
|
||||
nextVersion,
|
||||
new Date()
|
||||
);
|
||||
DocumentChunk update = new DocumentChunk();
|
||||
update.setId(chunkId);
|
||||
update.setContent(searchableContent);
|
||||
update.setOptions(options);
|
||||
update.setIndexSyncStatus(DocumentChunkSyncState.PENDING);
|
||||
update.setIndexSyncVersion(nextVersion);
|
||||
update.setIndexSyncErrorCode(null);
|
||||
update.setIndexSyncErrorMessage(null);
|
||||
if (documentChunkMapper.update(update) <= 0) {
|
||||
throw new BusinessException("分块更新失败");
|
||||
}
|
||||
touchDocument(current.getDocumentId(), null);
|
||||
current.setContent(searchableContent);
|
||||
current.setOptions(options);
|
||||
current.setIndexSyncStatus(DocumentChunkSyncState.PENDING);
|
||||
current.setIndexSyncVersion(nextVersion);
|
||||
current.setIndexSyncErrorCode(null);
|
||||
current.setIndexSyncErrorMessage(null);
|
||||
return new UpdateOutcome(current, task.getId());
|
||||
});
|
||||
if (outcome == null) {
|
||||
throw new BusinessException("分块更新失败");
|
||||
}
|
||||
return outcome;
|
||||
}
|
||||
|
||||
@Override
|
||||
public DocumentChunkDeleteResult deleteChunk(BigInteger knowledgeId, BigInteger chunkId) {
|
||||
DocumentChunk snapshot = requireChunk(knowledgeId, chunkId);
|
||||
DeleteOutcome outcome = withDocumentLock(
|
||||
snapshot.getDocumentId(),
|
||||
() -> deleteChunkLocked(knowledgeId, chunkId)
|
||||
);
|
||||
syncTaskAppService.dispatchBestEffort(outcome.taskId());
|
||||
return outcome.result();
|
||||
}
|
||||
|
||||
private <T> T withDocumentLock(BigInteger documentId, Supplier<T> action) {
|
||||
RedisLockExecutor.LockHandle handle;
|
||||
try {
|
||||
handle = redisLockExecutor.tryAcquire(
|
||||
DocumentChunkSyncState.parentLockKey(documentId),
|
||||
LOCK_WAIT_TIMEOUT,
|
||||
LOCK_LEASE_TIMEOUT
|
||||
);
|
||||
} catch (RuntimeException exception) {
|
||||
throw new BusinessException(
|
||||
503,
|
||||
DOCUMENT_CHUNK_LOCK_UNAVAILABLE,
|
||||
"分块正在处理中，请稍后重试",
|
||||
exception
|
||||
);
|
||||
}
|
||||
if (handle == null) {
|
||||
throw new BusinessException(
|
||||
503,
|
||||
DOCUMENT_CHUNK_LOCK_UNAVAILABLE,
|
||||
"分块正在处理中，请稍后重试"
|
||||
);
|
||||
}
|
||||
try (handle) {
|
||||
return action.get();
|
||||
}
|
||||
}
|
||||
|
||||
private DeleteOutcome deleteChunkLocked(BigInteger knowledgeId, BigInteger chunkId) {
|
||||
DocumentChunk current = requireChunk(knowledgeId, chunkId);
|
||||
DocumentCollection collection = requireKnowledge(knowledgeId);
|
||||
long nextVersion = current.getIndexSyncVersion() == null
|
||||
? 1L
|
||||
: current.getIndexSyncVersion() + 1L;
|
||||
DeleteOutcome outcome = transactionTemplate.execute(status -> {
|
||||
DocumentChunkSyncTask task = syncTaskAppService.createTask(
|
||||
current,
|
||||
collection,
|
||||
DocumentChunkSyncState.OPERATION_DELETE,
|
||||
nextVersion,
|
||||
new Date()
|
||||
);
|
||||
if (documentChunkMapper.deleteById(chunkId) <= 0) {
|
||||
throw new BusinessException("分块删除失败");
|
||||
}
|
||||
long remaining = documentChunkMapper.selectCountByQuery(
|
||||
QueryWrapper.create().eq(DocumentChunk::getDocumentId, current.getDocumentId())
|
||||
);
|
||||
touchDocument(current.getDocumentId(), remaining);
|
||||
return new DeleteOutcome(
|
||||
new DocumentChunkDeleteResult(chunkId, current.getDocumentId(), remaining),
|
||||
task.getId()
|
||||
);
|
||||
});
|
||||
if (outcome == null) {
|
||||
throw new BusinessException("分块删除失败");
|
||||
}
|
||||
return outcome;
|
||||
}
|
||||
|
||||
@Override
|
||||
public DocumentChunk retryIndexSync(
|
||||
BigInteger knowledgeId,
|
||||
BigInteger chunkId,
|
||||
long syncVersion
|
||||
) {
|
||||
requireChunk(knowledgeId, chunkId);
|
||||
return syncTaskAppService.retry(chunkId, syncVersion);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<DocumentChunkSyncStatus> listIndexSyncStatus(
|
||||
BigInteger knowledgeId,
|
||||
BigInteger documentId,
|
||||
List<BigInteger> chunkIds
|
||||
) {
|
||||
if (documentId == null || chunkIds == null || chunkIds.isEmpty()) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
BigInteger effectiveKnowledgeId = knowledgeId;
|
||||
if (effectiveKnowledgeId == null) {
|
||||
Document document = documentMapper.selectOneById(documentId);
|
||||
if (document == null || document.getCollectionId() == null) {
|
||||
throw new BusinessException("文档不存在");
|
||||
}
|
||||
effectiveKnowledgeId = document.getCollectionId();
|
||||
}
|
||||
if (chunkIds.size() > 200) {
|
||||
throw new BusinessException("单次最多查询 200 个分块状态");
|
||||
}
|
||||
List<DocumentChunk> chunks = documentChunkMapper.selectSyncStates(documentId, chunkIds);
|
||||
for (DocumentChunk chunk : chunks) {
|
||||
if (chunk.getDocumentCollectionId() == null
|
||||
|| chunk.getDocumentCollectionId().compareTo(effectiveKnowledgeId) != 0) {
|
||||
throw new BusinessException("分块不存在");
|
||||
}
|
||||
}
|
||||
return chunks.stream().map(chunk -> new DocumentChunkSyncStatus(
|
||||
chunk.getId(),
|
||||
chunk.getIndexSyncStatus(),
|
||||
chunk.getIndexSyncVersion(),
|
||||
chunk.getIndexSyncErrorCode(),
|
||||
chunk.getIndexSyncErrorMessage()
|
||||
)).toList();
|
||||
}
|
||||
|
||||
private DocumentChunk requireChunk(BigInteger knowledgeId, BigInteger chunkId) {
|
||||
if (knowledgeId == null || chunkId == null) {
|
||||
throw new BusinessException("分块不存在");
|
||||
}
|
||||
DocumentChunk chunk = documentChunkMapper.selectOneById(chunkId);
|
||||
if (chunk == null || chunk.getDocumentCollectionId() == null
|
||||
|| chunk.getDocumentCollectionId().compareTo(knowledgeId) != 0) {
|
||||
throw new BusinessException("分块不存在");
|
||||
}
|
||||
return chunk;
|
||||
}
|
||||
|
||||
private DocumentCollection requireKnowledge(BigInteger knowledgeId) {
|
||||
DocumentCollection knowledge = documentCollectionService.getById(knowledgeId);
|
||||
if (knowledge == null) {
|
||||
throw new BusinessException("知识库不存在");
|
||||
}
|
||||
return knowledge;
|
||||
}
|
||||
|
||||
private void touchDocument(BigInteger documentId, Long chunkCount) {
|
||||
if (documentId == null) {
|
||||
return;
|
||||
}
|
||||
Document update = new Document();
|
||||
update.setId(documentId);
|
||||
update.setModified(new Date());
|
||||
if (chunkCount != null) {
|
||||
int count = chunkCount > Integer.MAX_VALUE
|
||||
? Integer.MAX_VALUE
|
||||
: chunkCount.intValue();
|
||||
update.setTotalChunks(count);
|
||||
update.setCompletedChunks(count);
|
||||
update.setFailedChunks(0);
|
||||
update.setProgressPercent(100);
|
||||
}
|
||||
if (documentMapper.update(update) <= 0) {
|
||||
throw new BusinessException("文档状态更新失败");
|
||||
}
|
||||
}
|
||||
|
||||
private void assertWithinEmbeddingLimit(BigInteger chunkId, String content) {
|
||||
int tokenEstimate = BgeM3ChunkSafety.estimateContentTokens(content);
|
||||
if (tokenEstimate > RagDefaults.BGE_M3_HARD_CHUNK_TOKEN_LIMIT) {
|
||||
throw new BusinessException(
|
||||
"分块内容超过向量模型上下文上限，请缩短后保存：chunkId=" + chunkId
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private String toSearchableContent(String markdown) {
|
||||
return replaceHtmlImages(replaceMarkdownImages(markdown)).trim();
|
||||
}
|
||||
|
||||
private String replaceMarkdownImages(String markdown) {
|
||||
Matcher matcher = MARKDOWN_IMAGE_PATTERN.matcher(markdown);
|
||||
StringBuffer output = new StringBuffer();
|
||||
while (matcher.find()) {
|
||||
matcher.appendReplacement(
|
||||
output,
|
||||
Matcher.quoteReplacement(imageSearchText(matcher.group(1)))
|
||||
);
|
||||
}
|
||||
matcher.appendTail(output);
|
||||
return output.toString();
|
||||
}
|
||||
|
||||
private String replaceHtmlImages(String content) {
|
||||
Matcher matcher = HTML_IMAGE_PATTERN.matcher(content);
|
||||
StringBuffer output = new StringBuffer();
|
||||
while (matcher.find()) {
|
||||
Matcher altMatcher = HTML_ALT_PATTERN.matcher(matcher.group());
|
||||
String alt = altMatcher.find()
|
||||
? (altMatcher.group(1) == null ? altMatcher.group(2) : altMatcher.group(1))
|
||||
: null;
|
||||
matcher.appendReplacement(output, Matcher.quoteReplacement(imageSearchText(alt)));
|
||||
}
|
||||
matcher.appendTail(output);
|
||||
return output.toString();
|
||||
}
|
||||
|
||||
private String imageSearchText(String alt) {
|
||||
return alt == null || alt.trim().isEmpty() ? "图片" : alt.trim();
|
||||
}
|
||||
|
||||
private record UpdateOutcome(DocumentChunk chunk, BigInteger taskId) {
|
||||
}
|
||||
|
||||
private record DeleteOutcome(DocumentChunkDeleteResult result, BigInteger taskId) {
|
||||
}
|
||||
}
|
||||
|
||||
@@ -648,6 +648,7 @@ public class DocumentCollectionServiceImpl extends ServiceImpl<DocumentCollectio
|
||||
QueryWrapper chunkWrapper = QueryWrapper.create();
|
||||
chunkWrapper.in(DocumentChunk::getId, chunkIds);
|
||||
chunkWrapper.eq(DocumentChunk::getDocumentCollectionId, documentCollection.getId());
|
||||
chunkWrapper.eq(DocumentChunk::getIndexSyncStatus, "SYNCED");
|
||||
Map<String, DocumentChunk> chunkMap = documentChunkMapper.selectListByQuery(chunkWrapper).stream()
|
||||
.collect(Collectors.toMap(item -> item.getId().toString(), item -> item, (a, b) -> a));
|
||||
if (chunkMap.isEmpty()) {
|
||||
|
||||
@@ -31,6 +31,7 @@ import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.transaction.annotation.Transactional;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncState;
|
||||
import tech.easyflow.common.util.SearchKeywordUtil;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportDtos;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
@@ -40,6 +41,7 @@ import tech.easyflow.ai.documentimport.task.KnowledgeDocumentImportTaskAppServic
|
||||
import tech.easyflow.ai.entity.*;
|
||||
import tech.easyflow.ai.enums.DocumentProcessStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkSyncTaskMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentChunkService;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
@@ -47,6 +49,7 @@ import tech.easyflow.ai.service.DocumentService;
|
||||
import tech.easyflow.ai.service.ModelService;
|
||||
import tech.easyflow.ai.support.DocumentStoreLifecycleSupport;
|
||||
import tech.easyflow.common.ai.rag.ExcelDocumentSplitter;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.domain.Result;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
import tech.easyflow.common.util.FileUtil;
|
||||
@@ -58,8 +61,11 @@ import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.math.BigInteger;
|
||||
import java.math.BigDecimal;
|
||||
import java.time.Duration;
|
||||
import java.util.*;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import static tech.easyflow.ai.entity.DocumentCollection.KEY_CAN_UPDATE_EMBEDDING_MODEL;
|
||||
import static tech.easyflow.ai.entity.table.DocumentChunkTableDef.DOCUMENT_CHUNK;
|
||||
@@ -83,6 +89,9 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
@Resource
|
||||
private DocumentChunkMapper documentChunkMapper;
|
||||
|
||||
@Resource
|
||||
private DocumentChunkSyncTaskMapper documentChunkSyncTaskMapper;
|
||||
|
||||
@Resource
|
||||
private DocumentCollectionService knowledgeService;
|
||||
|
||||
@@ -110,6 +119,9 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
@Autowired
|
||||
private KnowledgeDocumentImportTaskAppService importTaskAppService;
|
||||
|
||||
@Autowired
|
||||
private RedisLockExecutor redisLockExecutor;
|
||||
|
||||
@Override
|
||||
public Page<Document> getDocumentList(String knowledgeId, int pageSize, int pageNum, String fileName) {
|
||||
return queryDocumentList(knowledgeId, pageSize, pageNum, fileName, null, false);
|
||||
@@ -202,6 +214,49 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
@Override
|
||||
@Transactional
|
||||
public boolean removeDoc(String id) {
|
||||
AtomicBoolean removalCompleted = new AtomicBoolean();
|
||||
AtomicReference<Boolean> removalResult = new AtomicReference<>();
|
||||
try {
|
||||
return redisLockExecutor.executeWithRenewingLock(
|
||||
DocumentChunkSyncState.parentLockKey(id),
|
||||
Duration.ofSeconds(5),
|
||||
Duration.ofSeconds(30),
|
||||
() -> {
|
||||
try {
|
||||
boolean result = removeDocLocked(id);
|
||||
removalResult.set(result);
|
||||
removalCompleted.set(true);
|
||||
return result;
|
||||
} catch (BusinessException exception) {
|
||||
throw exception;
|
||||
} catch (RuntimeException exception) {
|
||||
throw new DocumentRemovalExecutionException(exception);
|
||||
}
|
||||
}
|
||||
);
|
||||
} catch (BusinessException exception) {
|
||||
throw exception;
|
||||
} catch (DocumentRemovalExecutionException exception) {
|
||||
throw exception.getOriginalCause();
|
||||
} catch (RuntimeException exception) {
|
||||
if (removalCompleted.get()) {
|
||||
Log.warn(
|
||||
"文档删除已完成但分布式锁在收尾阶段失效，继续提交数据库事务: documentId={}",
|
||||
id,
|
||||
exception
|
||||
);
|
||||
return Boolean.TRUE.equals(removalResult.get());
|
||||
}
|
||||
throw new BusinessException(
|
||||
503,
|
||||
DocumentChunkServiceImpl.DOCUMENT_CHUNK_LOCK_UNAVAILABLE,
|
||||
"文档正在处理中，请稍后重试",
|
||||
exception
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private boolean removeDocLocked(String id) {
|
||||
// 查询该文档对应哪些分割的字段，先删除
|
||||
QueryWrapper queryWrapperDocument = QueryWrapper.create().eq(Document::getId, id);
|
||||
Document oneByQuery = documentMapper.selectOneByQuery(queryWrapperDocument);
|
||||
@@ -218,23 +273,38 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
|
||||
QueryWrapper queryWrapper = QueryWrapper.create()
|
||||
.select(DOCUMENT_CHUNK.ID).eq(DocumentChunk::getDocumentId, id);
|
||||
List<BigInteger> chunkIds = documentChunkMapper.selectListByQueryAs(
|
||||
queryWrapper,
|
||||
BigInteger.class
|
||||
List<BigInteger> chunkIds = new ArrayList<>(
|
||||
documentChunkMapper.selectListByQueryAs(queryWrapper, BigInteger.class)
|
||||
);
|
||||
chunkIds.sort(Comparator.naturalOrder());
|
||||
DocumentStore documentStore = null;
|
||||
try {
|
||||
Model model = null;
|
||||
if (!chunkIds.isEmpty()) {
|
||||
documentStore = knowledge.toDocumentStore();
|
||||
if (documentStore == null) {
|
||||
return false;
|
||||
throw new BusinessException("文档向量存储不可用");
|
||||
}
|
||||
Model model = modelService.getById(
|
||||
model = modelService.getById(
|
||||
knowledge.getVectorEmbedModelId()
|
||||
);
|
||||
if (model == null) {
|
||||
return false;
|
||||
throw new BusinessException("文档向量模型不存在");
|
||||
}
|
||||
}
|
||||
Date supersededAt = new Date();
|
||||
for (BigInteger chunkId : chunkIds) {
|
||||
redisLockExecutor.executeWithRenewingLock(
|
||||
DocumentChunkSyncState.syncLockKey(chunkId),
|
||||
Duration.ofSeconds(5),
|
||||
Duration.ofSeconds(30),
|
||||
() -> {
|
||||
documentChunkSyncTaskMapper.supersedeChunk(chunkId, supersededAt);
|
||||
return null;
|
||||
}
|
||||
);
|
||||
}
|
||||
if (!chunkIds.isEmpty()) {
|
||||
StoreOptions options = StoreOptions.ofCollectionName(
|
||||
knowledge.getVectorStoreCollection()
|
||||
);
|
||||
@@ -253,7 +323,16 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
// 删除搜索引擎中的数据
|
||||
DocumentSearcher searcher = searcherFactory.getSearcher();
|
||||
if (searcher != null) {
|
||||
chunkIds.forEach(searcher::deleteDocument);
|
||||
for (BigInteger chunkId : chunkIds) {
|
||||
if (!searcher.deleteDocument(chunkId)) {
|
||||
Log.error(
|
||||
"删除文档关键词索引失败: documentId={}, chunkId={}",
|
||||
id,
|
||||
chunkId
|
||||
);
|
||||
throw new BusinessException("文档关键词索引删除失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
int ck = documentChunkMapper.deleteByQuery(QueryWrapper.create().eq(DocumentChunk::getDocumentId, id));
|
||||
if (ck < 0) {
|
||||
@@ -1099,6 +1178,17 @@ public class DocumentServiceImpl extends ServiceImpl<DocumentMapper, Document> i
|
||||
}
|
||||
}
|
||||
|
||||
private static final class DocumentRemovalExecutionException extends RuntimeException {
|
||||
|
||||
private DocumentRemovalExecutionException(RuntimeException cause) {
|
||||
super(cause);
|
||||
}
|
||||
|
||||
private RuntimeException getOriginalCause() {
|
||||
return (RuntimeException) getCause();
|
||||
}
|
||||
}
|
||||
|
||||
public DocumentSplitter getDocumentSplitter(DocumentCollectionSplitParams params) {
|
||||
String splitterName = params.getSplitterName();
|
||||
int chunkSize = params.getChunkSize();
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
package tech.easyflow.ai.documentchunk;
|
||||
|
||||
import org.junit.Test;
|
||||
import org.mockito.InOrder;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.support.SimpleTransactionStatus;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
import tech.easyflow.ai.config.AiMilvusConfig;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentChunkSyncTask;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkSyncTaskMapper;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.ai.service.ModelService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
/**
|
||||
* 分块索引同步任务恢复与版本语义测试。
|
||||
*/
|
||||
public class DocumentChunkSyncTaskAppServiceTest {
|
||||
|
||||
@Test
|
||||
public void dispatchPendingShouldRecoverExpiredTasksAndSendClaimedRows() {
|
||||
Fixture fixture = fixture(1);
|
||||
Mockito.when(fixture.taskMapper.selectPendingDue(
|
||||
Mockito.any(), Mockito.any(), Mockito.anyInt()
|
||||
)).thenReturn(List.of(fixture.task));
|
||||
Mockito.when(fixture.taskMapper.markDispatched(
|
||||
Mockito.eq(fixture.taskId), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(1);
|
||||
|
||||
fixture.service.dispatchPendingTasks();
|
||||
|
||||
Mockito.verify(fixture.taskMapper).recoverExpired(Mockito.any());
|
||||
Mockito.verify(fixture.producer).send(fixture.taskId);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void outdatedTaskShouldBeSupersededWithoutTouchingIndexes() {
|
||||
Fixture fixture = fixture(1);
|
||||
fixture.chunk.setIndexSyncVersion(2L);
|
||||
|
||||
fixture.service.handleTask(fixture.taskId);
|
||||
|
||||
Mockito.verify(fixture.taskMapper).finishOwned(
|
||||
Mockito.eq(fixture.taskId), Mockito.anyString(),
|
||||
Mockito.eq(DocumentChunkSyncState.TASK_SUPERSEDED),
|
||||
Mockito.isNull(), Mockito.isNull(), Mockito.any()
|
||||
);
|
||||
Mockito.verifyNoInteractions(fixture.collectionService);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void transientFailureShouldReturnTaskToPending() {
|
||||
Fixture fixture = fixture(1);
|
||||
|
||||
fixture.service.handleTask(fixture.taskId);
|
||||
|
||||
InOrder order = Mockito.inOrder(
|
||||
fixture.taskMapper,
|
||||
fixture.collectionService
|
||||
);
|
||||
order.verify(fixture.taskMapper).lockChunkTasks(fixture.chunkId);
|
||||
order.verify(fixture.collectionService).getById(Mockito.any());
|
||||
Mockito.verify(fixture.taskMapper).failOrRetryOwned(
|
||||
Mockito.eq(fixture.taskId), Mockito.anyString(),
|
||||
Mockito.eq(DocumentChunkSyncState.PENDING), Mockito.any(),
|
||||
Mockito.eq("INDEX_SOURCE_MISSING"), Mockito.anyString(), Mockito.any()
|
||||
);
|
||||
Mockito.verify(fixture.chunkMapper, Mockito.never()).updateSyncState(
|
||||
Mockito.any(), Mockito.anyLong(), Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void workerShouldLockChunkTaskRangeBeforeReadingOwnedTask() {
|
||||
Fixture fixture = fixture(1);
|
||||
|
||||
fixture.service.handleTask(fixture.taskId);
|
||||
|
||||
InOrder order = Mockito.inOrder(fixture.taskMapper);
|
||||
order.verify(fixture.taskMapper).selectOneById(fixture.taskId);
|
||||
order.verify(fixture.taskMapper).claim(
|
||||
Mockito.eq(fixture.taskId), Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
);
|
||||
order.verify(fixture.taskMapper).lockChunkTasks(fixture.chunkId);
|
||||
order.verify(fixture.taskMapper).selectOneById(fixture.taskId);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void exhaustedFailureShouldExposeFailedStateOnCurrentVersion() {
|
||||
Fixture fixture = fixture(5);
|
||||
|
||||
fixture.service.handleTask(fixture.taskId);
|
||||
|
||||
Mockito.verify(fixture.taskMapper).failOrRetryOwned(
|
||||
Mockito.eq(fixture.taskId), Mockito.anyString(),
|
||||
Mockito.eq(DocumentChunkSyncState.FAILED), Mockito.any(),
|
||||
Mockito.eq("INDEX_SOURCE_MISSING"), Mockito.anyString(), Mockito.any()
|
||||
);
|
||||
Mockito.verify(fixture.chunkMapper).updateSyncState(
|
||||
fixture.chunkId, 1L, DocumentChunkSyncState.FAILED,
|
||||
"INDEX_SOURCE_MISSING", "分块或知识库不存在"
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void exhaustedDeleteFailureShouldRemainPendingForDurableCleanup() {
|
||||
Fixture fixture = fixture(5);
|
||||
fixture.task.setOperation(DocumentChunkSyncState.OPERATION_DELETE);
|
||||
Mockito.when(fixture.chunkMapper.selectOneById(fixture.chunkId)).thenReturn(null);
|
||||
|
||||
fixture.service.handleTask(fixture.taskId);
|
||||
|
||||
Mockito.verify(fixture.taskMapper).failOrRetryOwned(
|
||||
Mockito.eq(fixture.taskId), Mockito.anyString(),
|
||||
Mockito.eq(DocumentChunkSyncState.PENDING), Mockito.any(),
|
||||
Mockito.eq("INDEX_SYNC_FAILED"), Mockito.anyString(), Mockito.any()
|
||||
);
|
||||
Mockito.verify(fixture.chunkMapper, Mockito.never()).updateSyncState(
|
||||
Mockito.any(), Mockito.anyLong(), Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void manualRetryShouldResetTaskAndChunkBeforeDispatch() {
|
||||
Fixture fixture = fixture(5);
|
||||
fixture.chunk.setIndexSyncStatus(DocumentChunkSyncState.FAILED);
|
||||
fixture.task.setStatus(DocumentChunkSyncState.FAILED);
|
||||
Mockito.when(fixture.taskMapper.selectFailed(fixture.chunkId, 1L))
|
||||
.thenReturn(fixture.task);
|
||||
Mockito.when(fixture.taskMapper.retryFailed(
|
||||
Mockito.eq(fixture.taskId), Mockito.eq(fixture.chunkId),
|
||||
Mockito.eq(1L), Mockito.any()
|
||||
)).thenReturn(1);
|
||||
Mockito.when(fixture.chunkMapper.updateSyncState(
|
||||
fixture.chunkId, 1L, DocumentChunkSyncState.PENDING, null, null
|
||||
)).thenReturn(1);
|
||||
|
||||
fixture.service.retry(fixture.chunkId, 1L);
|
||||
|
||||
Mockito.verify(fixture.producer).send(fixture.taskId);
|
||||
}
|
||||
|
||||
private static Fixture fixture(int attemptCount) {
|
||||
BigInteger taskId = BigInteger.valueOf(10);
|
||||
BigInteger chunkId = BigInteger.valueOf(20);
|
||||
DocumentChunkSyncTask task = new DocumentChunkSyncTask();
|
||||
task.setId(taskId);
|
||||
task.setChunkId(chunkId);
|
||||
task.setDocumentId(BigInteger.valueOf(30));
|
||||
task.setDocumentCollectionId(BigInteger.valueOf(40));
|
||||
task.setVectorCollection("kb-test");
|
||||
task.setOperation(DocumentChunkSyncState.OPERATION_UPSERT);
|
||||
task.setSyncVersion(1L);
|
||||
task.setStatus(DocumentChunkSyncState.PENDING);
|
||||
task.setAttemptCount(attemptCount);
|
||||
|
||||
DocumentChunk chunk = new DocumentChunk();
|
||||
chunk.setId(chunkId);
|
||||
chunk.setDocumentCollectionId(task.getDocumentCollectionId());
|
||||
chunk.setIndexSyncVersion(1L);
|
||||
chunk.setIndexSyncStatus(DocumentChunkSyncState.PENDING);
|
||||
|
||||
DocumentChunkSyncTaskMapper taskMapper = Mockito.mock(DocumentChunkSyncTaskMapper.class);
|
||||
Mockito.when(taskMapper.selectOneById(taskId)).thenReturn(task);
|
||||
Mockito.when(taskMapper.claim(
|
||||
Mockito.eq(taskId), Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenAnswer(invocation -> {
|
||||
task.setStatus(DocumentChunkSyncState.TASK_RUNNING);
|
||||
task.setExecutionToken(invocation.getArgument(1));
|
||||
return 1;
|
||||
});
|
||||
Mockito.when(taskMapper.lockChunkTasks(chunkId)).thenReturn(List.of(taskId));
|
||||
Mockito.when(taskMapper.failOrRetryOwned(
|
||||
Mockito.any(), Mockito.anyString(), Mockito.anyString(), Mockito.any(),
|
||||
Mockito.any(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(1);
|
||||
Mockito.when(taskMapper.finishOwned(
|
||||
Mockito.any(), Mockito.anyString(), Mockito.anyString(), Mockito.any(),
|
||||
Mockito.any(), Mockito.any()
|
||||
)).thenReturn(1);
|
||||
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectOneById(chunkId)).thenReturn(chunk);
|
||||
DocumentCollectionService collectionService = Mockito.mock(DocumentCollectionService.class);
|
||||
ModelService modelService = Mockito.mock(ModelService.class);
|
||||
SearcherFactory searcherFactory = Mockito.mock(SearcherFactory.class);
|
||||
DocumentChunkSyncTaskProducer producer = Mockito.mock(DocumentChunkSyncTaskProducer.class);
|
||||
PlatformTransactionManager transactionManager = Mockito.mock(PlatformTransactionManager.class);
|
||||
Mockito.when(transactionManager.getTransaction(Mockito.any()))
|
||||
.thenReturn(new SimpleTransactionStatus());
|
||||
RedisLockExecutor redisLockExecutor = Mockito.mock(RedisLockExecutor.class);
|
||||
Mockito.doAnswer(invocation -> invocation.<Supplier<?>>getArgument(3).get())
|
||||
.when(redisLockExecutor).executeWithRenewingLock(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.<Supplier<?>>any()
|
||||
);
|
||||
|
||||
DocumentChunkSyncTaskAppService service = new DocumentChunkSyncTaskAppService(
|
||||
taskMapper,
|
||||
chunkMapper,
|
||||
collectionService,
|
||||
modelService,
|
||||
searcherFactory,
|
||||
producer,
|
||||
transactionManager,
|
||||
redisLockExecutor,
|
||||
Mockito.mock(AiMilvusConfig.class),
|
||||
Mockito.mock(ObjectProvider.class)
|
||||
);
|
||||
return new Fixture(service, taskMapper, chunkMapper, collectionService,
|
||||
producer, task, chunk, taskId, chunkId);
|
||||
}
|
||||
|
||||
private record Fixture(
|
||||
DocumentChunkSyncTaskAppService service,
|
||||
DocumentChunkSyncTaskMapper taskMapper,
|
||||
DocumentChunkMapper chunkMapper,
|
||||
DocumentCollectionService collectionService,
|
||||
DocumentChunkSyncTaskProducer producer,
|
||||
DocumentChunkSyncTask task,
|
||||
DocumentChunk chunk,
|
||||
BigInteger taskId,
|
||||
BigInteger chunkId
|
||||
) {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
package tech.easyflow.ai.service.impl;
|
||||
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.InOrder;
|
||||
import org.mockito.Mockito;
|
||||
import org.springframework.transaction.PlatformTransactionManager;
|
||||
import org.springframework.transaction.support.SimpleTransactionStatus;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncState;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncTaskAppService;
|
||||
import tech.easyflow.ai.documentimport.DocumentImportKeys;
|
||||
import tech.easyflow.ai.dto.DocumentChunkDeleteResult;
|
||||
import tech.easyflow.ai.entity.Document;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentChunkSyncTask;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.math.BigInteger;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* {@link DocumentChunkServiceImpl} 异步索引任务回归测试。
|
||||
*/
|
||||
public class DocumentChunkServiceImplTest {
|
||||
|
||||
@Test
|
||||
public void updateContentShouldPersistPendingVersionBeforeDispatch() {
|
||||
Fixture fixture = fixture();
|
||||
|
||||
DocumentChunk result = fixture.service.updateContent(
|
||||
fixture.knowledgeId,
|
||||
fixture.chunkId,
|
||||
"# 新标题\n\n\n正文"
|
||||
);
|
||||
|
||||
ArgumentCaptor<DocumentChunk> update = ArgumentCaptor.forClass(DocumentChunk.class);
|
||||
Mockito.verify(fixture.chunkMapper).update(update.capture());
|
||||
Assert.assertEquals("# 新标题\n\n说明\n正文", update.getValue().getContent());
|
||||
Assert.assertEquals(DocumentChunkSyncState.PENDING, update.getValue().getIndexSyncStatus());
|
||||
Assert.assertEquals(Long.valueOf(3L), update.getValue().getIndexSyncVersion());
|
||||
Assert.assertEquals(
|
||||
"# 新标题\n\n\n正文",
|
||||
update.getValue().getOptions().get(DocumentImportKeys.KEY_DOCUMENT_RENDER_MARKDOWN)
|
||||
);
|
||||
Assert.assertEquals(DocumentChunkSyncState.PENDING, result.getIndexSyncStatus());
|
||||
InOrder order = Mockito.inOrder(
|
||||
fixture.syncTaskAppService,
|
||||
fixture.chunkMapper
|
||||
);
|
||||
order.verify(fixture.syncTaskAppService).createTask(
|
||||
Mockito.eq(result),
|
||||
Mockito.eq(fixture.collection),
|
||||
Mockito.eq(DocumentChunkSyncState.OPERATION_UPSERT),
|
||||
Mockito.eq(3L),
|
||||
Mockito.any()
|
||||
);
|
||||
order.verify(fixture.chunkMapper).update(Mockito.any(DocumentChunk.class));
|
||||
Mockito.verify(fixture.syncTaskAppService).dispatchBestEffort(fixture.taskId);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void updateContentShouldRejectBlankWithoutSideEffects() {
|
||||
Fixture fixture = fixture();
|
||||
|
||||
try {
|
||||
fixture.service.updateContent(fixture.knowledgeId, fixture.chunkId, " \n ");
|
||||
Assert.fail("空内容必须转删除，不能进入更新链路");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals(
|
||||
DocumentChunkServiceImpl.DOCUMENT_CHUNK_EMPTY_REQUIRES_DELETE,
|
||||
expected.getErrorCode()
|
||||
);
|
||||
}
|
||||
|
||||
Mockito.verifyNoInteractions(fixture.chunkMapper, fixture.documentMapper,
|
||||
fixture.syncTaskAppService);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void deleteChunkShouldPersistCleanupTaskAndUpdateParentCount() {
|
||||
Fixture fixture = fixture();
|
||||
Mockito.when(fixture.chunkMapper.deleteById(fixture.chunkId)).thenReturn(1);
|
||||
Mockito.when(fixture.chunkMapper.selectCountByQuery(Mockito.any(QueryWrapper.class)))
|
||||
.thenReturn(0L);
|
||||
|
||||
DocumentChunkDeleteResult result = fixture.service.deleteChunk(
|
||||
fixture.knowledgeId,
|
||||
fixture.chunkId
|
||||
);
|
||||
|
||||
Assert.assertEquals(0L, result.remainingChunkCount());
|
||||
InOrder order = Mockito.inOrder(
|
||||
fixture.syncTaskAppService,
|
||||
fixture.chunkMapper,
|
||||
fixture.documentMapper
|
||||
);
|
||||
order.verify(fixture.syncTaskAppService).createTask(
|
||||
Mockito.any(), Mockito.eq(fixture.collection),
|
||||
Mockito.eq(DocumentChunkSyncState.OPERATION_DELETE),
|
||||
Mockito.eq(3L), Mockito.any()
|
||||
);
|
||||
order.verify(fixture.chunkMapper).deleteById(fixture.chunkId);
|
||||
order.verify(fixture.documentMapper).update(Mockito.any(Document.class));
|
||||
Mockito.verify(fixture.syncTaskAppService).dispatchBestEffort(fixture.taskId);
|
||||
ArgumentCaptor<Document> parent = ArgumentCaptor.forClass(Document.class);
|
||||
Mockito.verify(fixture.documentMapper).update(parent.capture());
|
||||
Assert.assertEquals(Integer.valueOf(0), parent.getValue().getTotalChunks());
|
||||
Assert.assertEquals(Integer.valueOf(100), parent.getValue().getProgressPercent());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void listStatusShouldRejectChunkFromAnotherKnowledge() {
|
||||
Fixture fixture = fixture();
|
||||
DocumentChunk foreign = new DocumentChunk();
|
||||
foreign.setId(BigInteger.valueOf(999));
|
||||
foreign.setDocumentId(fixture.documentId);
|
||||
foreign.setDocumentCollectionId(BigInteger.valueOf(888));
|
||||
Mockito.when(fixture.chunkMapper.selectSyncStates(
|
||||
fixture.documentId, List.of(foreign.getId())
|
||||
)).thenReturn(List.of(foreign));
|
||||
|
||||
try {
|
||||
fixture.service.listIndexSyncStatus(
|
||||
fixture.knowledgeId, fixture.documentId, List.of(foreign.getId())
|
||||
);
|
||||
Assert.fail("跨知识库状态查询必须被拒绝");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals("分块不存在", expected.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void listStatusShouldResolveKnowledgeFromDocumentForAdminPolling() {
|
||||
Fixture fixture = fixture();
|
||||
Document parent = new Document();
|
||||
parent.setId(fixture.documentId);
|
||||
parent.setCollectionId(fixture.knowledgeId);
|
||||
Mockito.when(fixture.documentMapper.selectOneById(fixture.documentId))
|
||||
.thenReturn(parent);
|
||||
DocumentChunk state = new DocumentChunk();
|
||||
state.setId(fixture.chunkId);
|
||||
state.setDocumentId(fixture.documentId);
|
||||
state.setDocumentCollectionId(fixture.knowledgeId);
|
||||
state.setIndexSyncStatus(DocumentChunkSyncState.PENDING);
|
||||
state.setIndexSyncVersion(3L);
|
||||
Mockito.when(fixture.chunkMapper.selectSyncStates(
|
||||
fixture.documentId, List.of(fixture.chunkId)
|
||||
)).thenReturn(List.of(state));
|
||||
|
||||
List<tech.easyflow.ai.dto.DocumentChunkSyncStatus> result =
|
||||
fixture.service.listIndexSyncStatus(
|
||||
null, fixture.documentId, List.of(fixture.chunkId)
|
||||
);
|
||||
|
||||
Assert.assertEquals(1, result.size());
|
||||
Assert.assertEquals(DocumentChunkSyncState.PENDING, result.get(0).indexSyncStatus());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void updateContentShouldExposeRetryableErrorWhenLockIsUnavailable() {
|
||||
Fixture fixture = fixture();
|
||||
Mockito.when(fixture.redisLockExecutor.tryAcquire(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(null);
|
||||
|
||||
try {
|
||||
fixture.service.updateContent(fixture.knowledgeId, fixture.chunkId, "更新内容");
|
||||
Assert.fail("锁不可用时必须返回可重试错误");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals(503, expected.getHttpStatus());
|
||||
Assert.assertEquals(
|
||||
DocumentChunkServiceImpl.DOCUMENT_CHUNK_LOCK_UNAVAILABLE,
|
||||
expected.getErrorCode()
|
||||
);
|
||||
}
|
||||
Mockito.verify(fixture.chunkMapper, Mockito.never()).update(Mockito.any());
|
||||
}
|
||||
|
||||
private static Fixture fixture() {
|
||||
BigInteger knowledgeId = BigInteger.valueOf(101);
|
||||
BigInteger documentId = BigInteger.valueOf(201);
|
||||
BigInteger chunkId = BigInteger.valueOf(301);
|
||||
BigInteger taskId = BigInteger.valueOf(401);
|
||||
|
||||
DocumentChunk current = new DocumentChunk();
|
||||
current.setId(chunkId);
|
||||
current.setDocumentId(documentId);
|
||||
current.setDocumentCollectionId(knowledgeId);
|
||||
current.setContent("旧内容");
|
||||
current.setIndexSyncStatus(DocumentChunkSyncState.SYNCED);
|
||||
current.setIndexSyncVersion(2L);
|
||||
HashMap<String, Object> options = new HashMap<>();
|
||||
options.put("existing", "保留");
|
||||
current.setOptions(options);
|
||||
|
||||
DocumentCollection collection = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(collection.getId()).thenReturn(knowledgeId);
|
||||
Mockito.when(collection.getVectorStoreCollection()).thenReturn("kb-test");
|
||||
DocumentCollectionService collectionService = Mockito.mock(DocumentCollectionService.class);
|
||||
Mockito.when(collectionService.getById(knowledgeId)).thenReturn(collection);
|
||||
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectOneById(chunkId)).thenReturn(current);
|
||||
Mockito.when(chunkMapper.update(Mockito.any(DocumentChunk.class))).thenReturn(1);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.update(Mockito.any(Document.class))).thenReturn(1);
|
||||
|
||||
DocumentChunkSyncTask task = new DocumentChunkSyncTask();
|
||||
task.setId(taskId);
|
||||
DocumentChunkSyncTaskAppService syncTaskAppService =
|
||||
Mockito.mock(DocumentChunkSyncTaskAppService.class);
|
||||
Mockito.when(syncTaskAppService.createTask(
|
||||
Mockito.any(), Mockito.any(), Mockito.anyString(), Mockito.anyLong(), Mockito.any()
|
||||
)).thenReturn(task);
|
||||
|
||||
PlatformTransactionManager transactionManager = Mockito.mock(PlatformTransactionManager.class);
|
||||
Mockito.when(transactionManager.getTransaction(Mockito.any()))
|
||||
.thenReturn(new SimpleTransactionStatus());
|
||||
RedisLockExecutor redisLockExecutor = Mockito.mock(RedisLockExecutor.class);
|
||||
RedisLockExecutor.LockHandle lockHandle = Mockito.mock(RedisLockExecutor.LockHandle.class);
|
||||
Mockito.when(redisLockExecutor.tryAcquire(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any()
|
||||
)).thenReturn(lockHandle);
|
||||
|
||||
DocumentChunkServiceImpl service = new DocumentChunkServiceImpl(
|
||||
chunkMapper,
|
||||
documentMapper,
|
||||
collectionService,
|
||||
syncTaskAppService,
|
||||
transactionManager,
|
||||
redisLockExecutor
|
||||
);
|
||||
return new Fixture(service, knowledgeId, documentId, chunkId, taskId,
|
||||
chunkMapper, documentMapper, collection, syncTaskAppService,
|
||||
redisLockExecutor);
|
||||
}
|
||||
|
||||
private record Fixture(
|
||||
DocumentChunkServiceImpl service,
|
||||
BigInteger knowledgeId,
|
||||
BigInteger documentId,
|
||||
BigInteger chunkId,
|
||||
BigInteger taskId,
|
||||
DocumentChunkMapper chunkMapper,
|
||||
DocumentMapper documentMapper,
|
||||
DocumentCollection collection,
|
||||
DocumentChunkSyncTaskAppService syncTaskAppService,
|
||||
RedisLockExecutor redisLockExecutor
|
||||
) {
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package tech.easyflow.ai.service.impl;
|
||||
import com.easyagents.core.document.Document;
|
||||
import com.easyagents.search.engine.service.DocumentSearcher;
|
||||
import com.easyagents.search.engine.service.KeywordSearchRequest;
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
@@ -20,6 +21,7 @@ import java.math.BigInteger;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import static tech.easyflow.ai.entity.DocumentCollection.KEY_DOC_RECALL_MAX_NUM;
|
||||
import static tech.easyflow.ai.entity.DocumentCollection.KEY_SIMILARITY_THRESHOLD;
|
||||
@@ -100,7 +102,9 @@ public class DocumentCollectionServiceImplTest {
|
||||
|
||||
DocumentCollectionServiceImpl service = new TestDocumentCollectionService(collection);
|
||||
setField(service, "searcherFactory", new SearcherFactory(new StaticObjectProvider<DocumentSearcher>(searcher)));
|
||||
setField(service, "documentChunkMapper", mockDocumentChunkMapper(completedChunk, indexingChunk));
|
||||
AtomicReference<QueryWrapper> chunkQuery = new AtomicReference<QueryWrapper>();
|
||||
setField(service, "documentChunkMapper",
|
||||
mockDocumentChunkMapper(chunkQuery, completedChunk, indexingChunk));
|
||||
setField(service, "documentMapper", mockDocumentMapper(completedDocument));
|
||||
|
||||
tech.easyflow.ai.rag.KnowledgeRetrievalRequest request = new tech.easyflow.ai.rag.KnowledgeRetrievalRequest();
|
||||
@@ -116,6 +120,10 @@ public class DocumentCollectionServiceImplTest {
|
||||
Assert.assertEquals(completedDocument.getTitle(), result.get(0).getTitle());
|
||||
Assert.assertEquals("completed chunk", result.get(0).getContent());
|
||||
Assert.assertEquals(String.valueOf(knowledgeId), searcher.lastKnowledgeId);
|
||||
Assert.assertTrue(
|
||||
"检索补齐必须过滤尚未同步的分块",
|
||||
chunkQuery.get().toSQL().toUpperCase().contains("INDEX_SYNC_STATUS")
|
||||
);
|
||||
Assert.assertEquals(
|
||||
tech.easyflow.ai.entity.DocumentCollection.TYPE_DOCUMENT,
|
||||
result.get(0).getMetadata("resultType")
|
||||
@@ -210,7 +218,10 @@ public class DocumentCollectionServiceImplTest {
|
||||
return document;
|
||||
}
|
||||
|
||||
private static DocumentChunkMapper mockDocumentChunkMapper(tech.easyflow.ai.entity.DocumentChunk... chunks) {
|
||||
private static DocumentChunkMapper mockDocumentChunkMapper(
|
||||
AtomicReference<QueryWrapper> query,
|
||||
tech.easyflow.ai.entity.DocumentChunk... chunks
|
||||
) {
|
||||
Map<String, tech.easyflow.ai.entity.DocumentChunk> chunkMap = new HashMap<String, tech.easyflow.ai.entity.DocumentChunk>();
|
||||
for (tech.easyflow.ai.entity.DocumentChunk chunk : chunks) {
|
||||
chunkMap.put(String.valueOf(chunk.getId()), chunk);
|
||||
@@ -220,6 +231,7 @@ public class DocumentCollectionServiceImplTest {
|
||||
new Class<?>[]{DocumentChunkMapper.class},
|
||||
(proxy, method, args) -> {
|
||||
if ("selectListByQuery".equals(method.getName())) {
|
||||
query.set((QueryWrapper) args[0]);
|
||||
return List.copyOf(chunkMap.values());
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
|
||||
@@ -3,22 +3,27 @@ package tech.easyflow.ai.service.impl;
|
||||
import com.easyagents.core.store.DocumentStore;
|
||||
import com.easyagents.core.store.StoreResult;
|
||||
import com.easyagents.rag.core.RagDefaults;
|
||||
import com.easyagents.search.engine.service.DocumentSearcher;
|
||||
import com.mybatisflex.core.query.QueryWrapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.InOrder;
|
||||
import org.mockito.Mockito;
|
||||
import tech.easyflow.ai.config.SearcherFactory;
|
||||
import tech.easyflow.ai.documentchunk.DocumentChunkSyncState;
|
||||
import tech.easyflow.ai.entity.Document;
|
||||
import tech.easyflow.ai.entity.DocumentChunk;
|
||||
import tech.easyflow.ai.entity.DocumentCollection;
|
||||
import tech.easyflow.ai.entity.Model;
|
||||
import tech.easyflow.ai.enums.DocumentProcessStatus;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentChunkSyncTaskMapper;
|
||||
import tech.easyflow.ai.mapper.DocumentMapper;
|
||||
import tech.easyflow.ai.service.DocumentCollectionService;
|
||||
import tech.easyflow.ai.service.ModelService;
|
||||
import tech.easyflow.common.filestorage.FileStorageService;
|
||||
import tech.easyflow.common.cache.RedisLockExecutor;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
@@ -28,6 +33,7 @@ import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.regex.Pattern;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
/**
|
||||
* {@link DocumentServiceImpl} 文档维护回归测试。
|
||||
@@ -155,9 +161,15 @@ public class DocumentServiceImplTest {
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
DocumentChunkSyncTaskMapper taskMapper = configureRemovalCoordination(service);
|
||||
|
||||
Assert.assertTrue(service.removeDoc(documentId.toString()));
|
||||
|
||||
InOrder order = Mockito.inOrder(taskMapper, documentStore);
|
||||
order.verify(taskMapper).supersedeChunk(
|
||||
Mockito.eq(BigInteger.valueOf(201)), Mockito.any()
|
||||
);
|
||||
order.verify(documentStore).delete(Mockito.anyList(), Mockito.any());
|
||||
Mockito.verify(documentMapper).deleteById(documentId);
|
||||
Mockito.verify(storageService).delete(document.getDocumentPath());
|
||||
}
|
||||
@@ -206,6 +218,7 @@ public class DocumentServiceImplTest {
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
configureRemovalCoordination(service);
|
||||
|
||||
Assert.assertTrue(service.removeDoc(documentId.toString()));
|
||||
|
||||
@@ -238,6 +251,7 @@ public class DocumentServiceImplTest {
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
configureRemovalCoordination(service);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
@@ -295,6 +309,7 @@ public class DocumentServiceImplTest {
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
configureRemovalCoordination(service);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
@@ -308,6 +323,176 @@ public class DocumentServiceImplTest {
|
||||
Mockito.verifyNoInteractions(storageService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证关键词索引拒绝删除时保留数据库记录，供用户安全重试。
|
||||
*
|
||||
* @throws Exception 反射注入异常
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldStopWhenKeywordDeleteFails() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(451);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(452);
|
||||
BigInteger modelId = BigInteger.valueOf(453);
|
||||
BigInteger chunkId = BigInteger.valueOf(454);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
DocumentStore documentStore = Mockito.mock(DocumentStore.class);
|
||||
Mockito.when(documentStore.delete(Mockito.anyList(), Mockito.any()))
|
||||
.thenReturn(StoreResult.success());
|
||||
DocumentCollection knowledge = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(knowledge.getVectorEmbedModelId()).thenReturn(modelId);
|
||||
Mockito.when(knowledge.getVectorStoreCollection()).thenReturn("kb-test");
|
||||
Mockito.when(knowledge.toDocumentStore()).thenReturn(documentStore);
|
||||
Model model = new Model();
|
||||
model.setModelName("embedding-test");
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(Mockito.any()))
|
||||
.thenReturn(document);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(), Mockito.eq(BigInteger.class)))
|
||||
.thenReturn(List.of(chunkId));
|
||||
DocumentCollectionService knowledgeService = Mockito.mock(
|
||||
DocumentCollectionService.class
|
||||
);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId)).thenReturn(knowledge);
|
||||
ModelService modelService = Mockito.mock(ModelService.class);
|
||||
Mockito.when(modelService.getById(modelId)).thenReturn(model);
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
DocumentSearcher searcher = Mockito.mock(DocumentSearcher.class);
|
||||
Mockito.when(searcher.deleteDocument(chunkId)).thenReturn(false);
|
||||
SearcherFactory searcherFactory = Mockito.mock(SearcherFactory.class);
|
||||
Mockito.when(searcherFactory.getSearcher()).thenReturn(searcher);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "modelService", modelService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", searcherFactory);
|
||||
configureRemovalCoordination(service);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
Assert.fail("关键词索引删除失败时应停止文档删除");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals("文档关键词索引删除失败", expected.getMessage());
|
||||
}
|
||||
|
||||
Mockito.verify(chunkMapper, Mockito.never()).deleteByQuery(Mockito.any());
|
||||
Mockito.verify(documentMapper, Mockito.never()).deleteById(Mockito.any());
|
||||
Mockito.verifyNoInteractions(storageService);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证外部与数据库删除已经执行完成时，父锁收尾阶段丢失不会回滚数据库事务。
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldCommitAfterPostExecutionLockLoss() throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(501);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(502);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setDocumentPath("storage://lock-lost.txt");
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(Mockito.any()))
|
||||
.thenReturn(document);
|
||||
Mockito.when(documentMapper.deleteById(documentId)).thenReturn(1);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(), Mockito.eq(BigInteger.class))).thenReturn(List.of());
|
||||
Mockito.when(chunkMapper.deleteByQuery(Mockito.any())).thenReturn(0);
|
||||
DocumentCollectionService knowledgeService = Mockito.mock(
|
||||
DocumentCollectionService.class
|
||||
);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId))
|
||||
.thenReturn(Mockito.mock(DocumentCollection.class));
|
||||
FileStorageService storageService = Mockito.mock(FileStorageService.class);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "storageService", storageService);
|
||||
setField(service, "searcherFactory", Mockito.mock(SearcherFactory.class));
|
||||
configureRemovalCoordination(service);
|
||||
RedisLockExecutor lockExecutor = getField(
|
||||
service, "redisLockExecutor", RedisLockExecutor.class
|
||||
);
|
||||
Mockito.doAnswer(invocation -> {
|
||||
String lockKey = invocation.getArgument(0);
|
||||
Object result = invocation.<Supplier<?>>getArgument(3).get();
|
||||
if (DocumentChunkSyncState.parentLockKey(documentId.toString())
|
||||
.equals(lockKey)) {
|
||||
throw new IllegalStateException("模拟 callback 完成后续租丢失");
|
||||
}
|
||||
return result;
|
||||
}).when(lockExecutor).executeWithRenewingLock(
|
||||
Mockito.anyString(), Mockito.any(), Mockito.any(),
|
||||
Mockito.<Supplier<?>>any()
|
||||
);
|
||||
|
||||
Assert.assertTrue(service.removeDoc(documentId.toString()));
|
||||
Mockito.verify(documentMapper).deleteById(documentId);
|
||||
Mockito.verify(storageService).delete(document.getDocumentPath());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证向量存储配置缺失时不会先废弃仍可恢复的分块同步任务。
|
||||
*/
|
||||
@Test
|
||||
public void removeDocShouldKeepSyncTasksWhenVectorStoreIsUnavailable()
|
||||
throws Exception {
|
||||
BigInteger documentId = BigInteger.valueOf(601);
|
||||
BigInteger knowledgeId = BigInteger.valueOf(602);
|
||||
BigInteger chunkId = BigInteger.valueOf(603);
|
||||
Document document = new Document();
|
||||
document.setId(documentId);
|
||||
document.setCollectionId(knowledgeId);
|
||||
document.setProcessStatus(DocumentProcessStatus.COMPLETED.name());
|
||||
|
||||
DocumentCollection knowledge = Mockito.mock(DocumentCollection.class);
|
||||
Mockito.when(knowledge.toDocumentStore()).thenReturn(null);
|
||||
DocumentMapper documentMapper = Mockito.mock(DocumentMapper.class);
|
||||
Mockito.when(documentMapper.selectOneByQuery(Mockito.any()))
|
||||
.thenReturn(document);
|
||||
DocumentChunkMapper chunkMapper = Mockito.mock(DocumentChunkMapper.class);
|
||||
Mockito.when(chunkMapper.selectListByQueryAs(
|
||||
Mockito.any(), Mockito.eq(BigInteger.class)))
|
||||
.thenReturn(List.of(chunkId));
|
||||
DocumentCollectionService knowledgeService = Mockito.mock(
|
||||
DocumentCollectionService.class
|
||||
);
|
||||
Mockito.when(knowledgeService.getById(knowledgeId)).thenReturn(knowledge);
|
||||
|
||||
DocumentServiceImpl service = new DocumentServiceImpl();
|
||||
setField(service, "documentMapper", documentMapper);
|
||||
setField(service, "documentChunkMapper", chunkMapper);
|
||||
setField(service, "knowledgeService", knowledgeService);
|
||||
setField(service, "modelService", Mockito.mock(ModelService.class));
|
||||
setField(service, "storageService", Mockito.mock(FileStorageService.class));
|
||||
setField(service, "searcherFactory", Mockito.mock(SearcherFactory.class));
|
||||
DocumentChunkSyncTaskMapper taskMapper = configureRemovalCoordination(service);
|
||||
|
||||
try {
|
||||
service.removeDoc(documentId.toString());
|
||||
Assert.fail("向量存储不可用时应停止文档删除");
|
||||
} catch (BusinessException expected) {
|
||||
Assert.assertEquals("文档向量存储不可用", expected.getMessage());
|
||||
}
|
||||
|
||||
Mockito.verify(taskMapper, Mockito.never())
|
||||
.supersedeChunk(Mockito.any(), Mockito.any());
|
||||
Mockito.verify(chunkMapper, Mockito.never()).deleteByQuery(Mockito.any());
|
||||
Mockito.verify(documentMapper, Mockito.never()).deleteById(Mockito.any());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证旧版向量化入口同样拒绝超过 BGE-M3 上下文预算的分块。
|
||||
*
|
||||
@@ -351,6 +536,33 @@ public class DocumentServiceImplTest {
|
||||
field.set(target, value);
|
||||
}
|
||||
|
||||
private static <T> T getField(Object target,
|
||||
String fieldName,
|
||||
Class<T> fieldType) throws Exception {
|
||||
Field field = DocumentServiceImpl.class.getDeclaredField(fieldName);
|
||||
field.setAccessible(true);
|
||||
return fieldType.cast(field.get(target));
|
||||
}
|
||||
|
||||
private static DocumentChunkSyncTaskMapper configureRemovalCoordination(
|
||||
DocumentServiceImpl service
|
||||
) throws Exception {
|
||||
DocumentChunkSyncTaskMapper taskMapper = Mockito.mock(
|
||||
DocumentChunkSyncTaskMapper.class
|
||||
);
|
||||
RedisLockExecutor lockExecutor = Mockito.mock(RedisLockExecutor.class);
|
||||
Mockito.doAnswer(invocation -> invocation.<Supplier<?>>getArgument(3).get())
|
||||
.when(lockExecutor).executeWithRenewingLock(
|
||||
Mockito.anyString(),
|
||||
Mockito.any(),
|
||||
Mockito.any(),
|
||||
Mockito.<Supplier<?>>any()
|
||||
);
|
||||
setField(service, "documentChunkSyncTaskMapper", taskMapper);
|
||||
setField(service, "redisLockExecutor", lockExecutor);
|
||||
return taskMapper;
|
||||
}
|
||||
|
||||
/**
|
||||
* 统一 SQL 文本格式，便于断言查询结构。
|
||||
*
|
||||
|
||||
Reference in New Issue
Block a user