mirror of
https://bitbucket.org/siakitem/my-pi.git
synced 2026-10-26 10:45:22 +00:00
feat: unify context compaction
This commit is contained in:
云服务部-叶林立
@@ -0,0 +1,381 @@
|
||||
import { createHash } from "node:crypto";
|
||||
import { estimateTokens } from "@earendil-works/pi-coding-agent";
|
||||
import {
|
||||
COMPACTION_CUSTOM_TYPE,
|
||||
COMPACTION_VERSION,
|
||||
type PiCompactionConfig,
|
||||
type ProjectionSnapshot,
|
||||
type SegmentSummaryRecord,
|
||||
type ToolHistoryRecord,
|
||||
type UsageBaseline,
|
||||
} from "./types.js";
|
||||
|
||||
function timestampOf(value: unknown): number {
|
||||
if (typeof value === "number" && Number.isFinite(value)) return value;
|
||||
if (typeof value === "string") {
|
||||
const parsed = Date.parse(value);
|
||||
if (Number.isFinite(parsed)) return parsed;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
export function sourceKey(toolCallId: string, timestamp: unknown): string {
|
||||
return `${toolCallId}@${timestampOf(timestamp)}`;
|
||||
}
|
||||
|
||||
export function contentText(content: unknown): string {
|
||||
if (typeof content === "string") return content;
|
||||
if (!Array.isArray(content)) return "";
|
||||
return content
|
||||
.flatMap((part) => part && typeof part === "object" && (part as any).type === "text" && typeof (part as any).text === "string"
|
||||
? [(part as any).text]
|
||||
: [])
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
function hash(value: string): string {
|
||||
return createHash("sha256").update(value).digest("hex").slice(0, 20);
|
||||
}
|
||||
|
||||
function json(value: unknown): string {
|
||||
try {
|
||||
return JSON.stringify(value) ?? "undefined";
|
||||
} catch {
|
||||
return "[unserializable]";
|
||||
}
|
||||
}
|
||||
|
||||
function boundedEvidence(text: string, maxChars: number): string {
|
||||
const normalized = text.replace(/\u001b\[[0-9;]*m/g, "").trim();
|
||||
if (normalized.length <= maxChars) return normalized;
|
||||
const marker = `\n… ${normalized.length - maxChars} chars omitted …\n`;
|
||||
const available = Math.max(0, maxChars - marker.length);
|
||||
const head = Math.ceil(available * 0.7);
|
||||
const tail = Math.max(0, available - head);
|
||||
return `${normalized.slice(0, head)}${marker}${tail ? normalized.slice(-tail) : ""}`;
|
||||
}
|
||||
|
||||
export function deterministicSummary(params: {
|
||||
toolName: string;
|
||||
args: unknown;
|
||||
rawText: string;
|
||||
isError: boolean;
|
||||
maxStubTokens: number;
|
||||
}): string {
|
||||
const maxChars = params.maxStubTokens * 4;
|
||||
const argsText = boundedEvidence(json(params.args), Math.min(600, Math.floor(maxChars * 0.25)));
|
||||
const evidenceBudget = Math.max(128, maxChars - argsText.length - 120);
|
||||
const evidence = boundedEvidence(params.rawText || "(no output)", evidenceBudget);
|
||||
return [
|
||||
`[pi-compaction tool=${params.toolName} status=${params.isError ? "error" : "ok"}]`,
|
||||
`Args: ${argsText}`,
|
||||
`Evidence:\n${evidence}`,
|
||||
].join("\n");
|
||||
}
|
||||
|
||||
function messageKey(message: any): string {
|
||||
const role = String(message?.role ?? "unknown");
|
||||
const timestamp = timestampOf(message?.timestamp);
|
||||
if (role === "toolResult") return `${role}:${message.toolCallId ?? ""}:${timestamp}`;
|
||||
if (role === "assistant") {
|
||||
const blocks = Array.isArray(message.content) ? message.content : [];
|
||||
const ids = blocks.flatMap((block: any) => block?.type === "toolCall" ? [block.id] : []).join(",");
|
||||
return `${role}:${timestamp}:${ids}:${blocks.length}`;
|
||||
}
|
||||
return `${role}:${timestamp}:${typeof message?.content === "string" ? message.content.length : (message?.content?.length ?? 0)}`;
|
||||
}
|
||||
|
||||
function isSummaryRecord(value: unknown): value is SegmentSummaryRecord {
|
||||
if (!value || typeof value !== "object") return false;
|
||||
const record = value as Partial<SegmentSummaryRecord>;
|
||||
return record.version === COMPACTION_VERSION
|
||||
&& record.kind === "segment-summary"
|
||||
&& typeof record.sourceKey === "string"
|
||||
&& typeof record.toolCallId === "string"
|
||||
&& typeof record.resultTimestamp === "number"
|
||||
&& typeof record.toolName === "string"
|
||||
&& typeof record.sourceHash === "string"
|
||||
&& typeof record.deterministicSummary === "string";
|
||||
}
|
||||
|
||||
export class SegmentIndex {
|
||||
private readonly records = new Map<string, ToolHistoryRecord>();
|
||||
private readonly ids = new Map<string, string[]>();
|
||||
private generation = 0;
|
||||
private semanticJobs = 0;
|
||||
private usageBaseline?: UsageBaseline;
|
||||
private sourceMessages: any[] = [];
|
||||
private sourceKeys: string[] = [];
|
||||
private projectedMessages: any[] = [];
|
||||
private rawEstimates: number[] = [];
|
||||
private effectiveEstimates: number[] = [];
|
||||
private cached?: ProjectionSnapshot;
|
||||
|
||||
readonly sessionId: string;
|
||||
|
||||
constructor(sessionId: string) {
|
||||
this.sessionId = sessionId;
|
||||
}
|
||||
|
||||
hydrate(branch: readonly any[], config: PiCompactionConfig): void {
|
||||
this.reset();
|
||||
if (!config.enabled) return;
|
||||
const summaries = new Map<string, SegmentSummaryRecord>();
|
||||
for (const entry of branch) {
|
||||
if (entry?.type === "custom" && entry.customType === COMPACTION_CUSTOM_TYPE && isSummaryRecord(entry.data)) {
|
||||
summaries.set(entry.data.sourceKey, entry.data);
|
||||
}
|
||||
}
|
||||
|
||||
const pendingCalls = new Map<string, { toolName: string; args: unknown }>();
|
||||
for (const entry of branch) {
|
||||
if (entry?.type !== "message") continue;
|
||||
const message = entry.message;
|
||||
if (message?.role === "assistant" && Array.isArray(message.content)) {
|
||||
for (const block of message.content) {
|
||||
if (block?.type === "toolCall" && typeof block.id === "string") {
|
||||
pendingCalls.set(block.id, {
|
||||
toolName: String(block.name ?? "unknown"),
|
||||
args: block.input ?? block.arguments ?? block.args ?? {},
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (message?.role === "toolResult" && typeof message.toolCallId === "string") {
|
||||
const call = pendingCalls.get(message.toolCallId) ?? {
|
||||
toolName: String(message.toolName ?? "unknown"),
|
||||
args: {},
|
||||
};
|
||||
const key = sourceKey(message.toolCallId, message.timestamp);
|
||||
const rawText = contentText(message.content);
|
||||
const persisted = summaries.get(key);
|
||||
const record: ToolHistoryRecord = {
|
||||
sourceKey: key,
|
||||
toolCallId: message.toolCallId,
|
||||
resultTimestamp: timestampOf(message.timestamp),
|
||||
toolName: call.toolName,
|
||||
args: call.args,
|
||||
isError: Boolean(message.isError),
|
||||
rawText,
|
||||
sourceHash: hash(`${call.toolName}\0${json(call.args)}\0${rawText}`),
|
||||
deterministicSummary: persisted?.deterministicSummary ?? deterministicSummary({
|
||||
toolName: call.toolName,
|
||||
args: call.args,
|
||||
rawText,
|
||||
isError: Boolean(message.isError),
|
||||
maxStubTokens: config.deterministic.maxStubTokens,
|
||||
}),
|
||||
semanticSummary: persisted?.semanticSummary,
|
||||
};
|
||||
this.put(record);
|
||||
pendingCalls.delete(message.toolCallId);
|
||||
}
|
||||
}
|
||||
this.generation++;
|
||||
}
|
||||
|
||||
reset(): void {
|
||||
this.records.clear();
|
||||
this.ids.clear();
|
||||
this.generation = 0;
|
||||
this.semanticJobs = 0;
|
||||
this.usageBaseline = undefined;
|
||||
this.sourceMessages = [];
|
||||
this.sourceKeys = [];
|
||||
this.projectedMessages = [];
|
||||
this.rawEstimates = [];
|
||||
this.effectiveEstimates = [];
|
||||
this.cached = undefined;
|
||||
}
|
||||
|
||||
private put(record: ToolHistoryRecord): void {
|
||||
this.records.set(record.sourceKey, record);
|
||||
const keys = this.ids.get(record.toolCallId) ?? [];
|
||||
if (!keys.includes(record.sourceKey)) keys.push(record.sourceKey);
|
||||
this.ids.set(record.toolCallId, keys);
|
||||
}
|
||||
|
||||
captureTurn(message: any, toolResults: any[], config: PiCompactionConfig): ToolHistoryRecord[] {
|
||||
if (!config.enabled || !config.deterministic.enabled || !Array.isArray(message?.content)) return [];
|
||||
const calls = new Map<string, any>();
|
||||
for (const block of message.content) {
|
||||
if (block?.type === "toolCall" && typeof block.id === "string") calls.set(block.id, block);
|
||||
}
|
||||
const captured: ToolHistoryRecord[] = [];
|
||||
for (const result of toolResults ?? []) {
|
||||
if (!result || typeof result.toolCallId !== "string") continue;
|
||||
const call = calls.get(result.toolCallId);
|
||||
if (!call) continue;
|
||||
const rawText = contentText(result.content);
|
||||
const toolName = String(call.name ?? result.toolName ?? "unknown");
|
||||
const args = call.input ?? call.arguments ?? call.args ?? {};
|
||||
const key = sourceKey(result.toolCallId, result.timestamp);
|
||||
const record: ToolHistoryRecord = {
|
||||
sourceKey: key,
|
||||
toolCallId: result.toolCallId,
|
||||
resultTimestamp: timestampOf(result.timestamp),
|
||||
toolName,
|
||||
args,
|
||||
isError: Boolean(result.isError),
|
||||
rawText,
|
||||
sourceHash: hash(`${toolName}\0${json(args)}\0${rawText}`),
|
||||
deterministicSummary: deterministicSummary({
|
||||
toolName,
|
||||
args,
|
||||
rawText,
|
||||
isError: Boolean(result.isError),
|
||||
maxStubTokens: config.deterministic.maxStubTokens,
|
||||
}),
|
||||
};
|
||||
this.put(record);
|
||||
captured.push(record);
|
||||
}
|
||||
if (captured.length) {
|
||||
this.generation++;
|
||||
this.cached = undefined;
|
||||
}
|
||||
return captured;
|
||||
}
|
||||
|
||||
applySemantic(sourceKeyValue: string, sourceHash: string, summary: string): SegmentSummaryRecord | undefined {
|
||||
const record = this.records.get(sourceKeyValue);
|
||||
if (!record || record.sourceHash !== sourceHash || !summary.trim()) return undefined;
|
||||
record.semanticSummary = summary.trim();
|
||||
this.generation++;
|
||||
this.cached = undefined;
|
||||
return this.toPersisted(record);
|
||||
}
|
||||
|
||||
toPersisted(record: ToolHistoryRecord): SegmentSummaryRecord {
|
||||
const effective = record.semanticSummary ?? record.deterministicSummary;
|
||||
return {
|
||||
version: COMPACTION_VERSION,
|
||||
kind: "segment-summary",
|
||||
sourceKey: record.sourceKey,
|
||||
toolCallId: record.toolCallId,
|
||||
resultTimestamp: record.resultTimestamp,
|
||||
toolName: record.toolName,
|
||||
sourceHash: record.sourceHash,
|
||||
deterministicSummary: record.deterministicSummary,
|
||||
...(record.semanticSummary ? { semanticSummary: record.semanticSummary } : {}),
|
||||
rawChars: record.rawText.length,
|
||||
effectiveChars: effective.length,
|
||||
};
|
||||
}
|
||||
|
||||
setSemanticJobs(count: number): void {
|
||||
this.semanticJobs = Math.max(0, count);
|
||||
if (this.cached) this.cached = { ...this.cached, semanticJobs: this.semanticJobs };
|
||||
}
|
||||
|
||||
recordUsage(tokens: number, messageEstimate: number, environmentHash: string): void {
|
||||
if (!Number.isFinite(tokens) || tokens <= 0) return;
|
||||
this.usageBaseline = { tokens, messageEstimate, environmentHash };
|
||||
this.cached = undefined;
|
||||
}
|
||||
|
||||
project(messages: readonly any[], params: {
|
||||
environmentHash: string;
|
||||
environmentTokens: number;
|
||||
contextWindow: number;
|
||||
}): ProjectionSnapshot {
|
||||
const sameTail = messages.length === this.sourceMessages.length
|
||||
&& messages.length > 0
|
||||
&& messageKey(messages[messages.length - 1]) === this.sourceKeys[this.sourceKeys.length - 1];
|
||||
if (this.cached && sameTail && this.cached.environmentHash === params.environmentHash) return this.cached;
|
||||
|
||||
const appendOnly = this.sourceMessages.length > 0
|
||||
&& messages.length >= this.sourceMessages.length
|
||||
&& messageKey(messages[this.sourceMessages.length - 1]) === this.sourceKeys[this.sourceKeys.length - 1]
|
||||
&& this.cached?.environmentHash === params.environmentHash;
|
||||
|
||||
let start = 0;
|
||||
if (!appendOnly) {
|
||||
this.sourceMessages = [];
|
||||
this.sourceKeys = [];
|
||||
this.projectedMessages = [];
|
||||
this.rawEstimates = [];
|
||||
this.effectiveEstimates = [];
|
||||
} else {
|
||||
start = this.sourceMessages.length;
|
||||
}
|
||||
|
||||
for (let index = start; index < messages.length; index++) {
|
||||
const message = messages[index];
|
||||
const projected = this.projectMessage(message);
|
||||
this.sourceMessages.push(message);
|
||||
this.sourceKeys.push(messageKey(message));
|
||||
this.projectedMessages.push(projected);
|
||||
this.rawEstimates.push(estimateTokens(message as any));
|
||||
this.effectiveEstimates.push(estimateTokens(projected as any));
|
||||
}
|
||||
|
||||
const rawMessageTokens = this.rawEstimates.reduce((sum, value) => sum + value, 0);
|
||||
const effectiveMessageTokens = this.effectiveEstimates.reduce((sum, value) => sum + value, 0);
|
||||
const rawTokens = rawMessageTokens + params.environmentTokens;
|
||||
const baseline = this.usageBaseline;
|
||||
const usageBacked = baseline?.environmentHash === params.environmentHash;
|
||||
const effectiveTokens = usageBacked
|
||||
? Math.max(0, baseline.tokens + effectiveMessageTokens - baseline.messageEstimate)
|
||||
: effectiveMessageTokens + params.environmentTokens;
|
||||
const confidence = usageBacked ? "usage-backed" as const : "estimated" as const;
|
||||
const lowerBound = Math.max(0, Math.floor(effectiveTokens * (usageBacked ? 0.98 : 0.9)));
|
||||
const upperBound = Math.ceil(effectiveTokens * (usageBacked ? 1.05 : 1.2));
|
||||
const fingerprint = hash(`${this.sessionId}:${this.generation}:${params.environmentHash}:${this.sourceKeys.join("|")}`);
|
||||
|
||||
this.cached = {
|
||||
sessionId: this.sessionId,
|
||||
generation: this.generation,
|
||||
fingerprint,
|
||||
environmentHash: params.environmentHash,
|
||||
messages: this.projectedMessages,
|
||||
messageEstimate: effectiveMessageTokens,
|
||||
rawTokens,
|
||||
effectiveTokens,
|
||||
reclaimedTokens: Math.max(0, rawTokens - effectiveTokens),
|
||||
contextWindow: params.contextWindow,
|
||||
confidence,
|
||||
lowerBound,
|
||||
upperBound,
|
||||
segmentCount: this.sourceMessages.length,
|
||||
semanticJobs: this.semanticJobs,
|
||||
};
|
||||
return this.cached;
|
||||
}
|
||||
|
||||
private projectMessage(message: any): any {
|
||||
if (message?.role !== "toolResult" || typeof message.toolCallId !== "string") return message;
|
||||
const record = this.records.get(sourceKey(message.toolCallId, message.timestamp));
|
||||
if (!record) return message;
|
||||
const summary = record.semanticSummary ?? record.deterministicSummary;
|
||||
return {
|
||||
...message,
|
||||
content: [{ type: "text", text: `${summary}\nSource: ${record.sourceKey}; recover with context_tree_query.` }],
|
||||
};
|
||||
}
|
||||
|
||||
query(ids: string[]): ToolHistoryRecord[] {
|
||||
const found: ToolHistoryRecord[] = [];
|
||||
const seen = new Set<string>();
|
||||
for (const id of ids) {
|
||||
const direct = this.records.get(id);
|
||||
const keys = direct ? [id] : this.ids.get(id) ?? [];
|
||||
for (const key of keys) {
|
||||
if (seen.has(key)) continue;
|
||||
const record = this.records.get(key);
|
||||
if (record) {
|
||||
found.push(record);
|
||||
seen.add(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
return found;
|
||||
}
|
||||
|
||||
allRecords(): ToolHistoryRecord[] {
|
||||
return [...this.records.values()];
|
||||
}
|
||||
|
||||
getSnapshot(): ProjectionSnapshot | undefined {
|
||||
return this.cached;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user