片段向量

概述

片段向量化调用知识库选择的远程向量模型，当前要求返回 1024 维。为了减少向量模型的调用次数，我们将向量模型的返回结果缓存到 moss_kb_embedding_cache 表中。以下是相关接口及代码的实现细节。

接口文档

批量处理文档向量

请求方法: POST

请求路径: /api/dataset/{id}/document/_bach

请求示例:

[
  {
    "name": "ICS111_31391_Miller_Syllabus_F24.pdf",
    "id": "file_id",
    "paragraphs": [
      {
        "title": "",
        "content": "ICS 111- Introduction to Computer Science I, 31391\nFall"
      }
    ]
  }
]

响应示例:

{
  "code": 200,
  "message": "success",
  "data": [
    {
      "create_time": "2024-11-05T14:24:12.315",
      "update_time": "2024-11-05T14:24:12.315",
      "id": "xxx",
      "name": "ICS111_31391_Miller_Syllabus_F24.pdf",
      "char_length": 20755,
      "status": "2",
      "is_active": true,
      "type": "0",
      "meta": {},
      "dataset_id": "xxx",
      "hit_handling_method": "optimization",
      "directly_return_similarity": 0.9,
      "paragraph_count": 6
    }
  ]
}

代码实现

控制器层

package nexus.io.mosskb.controller;

import java.util.List;

