init

fn-opensearch/src/main/java/net/yaoyi/pipeline/opensearch/OpenSearchHandlerV2.java

@@ -0,0 +1,359 @@
+package net.yaoyi.pipeline.opensearch;
+
+import com.aliyun.fc.runtime.Context;
+import com.aliyun.fc.runtime.PojoRequestHandler;
+import com.aliyun.ha3engine.vector.models.*;
+import com.fasterxml.jackson.databind.JsonNode;
+import lombok.extern.slf4j.Slf4j;
+import net.yaoyi.algorithm.domain.model.TaskImageDupDetail;
+import net.yaoyi.algorithm.domain.rabbitmq.TaskImageDupMsg;
+import net.yaoyi.pipeline.common.JacksonUtils;
+import net.yaoyi.pipeline.model.index.TaskImageIndex;
+import net.yaoyi.pipeline.model.rabbitmq.TaskImageMsg;
+import net.yaoyi.pipeline.model.rabbitmq.TaskImageResultMsg;
+
+import java.math.BigDecimal;
+import java.time.LocalDateTime;
+import java.time.format.DateTimeFormatter;
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+@Slf4j
+public class OpenSearchHandlerV2 implements PojoRequestHandler<List<TaskImageMsg>, TaskImageResultMsg> {
+
+    private static final int TOP_K = 10;
+    private static final float SCORE_THRESHOLD = 0.7f;
+
+    @Override
+    public TaskImageResultMsg handleRequest(List<TaskImageMsg> taskImageMsgList, Context context) {
+        log.info("收到请求消息：{}", JacksonUtils.toJson(taskImageMsgList));
+        if (taskImageMsgList.size() != 1) {
+            throw new RuntimeException("参数格式错误");
+        }
+
+        TaskImageMsg taskImageMsg = taskImageMsgList.get(0);
+        Long taskId = taskImageMsg.getTaskId();
+        String namespace = taskImageMsg.getNamespace();
+
+        // 1. 从数据到索引：将 image 数据转换为向量并写入 OpenSearch
+        saveByTask(taskImageMsg);
+
+        if (taskImageMsg.getStatus() < 2) {
+            return new TaskImageResultMsg();
+        }
+
+        List<TaskImageIndex> list = Collections.emptyList();
+        int times = 5;
+        while (times > 0) {
+            try {
+                Thread.sleep(2000);
+            } catch (Exception e) {
+                log.warn("等待被中断：" + e.getMessage());
+            }
+            list = fetchByTaskId(taskImageMsg.getIndexTableName(), namespace, taskId);
+            if (dataQuerySuccess(taskImageMsg.getImages(), list)) {
+                break;
+            } else {
+                times--;
+            }
+        }
+
+        // 2. 以 namespace + taskId 维度来完成数据更新，image相同的数据保持id不变
+        for (TaskImageIndex taskImageIndex : list) {
+            if (taskImageIndex.getVector() == null) {
+                throw new RuntimeException("尚未更新索引信息，id:" + taskImageIndex.getId());
+            }
+
+            // 维度一：namespace 维度
+            String fileterStartDateTime = LocalDateTime.now().minusYears(1).format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
+            String nsFilter = String.format("sourceSendPackageDeptId = %d AND taskId != %d AND status = 1 AND taskTime > \"%s\"",
+                    taskImageIndex.getSourceSendPackageDeptId(), taskImageIndex.getTaskId(), fileterStartDateTime);
+            List<TaskImageIndex.MultiResult> nsResult = knnSearch(taskImageMsg.getIndexTableName(), namespace, taskImageIndex.getVector(), nsFilter, "namespace");
+            List<TaskImageIndex.MultiResult> multiResultList = new ArrayList<>(nsResult);
+
+            // 维度二：namespace + taskId 维度
+            String nsTaskFilter = String.format("taskId = %d AND id != \"%s\"", taskImageIndex.getTaskId(), taskImageIndex.getId());
+            List<TaskImageIndex.MultiResult> nsTaskResult = knnSearch(taskImageMsg.getIndexTableName(), namespace, taskImageIndex.getVector(), nsTaskFilter, "task");
+
+            nsTaskResult = nsTaskResult.stream().filter(item -> !item.getImage().equals(taskImageIndex.getImage())).collect(Collectors.toList());
+            multiResultList.addAll(nsTaskResult);
+            taskImageIndex.setMultiResult(JacksonUtils.toJson(multiResultList));
+        }
+
+        //先落库
+        pushDocument(taskImageMsg.getIndexTableName(), list, Collections.emptyList());
+        //构建消息体
+        List<TaskImageDupMsg> msgList = new ArrayList<>();
+        for (TaskImageIndex taskImageIndex : list) {
+            msgList.add(buildImageMultiMsg(taskImageMsg, taskImageIndex));
+        }
+
+        TaskImageResultMsg taskImageResultMsg = new TaskImageResultMsg();
+        taskImageResultMsg.setQueueName(taskImageMsg.getResultQueueName());
+        taskImageResultMsg.setMsgList(msgList);
+        return taskImageResultMsg;
+    }
+
+    private boolean dataQuerySuccess(List<TaskImageMsg.TaskImage> images, List<TaskImageIndex> list) {
+        if (images.size() != list.size()) {
+            return false;
+        }
+
+        Set<String> set = images.stream().map(TaskImageMsg.TaskImage::getImage).collect(Collectors.toSet());
+        Set<String> set2 = list.stream().map(TaskImageIndex::getImage).collect(Collectors.toSet());
+        return set.removeAll(set2) && set.isEmpty();
+    }
+
+    private void saveByTask(TaskImageMsg taskImageMsg) {
+        final List<TaskImageIndex> addList = new ArrayList<>();
+        final List<TaskImageIndex> deleteList = new ArrayList<>();
+
+        List<TaskImageIndex> taskImageList = fetchByTaskId(taskImageMsg.getIndexTableName(), taskImageMsg.getNamespace(), taskImageMsg.getTaskId());
+        if (taskImageMsg.getStatus() == 0) {
+            deleteList.addAll(taskImageList);
+        } else {
+            Map<String, TaskImageIndex> map = taskImageList.stream().collect(Collectors.toMap(TaskImageIndex::getImage,
+                    Function.identity(), (a, b) -> {
+                        deleteList.add(b);
+                        return a;
+                    }));
+            List<TaskImageMsg.TaskImage> images = taskImageMsg.getImages();
+            for (TaskImageMsg.TaskImage imageInfo : images) {
+                TaskImageIndex exists = map.remove(imageInfo.getImage());
+                if (exists != null) {
+                    TaskImageIndex index = buildTaskImageIndex(exists.getId(), taskImageMsg, imageInfo);
+                    index.setVector(exists.getVector());
+                    index.setMultiResult(exists.getMultiResult());
+                    addList.add(index);
+                } else {
+                    String id = UUID.randomUUID().toString().replaceAll("-", "");
+                    addList.add(buildTaskImageIndex(id, taskImageMsg, imageInfo));
+                }
+            }
+
+            if (!map.isEmpty()) {
+                deleteList.addAll(map.values());
+            }
+        }
+
+        pushDocument(taskImageMsg.getIndexTableName(), addList, deleteList);
+    }
+
+    private void pushDocument(String tableName, List<TaskImageIndex> addList, List<TaskImageIndex> deleteList) {
+        if (addList.isEmpty() && deleteList.isEmpty()) {
+            return;
+        }
+
+        String fullTableName = Contents.INSTANCE_ID + "_" + tableName;
+
+        ArrayList<Map<String, ?>> documents = new ArrayList<>();
+        for (TaskImageIndex taskImageIndex : addList) {
+            documents.add(Map.of("cmd", "add", "fields", JacksonUtils.toMap(taskImageIndex)));
+        }
+        for (TaskImageIndex taskImageIndex : deleteList) {
+            documents.add(Map.of("cmd", "delete", "fields", Map.of("id", taskImageIndex.getId())));
+        }
+
+        PushDocumentsRequest request = new PushDocumentsRequest();
+        request.setBody(documents);
+        try {
+            PushDocumentsResponse response = OpenSearchUtils.getClient().pushDocuments(fullTableName, "id", request);
+            String responseBody = response.getBody();
+            JsonNode bodyNode = JacksonUtils.toJsonNode(responseBody);
+
+            if (!bodyNode.has("status") || !"OK".equals(bodyNode.get("status").asText())) {
+                log.error("push document with multiResult failed, request:{}, response:{}", JacksonUtils.toJson(request), responseBody);
+                throw new RuntimeException("push document failed");
+            }
+        } catch (Exception e) {
+            log.error("push document with multiResult exception, request:{}", JacksonUtils.toJson(request), e);
+            throw new RuntimeException(e);
+        }
+    }
+
+
+    private List<TaskImageIndex> fetchByTaskId(String tableName, String namespace, Long taskId) {
+        try {
+            FetchRequest request = new FetchRequest();
+            request.setTableName(tableName);
+            request.setFilter(String.format("taskId=%d AND namespace=\"%s\"", taskId, namespace));
+            request.setIncludeVector(true);
+
+            SearchResponse response = OpenSearchUtils.getClient().fetch(request);
+            JsonNode bodyNode = JacksonUtils.toJsonNode(response.getBody());
+            if (bodyNode.has("errorCode")) {
+                log.error("查询失败，body:{}", response.getBody());
+                throw new RuntimeException(String.format("查询失败,namespace:[%s], taskId:%d ", namespace, taskId));
+            }
+
+            int totalCount = bodyNode.get("totalCount").asInt();
+            if (totalCount == 0) {
+                return Collections.emptyList();
+            }
+            return parseToModelList(bodyNode.get("result"));
+        } catch (Exception e) {
+            throw new RuntimeException(String.format("查询异常,namespace:[%s], taskId:%d ", namespace, taskId), e);
+        }
+    }
+
+    private List<TaskImageIndex> parseToModelList(JsonNode result) {
+        List<TaskImageIndex> list = new ArrayList<>();
+        for (JsonNode jsonNode : result) {
+            list.add(parseToModel(jsonNode));
+        }
+        return list;
+    }
+
+    private TaskImageIndex parseToModel(JsonNode result) {
+        TaskImageIndex index = new TaskImageIndex();
+        if (result.has("id")) {
+            index.setId(result.get("id").asText());
+        }
+        if (result.has("namespace")) {
+            index.setNamespace(result.get("namespace").asText());
+        }
+        if (result.has("vector")) {
+            index.setVector(extractVector(result.get("vector")));
+        }
+
+        JsonNode fields = result.get("fields");
+        if (fields != null) {
+            if (fields.has("taskId")) {
+                index.setTaskId(fields.get("taskId").asLong());
+            }
+            if (fields.has("sourceSendPackageDeptId")) {
+                index.setSourceSendPackageDeptId(fields.get("sourceSendPackageDeptId").asLong());
+            }
+            if (fields.has("taskTypeId")) {
+                index.setTaskTypeId(fields.get("taskTypeId").asInt());
+            }
+            if (fields.has("userId")) {
+                index.setUserId(fields.get("userId").asLong());
+            }
+            if (fields.has("status")) {
+                index.setStatus(fields.get("status").asInt());
+            }
+            if (fields.has("image")) {
+                index.setImage(fields.get("image").asText());
+            }
+            if (fields.has("taskFiledValue")) {
+                index.setTaskFiledValue(fields.get("taskFiledValue").asText());
+            }
+            if (fields.has("taskTime")) {
+                index.setTaskTime(fields.get("taskTime").asText());
+            }
+            if (fields.has("multiResult")) {
+                index.setMultiResult(fields.get("multiResult").asText());
+            }
+        }
+        return index;
+    }
+
+    private List<TaskImageIndex.MultiResult> knnSearch(String tableName, String namespace, List<Float> vector, String filter, String type) {
+        try {
+            QueryRequest knn = new QueryRequest();
+            knn.setNamespace(namespace);
+            knn.setVector(vector);
+            knn.setTopK(TOP_K);
+            knn.setFilter(filter);
+            knn.setScoreThreshold(SCORE_THRESHOLD);
+
+            SearchRequest request = new SearchRequest();
+            request.setTableName(tableName);
+            request.setKnn(knn);
+            request.setOutputFields(List.of("taskId", "image", "taskFiledValue"));
+
+            SearchResponse response = OpenSearchUtils.getClient().search(request);
+            String responseBody = response.getBody();
+            log.info("knn search filter:{}, namespace:{}, response:{}", filter, namespace, responseBody);
+
+            JsonNode bodyNode = JacksonUtils.toJsonNode(responseBody);
+            if (bodyNode.has("errorCode")) {
+                log.error("knn search fail. filter:{}, namespace:{}, response:{}", filter, namespace, responseBody);
+                throw new RuntimeException("knn search failed");
+            }
+            return parseToMultiResultList(bodyNode, type);
+        } catch (Exception e) {
+            log.error("knnSearch exception, filter:{}", filter, e);
+            throw new RuntimeException("对比异常，filter:" + filter, e);
+        }
+    }
+
+    private List<TaskImageIndex.MultiResult> parseToMultiResultList(JsonNode bodyNode, String type) {
+        JsonNode resultArray;
+        if (bodyNode == null || !bodyNode.has("result") || !(resultArray = bodyNode.get("result")).isArray() || resultArray.isEmpty()) {
+            return Collections.emptyList();
+        }
+
+        List<TaskImageIndex.MultiResult> list = new ArrayList<>();
+        for (JsonNode jsonNode : resultArray) {
+            TaskImageIndex.MultiResult multiResult = new TaskImageIndex.MultiResult();
+            JsonNode fields = jsonNode.get("fields");
+
+            multiResult.setTaskId(fields.get("taskId").asLong());
+            multiResult.setImage(fields.get("image").asText());
+            multiResult.setTaskFiledValue(fields.get("taskFiledValue").asText());
+
+            if (jsonNode.has("score")) {
+                multiResult.setScore(jsonNode.get("score").asDouble());
+            } else if (jsonNode.has("similarityScore")) {
+                multiResult.setScore(jsonNode.get("similarityScore").asDouble());
+            }
+            multiResult.setType(type);
+            list.add(multiResult);
+        }
+        return list;
+    }
+
+    private TaskImageIndex buildTaskImageIndex(String id, TaskImageMsg msg, TaskImageMsg.TaskImage imageInfo) {
+        TaskImageIndex index = new TaskImageIndex();
+        index.setId(id);
+        index.setNamespace(msg.getNamespace());
+        index.setTaskId(msg.getTaskId());
+        index.setSourceSendPackageDeptId(msg.getSourceSendPackageDeptId());
+        index.setTaskTypeId(msg.getTaskTypeId());
+        index.setUserId(msg.getUserId());
+        index.setStatus(msg.getStatus());
+        index.setTaskTime(msg.getTaskTime());
+        index.setImage(imageInfo.getImage());
+        index.setTaskFiledValue(imageInfo.getTaskFiledValue());
+        return index;
+    }
+
+    private List<Float> extractVector(JsonNode vectorNode) {
+        if (vectorNode == null || !vectorNode.isArray() || vectorNode.isEmpty()) {
+            return null;
+        }
+        List<Float> vector = new ArrayList<>(vectorNode.size());
+        for (JsonNode v : vectorNode) {
+            vector.add((float) v.asDouble());
+        }
+        return vector;
+    }
+
+    private TaskImageDupMsg buildImageMultiMsg(TaskImageMsg taskImageMsg, TaskImageIndex taskImageIndex) {
+        TaskImageDupMsg msg = new TaskImageDupMsg();
+        msg.setAlgoConfigId(taskImageMsg.getAlgoConfigId());
+        msg.setTaskId(taskImageIndex.getTaskId());
+        msg.setImage(taskImageIndex.getImage());
+
+        if (taskImageIndex.getMultiResult() != null) {
+            TaskImageIndex.MultiResult[] results = JacksonUtils.parseObject(taskImageIndex.getMultiResult(), TaskImageIndex.MultiResult[].class);
+            if (results != null) {
+                List<TaskImageDupDetail> multiInfoList = new ArrayList<>();
+                for (TaskImageIndex.MultiResult r : results) {
+                    TaskImageDupDetail info = new TaskImageDupDetail();
+                    info.setTaskId(r.getTaskId());
+                    info.setImage(r.getImage());
+                    info.setTaskFiledValue(r.getTaskFiledValue());
+                    info.setScore(BigDecimal.valueOf(r.getScore()));
+                    multiInfoList.add(info);
+                }
+                msg.setMultiResult(multiInfoList);
+            }
+        }
+        return msg;
+    }
+
+}