mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
feat(workspace): keyword search across agent memory files
This commit is contained in:
matevip
parent
d90591fa35
commit
b9c0ead024
@ -9,12 +9,16 @@ import org.springframework.ai.tool.annotation.Tool;
|
||||
import org.springframework.ai.tool.annotation.ToolParam;
|
||||
import org.springframework.stereotype.Component;
|
||||
import vip.mate.memory.service.MemoryRecallTracker;
|
||||
import vip.mate.workspace.document.MemorySearchHit;
|
||||
import vip.mate.workspace.document.WorkspaceFileService;
|
||||
import vip.mate.workspace.document.model.WorkspaceFileEntity;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Comparator;
|
||||
import java.util.HashSet;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* 基于数据库工作区文件的长期记忆工具。
|
||||
@ -197,6 +201,88 @@ public class WorkspaceMemoryTool {
|
||||
return JSONUtil.toJsonPrettyStr(result);
|
||||
}
|
||||
|
||||
@Tool(description = """
|
||||
Search agent's workspace memory files (MEMORY.md, PROFILE.md, AGENTS.md, memory/*.md) \
|
||||
by keyword. Use this BEFORE read_workspace_memory_file when looking for a fact across \
|
||||
many memory entries. Returns ranked hits with filename, line number, and snippet \
|
||||
(matched terms wrapped in [[...]]).""")
|
||||
public String search_workspace_memory(
|
||||
@ToolParam(description = "当前 Agent 的 ID") Long agentId,
|
||||
@ToolParam(description = "关键词或短语，2-64 字符") String query,
|
||||
@ToolParam(description = "搜索范围：all（全部）/ memory（MEMORY.md 与 memory/）/ profile / persona，默认 all",
|
||||
required = false) String scope,
|
||||
@ToolParam(description = "返回的最大命中数，默认 10，上限 30", required = false) Integer limit) {
|
||||
|
||||
if (agentId == null) {
|
||||
return error("agentId 不能为空");
|
||||
}
|
||||
if (query == null || query.isBlank()) {
|
||||
return error("query 不能为空");
|
||||
}
|
||||
String trimmed = query.trim();
|
||||
if (trimmed.length() < 2) {
|
||||
return error("query 至少 2 个字符");
|
||||
}
|
||||
if (trimmed.length() > 64) {
|
||||
return error("query 不能超过 64 个字符");
|
||||
}
|
||||
|
||||
int effectiveLimit = limit == null ? 10 : Math.min(Math.max(limit, 1), 30);
|
||||
Set<String> prefixes = resolveScope(scope);
|
||||
|
||||
List<MemorySearchHit> hits = workspaceFileService.searchSnippets(
|
||||
agentId, trimmed, prefixes, effectiveLimit);
|
||||
|
||||
// Treat each unique file in the results as an active retrieval signal —
|
||||
// boosts that file's weight in the dream-consolidation ranker the same
|
||||
// way an explicit read_workspace_memory_file call would.
|
||||
Set<String> retrieved = new HashSet<>();
|
||||
for (MemorySearchHit hit : hits) {
|
||||
if (retrieved.add(hit.filename())) {
|
||||
WorkspaceFileEntity file = workspaceFileService.getFile(agentId, hit.filename());
|
||||
if (file != null && file.getContent() != null) {
|
||||
memoryRecallTracker.trackActiveRetrieval(agentId, hit.filename(), file.getContent());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
JSONArray hitsJson = new JSONArray();
|
||||
for (MemorySearchHit hit : hits) {
|
||||
JSONObject h = new JSONObject();
|
||||
h.set("filename", hit.filename());
|
||||
h.set("lineNumber", hit.lineNumber());
|
||||
h.set("snippet", hit.snippet());
|
||||
h.set("score", hit.score());
|
||||
hitsJson.add(h);
|
||||
}
|
||||
JSONObject result = new JSONObject();
|
||||
result.set("agentId", agentId);
|
||||
result.set("query", trimmed);
|
||||
result.set("scope", scope == null || scope.isBlank() ? "all" : scope);
|
||||
result.set("totalHits", hits.size());
|
||||
result.set("hits", hitsJson);
|
||||
if (!hits.isEmpty()) {
|
||||
result.set("hint", "Use read_workspace_memory_file to get full context of any hit.");
|
||||
}
|
||||
return JSONUtil.toJsonPrettyStr(result);
|
||||
}
|
||||
|
||||
/** Map the {@code scope} tool argument to a filename-prefix whitelist.
|
||||
* {@code "all"} (or null/blank) targets every memory-class file rather
|
||||
* than every workspace file the agent has, so a search doesn't surface
|
||||
* unrelated docs the user happens to store in the same workspace. */
|
||||
private static Set<String> resolveScope(String scope) {
|
||||
if (scope == null || scope.isBlank() || "all".equalsIgnoreCase(scope.trim())) {
|
||||
return new LinkedHashSet<>(List.of("memory/", "MEMORY.md", "PROFILE.md", "AGENTS.md"));
|
||||
}
|
||||
return switch (scope.trim().toLowerCase()) {
|
||||
case "memory" -> new LinkedHashSet<>(List.of("memory/", "MEMORY.md"));
|
||||
case "profile" -> new LinkedHashSet<>(List.of("PROFILE.md"));
|
||||
case "persona" -> new LinkedHashSet<>(List.of("AGENTS.md"));
|
||||
default -> new LinkedHashSet<>(List.of("memory/", "MEMORY.md", "PROFILE.md", "AGENTS.md"));
|
||||
};
|
||||
}
|
||||
|
||||
private String validate(Long agentId, String filename) {
|
||||
if (agentId == null) {
|
||||
return "agentId 不能为空";
|
||||
|
||||
@ -0,0 +1,17 @@
|
||||
package vip.mate.workspace.document;
|
||||
|
||||
/**
|
||||
* One snippet-level hit from {@link WorkspaceFileService#searchSnippets}.
|
||||
*
|
||||
* @param filename Workspace filename the hit was extracted from (e.g.
|
||||
* {@code MEMORY.md}, {@code memory/2026-05-10.md}).
|
||||
* @param lineNumber 1-based line number of the hit within {@code filename}.
|
||||
* @param snippet The original line clipped to {@code head(80) + match +
|
||||
* tail(80)} and with each matched term wrapped in
|
||||
* {@code [[...]]} for downstream highlighting.
|
||||
* @param score Composite relevance: number of terms matched on this line
|
||||
* weighted by per-file importance (see
|
||||
* {@code WorkspaceFileService#fileWeight}).
|
||||
*/
|
||||
public record MemorySearchHit(String filename, int lineNumber, String snippet, double score) {
|
||||
}
|
||||
@ -9,8 +9,12 @@ import vip.mate.workspace.document.model.WorkspaceFileEntity;
|
||||
import vip.mate.workspace.document.repository.WorkspaceFileMapper;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.Comparator;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
@ -128,6 +132,300 @@ public class WorkspaceFileService {
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Snippet search ====================
|
||||
|
||||
/** Hard cap on terms passed to the DB filter — keeps the {@code LIKE} chain
|
||||
* short and bounds the per-line scoring cost.
|
||||
* <p>
|
||||
* Sized so a typical CJK query of up to 12 chars (= 6 non-overlapping
|
||||
* 2-char windows) lands all of its tokens, not just the first four.
|
||||
* The trade-off: every extra term tightens the AND-LIKE filter, so a
|
||||
* user typing a long phrase whose words are spread across files will
|
||||
* match fewer candidates. Acceptable here because the typical search
|
||||
* is a short phrase ("用户跑步"), and the per-line scorer's term-hit
|
||||
* count keeps relevance ranking stable when only a subset matches.
|
||||
* Sliding-bigram tokenization is the next lever if real query data
|
||||
* shows a recall problem on long phrases. */
|
||||
private static final int MAX_TERMS = 6;
|
||||
|
||||
/** Minimum query length below which search is a no-op. Two characters
|
||||
* comfortably covers a single CJK word; one character would return the
|
||||
* bulk of every file. */
|
||||
private static final int MIN_QUERY_LENGTH = 2;
|
||||
|
||||
/** Candidate file cap. Picked large enough that scope filtering plus
|
||||
* AND-LIKE almost always leaves room for the relevant memory set, but
|
||||
* small enough that line-level scoring stays cheap. */
|
||||
private static final int CANDIDATE_FILE_CAP = 50;
|
||||
|
||||
/** Per-file hit cap. Prevents a single very long file (multi-week daily
|
||||
* ledger, copy-pasted log) from monopolising the result set. */
|
||||
private static final int PER_FILE_HIT_CAP = 5;
|
||||
|
||||
/** Snippet context window: characters before and after the first matched
|
||||
* term on the line. The original line is clipped to this window before
|
||||
* highlighting. */
|
||||
private static final int SNIPPET_CONTEXT = 80;
|
||||
|
||||
/**
|
||||
* Keyword-search the agent's workspace files and return line-level snippet
|
||||
* hits ranked by term-match count weighted by per-file importance.
|
||||
* <p>
|
||||
* Pipeline:
|
||||
* <ol>
|
||||
* <li>Tokenize {@code query} on whitespace and CJK/Latin boundaries.
|
||||
* CJK runs are split into non-overlapping 2-char windows; Latin /
|
||||
* digit runs are kept as whole tokens. Dedupe and cap at {@value
|
||||
* #MAX_TERMS}.</li>
|
||||
* <li>Pull candidate rows via {@link com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper}:
|
||||
* {@code agent_id} match, optional {@code filename LIKE prefix%} OR
|
||||
* group, and one {@code content LIKE %term%} per token (AND).</li>
|
||||
* <li>For each candidate file, scan lines and collect up to {@value
|
||||
* #PER_FILE_HIT_CAP} hits scored by {@code termsHit * fileWeight}.</li>
|
||||
* <li>Build a head/match/tail snippet around the first matched term on
|
||||
* the line, wrap each term in {@code [[...]]}, sort by score
|
||||
* descending and return the top {@code limit} results.</li>
|
||||
* </ol>
|
||||
*
|
||||
* @param agentId The agent whose workspace files are searched.
|
||||
* Required; null returns empty.
|
||||
* @param query The user-supplied search string. Returns empty
|
||||
* if {@code null}, blank, or shorter than {@value
|
||||
* #MIN_QUERY_LENGTH} characters after trim.
|
||||
* @param filenamePrefixes Optional filename prefix whitelist (each value
|
||||
* is a {@code likeRight} prefix). {@code null} or
|
||||
* empty means "all files for this agent".
|
||||
* @param limit Maximum number of hits returned across all
|
||||
* files. Capped at the candidate-file scan
|
||||
* output naturally; the caller is expected to
|
||||
* choose a sensible value (1-30 typical).
|
||||
*/
|
||||
public List<MemorySearchHit> searchSnippets(Long agentId, String query,
|
||||
Set<String> filenamePrefixes, int limit) {
|
||||
if (agentId == null || limit <= 0) {
|
||||
return List.of();
|
||||
}
|
||||
if (query == null || query.trim().length() < MIN_QUERY_LENGTH) {
|
||||
return List.of();
|
||||
}
|
||||
List<String> terms = tokenize(query);
|
||||
if (terms.isEmpty()) {
|
||||
return List.of();
|
||||
}
|
||||
|
||||
LambdaQueryWrapper<WorkspaceFileEntity> wrapper = new LambdaQueryWrapper<>();
|
||||
wrapper.eq(WorkspaceFileEntity::getAgentId, agentId);
|
||||
|
||||
if (filenamePrefixes != null && !filenamePrefixes.isEmpty()) {
|
||||
// Group prefix conditions inside a single AND-bracketed OR chain
|
||||
// so they don't interleave with the content-LIKE terms.
|
||||
wrapper.and(w -> {
|
||||
boolean[] first = {true};
|
||||
for (String prefix : filenamePrefixes) {
|
||||
if (prefix == null || prefix.isEmpty()) continue;
|
||||
if (first[0]) {
|
||||
w.likeRight(WorkspaceFileEntity::getFilename, prefix);
|
||||
first[0] = false;
|
||||
} else {
|
||||
w.or().likeRight(WorkspaceFileEntity::getFilename, prefix);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
for (String term : terms) {
|
||||
wrapper.like(WorkspaceFileEntity::getContent, term);
|
||||
}
|
||||
// Dialect-transparent LIMIT; @TableLogic appends deleted = 0 for us.
|
||||
wrapper.last("LIMIT " + CANDIDATE_FILE_CAP);
|
||||
|
||||
List<WorkspaceFileEntity> candidates = fileMapper.selectList(wrapper);
|
||||
if (candidates.isEmpty()) {
|
||||
return List.of();
|
||||
}
|
||||
|
||||
List<MemorySearchHit> all = new ArrayList<>();
|
||||
for (WorkspaceFileEntity file : candidates) {
|
||||
all.addAll(extractHitsFromFile(file, terms));
|
||||
}
|
||||
all.sort(Comparator.comparingDouble(MemorySearchHit::score).reversed()
|
||||
.thenComparing(MemorySearchHit::filename)
|
||||
.thenComparingInt(MemorySearchHit::lineNumber));
|
||||
if (all.size() > limit) {
|
||||
return new ArrayList<>(all.subList(0, limit));
|
||||
}
|
||||
return all;
|
||||
}
|
||||
|
||||
/** Split a query into searchable terms. Whitespace and non-letter/digit
|
||||
* separators end a token; CJK ↔ Latin/digit transitions end a token;
|
||||
* CJK runs of length ≥ 2 are further chunked into non-overlapping
|
||||
* 2-char windows (e.g. "用户喜欢" → ["用户", "喜欢"]). The result is
|
||||
* deduped (preserving order) and capped at {@value #MAX_TERMS}. */
|
||||
static List<String> tokenize(String query) {
|
||||
if (query == null) return List.of();
|
||||
List<String> raw = new ArrayList<>();
|
||||
StringBuilder cur = new StringBuilder();
|
||||
int prevType = -1; // 0 = CJK, 1 = Latin/digit, -1 = separator / not started
|
||||
for (int i = 0; i < query.length(); i++) {
|
||||
char c = query.charAt(i);
|
||||
int type = classify(c);
|
||||
if (type == -1) {
|
||||
flush(cur, raw);
|
||||
prevType = -1;
|
||||
continue;
|
||||
}
|
||||
if (prevType != -1 && type != prevType) {
|
||||
flush(cur, raw);
|
||||
}
|
||||
cur.append(c);
|
||||
prevType = type;
|
||||
}
|
||||
flush(cur, raw);
|
||||
|
||||
// Split long CJK runs into 2-char windows so a query like
|
||||
// "用户喜欢周日早上跑步" yields multiple usable terms instead of a
|
||||
// single substring rarely present verbatim in stored memory.
|
||||
List<String> exploded = new ArrayList<>();
|
||||
for (String token : raw) {
|
||||
if (!token.isEmpty() && classify(token.charAt(0)) == 0 && token.length() > 2) {
|
||||
for (int i = 0; i + 2 <= token.length(); i += 2) {
|
||||
exploded.add(token.substring(i, i + 2));
|
||||
}
|
||||
if (token.length() % 2 == 1) {
|
||||
// Trailing single char: pair it with the previous char as
|
||||
// an overlap so it can still match (better than dropping).
|
||||
exploded.add(token.substring(token.length() - 2));
|
||||
}
|
||||
} else {
|
||||
exploded.add(token);
|
||||
}
|
||||
}
|
||||
|
||||
LinkedHashSet<String> dedupe = new LinkedHashSet<>();
|
||||
for (String t : exploded) {
|
||||
if (!t.isEmpty()) dedupe.add(t);
|
||||
}
|
||||
if (dedupe.size() <= MAX_TERMS) {
|
||||
return List.copyOf(dedupe);
|
||||
}
|
||||
List<String> capped = new ArrayList<>(MAX_TERMS);
|
||||
int i = 0;
|
||||
for (String t : dedupe) {
|
||||
if (i++ >= MAX_TERMS) break;
|
||||
capped.add(t);
|
||||
}
|
||||
return Collections.unmodifiableList(capped);
|
||||
}
|
||||
|
||||
private static int classify(char c) {
|
||||
if (c >= 0x4E00 && c <= 0x9FFF) return 0; // CJK Unified Ideographs
|
||||
if (c >= 0x3400 && c <= 0x4DBF) return 0; // CJK Extension A
|
||||
if (Character.isLetterOrDigit(c)) return 1;
|
||||
return -1;
|
||||
}
|
||||
|
||||
private static void flush(StringBuilder cur, List<String> sink) {
|
||||
if (cur.length() > 0) {
|
||||
sink.add(cur.toString());
|
||||
cur.setLength(0);
|
||||
}
|
||||
}
|
||||
|
||||
private List<MemorySearchHit> extractHitsFromFile(WorkspaceFileEntity file, List<String> terms) {
|
||||
String content = file.getContent();
|
||||
if (content == null || content.isEmpty()) return List.of();
|
||||
String[] lines = content.split("\n", -1);
|
||||
double weight = fileWeight(file.getFilename());
|
||||
List<MemorySearchHit> hits = new ArrayList<>();
|
||||
for (int i = 0; i < lines.length; i++) {
|
||||
String line = lines[i];
|
||||
int termsHit = 0;
|
||||
for (String term : terms) {
|
||||
if (line.contains(term)) termsHit++;
|
||||
}
|
||||
if (termsHit == 0) continue;
|
||||
double score = termsHit * weight;
|
||||
hits.add(new MemorySearchHit(file.getFilename(), i + 1,
|
||||
buildSnippet(line, terms), score));
|
||||
if (hits.size() >= PER_FILE_HIT_CAP) break;
|
||||
}
|
||||
return hits;
|
||||
}
|
||||
|
||||
/** Per-file importance multiplier. {@code MEMORY.md} is the consolidated
|
||||
* long-term memory; daily ledger files under {@code memory/} are
|
||||
* high-recall but noisier; {@code PROFILE.md} is mostly static persona
|
||||
* facts; {@code AGENTS.md} is mostly tool / behavior config and rarely
|
||||
* the right answer to a memory query. */
|
||||
private double fileWeight(String filename) {
|
||||
if (filename == null) return 0.2;
|
||||
if ("MEMORY.md".equals(filename)) return 1.0;
|
||||
if (filename.startsWith("memory/")) return 0.7;
|
||||
if ("PROFILE.md".equals(filename)) return 0.5;
|
||||
if ("AGENTS.md".equals(filename)) return 0.3;
|
||||
return 0.2;
|
||||
}
|
||||
|
||||
/** Build a {@code head + match + tail} snippet around the first matched
|
||||
* term and wrap every non-overlapping match inside the window with
|
||||
* {@code [[...]]}. Term boundaries are preserved — adjacent matches from
|
||||
* two different terms render as {@code [[t1]][[t2]]}, not as one merged
|
||||
* bracket. Longer terms claim their span first so a shorter overlapping
|
||||
* term cannot fragment a longer match. */
|
||||
private String buildSnippet(String line, List<String> terms) {
|
||||
int firstMatch = Integer.MAX_VALUE;
|
||||
for (String term : terms) {
|
||||
if (term.isEmpty()) continue;
|
||||
int idx = line.indexOf(term);
|
||||
if (idx >= 0 && idx < firstMatch) firstMatch = idx;
|
||||
}
|
||||
if (firstMatch == Integer.MAX_VALUE) {
|
||||
return line;
|
||||
}
|
||||
|
||||
int start = Math.max(0, firstMatch - SNIPPET_CONTEXT);
|
||||
int end = Math.min(line.length(), firstMatch + SNIPPET_CONTEXT);
|
||||
String clipped = line.substring(start, end);
|
||||
boolean hasPrefix = start > 0;
|
||||
boolean hasSuffix = end < line.length();
|
||||
|
||||
List<String> longestFirst = new ArrayList<>(terms);
|
||||
longestFirst.sort(Comparator.comparingInt(String::length).reversed());
|
||||
|
||||
boolean[] covered = new boolean[clipped.length()];
|
||||
List<int[]> spans = new ArrayList<>();
|
||||
for (String term : longestFirst) {
|
||||
if (term.isEmpty()) continue;
|
||||
int from = 0;
|
||||
int idx;
|
||||
while ((idx = clipped.indexOf(term, from)) >= 0) {
|
||||
int spanEnd = idx + term.length();
|
||||
boolean overlap = false;
|
||||
for (int k = idx; k < spanEnd; k++) {
|
||||
if (covered[k]) { overlap = true; break; }
|
||||
}
|
||||
if (!overlap) {
|
||||
for (int k = idx; k < spanEnd; k++) covered[k] = true;
|
||||
spans.add(new int[]{idx, spanEnd});
|
||||
}
|
||||
from = spanEnd;
|
||||
}
|
||||
}
|
||||
spans.sort(Comparator.comparingInt(s -> s[0]));
|
||||
|
||||
StringBuilder sb = new StringBuilder(clipped.length() + spans.size() * 4 + 6);
|
||||
if (hasPrefix) sb.append("...");
|
||||
int pos = 0;
|
||||
for (int[] s : spans) {
|
||||
sb.append(clipped, pos, s[0]);
|
||||
sb.append("[[").append(clipped, s[0], s[1]).append("]]");
|
||||
pos = s[1];
|
||||
}
|
||||
sb.append(clipped, pos, clipped.length());
|
||||
if (hasSuffix) sb.append("...");
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
/**
|
||||
* 将启用的工作区文件拼接为系统提示词
|
||||
* <p>
|
||||
|
||||
@ -0,0 +1,280 @@
|
||||
package vip.mate.workspace.document;
|
||||
|
||||
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
|
||||
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
|
||||
import org.apache.ibatis.builder.MapperBuilderAssistant;
|
||||
import org.apache.ibatis.session.Configuration;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.DisplayName;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import vip.mate.workspace.document.model.WorkspaceFileEntity;
|
||||
import vip.mate.workspace.document.repository.WorkspaceFileMapper;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* Contract for {@link WorkspaceFileService#searchSnippets} — the back end
|
||||
* of the {@code search_workspace_memory} agent tool.
|
||||
* <p>
|
||||
* Tests run against a Mockito-stubbed {@link WorkspaceFileMapper}: the DB-side
|
||||
* AND-LIKE narrowing is MyBatis-Plus's problem, not this service's — what we
|
||||
* verify here is the post-fetch pipeline (tokenization, per-line extraction,
|
||||
* weighted scoring, snippet rendering) plus the wrapper construction so we
|
||||
* don't silently drop scope filters or term groups.
|
||||
*/
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class WorkspaceMemorySearchTest {
|
||||
|
||||
@Mock private WorkspaceFileMapper fileMapper;
|
||||
private WorkspaceFileService service;
|
||||
|
||||
@BeforeAll
|
||||
static void initMyBatisPlusCache() {
|
||||
// LambdaQueryWrapper resolves SFunction → column via TableInfoHelper.
|
||||
// Spring would init it during mapper scan; here we trigger it manually.
|
||||
TableInfoHelper.initTableInfo(
|
||||
new MapperBuilderAssistant(new Configuration(), ""),
|
||||
WorkspaceFileEntity.class);
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
service = new WorkspaceFileService(fileMapper);
|
||||
}
|
||||
|
||||
// ---------- tokenize ----------
|
||||
|
||||
@Test
|
||||
@DisplayName("tokenize: null / blank / single-char yields empty list")
|
||||
void tokenizeShortInputs() {
|
||||
assertThat(WorkspaceFileService.tokenize(null)).isEmpty();
|
||||
assertThat(WorkspaceFileService.tokenize("")).isEmpty();
|
||||
assertThat(WorkspaceFileService.tokenize(" ")).isEmpty();
|
||||
// A single CJK char DOES produce one token — query-length guard at
|
||||
// searchSnippets level is what enforces the 2-char minimum.
|
||||
assertThat(WorkspaceFileService.tokenize("好")).containsExactly("好");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("tokenize: long CJK runs split into non-overlapping 2-char windows (within cap)")
|
||||
void tokenizeCjkPairs() {
|
||||
// 10-char CJK run → 5 windows; all five survive under MAX_TERMS = 6.
|
||||
assertThat(WorkspaceFileService.tokenize("用户喜欢周日早上跑步"))
|
||||
.containsExactly("用户", "喜欢", "周日", "早上", "跑步");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("tokenize: CJK/Latin boundary splits, Latin runs stay intact")
|
||||
void tokenizeMixedScripts() {
|
||||
// "用户Foo123跑步" → CJK run "用户" (2-char window) + Latin/digit run
|
||||
// "Foo123" + CJK run "跑步".
|
||||
assertThat(WorkspaceFileService.tokenize("用户Foo123跑步"))
|
||||
.containsExactly("用户", "Foo123", "跑步");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("tokenize: dedupe + cap at 6")
|
||||
void tokenizeDedupeAndCap() {
|
||||
// Whitespace splits; "abc" repeats are dedup'd; only first 6 kept.
|
||||
// The 7th and 8th tokens ("ppp" / "qqq") get trimmed by the cap.
|
||||
assertThat(WorkspaceFileService.tokenize("abc abc def ghi jkl mno opq ppp qqq"))
|
||||
.containsExactly("abc", "def", "ghi", "jkl", "mno", "opq");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("tokenize: a 12-char CJK run yields 6 windows and stops at the cap")
|
||||
void tokenizeLongCjkFitsCap() {
|
||||
// 12-char run → 6 windows; the 7th and 8th would be "测一" / "下下"
|
||||
// if we extended the query, but the cap stops at 6 either way.
|
||||
assertThat(WorkspaceFileService.tokenize("用户喜欢周日早上跑步公园"))
|
||||
.containsExactly("用户", "喜欢", "周日", "早上", "跑步", "公园");
|
||||
}
|
||||
|
||||
// ---------- searchSnippets contract ----------
|
||||
|
||||
@Test
|
||||
@DisplayName("Short query: returns empty without touching the mapper")
|
||||
void shortQueryReturnsEmpty() {
|
||||
assertThat(service.searchSnippets(1L, "a", null, 10)).isEmpty();
|
||||
assertThat(service.searchSnippets(1L, " ", null, 10)).isEmpty();
|
||||
assertThat(service.searchSnippets(1L, "", null, 10)).isEmpty();
|
||||
assertThat(service.searchSnippets(1L, null, null, 10)).isEmpty();
|
||||
// limit <= 0 short-circuits.
|
||||
assertThat(service.searchSnippets(1L, "running", null, 0)).isEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("No candidate files: returns empty list")
|
||||
void noCandidateRowsReturnsEmpty() {
|
||||
when(fileMapper.selectList(any())).thenReturn(List.of());
|
||||
assertThat(service.searchSnippets(1L, "running", null, 10)).isEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Score ordering: MEMORY > memory/* > PROFILE > AGENTS for identical term-hit counts")
|
||||
void scoreOrdering() {
|
||||
List<WorkspaceFileEntity> candidates = List.of(
|
||||
file("AGENTS.md", "Line about running.\n"),
|
||||
file("MEMORY.md", "Line about running.\n"),
|
||||
file("PROFILE.md", "Line about running.\n"),
|
||||
file("memory/2026-05-10.md", "Line about running.\n"));
|
||||
when(fileMapper.selectList(any())).thenReturn(candidates);
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "running", null, 10);
|
||||
|
||||
assertThat(hits).extracting(MemorySearchHit::filename)
|
||||
.containsExactly("MEMORY.md", "memory/2026-05-10.md", "PROFILE.md", "AGENTS.md");
|
||||
assertThat(hits.get(0).score()).isGreaterThan(hits.get(1).score());
|
||||
assertThat(hits.get(1).score()).isGreaterThan(hits.get(2).score());
|
||||
assertThat(hits.get(2).score()).isGreaterThan(hits.get(3).score());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Term-hit count: line matching 2 terms scores 2× the per-file weight")
|
||||
void termHitCountAffectsScore() {
|
||||
List<WorkspaceFileEntity> candidates = List.of(
|
||||
file("MEMORY.md",
|
||||
"Line with running and shoes.\n" +
|
||||
"Line with only running.\n"));
|
||||
when(fileMapper.selectList(any())).thenReturn(candidates);
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "running shoes", null, 10);
|
||||
|
||||
// Both lines surface; the 2-term line ranks first with score 2.0, the
|
||||
// 1-term line ranks second with score 1.0 (MEMORY weight = 1.0).
|
||||
assertThat(hits).hasSize(2);
|
||||
assertThat(hits.get(0).score()).isEqualTo(2.0);
|
||||
assertThat(hits.get(0).snippet()).contains("[[running]]").contains("[[shoes]]");
|
||||
assertThat(hits.get(1).score()).isEqualTo(1.0);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Per-file hit cap: a flood of matching lines in one file caps at 5")
|
||||
void perFileHitCap() {
|
||||
StringBuilder content = new StringBuilder();
|
||||
for (int i = 0; i < 12; i++) content.append("line ").append(i).append(" running\n");
|
||||
List<WorkspaceFileEntity> candidates = List.of(file("MEMORY.md", content.toString()));
|
||||
when(fileMapper.selectList(any())).thenReturn(candidates);
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "running", null, 30);
|
||||
|
||||
assertThat(hits).hasSize(5);
|
||||
// First 5 lines only.
|
||||
assertThat(hits).extracting(MemorySearchHit::lineNumber)
|
||||
.containsExactly(1, 2, 3, 4, 5);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Limit caps the total result count after global ranking")
|
||||
void limitClampsResults() {
|
||||
StringBuilder mem = new StringBuilder();
|
||||
for (int i = 0; i < 5; i++) mem.append("MEM running\n");
|
||||
StringBuilder daily = new StringBuilder();
|
||||
for (int i = 0; i < 5; i++) daily.append("DAILY running\n");
|
||||
when(fileMapper.selectList(any())).thenReturn(List.of(
|
||||
file("MEMORY.md", mem.toString()),
|
||||
file("memory/2026-05-10.md", daily.toString())));
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "running", null, 3);
|
||||
|
||||
assertThat(hits).hasSize(3);
|
||||
// Highest-weight hits come first — all from MEMORY.md.
|
||||
assertThat(hits).extracting(MemorySearchHit::filename)
|
||||
.containsExactly("MEMORY.md", "MEMORY.md", "MEMORY.md");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Snippet: each matched term is wrapped in [[...]] with term boundaries preserved")
|
||||
void snippetHighlightingTermBoundaries() {
|
||||
List<WorkspaceFileEntity> candidates = List.of(file("MEMORY.md",
|
||||
"用户喜欢在公园跑步\n"));
|
||||
when(fileMapper.selectList(any())).thenReturn(candidates);
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "用户喜欢跑步", null, 10);
|
||||
// tokenize("用户喜欢跑步") → ["用户","喜欢","跑步"]; line contains all three.
|
||||
assertThat(hits).hasSize(1);
|
||||
assertThat(hits.get(0).snippet())
|
||||
.contains("[[用户]]")
|
||||
.contains("[[喜欢]]")
|
||||
.contains("[[跑步]]")
|
||||
// Adjacent term matches keep their boundary, not merged into one bracket.
|
||||
.contains("[[用户]][[喜欢]]");
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Snippet: long line is clipped to ±80 chars around the first match with ellipses")
|
||||
void snippetTruncation() {
|
||||
String lead = "x".repeat(200);
|
||||
String tail = "y".repeat(200);
|
||||
String line = lead + " running " + tail;
|
||||
when(fileMapper.selectList(any())).thenReturn(List.of(file("MEMORY.md", line + "\n")));
|
||||
|
||||
List<MemorySearchHit> hits = service.searchSnippets(1L, "running", null, 10);
|
||||
|
||||
assertThat(hits).hasSize(1);
|
||||
String snippet = hits.get(0).snippet();
|
||||
assertThat(snippet).startsWith("...");
|
||||
assertThat(snippet).endsWith("...");
|
||||
assertThat(snippet).contains("[[running]]");
|
||||
// Clip window is 80 chars on each side plus the match (7 chars) plus
|
||||
// two "..." markers (6 chars). Should be << original length.
|
||||
assertThat(snippet.length()).isLessThan(line.length());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Wrapper carries one content-LIKE per token plus the prefix group and LIMIT 50")
|
||||
void wrapperContainsTermsAndPrefixes() {
|
||||
when(fileMapper.selectList(any())).thenReturn(List.of());
|
||||
Set<String> prefixes = new LinkedHashSet<>(List.of("memory/", "MEMORY.md"));
|
||||
service.searchSnippets(42L, "running shoes 跑步", prefixes, 10);
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<LambdaQueryWrapper<WorkspaceFileEntity>> captor =
|
||||
ArgumentCaptor.forClass(LambdaQueryWrapper.class);
|
||||
org.mockito.Mockito.verify(fileMapper).selectList(captor.capture());
|
||||
LambdaQueryWrapper<WorkspaceFileEntity> wrapper = captor.getValue();
|
||||
|
||||
// Force SQL rendering so paramNameValuePairs gets populated; assert
|
||||
// on the SQL shape (placeholders only — column-name casing depends
|
||||
// on global underscore-camelCase config not present in this unit
|
||||
// test) and on the bound literal values.
|
||||
String sql = wrapper.getTargetSql();
|
||||
// One equality on agent + two prefix LIKEs in an OR group + three
|
||||
// content LIKEs, ANDed together, suffixed with the candidate cap.
|
||||
assertThat(sql).contains("LIKE ? OR")
|
||||
.contains("AND content")
|
||||
.contains("LIMIT 50");
|
||||
assertThat(sql.chars().filter(ch -> ch == '?').count())
|
||||
.as("one agentId + two prefix LIKEs + three content LIKEs = 6 bind params")
|
||||
.isEqualTo(6);
|
||||
|
||||
List<Object> values = new ArrayList<>(wrapper.getParamNameValuePairs().values());
|
||||
assertThat(values).contains(42L);
|
||||
// Each content-LIKE term gets %term% by MyBatis-Plus's like().
|
||||
assertThat(values).contains("%running%", "%shoes%", "%跑步%");
|
||||
// likeRight produces "prefix%" — confirms both prefixes were bound.
|
||||
assertThat(values).contains("memory/%", "MEMORY.md%");
|
||||
}
|
||||
|
||||
// ---------- helpers ----------
|
||||
|
||||
private static WorkspaceFileEntity file(String filename, String content) {
|
||||
WorkspaceFileEntity e = new WorkspaceFileEntity();
|
||||
e.setAgentId(1L);
|
||||
e.setFilename(filename);
|
||||
e.setContent(content);
|
||||
return e;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user