import nexus.io.annotation.Delete;
import nexus.io.annotation.Get;
import nexus.io.annotation.Post;
import nexus.io.annotation.Put;
import nexus.io.annotation.RequestPath;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.service.kb.MossKbDatasetHitTestService;
import nexus.io.mosskb.service.kb.MossKbDatasetService;
import nexus.io.mosskb.service.kb.MossKbDocumentService;
import nexus.io.mosskb.service.kb.MossKbParagraphServcie;
import nexus.io.mosskb.service.kb.MossKbParagraphSplitService;
import nexus.io.mosskb.service.kb.MossKbProblemService;
import nexus.io.mosskb.vo.KbDatasetModel;
import nexus.io.mosskb.vo.MossKbUpdateDocumentRequestVo;
import nexus.io.mosskb.vo.Paragraph;
import nexus.io.mosskb.vo.ParagraphBatchVo;
import nexus.io.mosskb.vo.ProbrolemCreateBatch;
import nexus.io.model.result.ResultVo;
import nexus.io.tio.boot.http.TioRequestContext;
import nexus.io.tio.http.common.HttpRequest;
import nexus.io.tio.utils.json.JsonUtils;

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

  @Get("")
  public ResultVo list() {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetService.class).list(userId);
  }

  @Get("/{pageNo}/{pageSize}")
  public ResultVo page(Integer pageNo, Integer pageSize, String name) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetService.class).page(userId, pageNo, pageSize, name);
  }

  @Post("")
  public ResultVo save(HttpRequest request) {
    String bodyString = request.getBodyString();
    KbDatasetModel kbDatasetModel = JsonUtils.parse(bodyString, KbDatasetModel.class);
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetService.class).save(userId, kbDatasetModel);
  }

  @Get("/{id}")
  public ResultVo get(Long id) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetService.class).get(userId, id);
  }

  @Put("/{id}")
  public ResultVo updateDataset(Long id, HttpRequest request) {
    KbDatasetModel model = JsonUtils.parse(request.getBodyString(), KbDatasetModel.class);
    model.setId(id);
    return Aop.get(MossKbDatasetService.class).save(TioRequestContext.getUserIdLong(), model);
  }

  @Put("/{datasetId}/document/{documentId}/paragraph/{paragraphId}")
  public ResultVo updateParagraph(Long datasetId, Long documentId, Long paragraphId, HttpRequest request) {
    return Aop.get(MossKbParagraphServcie.class).update(TioRequestContext.getUserIdLong(), datasetId, documentId, paragraphId,
        JsonUtils.parse(request.getBodyString(), Paragraph.class));
  }

  @Delete("/{id}")
  public ResultVo delete(Long id) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDatasetService.class).delete(userId, id);
  }

  @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);
  }

  @Post("/{id}/document/_bach")
  public ResultVo batch(Long id, HttpRequest request) {
    String bodyString = request.getBodyString();
    Long userId = TioRequestContext.getUserIdLong();
    List<ParagraphBatchVo> list = JsonUtils.parseArray(bodyString, ParagraphBatchVo.class);
    return Aop.get(MossKbParagraphSplitService.class).batch(userId, id, list);
  }

  @Get("/{id}/document/{pageNo}/{pageSize}")
  public ResultVo pageDocument(Long id, Integer pageNo, Integer pageSize) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDocumentService.class).page(userId, id, pageNo, pageSize);
  }

  @Get("/{datasetId}/document")
  public ResultVo documentList(Long datasetId) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDocumentService.class).list(userId, datasetId);
  }

  @Get("/{datasetId}/document/{documentId}")
  public ResultVo getDocument(Long datasetId, Long documentId) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDocumentService.class).get(userId, datasetId, documentId);
  }

  @Delete("/{datasetId}/document/{documentId}")
  public ResultVo deleteDocument(Long datasetId, Long documentId) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbDocumentService.class).delete(userId, datasetId, documentId);
  }

  @Put("/{datasetId}/document/{documentId}")
  public ResultVo updateDocument(Long datasetId, Long documentId, HttpRequest httpRequest) {
    Long userId = TioRequestContext.getUserIdLong();
    String bodyString = httpRequest.getBodyString();
    if (bodyString == null) {
      return ResultVo.fail("request body can not be empty");
    }
    MossKbUpdateDocumentRequestVo vo = JsonUtils.parse(bodyString, MossKbUpdateDocumentRequestVo.class);
    return Aop.get(MossKbDocumentService.class).update(userId, datasetId, documentId, vo);
  }

  @Get("/{datasetId}/application")
  public ResultVo getApplicationByDatasetId(Long datasetId) {
    return Aop.get(MossKbDatasetService.class).getApplicationByDatasetId(datasetId);
  }

  @Post("/{datasetId}/document/{documentId}/paragraph")
  public ResultVo createParagraph(Long datasetId, Long documentId, HttpRequest request) {
    String bodyString = request.getBodyString();
    Paragraph paragraph = JsonUtils.parse(bodyString, Paragraph.class);
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbParagraphServcie.class).create(userId, datasetId, documentId, paragraph);
  }

  // paragraph
  @Get("/{datasetId}/document/{documentId}/paragraph/{pageNo}/{pageSize}")
  public ResultVo listParagraph(Long datasetId, Long documentId, Integer pageNo, Integer pageSize) {
    Long userId = TioRequestContext.getUserIdLong();
    return Aop.get(MossKbParagraphServcie.class).page(userId, datasetId, documentId, pageNo, pageSize);
  }

  @Get("/{datasetId}/document/{documentId}/paragraph/{paragraphId}/problem")
  public ResultVo listProblem(Long datasetId, Long documentId, Long paragraphId) {
    return Aop.get(MossKbParagraphServcie.class).listProblemByParagraphId(datasetId, documentId, paragraphId);
  }

  @Post("/{datasetId}/document/{documentId}/paragraph/{paragraphId}/problem")
  public ResultVo addProblems(Long datasetId, Long documentId, Long paragraphId, HttpRequest request) {
    String bodyString = request.getBodyString();
    String content = JsonUtils.parseToMap(bodyString, String.class, String.class).get("content");
    return Aop.get(MossKbParagraphServcie.class).addProblemById(datasetId, documentId, paragraphId, content);
  }

  @Post("/{datasetId}/problem")
  public ResultVo createProblem(Long datasetId, HttpRequest request) {
    String bodyString = request.getBodyString();
    List<String> problems = JsonUtils.parseArray(bodyString, String.class);
    return Aop.get(MossKbProblemService.class).create(datasetId, problems);
  }


  @Get("/{datasetId}/problem/{pageNo}/{pageSize}")
  public ResultVo pageDatasetProbleam(Long datasetId, Integer pageNo, Integer pageSize) {
    return Aop.get(MossKbProblemService.class).page(datasetId, pageNo, pageSize);
  }

  // api/dataset/443309276048408576/problem/444734959665311744/paragraph
  @Get("/{datasetId}/problem/{problemId}/paragraph")
  public ResultVo listParagraphByProblemId(Long datasetId, Long problemId) {
    return Aop.get(MossKbProblemService.class).listParagraphByProblemId(datasetId, problemId);
  }

  @Put("/{datasetId}/document/{documentId}/paragraph/{paragraphId}/problem/{problemId}/association")
  public ResultVo association(Long datasetId, Long documentId, Long paragraphId, Long problemId) {
    return Aop.get(MossKbProblemService.class).association(datasetId, documentId, paragraphId, problemId);
  }

  @Put("/{datasetId}/document/{documentId}/paragraph/{paragraphId}/problem/{problemId}/un_association")
  public ResultVo unAssociation(Long datasetId, Long documentId, Long paragraphId, Long problemId) {
    return Aop.get(MossKbProblemService.class).unAssociation(datasetId, documentId, paragraphId, problemId);
  }

  @Post("/{datasetId}/problem/_batch")
  public ResultVo addProbrolems(Long datasetId, HttpRequest request) {
    String bodyString = request.getBodyString();
    ProbrolemCreateBatch batchReequest = JsonUtils.parse(bodyString, ProbrolemCreateBatch.class);
    return Aop.get(MossKbProblemService.class).addProbrolems(datasetId, batchReequest);
  }

  @Delete("/{datasetId}/problem/{problemId}")
  public ResultVo deleteProblem(Long datasetId, Long problemId) {
    return Aop.get(MossKbProblemService.class).delete(datasetId, problemId);
  }
  

  @Delete("/{datasetId}/problem/_batch")
  public ResultVo deleteProblemBatch(Long datasetId, HttpRequest request) {
    String bodyString = request.getBodyString();
    List<Long> problems = JsonUtils.parseArray(bodyString, Long.class);
    return Aop.get(MossKbProblemService.class).delete(datasetId, problems);
  }

}

