RAG/src/main/java/com/wok/supportbot/controller/AiController.java

package com.wok.supportbot.controller;

import com.wok.supportbot.app.AssistantApp;
import com.wok.supportbot.config.RoleAccessConfig;
import com.wok.supportbot.service.ConversationService;
import com.wok.supportbot.service.CustomerAccountService;
import com.wok.supportbot.service.CustomerAccountService.AccountScope;
import com.wok.supportbot.service.CustomerServiceRoleService;
import com.wok.supportbot.service.CustomerServiceRoleService.RoleScope;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;


@RestController
@RequestMapping("/ai")
public class AiController {

    @Resource
    private AssistantApp assistantApp;
    @Resource
    private CustomerServiceRoleService customerServiceRoleService;
    @Resource
    private CustomerAccountService customerAccountService;
    @Resource
    private ConversationService conversationService;
    @Resource
    private RoleAccessConfig roleAccessConfig;


    /**
     * 同步调用 AI 智能客服应用
     *
     * @param message      用户消息
     * @param chatId       会话ID
     * @param roleId       客服角色ID（可选）；命中角色时由服务端强制套用该角色人设
     * @param systemPrompt 角色人设兜底（仅在未命中角色时生效）
     * @return AI 回答
     */
    @GetMapping("/assistant_app/chat/sync")
    public String doChatWithAssistantAppSync(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        return assistantApp.doChat(message, chatId, resolveSystemPrompt(scope, systemPrompt));
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 返回Flux 响应式؜对象，并且添加 SSE 对应的 MediaType
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> doChatWithLoveAppSSE(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        return assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt));
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 返回 Flux 对象，并且؜设置泛型为 ServerSentEvent。使用这种方式可以省略 MediaType
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/server_sent_event")
    public Flux<ServerSentEvent<String>> doChatWithAssistantAppServerSentEvent(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        return assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt))
                .map(chunk -> ServerSentEvent.<String>builder()
                        .data(chunk)
                        .build());
    }

    /**
     * SSE 流式调用 AI 智能客服应用
     * 使用 SSEEmiter，؜通过 send 方法持续向 SseEmitter 发送消息
     *
     * @param message
     * @param chatId
     * @return
     */
    @GetMapping(value = "/assistant_app/chat/sse_emitter")
    public SseEmitter doChatWithAssistantAppServerSseEmitter(String message, String chatId, Long roleId, String accountId, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        // 创建一个超时时间较长的 SseEmitter
        SseEmitter sseEmitter = new SseEmitter(180000L); // 3 分钟超时
        // 获取 Flux 响应式数据流并且直接通过订阅推送给 SseEmitter
        assistantApp.doChatByStream(message, chatId, resolveSystemPrompt(scope, systemPrompt))
                .subscribe(chunk -> {
                    try {
                        sseEmitter.send(chunk);
                    } catch (IOException e) {
                        sseEmitter.completeWithError(e);
                    }
                }, sseEmitter::completeWithError, sseEmitter::complete);
        // 返回
        return sseEmitter;
    }

    /**
     * RAG 知识库同步对话（支持查询重写策略）
     *
     * @param message         用户消息
     * @param chatId          会话ID
     * @param rewriteStrategy 查询重写策略（可选）：NONE/REWRITE/TRANSLATION/COMPRESSION/MULTI_QUERY，默认为 MULTI_QUERY
     * @param roleId          客服角色ID（可选）；命中角色时由服务端强制限定检索范围与人设，客户端无法跨域
     * @param categoryId      单个分类过滤（仅未命中角色时生效）
     * @param categoryIds     多个分类过滤，逗号分隔（仅未命中角色时生效）
     * @param systemPrompt    角色人设兜底（仅未命中角色时生效）
     * @return AI 回答
     */
    @GetMapping("/assistant_app/chat/rag/sync")
    public String doChatWithRagSync(String message, String chatId, String rewriteStrategy, Long roleId, String accountId, Long categoryId, String categoryIds, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        String sys = resolveSystemPrompt(scope, systemPrompt);
        // 严格隔离：未授权任何知识库的角色退回普通对话，绝不检索 KB
        if (isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return assistantApp.doChat(message, chatId, sys);
        }
        return assistantApp.doChatWithRagStrategy(
                message, chatId, normalizeStrategy(rewriteStrategy),
                resolveCategoryIds(scope, categoryId, categoryIds), sys);
    }

    /**
     * RAG 知识库流式对话（SSE，支持查询重写策略）
     *
     * @param message         用户消息
     * @param chatId          会话ID
     * @param rewriteStrategy 查询重写策略（可选）：NONE/REWRITE/TRANSLATION/COMPRESSION/MULTI_QUERY，默认为 MULTI_QUERY
     * @param roleId          客服角色ID（可选）；命中角色时由服务端强制限定检索范围与人设，客户端无法跨域
     * @param categoryId      单个分类过滤（仅未命中角色时生效）
     * @param categoryIds     多个分类过滤，逗号分隔（仅未命中角色时生效）
     * @param systemPrompt    角色人设兜底（仅未命中角色时生效）
     * @return 流式 AI 回答
     */
    @GetMapping(value = "/assistant_app/chat/rag/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> doChatWithRagSSE(String message, String chatId, String rewriteStrategy, Long roleId, String accountId, Long categoryId, String categoryIds, String systemPrompt) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        bindConversation(chatId, context);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        String sys = resolveSystemPrompt(scope, systemPrompt);
        // 严格隔离：未授权任何知识库的角色退回普通流式对话，绝不检索 KB
        if (isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return assistantApp.doChatByStream(message, chatId, sys);
        }
        return assistantApp.doChatWithRagStrategyByStream(
                message, chatId, normalizeStrategy(rewriteStrategy),
                resolveCategoryIds(scope, categoryId, categoryIds), sys);
    }

    /**
     * RAG 引用来源：返回本次问题命中的知识库片段（不生成回答），供回答下方展示来源。
     * 复用与 RAG 回答相同的角色范围、分类过滤与查询改写策略，确保来源即答案所依据的片段。
     */
    @GetMapping("/assistant_app/rag/sources")
    public Map<String, Object> getRagSources(String message, String chatId, String rewriteStrategy,
                                             Long roleId, String accountId, Long categoryId, String categoryIds) {
        AccountRoleContext context = resolveAccountRole(accountId, roleId);
        RoleScope scope = customerServiceRoleService.getRoleScope(context.roleId());
        if (message == null || message.isBlank() || isKbDenied(scope) || shouldBypassKnowledgeRetrieval(message)) {
            return Map.of("success", true, "data", List.of());
        }
        List<Long> cats = resolveCategoryIds(scope, categoryId, categoryIds);
        List<Document> docs = assistantApp.retrieveRagSources(message, chatId, normalizeStrategy(rewriteStrategy), cats);
        List<Map<String, Object>> out = new ArrayList<>();
        for (Document doc : docs) {
            Map<String, Object> meta = doc.getMetadata();
            Map<String, Object> item = new LinkedHashMap<>();
            item.put("documentId", meta.get("documentId"));
            item.put("title", meta.get("title"));
            item.put("sourceName", meta.get("sourceName"));
            item.put("chunkIndex", meta.get("chunkIndex"));
            item.put("score", meta.get("distance"));
            String text = doc.getText();
            item.put("snippet", text != null && text.length() > 160 ? text.substring(0, 160) + "…" : text);
            out.add(item);
        }
        return Map.of("success", true, "data", out);
    }

    /**
     * 严格隔离模式下，命中角色但未绑定任何知识库分类 ⇒ 拒绝检索任何 KB。
     * 非严格模式下永远返回 false（未绑定 = 可检索全部，沿用通用角色语义）。
     */
    private boolean isKbDenied(RoleScope scope) {
        return roleAccessConfig.isStrictIsolation() && scope.hasRole() && scope.categoryIds().isEmpty();
    }

    /**
     * 寒暄、感谢、告别这类短消息没有知识库检索意图，直接走普通对话。
     * 否则向量库会被迫召回一个“最相似但实际无关”的片段，导致回答下方出现误导性来源。
     */
    private boolean shouldBypassKnowledgeRetrieval(String message) {
        if (!StringUtils.hasText(message)) {
            return true;
        }
        String normalized = message.trim()
                .toLowerCase(Locale.ROOT)
                .replaceAll("[\\s，。！？!?,.;；：:、~～…]+", "");
        if (normalized.isEmpty()) {
            return true;
        }
        if (normalized.length() > 12) {
            return false;
        }
        return List.of(
                "你好", "您好", "hello", "hi", "哈喽", "嗨", "在吗",
                "早上好", "上午好", "中午好", "下午好", "晚上好",
                "谢谢", "感谢", "多谢", "好的", "好", "嗯", "嗯嗯",
                "再见", "拜拜", "辛苦了"
        ).contains(normalized);
    }

    /** 未指定策略时默认 MULTI_QUERY（多路扩展）。 */
    private String normalizeStrategy(String rewriteStrategy) {
        return (rewriteStrategy != null && !rewriteStrategy.isEmpty()) ? rewriteStrategy : "MULTI_QUERY";
    }

    /** 命中角色时用角色人设，否则用客户端兜底人设。 */
    private AccountRoleContext resolveAccountRole(String accountId, Long fallbackRoleId) {
        AccountScope accountScope = customerAccountService.getOrCreateAccountScope(accountId, fallbackRoleId);
        Long effectiveRoleId = accountScope.hasAccount() && accountScope.roleId() != null
                ? accountScope.roleId()
                : fallbackRoleId;
        Long effectiveAccountId = accountScope.hasAccount() ? accountScope.accountId() : null;
        return new AccountRoleContext(effectiveAccountId, effectiveRoleId);
    }

    private void bindConversation(String chatId, AccountRoleContext context) {
        conversationService.bindConversation(chatId, context.accountId(), context.roleId());
    }

    private String resolveSystemPrompt(RoleScope scope, String fallbackSystemPrompt) {
        if (scope.hasRole() && StringUtils.hasText(scope.systemPrompt())) {
            return scope.systemPrompt();
        }
        return fallbackSystemPrompt;
    }

    /**
     * 解析最终的检索分类范围。
     * 命中角色时，强制使用角色绑定的分类（忽略客户端传入），防止越权跨域；
     * 未命中角色时，沿用客户端传入的分类（便于直接调试/无角色调用）。
     */
    private List<Long> resolveCategoryIds(RoleScope scope, Long categoryId, String categoryIds) {
        if (scope.hasRole()) {
            return scope.categoryIds();
        }
        List<Long> ids = parseCategoryIds(categoryIds);
        if (ids.isEmpty() && categoryId != null) {
            return List.of(categoryId);
        }
        return ids;
    }

    private List<Long> parseCategoryIds(String categoryIds) {
        if (categoryIds == null || categoryIds.isBlank()) {
            return Collections.emptyList();
        }
        return Arrays.stream(categoryIds.split(","))
                .map(String::trim)
                .filter(item -> !item.isEmpty())
                .map(Long::valueOf)
                .filter(id -> id > 0)
                .distinct()
                .toList();
    }

    private record AccountRoleContext(Long accountId, Long roleId) {
    }

}