检索性能优化

package nexus.io.mosskb.service.kb;

import java.util.Arrays;
import java.util.concurrent.locks.Lock;
import org.postgresql.util.PGobject;
import com.google.common.util.concurrent.Striped;
import nexus.io.chat.PlatformInput;
import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.db.utils.PgVectorUtils;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.tio.utils.snowflake.SnowflakeIdUtils;

/** Cache keys isolate model instances; documents and queries use the same model. */
public class KbEmbeddingService {
  private static final Striped<Lock> LOCKS = Striped.lock(256);

  public Long getVectorId(String text) { return cached(text).getLong("id"); }
  public PGobject getVector(String text) { return cached(text).get("v"); }
  public PGobject getVectorForModel(String text, Long modelId) { return cached(text, modelId).get("v"); }
  public Long getVectorIdForModel(String text, Long modelId) { return cached(text, modelId).getLong("id"); }

  private Row cached(String text) { return cached(text, 1002L); }

  private Row cached(String text, Long modelId) {
    if (text == null || text.isBlank()) {
      throw new IllegalArgumentException("Embedding text must not be empty");
    }
    String hash = Md5Utils.md5Hex(text);
    long resolvedId = modelId == null ? 1002L : modelId;
    Row selected = Db.findById("moss_kb_model", resolvedId);
    if (selected == null || !"EMBEDDING".equals(selected.getStr("model_type"))) {
      throw new IllegalArgumentException("知识库配置的向量模型不存在");
    }
    nexus.io.mosskb.vo.CredentialVo credential = ModelCatalogService.credential(selected);
    String base = ModelCatalogService.baseUrl(credential.getApi_base());
    String model = resolvedId == 1002L ? KnowledgeModelService.embeddingModel() + ":1024"
        : resolvedId + ":" + Md5Utils.md5Hex(base + "|" + selected.getStr("model_name")) + ":1024";
    Lock lock = LOCKS.get(model + hash);
    lock.lock();
    try {
      Row row = Db.findFirst("select id,v from moss_kb_embedding_cache where md5=? and m=?", hash, model);
      if (row != null) {
        return row;
      }
      // 未命中：在锁内调用远程 embeddings 接口，校验 1024 维与有限数值后写回缓存
      return embedAndCache(text, hash, model, base, credential, selected);
    } finally {
      lock.unlock();
    }
  }
}

package nexus.io.mosskb.service.kb;

import java.util.ArrayList;
import java.util.List;
import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.vo.MossKbRetrieveResult;
import nexus.io.mosskb.vo.ParagraphSearchResultVo;

public class MossKbParagraphRetrieveService {