服务层

MossKbParagraphSplitService（文档分段与向量保存）

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) {
    Row dataset = Db.findById("moss_kb_dataset", datasetId);
    if (userId == null || dataset == null || (userId != 1L && !userId.equals(dataset.getLong("user_id")))) {
      return ResultVo.fail("知识库不存在或无权访问");
    }
    if (documents == null || documents.isEmpty()) {
      return ResultVo.fail("请选择需要导入的文档");
    }
    List<Kv> result = new ArrayList<>();
    for (ParagraphBatchVo input : documents) {
      if (input.getId() == null || input.getParagraphs() == null || input.getParagraphs().isEmpty()) {
        return ResultVo.fail("文档需要文件ID和至少一个有效分段");
      }
      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());
      // Remote requests finish before publishing a document or replacing existing paragraphs.
      for (Paragraph p : input.getParagraphs()) {
        if (p.getContent() == null || p.getContent().isBlank()) {
          return ResultVo.fail("文档分段不能为空");
        }
        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"))));
      }
      Row document = Row.by("id", documentId).set("file_id", input.getId()).set("user_id", userId)
          .set("name", input.getName()).set("char_length", charLength).set("status", "2")
          .set("is_active", true).set("type", type).set("dataset_id", datasetId)
          .set("paragraph_count", paragraphs.size()).set("hit_handling_method", "optimization")
          .set("directly_return_similarity", 0.9);
      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);
  }
}

