agent-management/backend/src/main/java/com/agent/management/service/impl/EmbeddingBridgeClient.java

package com.agent.management.service.impl;

import com.agent.management.common.LogRedactor;
import com.agent.management.common.exception.BusinessException;
import com.agent.management.config.EmbeddingBridgeProperties;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.*;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestClientResponseException;
import org.springframework.web.client.RestTemplate;

import java.util.List;
import java.util.Map;

/**
 * 本地 Embedding Bridge HTTP 客户端
 */
@Slf4j
@Component
public class EmbeddingBridgeClient {

    /**
     * LLM/智能体调用专用 logger：
     * 由 logback-spring.xml 中 LLM_HTTP logger 配置为 DEBUG 级别，仅写入独立日志文件，不输出到控制台。
     */
    private static final Logger LLM_HTTP = LoggerFactory.getLogger("LLM_HTTP");

    private final EmbeddingBridgeProperties props;
    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;

    public EmbeddingBridgeClient(EmbeddingBridgeProperties props) {
        this.props = props;
        SimpleClientHttpRequestFactory factory = new SimpleClientHttpRequestFactory();
        int timeout = Math.max(props.getHttpTimeout(), 60000);
        factory.setConnectTimeout(5000);
        factory.setReadTimeout(timeout);
        this.restTemplate = new RestTemplate(factory);
        this.objectMapper = new ObjectMapper();
    }

    private String baseUrl() {
        return "http://" + props.getHost() + ":" + props.getPort();
    }

    private HttpHeaders headers() {
        HttpHeaders h = new HttpHeaders();
        h.setContentType(MediaType.APPLICATION_JSON);
        if (props.getAuthToken() != null && !props.getAuthToken().isBlank()) {
            h.set("X-Bridge-Token", props.getAuthToken());
        }
        return h;
    }

    /**
     * 健康检查
     */
    public boolean health() {
        try {
            ResponseEntity<String> resp = restTemplate.exchange(
                    baseUrl() + "/health", HttpMethod.GET, new HttpEntity<>(headers()), String.class);
            return resp.getStatusCode().is2xxSuccessful();
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 对文档全文执行三级分块 + 向量化
     */
    @SuppressWarnings("unchecked")
    public Map<String, Object> indexDocument(long documentId, String text, String filename,
                                             String fileType, String filePath, Long categoryId) {
        Map<String, Object> body = Map.of(
                "document_id", documentId,
                "text", text,
                "filename", filename,
                "file_type", fileType,
                "file_path", filePath == null ? "" : filePath,
                "page_number", 0,
                "chunk_size", props.getChunkSize(),
                "chunk_overlap", props.getChunkOverlap(),
                "category_id", categoryId == null ? 0L : categoryId
        );
        return postJson("/index", body, new TypeReference<>() {});
    }

    /**
     * 批量向量化指定 chunk（仅用于补向量）
     */
    @SuppressWarnings("unchecked")
    public List<String> vectorizeChunks(long documentId, String filename, String fileType,
                                        String filePath, List<Map<String, Object>> chunks) {
        Map<String, Object> body = Map.of(
                "document_id", documentId,
                "filename", filename,
                "file_type", fileType,
                "file_path", filePath,
                "chunks", chunks
        );
        Map<String, Object> resp = postJson("/vectorize", body, new TypeReference<>() {});
        Object ids = resp.get("vector_ids");
        if (ids instanceof List<?> list) {
            return list.stream().map(String::valueOf).toList();
        }
        return List.of();
    }

    /**
     * 按文档 ID 删除 Milvus 向量
     */
    public int deleteByDocument(long documentId) {
        Map<String, Object> body = Map.of("document_id", documentId);
        Map<String, Object> resp = postJson("/delete_by_document", body, new TypeReference<>() {});
        Object deleted = resp.get("deleted");
        return deleted instanceof Number n ? n.intValue() : 0;
    }

    /**
     * 按 vector_id 删除 Milvus 向量
     */
    public int deleteByVectorIds(List<String> vectorIds) {
        Map<String, Object> body = Map.of("vector_ids", vectorIds);
        Map<String, Object> resp = postJson("/delete_by_vector_ids", body, new TypeReference<>() {});
        Object deleted = resp.get("deleted");
        return deleted instanceof Number n ? n.intValue() : 0;
    }

    public Map<String, Object> retrieve(String query, int topK, String mode, String filterExpr) {
        Map<String, Object> body = Map.of(
                "query", query, "top_k", topK,
                "mode", mode == null || mode.isBlank() ? "hybrid" : mode,
                "filter_expr", filterExpr == null ? "" : filterExpr);
        return postJson("/retrieve", body, new TypeReference<>() {});
    }

    @SuppressWarnings("unchecked")
    public List<List<Double>> embed(List<String> texts) {
        Map<String, Object> response = postJson("/embed", Map.of("texts", texts), new TypeReference<>() {});
        Object vectors = response.get("vectors");
        if (!(vectors instanceof List<?> rows)) return List.of();
        return rows.stream().map(row -> ((List<?>) row).stream()
                .map(value -> ((Number) value).doubleValue()).toList()).toList();
    }

    private <T> T postJson(String path, Object body, TypeReference<T> typeRef) {
        try {
            String jsonBody = objectMapper.writeValueAsString(body);
            HttpEntity<String> entity = new HttpEntity<>(jsonBody, headers());
            String fullUrl = baseUrl() + path;

            // === 请求 DEBUG 日志（仅写入 LLM 日志文件，不输出到控制台） ===
            LLM_HTTP.debug("[Embed-REQ] ========== 请求开始 ==========");
            LLM_HTTP.debug("[Embed-REQ] POST {}", fullUrl);
            LLM_HTTP.debug("[Embed-REQ] Body ({} bytes): {}", jsonBody.length(), jsonBody);

            long startTime = System.currentTimeMillis();
            ResponseEntity<String> resp = restTemplate.exchange(
                    fullUrl, HttpMethod.POST, entity, String.class);
            long elapsed = System.currentTimeMillis() - startTime;

            String respBody = resp.getBody();
            // === 响应 DEBUG 日志 ===
            LLM_HTTP.debug("[Embed-RESP] HTTP {} ({}ms), body 长度={}",
                    resp.getStatusCode(), elapsed, respBody == null ? 0 : respBody.length());
            LLM_HTTP.debug("[Embed-RESP] Body: {}", respBody == null ? null : LogRedactor.redactVectors(respBody));
            LLM_HTTP.debug("[Embed-RESP] ========== 请求结束 ==========");

            if (!resp.getStatusCode().is2xxSuccessful() || respBody == null) {
                throw new BusinessException("Embedding Bridge 调用失败: HTTP " + resp.getStatusCode());
            }
            return objectMapper.readValue(respBody, typeRef);
        } catch (RestClientResponseException e) {
            log.error("Embedding Bridge 调用失败: {}", e.getResponseBodyAsString(), e);
            throw new BusinessException("Embedding Bridge 调用失败: " + e.getResponseBodyAsString());
        } catch (Exception e) {
            log.error("Embedding Bridge 调用失败: {}", e.getMessage(), e);
            throw new BusinessException("Embedding Bridge 调用失败: " + e.getMessage());
        }
    }
}