  public List<Row> searchRows(Long[] datasets, Float threshold, Integer limit, String question, String mode) {
    if (datasets == null || datasets.length == 0) {
      return List.of();
    }
    if (question == null || question.isBlank()) {
      throw new IllegalArgumentException("检索问题不能为空");
    }
    String searchMode = mode == null ? "embedding" : mode;
    if (!List.of("embedding", "keywords", "blend").contains(searchMode)) {
      throw new IllegalArgumentException("未知的检索模式");
    }
    // keywords 模式不需要向量，直接单次查询
    if (!searchMode.equals("keywords")) {
      java.util.Map<Long, List<Long>> groups = new java.util.LinkedHashMap<>();
      for (Row dataset : Db.find("select id,coalesce(embedding_mode_id,1002) as model_id from moss_kb_dataset where id=any(?) and deleted=0", (Object) datasets)) {
        groups.computeIfAbsent(dataset.getLong("model_id"), key -> new ArrayList<>()).add(dataset.getLong("id"));
      }
      List<Row> all = new ArrayList<>();
      for (java.util.Map.Entry<Long, List<Long>> group : groups.entrySet()) {
        Long vectorId = Aop.get(KbEmbeddingService.class).getVectorIdForModel(question, group.getKey());
        all.addAll(searchGroup(group.getValue().toArray(new Long[0]), threshold, limit, question, searchMode, vectorId));
      }
      all.sort(java.util.Comparator.<Row>comparingDouble(row -> ((Number) row.get("similarity")).doubleValue()).reversed().thenComparing(row -> row.getLong("id")));
      return new ArrayList<>(all.subList(0, Math.min(all.size(), limit == null ? 10 : Math.max(1, Math.min(50, limit)))));
    }
    return searchGroup(datasets, threshold, limit, question, searchMode, null);
  }
}

  private List<Row> searchGroup(Long[] datasets, Float threshold, Integer limit, String question, String searchMode, Long vectorId) {
    List<Object> args = new ArrayList<>();
    String score;
    String cacheJoin = "";
    if (searchMode.equals("keywords")) {
      score = "word_similarity(?,p.content)";
      args.add(question);
    } else {
      score = "1-(p.embedding <=> c.v)";
      if (searchMode.equals("blend")) {
        score = "(" + score + ")+word_similarity(?,p.content)";
        args.add(question);
      }
      cacheJoin = " join moss_kb_embedding_cache c on c.id=?";
      args.add(vectorId);
    }
    String sql = "select sub.*,similarity as comprehensive_score from (select p.id,p.content,p.title,p.status,p.hit_num,p.is_active,p.dataset_id,p.document_id,d.name as document_name,ds.name as dataset_name," + score
        + " as similarity from moss_kb_paragraph p join moss_kb_document d on d.id=p.document_id join moss_kb_dataset ds on ds.id=p.dataset_id" + cacheJoin
        + " where p.is_active=true and d.is_active=true and p.deleted=0 and d.deleted=0 and ds.deleted=0 and p.dataset_id=any(?)) sub where similarity>=? order by similarity desc,id limit ?";
    args.add(datasets);
    args.add(threshold == null ? 0 : Math.max(0, threshold));
    args.add(limit == null ? 10 : Math.max(1, Math.min(50, limit)));
    return Db.find(sql, args.toArray());
  }

package nexus.io.mosskb.service.kb;

import java.util.ArrayList;
import java.util.List;
import com.jfinal.kit.Kv;
import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.vo.Paragraph;
import nexus.io.mosskb.vo.ParagraphBatchVo;
import nexus.io.model.result.ResultVo;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.tio.utils.hutool.FilenameUtils;
import nexus.io.tio.utils.snowflake.SnowflakeIdUtils;

public class MossKbParagraphSplitService {

  public ResultVo batch(Long userId, Long datasetId, List<ParagraphBatchVo> documents) {
    // 权限与参数校验（略）
    List<Kv> result = new ArrayList<>();
    for (ParagraphBatchVo input : documents) {
      Long existingId = Db.queryLong("select id from moss_kb_document where user_id=? and file_id=? and dataset_id=?", userId, input.getId(), datasetId);
      long documentId = existingId == null ? SnowflakeIdUtils.id() : existingId;
      List<Row> paragraphs = new ArrayList<>();
      int charLength = 0;
      String type = FilenameUtils.getSuffix(input.getName());
      // 远程请求全部完成后再发布文档或替换已有分段
      for (Paragraph p : input.getParagraphs()) {
        charLength += p.getContent().length();
        paragraphs.add(Row.by("id", SnowflakeIdUtils.id()).set("source_id", input.getId())
            .set("source_type", type).set("title", p.getTitle()).set("content", p.getContent())
            .set("md5", Md5Utils.md5Hex(p.getContent())).set("status", "2").set("hit_num", 0)
            .set("is_active", true).set("dataset_id", datasetId).set("document_id", documentId)
            .set("embedding", Aop.get(KbEmbeddingService.class).getVectorForModel(p.getContent(), dataset.getLong("embedding_mode_id"))));
      }
      boolean saved = Db.tx(() -> {
        if (existingId == null) {
          Db.save("moss_kb_document", document);
        } else {
          Db.update("moss_kb_document", document);
          Db.update("delete from moss_kb_paragraph where document_id=?", documentId);
        }
        Db.batchSave("moss_kb_paragraph", paragraphs, 2000);
        return true;
      });
      if (!saved) {
        return ResultVo.fail("文档保存失败，请重试");
      }
      result.add(document.toKv());
    }
    return ResultVo.ok(result);
  }
}