KbEmbeddingService（按模型实例隔离缓存）

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 Long getVectorId(String text, String model) { return getVectorId(text); }
  public Long getVectorId(String text, int areaCode, String model) { return getVectorId(text); }
  public Long getVectorId(String text, PlatformInput input) { return getVectorId(text); }
  public PGobject getVector(String text) { return cached(text).get("v"); }
  public PGobject getVector(String text, String model) { return getVector(text); }
  public PGobject getVector(String text, PlatformInput input) { return getVector(text); }
  public float[] embedding(String text, PlatformInput input) { return KnowledgeModelService.embedding(text); }
  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;
      }
      com.alibaba.fastjson2.JSONObject body = new com.alibaba.fastjson2.JSONObject();
      body.put("model", selected.getStr("model_name")); body.put("input", java.util.List.of(text)); body.put("dimensions", 1024);
      float[] vector;
      try (okhttp3.Response response = nexus.io.openai.client.OpenAiClient.embeddings(base, credential.getApi_key(), body.toJSONString())) {
        if (!response.isSuccessful() || response.body() == null) {
          throw new IllegalStateException("向量请求失败，HTTP " + response.code());
        }
        vector = com.alibaba.fastjson2.JSON.parseObject(response.body().string()).getJSONArray("data").getJSONObject(0).getObject("embedding", float[].class);
      } catch (java.io.IOException e) {
        throw new IllegalStateException("向量服务连接失败", e);
      }
      if (vector == null || vector.length != 1024) {
        throw new IllegalStateException("向量模型必须返回 1024 维");
      }
      for (float value : vector) {
        if (!Float.isFinite(value)) {
          throw new IllegalStateException("向量模型返回非有限数值");
        }
      }
      row = Row.by("id", SnowflakeIdUtils.id()).set("t", text).set("md5", hash).set("m", model)
          .set("v", PgVectorUtils.getPgVector(Arrays.toString(vector)));
      Db.save("moss_kb_embedding_cache", row); return row;
    } finally { lock.unlock(); }
  }
}

说明

  1. 缓存机制: 为了优化性能，KbEmbeddingService 在生成向量时会先查询 moss_kb_embedding_cache 表。如果缓存中存在对应的向量，则直接使用缓存；否则，调用所选模型的 OpenAI 兼容 embeddings 接口生成向量，并将结果缓存到数据库中。

  2. 线程安全: 在 KbEmbeddingService 中，按模型与文本摘要选择的 Striped<Lock> 在当前 JVM 内协调向量生成和缓存写入；它不等同于分布式锁。

  3. 事务管理: 在 MossKbParagraphSplitService 中，批量保存段落记录时使用事务管理，确保数据的一致性。

  4. 唯一标识: 使用 SnowflakeIdUtils 生成唯一的 ID，确保记录的唯一性。

  5. MD5 校验: 使用 MD5 对段落内容进行校验，以便快速查询和缓存。

总结

通过以上实现，我们完成了片段向量化的功能，利用缓存机制有效减少了对向量模型的调用次数，提升了系统的整体性能和响应速度。后续可以根据需求进一步优化缓存策略和并发处理机制。

当前接口与向量空间说明

请求中的 id 必须使用解析预览返回的文件 ID。服务先完成远程向量请求，再在事务中保存文档与分段。文档导入、分段修改和查询使用同一知识库的向量模型，不能在已有分段时直接更换模型空间。

当前使用远程 embeddings 接口，默认 Gitee Qwen3-Embedding-8B，存储维度为 1024。知识库通过 embedding_mode_id 关联用户配置模型；用户可以选择其他兼容服务，但创建模型时必须通过 1024 维返回校验。

文档导入、分段创建或修改通过 KbEmbeddingService 的模型实例入口生成向量。查询按知识库的模型分组，在同一模型空间比较余弦相似度。向量缓存由模型实例或默认模型标识、API 地址/模型名称的摘要、维度与内容摘要区分，避免不同模型共用缓存。

已有知识库不能直接修改向量空间。被引用的向量模型也不能更换模型 ID、类型或 API 地址；更换密钥可以保留模型空间。迁移到新向量模型时，应新建知识库并重新导入文档。

参见数据库初始化、支持自定义模型。