package com.gzzm.lobster.tool.builtin;
import com.gzzm.lobster.common.MemoryCategory;
import com.gzzm.lobster.common.ToolCategory;
import com.gzzm.lobster.common.ToolRiskLevel;
import com.gzzm.lobster.memory.MemoryService;
import com.gzzm.lobster.memory.PersonalMemory;
import com.gzzm.lobster.tool.*;
import net.cyan.nest.annotation.Inject;
import java.util.*;
/**
* MemoryTools —— 个人记忆工具（Claude Code 风格的"索引 + 正文"两段式）。
*
*
* - {@code memory_search} —— 返回「索引行」（不含 content），轻量
* - {@code memory_read} —— 按 memoryId 拉取完整正文
* - {@code memory_write} —— 写入新记忆，强制 name/description/category
*
*/
public class MemoryTools {
@Inject private MemoryService memoryService;
public void registerTo(ToolRegistry registry) {
registry.register(searchDef(), this::search);
registry.register(readDef(), this::read);
registry.register(writeDef(), this::write);
}
// ---------------- memory_search ----------------
private BuiltinToolDefinition searchDef() {
return BuiltinToolDefinition.builder()
.name("memory_search")
.displayName("搜索个人记忆索引")
.description("按关键词检索当前用户的个人记忆索引，返回候选条目的 memoryId + name + description + category。"
+ "正文需另行通过 memory_read 工具按 memoryId 拉取。空 query 返回最近若干条索引。"
+ "严格个人隔离，仅返回当前用户数据。")
.category(ToolCategory.MEMORY)
.risk(ToolRiskLevel.READ_ONLY)
.inputSchema(SchemaBuilder.obj()
.prop("query", "string", "关键词（可留空）")
.propEnum("category", "分类筛选（可选）", "user", "feedback", "project", "reference", "ALL")
.propInt("maxResults", "最大返回数，默认 10，上限 50")
.build())
.build();
}
private ToolResult search(ToolContext ctx, Map args) throws Exception {
String query = asStr(args.get("query"));
String catStr = asStr(args.get("category"));
int maxResults = clamp(asInt(args.get("maxResults"), 10), 1, 50);
MemoryCategory category = null;
if (catStr != null && !catStr.isEmpty() && !"ALL".equalsIgnoreCase(catStr)) {
try {
category = MemoryCategory.valueOf(catStr);
} catch (IllegalArgumentException e) {
return ToolResult.error("memory.invalid_category: " + catStr
+ "（可选：user / feedback / project / reference / ALL）");
}
}
// MemoryService.search 统一处理：空/短 query + category → listByCategory；
// 有 query + category → bigram 检索后再用 category 过滤。
List hits = memoryService.search(
ctx.getUserContext(), query, category, maxResults);
List> out = new ArrayList<>();
for (PersonalMemory m : hits) out.add(toIndexRow(m));
Map data = new LinkedHashMap<>();
data.put("results", out);
return ToolResult.okData(data);
}
// ---------------- memory_read ----------------
private BuiltinToolDefinition readDef() {
return BuiltinToolDefinition.builder()
.name("memory_read")
.displayName("读取个人记忆正文")
.description("按 memoryId 拉取完整记忆正文（包括 content）。先用 memory_search 找到候选 id，"
+ "再用本工具展开。仅允许读取当前用户自己的记忆。")
.category(ToolCategory.MEMORY)
.risk(ToolRiskLevel.READ_ONLY)
.inputSchema(SchemaBuilder.obj()
.prop("memoryId", "string", "记忆 ID")
.required("memoryId")
.build())
.build();
}
private ToolResult read(ToolContext ctx, Map args) throws Exception {
String memoryId = asStr(args.get("memoryId"));
PersonalMemory m = memoryService.read(ctx.getUserContext(), memoryId);
Map data = toIndexRow(m);
data.put("content", m.getContent());
data.put("sourceType", m.getSourceType());
data.put("updateTime", m.getUpdateTime());
return ToolResult.okData(data);
}
// ---------------- memory_write ----------------
private BuiltinToolDefinition writeDef() {
return BuiltinToolDefinition.builder()
.name("memory_write")
.displayName("写入个人记忆")
.description("写入一条个人记忆，按 Claude Code 风格需提供：\n"
+ "- name：短标题（≤80），名词短语且具体可辨识\n"
+ "- description：一句话钩子（≤200），用于未来对话判断是否相关——这是记忆最重要的字段\n"
+ "- category：user / feedback / project / reference，按「读取时机」分类\n"
+ "- content：可选的完整正文（≤2000）。feedback/project 类推荐结构化为：规则/事实 + Why + How to apply\n"
+ "同 name 视为更新：自动覆盖 description + content 并刷新 updateTime。单用户 20 次/分钟 限流。\n"
+ "不要用于：代码模式、调用路径、git 历史、调试 fix、CLAUDE.md 已覆盖内容——这些不属于记忆。")
.category(ToolCategory.MEMORY)
.risk(ToolRiskLevel.WRITE)
.inputSchema(SchemaBuilder.obj()
.prop("name", "string", "短标题，≤80 字")
.prop("description", "string", "一句话钩子，≤200 字")
.propEnum("category", "记忆分类", "user", "feedback", "project", "reference")
.prop("content", "string", "完整正文（可选），≤2000 字")
.propEnum("sourceType", "来源",
"explicit_user_write", "assistant_confirmed_write",
"imported_preference", "derived_summary")
.required("name", "description", "category")
.build())
.build();
}
private ToolResult write(ToolContext ctx, Map args) throws Exception {
String name = asStr(args.get("name"));
String desc = asStr(args.get("description"));
String content = asStr(args.get("content"));
String catStr = asStr(args.get("category"));
String source = asStr(args.get("sourceType"));
if (catStr == null || catStr.isEmpty()) {
return ToolResult.error("memory.invalid_category: category 必填（user / feedback / project / reference）");
}
MemoryCategory category;
try {
category = MemoryCategory.valueOf(catStr);
} catch (IllegalArgumentException e) {
// 静默回落会让错类别的 LLM 写入蒙混过关，索引长期污染。直接拒绝，让模型重写。
return ToolResult.error("memory.invalid_category: " + catStr
+ "（可选：user / feedback / project / reference）");
}
PersonalMemory m = memoryService.write(ctx.getUserContext(), name, desc, content, category, source);
Map data = toIndexRow(m);
data.put("status", "SAVED");
return ToolResult.ok("记忆已写入", data);
}
// ---------------- shared ----------------
private Map toIndexRow(PersonalMemory m) {
Map row = new LinkedHashMap<>();
row.put("memoryId", m.getMemoryId());
row.put("name", m.getName());
row.put("description", m.getDescription());
row.put("category", m.getCategory() == null ? null : m.getCategory().name());
row.put("createTime", m.getCreateTime());
return row;
}
private String asStr(Object o) { return o == null ? null : String.valueOf(o); }
private int asInt(Object o, int def) {
if (o == null) return def;
if (o instanceof Number) return ((Number) o).intValue();
try { return Integer.parseInt(String.valueOf(o)); } catch (Exception e) { return def; }
}
private int clamp(int v, int min, int max) { return Math.max(min, Math.min(max, v)); }
}