package nexus.io.mosskb.service.kb;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import org.postgresql.util.PGobject;
import lombok.extern.slf4j.Slf4j;
import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.constant.MossKbTableNames;
import nexus.io.mosskb.model.MossKbSentence;
import nexus.io.mosskb.utils.ExecutorServiceUtils;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.tio.utils.snowflake.SnowflakeIdUtils;

@Slf4j
public class MossKbSentenceService {
  KbEmbeddingService mossKbEmbeddingService = Aop.get(KbEmbeddingService.class);
  CompletionService<Row> completionServiceRow = new ExecutorCompletionService<>(ExecutorServiceUtils.getExecutorService());

  public boolean splitToSentenceAndSave(Long dataset_id, String modelName, List<Row> paragraphRecords, Long documentIdFinal) {
    List<MossKbSentence> sentences = new ArrayList<>();
    // 递归分句并构造 MossKbSentence（略）
    List<Future<Row>> futures = new ArrayList<>(sentences.size());
    for (MossKbSentence sentence : sentences) {
      futures.add(completionServiceRow.submit(() -> {
        PGobject vector = mossKbEmbeddingService.getVector(sentence.getContent(), modelName);
        Row record = sentence.toRow();
        record.set("embedding", vector);
        return record;
      }));
    }
    List<Row> sentenceRows = new ArrayList<>();
    for (int i = 0; i < sentences.size(); i++) {
      try {
        Row record = completionServiceRow.take().get();
        if (record != null) {
          sentenceRows.add(record);
        }
      } catch (Exception e) {
        log.error(e.getMessage(), e);
      }
    }
    return Db.tx(() -> {
      Db.deleteById(MossKbTableNames.moss_kb_sentence, "document_id", documentIdFinal);
      Db.batchSave(MossKbTableNames.moss_kb_sentence, sentenceRows, 2000);
      return true;
    });
  }
}

package nexus.io.mosskb.service.kb;

import java.util.Locale;

  private String normalize(String query) {
    return query.toLowerCase(Locale.ROOT).replaceAll("[\\s\\p{P}]+", "");
  }

package nexus.io.mosskb.controller;

import nexus.io.annotation.Get;
import nexus.io.annotation.RequestPath;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.service.kb.MossKbDatasetHitTestService;
import nexus.io.model.result.ResultVo;
import nexus.io.tio.boot.http.TioRequestContext;
import nexus.io.tio.http.common.HttpRequest;

@RequestPath("/api/dataset")
public class ApiDatasetController {

  @Get("/{id}/hit_test")
  public ResultVo hitTest(Long id, HttpRequest request) {
    String query_text = request.getParam("query_text");
    Double similarity = request.getDouble("similarity");
    Integer top_number = request.getInt("top_number");
    String search_mode = request.getParam("search_mode");
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetHitTestService.class).hitTest(userId, id, query_text, similarity, top_number,
        search_mode);
  }

package nexus.io.mosskb.vo;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.experimental.Accessors;

@Data
@NoArgsConstructor
@AllArgsConstructor
@Accessors(chain = true)
public class MossKbDatasetSettingVo {
  private Integer top_n;
  private Float similarity;
  private Integer max_paragraph_char_number;
  private String search_mode;
  private MossKbApplicationNoReferencesSetting no_references_setting;
}

      MossKbDatasetSettingVo dataset_setting = applicationVo.getDataset_setting();
      Float similarity = 0.0f;
      Integer top_n = 10;
      if (dataset_setting != null) {
        if (dataset_setting.getSimilarity() != null) {
          similarity = dataset_setting.getSimilarity();
        }
        if (dataset_setting.getTop_n() != null) {
          top_n = dataset_setting.getTop_n();
        }
      }