feat(wiki): per-document entity extraction action (#673)

- Add an Extract entities action to completed documents in the knowledge base document list when entity extraction is enabled. - Add a per-document extraction endpoint with workspace, ownership and processing checks. - Rename the KB-level button to Re-extract whole KB.
2026-11-02 14:11:49 +08:00 · 2026-10-29 20:29:54 +08:00 · 2026-10-29 20:29:54 +08:00 · 2d045b252e
commit 2d045b252e
parent 9aa4619ec6
9 changed files with 251 additions and 5 deletions
--- a/mateclaw-server/src/main/java/vip/mate/wiki/controller/WikiEntityController.java
+++ b/mateclaw-server/src/main/java/vip/mate/wiki/controller/WikiEntityController.java
@ -13,10 +13,12 @@ import vip.mate.exception.MateClawException;
 import vip.mate.wiki.dto.WikiEntityGraphView;
 import vip.mate.wiki.dto.WikiEntityView;
 import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
+import vip.mate.wiki.model.WikiRawMaterialEntity;
 import vip.mate.wiki.service.WikiEntityExtractionService;
 import vip.mate.wiki.service.WikiEntityGraphService;
 import vip.mate.wiki.service.WikiKnowledgeBaseService;
 import vip.mate.wiki.service.WikiProcessingService;
+import vip.mate.wiki.service.WikiRawMaterialService;
 import vip.mate.workspace.core.annotation.RequireWorkspaceRole;

 import java.util.List;
@ -36,6 +38,7 @@ public class WikiEntityController {
    private final WikiEntityGraphService graphService;
    private final WikiEntityExtractionService extractionService;
    private final WikiKnowledgeBaseService kbService;
+    private final WikiRawMaterialService rawMaterialService;

    /** List entities in a KB, optionally filtered by type, ranked by salience. */
    @RequireWorkspaceRole("viewer")
@ -92,6 +95,40 @@ public class WikiEntityController {
        return Map.of("status", "started", "kbId", kbId);
    }

+    /**
+     * Manually trigger entity extraction for one document. Runs on the wiki
+     * executor so the request returns immediately.
+     *
+     * @param force when true, re-extract chunks that already have mentions
+     */
+    @RequireWorkspaceRole("member")
+    @PostMapping("/kb/{kbId}/raw/{rawId}/entities/extract")
+    public Map<String, Object> extractRaw(@PathVariable Long kbId,
+                                          @PathVariable Long rawId,
+                                          @RequestParam(defaultValue = "false") boolean force,
+                                          @RequestHeader(value = "X-Workspace-Id", required = false) Long workspaceId) {
+        verifyKBWorkspace(kbId, workspaceId);
+        WikiRawMaterialEntity raw = rawMaterialService.getById(rawId);
+        if (raw == null || !kbId.equals(raw.getKbId())) {
+            throw new MateClawException(404, "Raw material not found");
+        }
+        // Extraction reads the document's chunks, which exist only after processing completes.
+        if (!"completed".equals(raw.getProcessingStatus())) {
+            throw new MateClawException(400, "Raw material has not finished processing");
+        }
+        WikiProcessingService.WIKI_EXECUTOR.submit(() -> {
+            try {
+                int count = extractionService.extractForRaw(kbId, rawId, force);
+                log.info("[WikiEntity] Manual document extraction completed: kbId={}, rawId={}, entities={}",
+                        kbId, rawId, count);
+            } catch (Exception e) {
+                log.warn("[WikiEntity] Manual document extraction failed for kbId={}, rawId={}: {}",
+                        kbId, rawId, e.getMessage());
+            }
+        });
+        return Map.of("status", "started", "kbId", kbId, "rawId", rawId);
+    }
+
    // ==================== Workspace Verification ====================

    private void verifyKBWorkspace(Long kbId, Long headerWorkspaceId) {
--- a/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiEntityExtractionService.java
+++ b/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiEntityExtractionService.java
@ -81,7 +81,23 @@ public class WikiEntityExtractionService {

    /** Extract entities from every not-yet-processed chunk of one raw material. */
    public int extractForRaw(Long kbId, Long rawId) {
-        return extract(kbId, chunkService.listByRawId(rawId), false);
+        return extractForRaw(kbId, rawId, false);
+    }
+
+    /**
+     * Extract entities from one raw material's chunks.
+     *
+     * @param force when {@code true}, re-extract chunks that already have
+     *              mentions (manual per-document rebuild); otherwise skip them
+     */
+    public int extractForRaw(Long kbId, Long rawId, boolean force) {
+        int touched = extract(kbId, chunkService.listByRawId(rawId), force);
+        if (force && touched > 0) {
+            // Re-extraction may drop entities that only this document mentioned.
+            // Skip pruning when nothing resolved so a model failure keeps the graph.
+            pruneOrphanEntities(kbId);
+        }
+        return touched;
    }

    /**
--- a/mateclaw-server/src/test/java/vip/mate/wiki/controller/WikiEntityControllerIdorTest.java
+++ b/mateclaw-server/src/test/java/vip/mate/wiki/controller/WikiEntityControllerIdorTest.java
@ -8,6 +8,7 @@ import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
 import vip.mate.wiki.service.WikiEntityExtractionService;
 import vip.mate.wiki.service.WikiEntityGraphService;
 import vip.mate.wiki.service.WikiKnowledgeBaseService;
+import vip.mate.wiki.service.WikiRawMaterialService;

 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@ -31,7 +32,8 @@ class WikiEntityControllerIdorTest {
        controller = new WikiEntityController(
                mock(WikiEntityGraphService.class),
                mock(WikiEntityExtractionService.class),
-                kbService);
+                kbService,
+                mock(WikiRawMaterialService.class));
    }

    @Test
@ -70,6 +72,15 @@ class WikiEntityControllerIdorTest {
                .isInstanceOf(MateClawException.class);
    }

+    @Test
+    @DisplayName("extractRaw (member-level write) on another workspace's KB → 403")
+    void extractRawCrossWorkspaceRejected() {
+        when(kbService.getById(10L)).thenReturn(kb(10L, 2L));
+
+        assertThatThrownBy(() -> controller.extractRaw(10L, 5L, false, 1L))
+                .isInstanceOf(MateClawException.class);
+    }
+
    @Test
    @DisplayName("unknown kbId → 404")
    void unknownKbReturns404() {
--- a/mateclaw-server/src/test/java/vip/mate/wiki/controller/WikiEntityControllerRawExtractionTest.java
+++ b/mateclaw-server/src/test/java/vip/mate/wiki/controller/WikiEntityControllerRawExtractionTest.java
@ -0,0 +1,96 @@
+package vip.mate.wiki.controller;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import vip.mate.exception.MateClawException;
+import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
+import vip.mate.wiki.model.WikiRawMaterialEntity;
+import vip.mate.wiki.service.WikiEntityExtractionService;
+import vip.mate.wiki.service.WikiEntityGraphService;
+import vip.mate.wiki.service.WikiKnowledgeBaseService;
+import vip.mate.wiki.service.WikiRawMaterialService;
+
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/** Per-document entity extraction endpoint: ownership, readiness and dispatch. */
+class WikiEntityControllerRawExtractionTest {
+
+    private static final long KB_ID = 10L;
+    private static final long RAW_ID = 2055137662148763649L;
+
+    private WikiKnowledgeBaseService kbService;
+    private WikiRawMaterialService rawService;
+    private WikiEntityExtractionService extractionService;
+    private WikiEntityController controller;
+
+    @BeforeEach
+    void setUp() {
+        kbService = mock(WikiKnowledgeBaseService.class);
+        rawService = mock(WikiRawMaterialService.class);
+        extractionService = mock(WikiEntityExtractionService.class);
+        controller = new WikiEntityController(mock(WikiEntityGraphService.class), extractionService, kbService, rawService);
+        WikiKnowledgeBaseEntity kb = new WikiKnowledgeBaseEntity();
+        kb.setId(KB_ID);
+        kb.setWorkspaceId(1L);
+        when(kbService.getById(KB_ID)).thenReturn(kb);
+    }
+
+    @Test
+    @DisplayName("completed document in the KB → extraction dispatched with force flag")
+    void completedDocumentStartsExtraction() {
+        when(rawService.getById(RAW_ID)).thenReturn(raw(KB_ID, "completed"));
+
+        Map<String, Object> result = controller.extractRaw(KB_ID, RAW_ID, true, 1L);
+
+        assertThat(result).containsEntry("status", "started").containsEntry("rawId", RAW_ID);
+        verify(extractionService, timeout(2000)).extractForRaw(KB_ID, RAW_ID, true);
+    }
+
+    @Test
+    @DisplayName("document from another KB → 404, nothing dispatched")
+    void documentFromAnotherKbRejected() {
+        when(rawService.getById(RAW_ID)).thenReturn(raw(99L, "completed"));
+
+        assertThatThrownBy(() -> controller.extractRaw(KB_ID, RAW_ID, false, 1L))
+                .isInstanceOf(MateClawException.class)
+                .hasMessageContaining("not found");
+        verify(extractionService, never()).extractForRaw(anyLong(), anyLong(), anyBoolean());
+    }
+
+    @Test
+    @DisplayName("missing document → 404")
+    void missingDocumentRejected() {
+        assertThatThrownBy(() -> controller.extractRaw(KB_ID, RAW_ID, false, 1L))
+                .isInstanceOf(MateClawException.class)
+                .hasMessageContaining("not found");
+    }
+
+    @Test
+    @DisplayName("document still processing → rejected, nothing dispatched")
+    void unfinishedDocumentRejected() {
+        when(rawService.getById(RAW_ID)).thenReturn(raw(KB_ID, "processing"));
+
+        assertThatThrownBy(() -> controller.extractRaw(KB_ID, RAW_ID, false, 1L))
+                .isInstanceOf(MateClawException.class);
+        verify(extractionService, never()).extractForRaw(anyLong(), anyLong(), anyBoolean());
+    }
+
+    private static WikiRawMaterialEntity raw(long kbId, String status) {
+        WikiRawMaterialEntity raw = new WikiRawMaterialEntity();
+        raw.setId(RAW_ID);
+        raw.setKbId(kbId);
+        raw.setProcessingStatus(status);
+        return raw;
+    }
+}
--- a/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiEntityExtractionServiceTest.java
+++ b/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiEntityExtractionServiceTest.java
@ -224,6 +224,38 @@ class WikiEntityExtractionServiceTest {
        verify(entityMapper, times(0)).insert(any(WikiEntityEntity.class));
    }

+    @Test
+    @DisplayName("extractForRaw(force): re-extracts processed chunks and prunes orphans")
+    void extractForRaw_forceReextractsAndPrunes() {
+        when(chunkService.listByRawId(RAW_ID)).thenReturn(List.of(chunk(1L, "Alice works at Acme.")));
+        when(mentionMapper.selectCount(any(Wrapper.class))).thenReturn(1L);
+        when(mentionMapper.selectList(any())).thenReturn(Collections.emptyList());
+
+        int touched = service.extractForRaw(KB_ID, RAW_ID, true);
+
+        assertEquals(2, touched);
+        verify(mentionMapper, times(1)).delete(any());
+        verify(mentionMapper, times(2)).insert(any(WikiEntityMentionEntity.class));
+        // Orphan sweep runs after a successful forced rebuild.
+        verify(entityMapper, times(2)).selectList(any());
+    }
+
+    @Test
+    @DisplayName("extractForRaw(force): model failure keeps the existing graph")
+    void extractForRaw_forceWithoutResultKeepsGraph() {
+        when(chunkService.listByRawId(RAW_ID)).thenReturn(List.of(chunk(1L, "Alice works at Acme.")));
+        when(mentionMapper.selectCount(any(Wrapper.class))).thenReturn(1L);
+        ChatModel failing = failingModel();
+        when(routingService.buildChatModel(7L)).thenReturn(failing);
+
+        int touched = service.extractForRaw(KB_ID, RAW_ID, true);
+
+        assertEquals(0, touched);
+        verify(mentionMapper, times(0)).delete(any());
+        verify(relationMapper, times(0)).delete(any());
+        verify(entityMapper, times(0)).delete(any());
+    }
+
    private WikiChunkEntity chunk(Long id, String content) {
        WikiChunkEntity c = new WikiChunkEntity();
        c.setId(id);
@ -233,6 +265,12 @@ class WikiEntityExtractionServiceTest {
        return c;
    }

+    private ChatModel failingModel() {
+        ChatModel model = mock(ChatModel.class);
+        when(model.call(any(Prompt.class))).thenThrow(new IllegalStateException("model unavailable"));
+        return model;
+    }
+
    private ChatModel cannedModel(String body) {
        ChatModel model = mock(ChatModel.class);
        when(model.call(any(Prompt.class))).thenReturn(new ChatResponse(List.of(
--- a/mateclaw-ui/src/api/index.ts
+++ b/mateclaw-ui/src/api/index.ts
@ -1304,6 +1304,9 @@ export const wikiApi = {
    http.get(`/wiki/kb/${kbId}/entities/${entityId}/graph`, { params: { limit } }),
  extractEntities: (kbId: number | string, force = false) =>
    http.post(`/wiki/kb/${kbId}/entities/extract`, null, { params: { force } }),
+  // Entity extraction for a single processed document.
+  extractRawEntities: (kbId: number | string, rawId: number | string, force = false) =>
+    http.post(`/wiki/kb/${kbId}/raw/${rawId}/entities/extract`, null, { params: { force } }),

  // RFC-030: Jobs
  getWikiJobs: (kbId: number, rawId: number) =>
--- a/mateclaw-ui/src/i18n/locales/en-US.ts
+++ b/mateclaw-ui/src/i18n/locales/en-US.ts
@ -2968,6 +2968,8 @@ export default {
    cancelling: 'Cancelling…',
    cancelledHint: 'Cancelled by user',
    download: 'Download original file',
+    extractEntities: 'Extract entities',
+    extractEntitiesStarted: 'Started extracting entities from "{title}". Check the entity graph when it finishes.',
    downloadFailed: 'Download failed',
    processAll: 'Process All Pending',
    materialTitle: 'Title',
@ -3098,7 +3100,7 @@ export default {
      entityExtraction: 'Entity extraction',
      entityExtractionHint: 'Extract named entities (people, organizations, locations, ...) and their relations from raw materials to build an entity-level knowledge graph. Adds processing time and token cost; off by default.',
      entityExtractionEnable: 'Enable entity extraction',
-      entityExtractionRun: 'Extract now',
+      entityExtractionRun: 'Re-extract whole KB',
      entityExtractionRunning: 'Extracting…',
      entityTypesLabel: 'Entity types to extract',
      entityTypesPlaceholder: 'Select or type an entity type, Enter to add',
--- a/mateclaw-ui/src/i18n/locales/zh-CN.ts
+++ b/mateclaw-ui/src/i18n/locales/zh-CN.ts
@ -2980,6 +2980,8 @@ export default {
    cancelling: '正在取消…',
    cancelledHint: '用户已取消处理',
    download: '下载原始文件',
+    extractEntities: '实体抽取',
+    extractEntitiesStarted: '已开始抽取「{title}」的实体，完成后可在实体图谱中查看',
    downloadFailed: '下载失败',
    processAll: '处理所有待处理材料',
    materialTitle: '标题',
@ -3110,7 +3112,7 @@ export default {
      entityExtraction: '实体抽取',
      entityExtractionHint: '从原始材料中抽取命名实体（人物、组织、地点等）及其关系，构建实体级知识图谱。会增加处理时间与 Token 消耗，默认关闭。',
      entityExtractionEnable: '开启实体抽取',
-      entityExtractionRun: '立即抽取',
+      entityExtractionRun: '全库重新抽取',
      entityExtractionRunning: '抽取中…',
      entityTypesLabel: '抽取的实体类型',
      entityTypesPlaceholder: '选择或输入实体类型，回车添加',
--- a/mateclaw-ui/src/views/Wiki/components/RawMaterialPanel.vue
+++ b/mateclaw-ui/src/views/Wiki/components/RawMaterialPanel.vue
@ -290,6 +290,14 @@
                <path d="M20.49 15a9 9 0 1 1-2.12-9.36L23 10"/>
              </svg>
            </button>
+            <button
+              v-if="entityExtractionEnabled && raw.processingStatus === 'completed'"
+              class="btn-icon" :title="t('wiki.extractEntities')"
+              :disabled="extractingIds.has(String(raw.id))"
+              @click.stop="extractRawEntities(raw)"
+            >
+              <el-icon :size="14"><Connection /></el-icon>
+            </button>
            <button
              v-if="raw.processingStatus !== 'uploading'"
              class="btn-icon" :title="t('wiki.download')"
@ -378,7 +386,7 @@ import { useI18n } from 'vue-i18n'
 import { mcToast } from '@/composables/useMcToast'
 import { mcConfirm } from '@/components/common/useConfirm'
 import { useFileDrop } from '@/composables/useFileDrop'
-import { Download } from '@element-plus/icons-vue'
+import { Connection, Download } from '@element-plus/icons-vue'
 import { useWikiStore } from '@/stores/useWikiStore'
 import { useWorkspaceStore } from '@/stores/useWorkspaceStore'
 import { wikiApi } from '@/api/index'
@ -906,6 +914,39 @@ async function handleAddText() {
  textContent.value = ''
 }

+// Per-document entity extraction is offered only once the KB has entity
+// extraction switched on in its processing config.
+const entityExtractionEnabled = computed(() => {
+  const content = store.currentKB?.configContent
+  if (!content) return false
+  try {
+    return JSON.parse(content)?.entityExtractionEnabled === true
+  } catch {
+    return false
+  }
+})
+// Keyed by String(id): raw ids are Snowflake values and must stay strings.
+const extractingIds = ref<Set<string>>(new Set())
+
+async function extractRawEntities(raw: { id: number | string; title: string }) {
+  if (!store.currentKB) return
+  const key = String(raw.id)
+  if (extractingIds.value.has(key)) return
+  extractingIds.value = new Set(extractingIds.value).add(key)
+  try {
+    // force=true re-extracts this document only, so the latest entity-type
+    // config also applies to documents extracted before it changed.
+    await wikiApi.extractRawEntities(store.currentKB.id, raw.id, true)
+    mcToast.success(t('wiki.extractEntitiesStarted', { title: raw.title }))
+  } catch (e) {
+    mcToast.error(e instanceof Error ? e.message : String(e))
+  } finally {
+    const next = new Set(extractingIds.value)
+    next.delete(key)
+    extractingIds.value = next
+  }
+}
+
 async function reprocess(rawId: number) {
  if (!store.currentKB) return
  const kbId = store.currentKB.id