Skip to content

perf: bound graph-memory recall token injection #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Liyuk wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: bound graph-memory recall token injection #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion dsh.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,10 @@ export interface Config {
recallEnabled?: boolean;
recallMaxNodes?: number;
recallMaxDepth?: number;
/** 每轮召回注入的 token 预算（0 = 不限制） */
recallTokenBudget?: number;
/** 单节点 content 注入的最大字符数（0 = 不截断） */
contentMaxChars?: number;
maintenanceInterval?: number;
llmProvider?: string;
llmModel?: string;
Expand Down Expand Up @@ -151,6 +155,8 @@ export function apply(ctx: DshContext, input: Config = {}): void {
compactTurnCount: input.maintenanceInterval ?? DEFAULT_CONFIG.compactTurnCount,
recallMaxNodes: input.recallMaxNodes ?? DEFAULT_CONFIG.recallMaxNodes,
recallMaxDepth: input.recallMaxDepth ?? DEFAULT_CONFIG.recallMaxDepth,
recallTokenBudget: input.recallTokenBudget ?? DEFAULT_CONFIG.recallTokenBudget,
contentMaxChars: input.contentMaxChars ?? DEFAULT_CONFIG.contentMaxChars,
embedding,
};
const extractionEnabled = input.extractionEnabled ?? true;
Expand Down Expand Up @@ -354,7 +360,8 @@ export function apply(ctx: DshContext, input: Config = {}): void {
const activeIds = new Set(activeNodes.map((node) => node.id));
const activeEdges = allEdges(db).filter((edge) => activeIds.has(edge.fromId) && activeIds.has(edge.toId));
const built = assembleContext(db, {
tokenBudget: 0,
tokenBudget: config.recallTokenBudget,
contentMaxChars: config.contentMaxChars,
activeNodes,
activeEdges,
recalledNodes: recalled.nodes,
Expand Down
3 changes: 2 additions & 1 deletion index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,7 +366,8 @@ const graphMemoryPlugin = {

// ── 2. 图谱 + 溯源 ─────────────────────────────
const { xml, systemPrompt, tokens: gmTokens, episodicXml, episodicTokens } = assembleContext(db, {
tokenBudget: 0,
tokenBudget: cfg.recallTokenBudget,
contentMaxChars: cfg.contentMaxChars,
activeNodes,
activeEdges,
recalledNodes: rec.nodes,
Expand Down
120 changes: 95 additions & 25 deletions src/format/assemble.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,65 @@ import { getCommunitySummary, getEpisodicMessages } from "../store/store.ts";

const CHARS_PER_TOKEN = 3;

/**
* 单节点 content 注入的最大字符数（0 = 不截断）。
* 完整 content 始终可通过 gm_search 按需取回，注入端只需保留核心步骤/结论。
*/
export const DEFAULT_CONTENT_MAX_CHARS = 400;

/** 溯源（episodic）注入的收紧参数：top 节点数 / 单节点总字符 / 单条消息字符。 */
const EPISODIC_TOP_NODES = 2;
const EPISODIC_MAX_CHARS = 300;
const EPISODIC_MESSAGE_CHARS = 150;

/** XML 标签/属性等固定开销的粗略字符数，纳入节点 token 估算。 */
const NODE_XML_OVERHEAD_CHARS = 80;

/**
* 估算单节点注入后的 token 数（name + description + content + XML 固定开销）。
*/
function estimateNodeTokens(n: GmNode): number {
return Math.ceil(
(n.name.length + n.description.length + n.content.length + NODE_XML_OVERHEAD_CHARS) / CHARS_PER_TOKEN,
);
}

/**
* 按字符上限截断 content 正文；截断时追加省略号以提示内容被裁剪。
*/
function clipBody(content: string, maxChars: number): string {
const trimmed = content.trim();
if (maxChars > 0 && trimmed.length > maxChars) {
return `${trimmed.slice(0, maxChars)}…`;
}
return trimmed;
}

/**
* 去掉 content 首行与节点 name 重复的内容。
*
* 提取模板让 content 首行以 `[name]`（或裸 name）开头，而 XML 标签已有 name 属性，
* 首行是纯冗余。去掉后零信息损失，且对存量节点同样生效（无需迁移数据）。
*/
function stripLeadingName(content: string, name: string): string {
const lines = content.split("\n");
if (!lines.length) return content;
const first = lines[0].trim();
const bare = first.replace(/^\[|\]$/g, "").trim();
if (first === name || bare === name) {
return lines.slice(1).join("\n").trim();
}
return content;
}

/**
* 生成节点正文：先去掉首行 name 冗余，再按上限截断，最后 XML 转义。
*/
function renderBody(n: GmNode, contentMaxChars: number): string {
const stripped = stripLeadingName(n.content, n.name);
return escapeXml(clipBody(stripped, contentMaxChars));
}

/**
* 构建知识图谱的 system prompt 引导文字
*/
Expand All @@ -33,41 +92,34 @@ export function buildSystemPromptAddition(params: {
sections.push(
"## Graph Memory — 知识图谱记忆",
"",
"Below `<knowledge_graph>` is your accumulated experience from past conversations.",
"It contains structured knowledge — NOT raw conversation history.",
"Below `<knowledge_graph>` is your accumulated experience from past conversations — structured knowledge, NOT raw history.",
"",
`Current graph: ${skillCount} skills, ${eventCount} events, ${taskCount} tasks, ${edgeCount} relationships.`,
);

if (hasRecalled) {
sections.push(
"",
`**${recalledCount} nodes recalled from OTHER conversations** — these are proven solutions that worked before.`,
"Apply them directly when the current situation matches their trigger conditions.",
`**${recalledCount} nodes recalled from OTHER conversations** — proven solutions; apply directly when the situation matches their trigger conditions.`,
);
}

sections.push(
"",
"## Recalled context for this query",
"",
"This is a context engine. The following was retrieved by semantic search for the current message:",
"Retrieved by semantic search for the current message:",
"",
"- **`<episodic_context>`** — Trimmed conversation traces from sessions that produced the knowledge nodes, ordered by time.",
"- **`<episodic_context>`** — Trimmed conversation traces from sessions that produced the knowledge nodes.",
"- **`<knowledge_graph>`** — Relevant triples (TASK/SKILL/EVENT) and edges, grouped by community.",
"- **Recent 5 turns** — Last turn in full, previous 4 turns as user+assistant text only.",
"",
"Read this context first. Use `gm_search` only if insufficient. Use `gm_record` to save new knowledge.",
"Read this first. Use `gm_search` if insufficient, `gm_record` to save new knowledge.",
);

if (isRich) {
sections.push(
"",
"**Graph navigation:** Edges show how knowledge connects:",
"- `SOLVED_BY`: an EVENT was fixed by a SKILL — apply the skill when you see similar errors",
"- `USED_SKILL`: a TASK used a SKILL — reuse the same approach for similar tasks",
"- `PATCHES`: a newer SKILL corrects an older one — prefer the newer version",
"- `CONFLICTS_WITH`: two SKILLs are mutually exclusive — check conditions before choosing",
"**Graph navigation:** `SOLVED_BY`=EVENT fixed by SKILL · `USED_SKILL`=TASK used SKILL · `PATCHES`=newer SKILL corrects older · `CONFLICTS_WITH`=mutually exclusive SKILLs.",
);
}

Expand All @@ -76,6 +128,9 @@ export function buildSystemPromptAddition(params: {

/**
* 组装知识图谱为 XML context
*
* tokenBudget > 0 时按预算贪心裁剪节点（已排序，优先级最高者先保留，至少保留 1 个）；
* tokenBudget <= 0 时全量放入（向后兼容）。
*/
export function assembleContext(
db: DatabaseSyncInstance,
Expand All @@ -85,9 +140,11 @@ export function assembleContext(
activeEdges: GmEdge[];
recalledNodes: GmNode[];
recalledEdges: GmEdge[];
contentMaxChars?: number;
},
): { xml: string | null; systemPrompt: string; tokens: number; episodicXml: string; episodicTokens: number } {
// recall 返回多少节点就放多少，不截断
const contentMaxChars = params.contentMaxChars ?? DEFAULT_CONTENT_MAX_CHARS;

const map = new Map<string, GmNode & { src: "active" | "recalled" }>();
for (const n of params.recalledNodes) map.set(n.id, { ...n, src: "recalled" });
for (const n of params.activeNodes) map.set(n.id, { ...n, src: "active" });
Expand All @@ -103,8 +160,21 @@ export function assembleContext(
b.pagerank - a.pagerank
);

// recall 返回的已经是 PPR 排序过的，全量放入
const selected = sorted;
// ── 按 token 预算裁剪（0 = 不限制，向后兼容）────────────
let selected = sorted;
if (params.tokenBudget > 0) {
const kept: typeof sorted = [];
let used = 0;
for (const n of sorted) {
const cost = estimateNodeTokens(n);
if (kept.length > 0 && used + cost > params.tokenBudget) break;
kept.push(n);
used += cost;
}
// 至少保留优先级最高的 1 个节点，即使单节点就超预算
if (kept.length === 0 && sorted.length) kept.push(sorted[0]);
selected = kept;
}

if (!selected.length) return { xml: null, systemPrompt: "", tokens: 0, episodicXml: "", episodicTokens: 0 };

Expand Down Expand Up @@ -140,8 +210,8 @@ export function assembleContext(
for (const n of members) {
const tag = n.type.toLowerCase();
const srcAttr = n.src === "recalled" ? ` source="recalled"` : "";
const timeAttr = ` updated="${new Date(n.updatedAt).toISOString().slice(0, 10)}"`;
xmlParts.push(` <${tag} name="${n.name}" desc="${escapeXml(n.description)}"${srcAttr}${timeAttr}>\n${n.content.trim()}\n </${tag}>`);
const body = renderBody(n, contentMaxChars);
xmlParts.push(` <${tag} name="${n.name}" desc="${escapeXml(n.description)}"${srcAttr}>\n${body}\n </${tag}>`);
}
xmlParts.push(` </community>`);
}
Expand All @@ -150,8 +220,8 @@ export function assembleContext(
for (const n of noCommunity) {
const tag = n.type.toLowerCase();
const srcAttr = n.src === "recalled" ? ` source="recalled"` : "";
const timeAttr = ` updated="${new Date(n.updatedAt).toISOString().slice(0, 10)}"`;
xmlParts.push(` <${tag} name="${n.name}" desc="${escapeXml(n.description)}"${srcAttr}${timeAttr}>\n${n.content.trim()}\n </${tag}>`);
const body = renderBody(n, contentMaxChars);
xmlParts.push(` <${tag} name="${n.name}" desc="${escapeXml(n.description)}"${srcAttr}>\n${body}\n </${tag}>`);
}

const nodesXml = xmlParts.join("\n");
Expand All @@ -172,19 +242,19 @@ export function assembleContext(
edgeCount: edges.length,
});

// ── 溯源选拉：PPR top 3 节点 → 拉原始 user/assistant 对话 ──
const topNodes = selected.slice(0, 3);
// ── 溯源选拉：PPR top N 节点 → 拉原始 user/assistant 对话 ──
const topNodes = selected.slice(0, EPISODIC_TOP_NODES);
const episodicParts: string[] = [];

for (const node of topNodes) {
if (!node.sourceSessions?.length) continue;
// 取最近的 2 个 session
const recentSessions = node.sourceSessions.slice(-2);
const msgs = getEpisodicMessages(db, recentSessions, node.updatedAt, 500);
const msgs = getEpisodicMessages(db, recentSessions, node.updatedAt, EPISODIC_MAX_CHARS);
if (!msgs.length) continue;

const lines = msgs.map(m =>
` [${m.role.toUpperCase()}] ${escapeXml(m.text.slice(0, 200))}`
` [${m.role.toUpperCase()}] ${escapeXml(m.text.slice(0, EPISODIC_MESSAGE_CHARS))}`
).join("\n");
episodicParts.push(` <trace node="${node.name}">\n${lines}\n </trace>`);
}
Expand All @@ -205,4 +275,4 @@ export function assembleContext(

function escapeXml(s: string): string {
return s.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;").replace(/"/g, "&quot;");
}
}
25 changes: 18 additions & 7 deletions src/recaller/recall.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ export class Recaller {
const precise = await this.recallPrecise(query, limit);
const generalized = await this.recallGeneralized(query, limit);

// ── 合并去重（全部保留，只去重复节点） ────────────────
const merged = this.mergeResults(precise, generalized);
// ── 合并：精确路径优先，泛化路径只补精确未覆盖的社区，总量封顶 limit ──
const merged = this.mergeResults(precise, generalized, limit);

return merged;
}
Expand Down Expand Up @@ -179,19 +179,30 @@ export class Recaller {
}

/**
* 合并两条路径的结果：全部保留，只去重复节点
* 合并两条路径的结果：精确路径优先，泛化路径只补充精确路径未覆盖的社区，总量封顶 limit。
*
* 之前两条路径各自跑满 limit 后全量合并，节点数可能翻倍；且泛化路径常补进与精确
* 路径同社区的高度重复节点。现在泛化路径只在「目标节点属于精确路径未覆盖社区」时才
* 补入，既保留跨领域概览，又避免同社区冗余，同时把总节点数稳定封顶在 limit。
*/
private mergeResults(precise: RecallResult, generalized: RecallResult): RecallResult {
private mergeResults(precise: RecallResult, generalized: RecallResult, limit: number): RecallResult {
const nodeMap = new Map<string, GmNode>();
const edgeMap = new Map<string, GmEdge>();

// 精确路径全部入场
// 精确路径全部入场（已按 PPR 排序并 slice 到 limit）
for (const n of precise.nodes) nodeMap.set(n.id, n);
for (const e of precise.edges) edgeMap.set(e.id, e);

// 泛化路径去重后全部入场
// 泛化路径：只补精确路径未覆盖的社区，直到达到总配额
const preciseCommunityIds = new Set(
precise.nodes.map(n => n.communityId).filter((cid): cid is string => Boolean(cid)),
);
for (const n of generalized.nodes) {
if (!nodeMap.has(n.id)) nodeMap.set(n.id, n);
if (nodeMap.size >= limit) break;
if (nodeMap.has(n.id)) continue;
// 泛化节点若属于精确路径已覆盖的社区，视为冗余，跳过
if (n.communityId && preciseCommunityIds.has(n.communityId)) continue;
nodeMap.set(n.id, n);
}

// 合并边：两端都在最终节点集中的边才保留
Expand Down
6 changes: 6 additions & 0 deletions src/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,10 @@ export interface GmConfig {
compactTurnCount: number;
recallMaxNodes: number;
recallMaxDepth: number;
/** 每轮召回注入的 token 预算（0 = 不限制，向后兼容） */
recallTokenBudget: number;
/** 单节点 content 注入的最大字符数（0 = 不截断） */
contentMaxChars: number;
freshTailCount: number;
embedding?: EmbeddingConfig;
llm?: {
Expand All @@ -152,6 +156,8 @@ export const DEFAULT_CONFIG: GmConfig = {
compactTurnCount: 6,
recallMaxNodes: 6,
recallMaxDepth: 2,
recallTokenBudget: 2000,
contentMaxChars: 400,
freshTailCount: 10,
dedupThreshold: 0.90,
pagerankDamping: 0.85,
Expand Down
Loading