AIyaodianzhushou/backend-java/src/main/java/com/pharmacopoeia/nihaisha/service/NihQueryRewriterService.java

package com.pharmacopoeia.nihaisha.service;

import com.pharmacopoeia.nihaisha.config.NihaishaProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.netty.http.client.HttpClient;

import java.time.Duration;
import java.util.*;

/**
 * 查询改写服务 — 本地 Ollama 主路径 + 百炼 API 降级
 */
@Service
public class NihQueryRewriterService {

    private static final Logger log = LoggerFactory.getLogger(NihQueryRewriterService.class);

    private final NihaishaProperties properties;
    private final WebClient webClient;

    public NihQueryRewriterService(NihaishaProperties properties) {
        this.properties = properties;
        this.webClient = WebClient.builder()
                .clientConnector(new ReactorClientHttpConnector(
                        HttpClient.create().responseTimeout(Duration.ofSeconds(30))))
                .build();
    }

    /**
     * 改写查询 — 调用 LLM 生成多个变体
     * 主路径：本地 Ollama，降级：百炼 API
     */
    public List<String> rewriteQuery(String originalQuery) {
        var primary = properties.getLlm().getPrimary();
        var fallback = properties.getLlm().getFallback();

        // 尝试主路径（本地 Ollama）
        if (primary.getUrl() != null && !primary.getUrl().isEmpty()) {
            try {
                return callLlm(primary.getUrl(), primary.getApiKey(), primary.getModel(), originalQuery);
            } catch (Exception e) {
                log.warn("主 LLM (Ollama) 不可用: {}, 降级到百炼 API", e.getMessage());
            }
        }

        // 降级路径（百炼 API）
        if (fallback.getUrl() != null && !fallback.getUrl().isEmpty()
                && fallback.getApiKey() != null && !fallback.getApiKey().isEmpty()) {
            try {
                return callLlm(fallback.getUrl(), fallback.getApiKey(), fallback.getModel(), originalQuery);
            } catch (Exception e) {
                log.warn("降级 LLM (百炼) 也不可用: {}, 返回原始查询", e.getMessage());
            }
        }

        return List.of(originalQuery);
    }

    @SuppressWarnings("unchecked")
    private List<String> callLlm(String url, String apiKey, String model, String query) {
        Map<String, Object> body = buildRequestBody(model, query);

        Map<String, Object> response;
        try {
            var request = webClient.post()
                    .uri(url + "/chat/completions")
                    .header("Content-Type", "application/json");

            if (apiKey != null && !apiKey.isEmpty()) {
                request.header("Authorization", "Bearer " + apiKey);
            }

            response = request.bodyValue(body)
                    .retrieve()
                    .bodyToMono(Map.class)
                    .block(properties.getLlm().getTimeout());
        } catch (WebClientResponseException e) {
            throw new RuntimeException("LLM API 调用失败: " + e.getStatusCode() + " " + e.getResponseBodyAsString());
        } catch (Exception e) {
            throw new RuntimeException("LLM API 调用失败: " + e.getMessage());
        }

        if (response == null || !response.containsKey("choices")) {
            return List.of(query);
        }

        List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");
        if (choices.isEmpty()) return List.of(query);

        Map<String, Object> choice = choices.get(0);
        Map<String, Object> message = (Map<String, Object>) choice.get("message");
        String content = (String) message.get("content");

        // 解析 LLM 返回的多个查询变体（按行分割）
        List<String> rewrites = new ArrayList<>();
        rewrites.add(query);  // 始终包含原始查询
        if (content != null) {
            for (String line : content.split("\n")) {
                String trimmed = line.replaceAll("^[\\d\\.\\-\\* ]+", "").trim();
                if (!trimmed.isEmpty() && !trimmed.equals(query) && trimmed.length() < 2000) {
                    rewrites.add(trimmed);
                }
            }
        }

        log.debug("查询改写: {} -> {} 个变体", query.length() > 30 ? query.substring(0, 30) + "..." : query, rewrites.size());
        return rewrites;
    }

    private Map<String, Object> buildRequestBody(String model, String query) {
        Map<String, Object> body = new HashMap<>();
        body.put("model", model);

        List<Map<String, Object>> messages = new ArrayList<>();

        Map<String, Object> systemMsg = new HashMap<>();
        systemMsg.put("role", "system");
        systemMsg.put("content", """
            你是中医检索专家。请将用户的问题改写为2-3个语义相同但表述不同的查询，
            用于提高检索召回率。每个改写一行，不要编号，不要解释，只输出查询文本。""");
        messages.add(systemMsg);

        Map<String, Object> userMsg = new HashMap<>();
        userMsg.put("role", "user");
        userMsg.put("content", query);
        messages.add(userMsg);

        body.put("messages", messages);
        body.put("temperature", 0.3);
        body.put("max_tokens", 200);

        return body;
    }
}