mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
450 lines
18 KiB
Java
450 lines
18 KiB
Java
package vip.mate.wiki.service;
|
||
|
||
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
|
||
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
|
||
import lombok.RequiredArgsConstructor;
|
||
import lombok.extern.slf4j.Slf4j;
|
||
import org.springframework.ai.embedding.EmbeddingModel;
|
||
import org.springframework.ai.embedding.EmbeddingRequest;
|
||
import org.springframework.ai.embedding.EmbeddingResponse;
|
||
import org.springframework.stereotype.Service;
|
||
import vip.mate.llm.embedding.EmbeddingModelFactory;
|
||
import vip.mate.llm.model.ModelConfigEntity;
|
||
import vip.mate.llm.service.ModelConfigService;
|
||
import vip.mate.system.model.SystemSettingEntity;
|
||
import vip.mate.system.repository.SystemSettingMapper;
|
||
import vip.mate.wiki.WikiProperties;
|
||
import vip.mate.wiki.model.WikiChunkEntity;
|
||
import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
|
||
import vip.mate.wiki.repository.WikiChunkMapper;
|
||
|
||
import java.nio.ByteBuffer;
|
||
import java.nio.ByteOrder;
|
||
import java.util.ArrayList;
|
||
import java.util.List;
|
||
|
||
/**
|
||
* RFC-011 + Embedding-UI-Config: Wiki 嵌入服务
|
||
* <p>
|
||
* 按知识库（KB）动态解析应使用的 Embedding 模型，解析优先级：
|
||
* <ol>
|
||
* <li>KB 级绑定：{@link WikiKnowledgeBaseEntity#getEmbeddingModelId()}</li>
|
||
* <li>系统默认：{@code mate_system_setting.setting_key = 'embedding.default.model.id'}</li>
|
||
* <li>任意 enabled 的 embedding 模型（取第一个）</li>
|
||
* <li>全无 → 返回不可用，上层降级（语义搜索返回空，关键词搜索仍可用）</li>
|
||
* </ol>
|
||
*
|
||
* @author MateClaw Team
|
||
*/
|
||
@Slf4j
|
||
@Service
|
||
@RequiredArgsConstructor
|
||
public class WikiEmbeddingService {
|
||
|
||
private final WikiChunkMapper chunkMapper;
|
||
private final WikiProperties properties;
|
||
private final EmbeddingModelFactory factory;
|
||
private final ModelConfigService modelConfigService;
|
||
private final WikiKnowledgeBaseService kbService;
|
||
private final SystemSettingMapper systemSettingMapper;
|
||
|
||
/** 系统默认 embedding 模型的 mate_system_setting key */
|
||
public static final String SYSTEM_SETTING_DEFAULT_EMBEDDING_ID = "embedding.default.model.id";
|
||
|
||
/**
|
||
* 判断全局是否有可用的 embedding 能力（任何 enabled 的 embedding 模型配置）
|
||
*/
|
||
public boolean isAvailable() {
|
||
try {
|
||
ModelConfigEntity fallback = modelConfigService.findFirstEnabledEmbedding();
|
||
return fallback != null;
|
||
} catch (Exception e) {
|
||
return false;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 解析指定 KB 应使用的 embedding 模型与模型实例
|
||
*/
|
||
public Resolved resolveForKb(Long kbId) {
|
||
// 优先级 1：KB 级绑定
|
||
try {
|
||
WikiKnowledgeBaseEntity kb = kbService.getById(kbId);
|
||
if (kb != null && kb.getEmbeddingModelId() != null) {
|
||
ModelConfigEntity model = safeGetModel(kb.getEmbeddingModelId());
|
||
if (isUsable(model)) {
|
||
return new Resolved(factory.build(model), model.getModelName());
|
||
}
|
||
log.warn("[WikiEmbedding] KB {} bound embedding model {} is unusable, falling back",
|
||
kbId, kb.getEmbeddingModelId());
|
||
}
|
||
} catch (Exception e) {
|
||
log.debug("[WikiEmbedding] KB binding resolve failed for kbId={}: {}", kbId, e.getMessage());
|
||
}
|
||
|
||
// 优先级 2：系统默认
|
||
Long defaultId = readSystemDefaultEmbeddingId();
|
||
if (defaultId != null) {
|
||
ModelConfigEntity model = safeGetModel(defaultId);
|
||
if (isUsable(model)) {
|
||
return new Resolved(factory.build(model), model.getModelName());
|
||
}
|
||
log.warn("[WikiEmbedding] System default embedding model {} is unusable, falling back", defaultId);
|
||
}
|
||
|
||
// 优先级 3：任意 enabled
|
||
ModelConfigEntity anyEnabled = modelConfigService.findFirstEnabledEmbedding();
|
||
if (isUsable(anyEnabled)) {
|
||
return new Resolved(factory.build(anyEnabled), anyEnabled.getModelName());
|
||
}
|
||
|
||
log.warn("[WikiEmbedding] No usable embedding model configured. "
|
||
+ "Configure one under Settings → Models → Embedding tab.");
|
||
return null;
|
||
}
|
||
|
||
/**
|
||
* 批量嵌入指定 KB 中缺失 embedding 的 chunk。
|
||
* <p>
|
||
* 只嵌入 embedding 为 NULL 或 embeddingModel 与当前解析出的模型不一致的 chunk。
|
||
* 模型切换时自动触发全量重嵌（通过 embedding_model 字段比对）。
|
||
*/
|
||
public int embedMissingChunks(Long kbId) {
|
||
Resolved r = resolveForKb(kbId);
|
||
if (r == null) {
|
||
log.debug("[WikiEmbedding] Skipping kbId={} — no embedding model available", kbId);
|
||
return 0;
|
||
}
|
||
|
||
String modelName = r.modelName();
|
||
List<WikiChunkEntity> pending = chunkMapper.selectList(
|
||
new LambdaQueryWrapper<WikiChunkEntity>()
|
||
.eq(WikiChunkEntity::getKbId, kbId)
|
||
.and(w -> w.isNull(WikiChunkEntity::getEmbedding)
|
||
.or().ne(WikiChunkEntity::getEmbeddingModel, modelName)));
|
||
|
||
if (pending.isEmpty()) {
|
||
log.debug("[WikiEmbedding] No chunks need embedding for kbId={}", kbId);
|
||
return 0;
|
||
}
|
||
|
||
int batchSize = Math.max(1, properties.getEmbeddingBatchSize());
|
||
int maxChars = Math.max(500, properties.getEmbeddingMaxChars());
|
||
int total = 0;
|
||
|
||
for (int offset = 0; offset < pending.size(); offset += batchSize) {
|
||
List<WikiChunkEntity> batch = pending.subList(offset, Math.min(offset + batchSize, pending.size()));
|
||
|
||
// Split the batch into short chunks (direct batch embed) and long chunks
|
||
// (split into sub-segments, embed each, then mean-pool into a single vector)
|
||
List<WikiChunkEntity> shortBatch = new ArrayList<>();
|
||
List<WikiChunkEntity> longChunks = new ArrayList<>();
|
||
for (WikiChunkEntity c : batch) {
|
||
if (c.getContent() == null || c.getContent().isBlank()) continue;
|
||
if (c.getContent().length() <= maxChars) {
|
||
shortBatch.add(c);
|
||
} else {
|
||
longChunks.add(c);
|
||
}
|
||
}
|
||
|
||
// Short chunks: existing batch path
|
||
if (!shortBatch.isEmpty()) {
|
||
total += embedShortBatch(shortBatch, r.model(), modelName, kbId);
|
||
}
|
||
|
||
// Long chunks: each goes through sub-segment split + mean pool
|
||
for (WikiChunkEntity longChunk : longChunks) {
|
||
if (embedLongChunk(longChunk, r.model(), modelName, maxChars)) {
|
||
total++;
|
||
}
|
||
}
|
||
}
|
||
|
||
if (total == 0 && !pending.isEmpty()) {
|
||
log.warn("[WikiEmbedding] ALL {} chunks failed for kbId={} model={} — check API key / model availability",
|
||
pending.size(), kbId, modelName);
|
||
} else {
|
||
log.info("[WikiEmbedding] Embedded {}/{} chunks for kbId={}, model={}",
|
||
total, pending.size(), kbId, modelName);
|
||
}
|
||
return total;
|
||
}
|
||
|
||
/**
|
||
* Embed a batch of chunks whose content fits within the per-segment char limit.
|
||
* One API call per batch; individual results are persisted independently.
|
||
* Returns the number of chunks that were successfully embedded and persisted.
|
||
*/
|
||
private int embedShortBatch(List<WikiChunkEntity> batch, EmbeddingModel model,
|
||
String modelName, Long kbId) {
|
||
try {
|
||
List<String> inputs = batch.stream().map(WikiChunkEntity::getContent).toList();
|
||
EmbeddingResponse resp = model.call(new EmbeddingRequest(inputs, null));
|
||
for (int i = 0; i < batch.size(); i++) {
|
||
float[] vec = resp.getResults().get(i).getOutput();
|
||
WikiChunkEntity chunk = batch.get(i);
|
||
chunk.setEmbedding(floatsToBytes(vec));
|
||
chunk.setEmbeddingModel(modelName);
|
||
chunkMapper.updateById(chunk);
|
||
}
|
||
return batch.size();
|
||
} catch (Exception e) {
|
||
log.error("[WikiEmbedding] Short-batch embedding failed (kbId={}, batchSize={}, model={}): {}",
|
||
kbId, batch.size(), modelName, e.getMessage());
|
||
return 0;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* Embed a single chunk whose content exceeds the per-segment char limit:
|
||
* 1. Split into sub-segments (each ≤ maxChars) along sentence boundaries
|
||
* 2. Batch-embed all sub-segments in one API call
|
||
* 3. Fall back to per-segment retry if the batch fails (partial recovery)
|
||
* 4. Mean-pool successful vectors and re-normalize (L2) to restore unit length
|
||
* 5. Store the single pooled vector against this chunk's id
|
||
* <p>
|
||
* Returns true if at least one sub-segment succeeded and the chunk was persisted.
|
||
*/
|
||
private boolean embedLongChunk(WikiChunkEntity chunk, EmbeddingModel model,
|
||
String modelName, int maxChars) {
|
||
List<String> segments = splitForEmbedding(chunk.getContent(), maxChars);
|
||
if (segments.isEmpty()) {
|
||
log.warn("[WikiEmbedding] Chunk {} produced no embeddable segments after split", chunk.getId());
|
||
return false;
|
||
}
|
||
log.info("[WikiEmbedding] Chunk {} ({} chars) split into {} sub-segments",
|
||
chunk.getId(), chunk.getContent().length(), segments.size());
|
||
|
||
List<float[]> vectors = new ArrayList<>();
|
||
try {
|
||
EmbeddingResponse resp = model.call(new EmbeddingRequest(segments, null));
|
||
for (int i = 0; i < resp.getResults().size(); i++) {
|
||
vectors.add(resp.getResults().get(i).getOutput());
|
||
}
|
||
} catch (Exception e) {
|
||
// Batch failed — degrade to per-segment retry; whatever succeeds is still usable
|
||
log.warn("[WikiEmbedding] Batch of {} sub-segments failed for chunk {}: {} — retrying one-by-one",
|
||
segments.size(), chunk.getId(), e.getMessage());
|
||
vectors.clear();
|
||
for (String seg : segments) {
|
||
try {
|
||
EmbeddingResponse single = model.call(new EmbeddingRequest(List.of(seg), null));
|
||
vectors.add(single.getResults().get(0).getOutput());
|
||
} catch (Exception ignored) {
|
||
// Skip this segment; proceed with remaining
|
||
}
|
||
}
|
||
}
|
||
|
||
if (vectors.isEmpty()) {
|
||
log.error("[WikiEmbedding] All {} sub-segments failed for chunk {}", segments.size(), chunk.getId());
|
||
return false;
|
||
}
|
||
|
||
float[] pooled = averageAndNormalize(vectors);
|
||
chunk.setEmbedding(floatsToBytes(pooled));
|
||
chunk.setEmbeddingModel(modelName);
|
||
chunkMapper.updateById(chunk);
|
||
return true;
|
||
}
|
||
|
||
/**
|
||
* Split a long text into sub-segments of at most {@code maxChars} characters,
|
||
* respecting sentence boundaries when possible.
|
||
* Boundary priority: double-newline > Chinese period > English period > newline > space.
|
||
* Hard-truncation fallback is applied when no boundary is found (e.g. a single
|
||
* run-on passage with no punctuation).
|
||
*/
|
||
private List<String> splitForEmbedding(String text, int maxChars) {
|
||
if (text == null || text.isBlank()) return List.of();
|
||
if (text.length() <= maxChars) return List.of(text);
|
||
|
||
List<String> segments = new ArrayList<>();
|
||
int start = 0;
|
||
while (start < text.length()) {
|
||
int end = Math.min(start + maxChars, text.length());
|
||
if (end < text.length()) {
|
||
int boundary = findEmbeddingBoundary(text, start, end, maxChars);
|
||
if (boundary > start) end = boundary;
|
||
}
|
||
String seg = text.substring(start, end).trim();
|
||
if (!seg.isBlank()) {
|
||
// Hard-truncation fallback: a boundary beyond maxChars shouldn't happen
|
||
// with the logic above, but guard against edge cases defensively.
|
||
if (seg.length() > maxChars) {
|
||
seg = seg.substring(0, maxChars);
|
||
}
|
||
segments.add(seg);
|
||
}
|
||
int nextStart = end;
|
||
if (nextStart <= start) nextStart = start + maxChars; // prevent infinite loop
|
||
start = nextStart;
|
||
}
|
||
return segments;
|
||
}
|
||
|
||
/**
|
||
* Find a sentence boundary within [start, end] for embedding sub-segmentation.
|
||
* Only returns boundaries past the midpoint so we don't produce tiny segments.
|
||
*/
|
||
private int findEmbeddingBoundary(String text, int start, int end, int maxChars) {
|
||
int halfChunk = start + maxChars / 2;
|
||
|
||
int lastPara = text.lastIndexOf("\n\n", end);
|
||
if (lastPara > halfChunk) return lastPara + 2;
|
||
|
||
int lastChinese = text.lastIndexOf("。", end);
|
||
if (lastChinese > halfChunk) return lastChinese + 1;
|
||
|
||
for (int i = end - 1; i > halfChunk; i--) {
|
||
if (text.charAt(i) == '.' && i + 1 < text.length()
|
||
&& (text.charAt(i + 1) == ' ' || text.charAt(i + 1) == '\n')) {
|
||
return i + 1;
|
||
}
|
||
}
|
||
|
||
int lastNewline = text.lastIndexOf("\n", end);
|
||
if (lastNewline > halfChunk) return lastNewline + 1;
|
||
|
||
int lastSpace = text.lastIndexOf(" ", end);
|
||
if (lastSpace > halfChunk) return lastSpace + 1;
|
||
|
||
return end; // hard cut
|
||
}
|
||
|
||
/**
|
||
* 查询向量化（混合搜索时调用，需指定 KB 以便解析对应模型）
|
||
*/
|
||
public float[] embedQuery(Long kbId, String query) {
|
||
Resolved r = resolveForKb(kbId);
|
||
if (r == null) return null;
|
||
// Defensive truncation: user queries are usually short, but guard against
|
||
// callers that accidentally pass document-sized text as a query.
|
||
int maxChars = Math.max(500, properties.getEmbeddingMaxChars());
|
||
String safeQuery = (query != null && query.length() > maxChars)
|
||
? query.substring(0, maxChars) : query;
|
||
try {
|
||
EmbeddingResponse resp = r.model().call(new EmbeddingRequest(List.of(safeQuery), null));
|
||
return resp.getResults().get(0).getOutput();
|
||
} catch (Exception e) {
|
||
log.error("[WikiEmbedding] Query embedding failed for kbId={}: {}", kbId, e.getMessage());
|
||
return null;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 清空指定 KB 的所有 embedding（模型切换时调用）
|
||
*/
|
||
public void clearEmbeddings(Long kbId) {
|
||
chunkMapper.update(null, new LambdaUpdateWrapper<WikiChunkEntity>()
|
||
.eq(WikiChunkEntity::getKbId, kbId)
|
||
.set(WikiChunkEntity::getEmbedding, null)
|
||
.set(WikiChunkEntity::getEmbeddingModel, null));
|
||
log.info("[WikiEmbedding] Cleared all embeddings for kbId={}", kbId);
|
||
}
|
||
|
||
// ==================== 私有 helper ====================
|
||
|
||
private ModelConfigEntity safeGetModel(Long id) {
|
||
try {
|
||
return modelConfigService.getModel(id);
|
||
} catch (Exception e) {
|
||
return null;
|
||
}
|
||
}
|
||
|
||
private boolean isUsable(ModelConfigEntity model) {
|
||
return model != null
|
||
&& Boolean.TRUE.equals(model.getEnabled())
|
||
&& "embedding".equals(model.getModelType());
|
||
}
|
||
|
||
private Long readSystemDefaultEmbeddingId() {
|
||
try {
|
||
SystemSettingEntity entity = systemSettingMapper.selectOne(
|
||
new LambdaQueryWrapper<SystemSettingEntity>()
|
||
.eq(SystemSettingEntity::getSettingKey, SYSTEM_SETTING_DEFAULT_EMBEDDING_ID)
|
||
.last("LIMIT 1"));
|
||
if (entity == null || entity.getSettingValue() == null || entity.getSettingValue().isBlank()) {
|
||
return null;
|
||
}
|
||
return Long.parseLong(entity.getSettingValue().trim());
|
||
} catch (Exception e) {
|
||
log.debug("[WikiEmbedding] Failed to read system default embedding id: {}", e.getMessage());
|
||
return null;
|
||
}
|
||
}
|
||
|
||
// ==================== 向量序列化 ====================
|
||
|
||
public static byte[] floatsToBytes(float[] vec) {
|
||
ByteBuffer buf = ByteBuffer.allocate(vec.length * 4).order(ByteOrder.LITTLE_ENDIAN);
|
||
for (float v : vec) buf.putFloat(v);
|
||
return buf.array();
|
||
}
|
||
|
||
public static float[] bytesToFloats(byte[] bytes) {
|
||
ByteBuffer buf = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
|
||
float[] vec = new float[bytes.length / 4];
|
||
for (int i = 0; i < vec.length; i++) vec[i] = buf.getFloat();
|
||
return vec;
|
||
}
|
||
|
||
/**
|
||
* Arithmetic mean of N equal-length vectors, followed by L2 normalization.
|
||
* <p>
|
||
* Individual embedding outputs are usually unit vectors, but the arithmetic mean
|
||
* of multiple unit vectors is generally not unit length (||v̄|| < 1 unless all
|
||
* inputs are identical). Re-normalizing to unit length preserves the cosine
|
||
* similarity semantics used by downstream retrievers.
|
||
*/
|
||
public static float[] averageAndNormalize(List<float[]> vectors) {
|
||
if (vectors == null || vectors.isEmpty()) {
|
||
throw new IllegalArgumentException("vectors must not be empty");
|
||
}
|
||
int dim = vectors.get(0).length;
|
||
float[] avg = new float[dim];
|
||
for (float[] v : vectors) {
|
||
if (v.length != dim) {
|
||
throw new IllegalArgumentException("dimension mismatch: expected " + dim + " got " + v.length);
|
||
}
|
||
for (int i = 0; i < dim; i++) avg[i] += v[i];
|
||
}
|
||
float n = vectors.size();
|
||
for (int i = 0; i < dim; i++) avg[i] /= n;
|
||
|
||
double norm = 0;
|
||
for (float x : avg) norm += x * x;
|
||
norm = Math.sqrt(norm);
|
||
if (norm > 0) {
|
||
float invNorm = (float) (1.0 / norm);
|
||
for (int i = 0; i < dim; i++) avg[i] *= invNorm;
|
||
}
|
||
return avg;
|
||
}
|
||
|
||
/** 余弦相似度 */
|
||
public static float cosine(float[] a, float[] b) {
|
||
if (a.length != b.length) return 0f;
|
||
float dot = 0, normA = 0, normB = 0;
|
||
for (int i = 0; i < a.length; i++) {
|
||
dot += a[i] * b[i];
|
||
normA += a[i] * a[i];
|
||
normB += b[i] * b[i];
|
||
}
|
||
float denom = (float) (Math.sqrt(normA) * Math.sqrt(normB));
|
||
return denom == 0 ? 0f : dot / denom;
|
||
}
|
||
|
||
/** 解析结果 DTO */
|
||
public record Resolved(EmbeddingModel model, String modelName) {}
|
||
|
||
/**
|
||
* 暴露 factory 给外部（如测试连通性 API）
|
||
*/
|
||
public EmbeddingModelFactory getFactory() {
|
||
return factory;
|
||
}
|
||
}
|