SuperBizAgent-java/src/main/java/com/superbiz/agent/service/DocumentFieldEnricher.java

package com.superbiz.agent.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.dto.Frontmatter;
import com.superbiz.agent.dto.KnowledgeEntry;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Service;

import jakarta.annotation.PostConstruct;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 文档字段补全服务
 * 上传时调用 LLM 生成 covers 和 whenToRetrieve
 */
@Slf4j
@Service
public class DocumentFieldEnricher {

    @Autowired
    private ChatModel chatModel;

    @Autowired
    private ObjectMapper objectMapper;

    @Autowired
    private KnowledgeIndexService knowledgeIndexService;

    private String promptTemplate;

    @PostConstruct
    public void init() {
        try {
            promptTemplate = new String(
                new ClassPathResource("prompts/doc-field-enricher-prompt.md").getInputStream().readAllBytes(),
                StandardCharsets.UTF_8);
            log.info("DocumentFieldEnricher prompt 加载成功");
        } catch (IOException e) {
            log.error("加载 doc-field-enricher-prompt.md 失败", e);
            throw new RuntimeException("Failed to load doc-field-enricher prompt", e);
        }
    }

    public void enrich(Frontmatter frontmatter, String bodyText) {
        enrich(frontmatter, bodyText, null);
    }

    /**
     * 为 Frontmatter 补全 covers 和 whenToRetrieve
     * 若已有值则跳过；LLM 失败时降级，不阻断主流程
     *
     * @param frontmatter 待补全的 frontmatter
     * @param bodyText    文档正文
     * @param category    文档所属域（用于查找同域其他文档）
     */
    public void enrich(Frontmatter frontmatter, String bodyText, String category) {
        if (frontmatter == null) return;

        boolean needsCovers = frontmatter.getCovers() == null || frontmatter.getCovers().isEmpty();
        boolean needsWhen = frontmatter.getWhenToRetrieve() == null || frontmatter.getWhenToRetrieve().isBlank();

        if (!needsCovers && !needsWhen) {
            log.debug("covers 和 whenToRetrieve 已存在，跳过 LLM 生成");
            return;
        }

        try {
            String snippet = bodyText != null && bodyText.length() > 1000
                    ? bodyText.substring(0, 1000) : (bodyText != null ? bodyText : "");

            String sameDomainDocs = buildSameDomainDocs(frontmatter.getTitle(), category);

            String promptText = String.format(promptTemplate,
                    frontmatter.getTitle(),
                    frontmatter.getSummary(),
                    sameDomainDocs,
                    snippet);

            String response = chatModel.call(new Prompt(promptText))
                    .getResult().getOutput().getText();

            // 提取 JSON 部分（防止模型输出多余文本）
            String json = extractJson(response);
            JsonNode node = objectMapper.readTree(json);

            if (needsCovers && node.has("covers")) {
                List<String> covers = new ArrayList<>();
                node.get("covers").forEach(n -> covers.add(n.asText()));
                frontmatter.setCovers(covers);
                log.debug("LLM 生成 covers: {}", covers);
            }

            if (needsWhen && node.has("whenToRetrieve")) {
                frontmatter.setWhenToRetrieve(node.get("whenToRetrieve").asText());
                log.debug("LLM 生成 whenToRetrieve: {}", frontmatter.getWhenToRetrieve());
            }

        } catch (Exception e) {
            log.warn("LLM 生成文档字段失败，降级处理: title={}", frontmatter.getTitle(), e);
            if (needsCovers) frontmatter.setCovers(List.of());
            if (needsWhen) frontmatter.setWhenToRetrieve(frontmatter.getSummary());
        }
    }

    private String extractJson(String text) {
        if (text == null) return "{}";
        int start = text.indexOf('{');
        int end = text.lastIndexOf('}');
        if (start == -1 || end == -1 || end <= start) return "{}";
        return text.substring(start, end + 1);
    }

    /**
     * 构建同域其他文档标题列表（供 LLM 做排除判断）
     */
    private String buildSameDomainDocs(String currentTitle, String category) {
        if (category == null || category.isBlank()) return "（无同域文档信息）";
        List<String> otherTitles = knowledgeIndexService.getAllEntries().stream()
                .filter(e -> category.equals(e.getCategory()))
                .map(KnowledgeEntry::getTitle)
                .filter(t -> t != null && !t.equals(currentTitle))
                .collect(Collectors.toList());
        if (otherTitles.isEmpty()) return "（无同域其他文档）";
        return String.join("、", otherTitles);
    }
}