fable-agent/src/tier1-foundation/context-manager.ts

206 lines
6.2 KiB
TypeScript

import type { ContextEntry, ContextEntryPriority, ContextWindow } from "../core/types.js";
const PRIORITY_WEIGHT: Record<ContextEntryPriority, number> = {
critical: 100,
high: 50,
normal: 10,
low: 1,
};
/**
* Step 2: Context Manager
* Manages the agent's context window with sliding-window summarization,
* priority-based pruning, and token budget enforcement.
*
* "Massive context utilization" — Fable 5's headline capability.
*/
export class ContextManager {
private window: ContextWindow;
constructor(budget: number = 128_000) {
this.window = {
entries: [],
totalTokens: 0,
budget,
};
}
// ── Entry Management ──────────────────────────────────────
/** Add an entry and enforce budget */
add(entry: Omit<ContextEntry, "id" | "timestamp">): ContextEntry {
const full: ContextEntry = {
id: cryptoUid(),
timestamp: new Date().toISOString(),
...entry,
};
this.window.entries.push(full);
this.window.totalTokens += full.tokenCount;
this.enforceBudget();
return full;
}
/** Add a batch of entries at once */
addBatch(entries: Array<Omit<ContextEntry, "id" | "timestamp">>): ContextEntry[] {
return entries.map((e) => this.add(e));
}
// ── Budget Enforcement ────────────────────────────────────
/**
* When the budget is exceeded:
* 1. Summarize the lowest-priority entries (oldest 30% by score)
* 2. Drop entries that have already been summarized and are low-priority
*/
private enforceBudget(): void {
if (this.window.totalTokens <= this.window.budget) return;
// Sort by priority weight asc, then by timestamp asc (oldest first)
const sorted = [...this.window.entries].sort((a, b) => {
const pa = PRIORITY_WEIGHT[a.priority];
const pb = PRIORITY_WEIGHT[b.priority];
if (pa !== pb) return pa - pb;
return new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime();
});
let overage = this.window.totalTokens - this.window.budget;
const toRemove: Set<string> = new Set();
for (const entry of sorted) {
if (overage <= 0) break;
if (entry.priority === "critical") continue; // never drop critical
if (entry.summary) {
// Already summarized — safe to drop
toRemove.add(entry.id);
this.window.totalTokens -= entry.tokenCount;
overage -= entry.tokenCount;
} else {
// Summarize first, then mark for removal on next pass
entry.summary = this.summarize(entry.content);
// Reduce token count to just the summary
const saved = entry.tokenCount - entry.summary.length;
this.window.totalTokens -= saved;
overage -= saved;
entry.tokenCount = entry.summary.length;
}
}
// Remove entries marked for removal
if (toRemove.size > 0) {
this.window.entries = this.window.entries.filter(
(e) => !toRemove.has(e.id)
);
}
}
// ── Summarization ─────────────────────────────────────────
/** Produce a condensed summary of content (character-count based) */
private summarize(content: string): string {
const words = content.split(/\s+/);
if (words.length <= 50) return content;
// Extract first 2 sentences + key terms
const sentences = content.match(/[^.!?\n]+[.!?]/g) || [content];
const keyTerms = this.extractKeyTerms(content, 10);
let summary = sentences.slice(0, 2).join(" ").trim();
if (keyTerms.length > 0) {
summary += ` [Key: ${keyTerms.join(", ")}]`;
}
if (summary.length > 500) {
summary = summary.slice(0, 497) + "...";
}
return summary;
}
/** Extract significant terms (capitalized words, frequent terms) */
private extractKeyTerms(content: string, max: number): string[] {
const words = content.split(/\s+/);
const freq = new Map<string, number>();
for (const w of words) {
const clean = w.replace(/[^a-zA-Z0-9_-]/g, "");
if (clean.length < 4) continue;
// Prefer capitalized/technical terms
if (/^[A-Z]/.test(clean) || /[_-]/.test(clean)) {
freq.set(clean, (freq.get(clean) || 0) + 1);
}
}
return [...freq.entries()]
.sort((a, b) => b[1] - a[1])
.slice(0, max)
.map(([term]) => term);
}
// ── Queries ───────────────────────────────────────────────
getWindow(): ContextWindow {
return { ...this.window, entries: [...this.window.entries] };
}
getEntries(options?: {
priority?: ContextEntryPriority;
source?: string;
limit?: number;
}): ContextEntry[] {
let filtered = this.window.entries;
if (options?.priority) {
filtered = filtered.filter((e) => e.priority === options.priority);
}
if (options?.source) {
filtered = filtered.filter((e) => e.source === options.source);
}
if (options?.limit) {
filtered = filtered.slice(0, options.limit);
}
return filtered;
}
/** Build a single text block from all entries for LLM context injection */
render(): string {
return this.window.entries
.sort((a, b) => {
const pa = PRIORITY_WEIGHT[a.priority];
const pb = PRIORITY_WEIGHT[b.priority];
if (pa !== pb) return pb - pa; // higher priority first
return (
new Date(a.timestamp).getTime() -
new Date(b.timestamp).getTime()
);
})
.map((e) => {
const content = e.summary || e.content;
return `[${e.source}] ${content}`;
})
.join("\n\n");
}
getUsage(): { used: number; budget: number; pct: number } {
return {
used: this.window.totalTokens,
budget: this.window.budget,
pct: Math.round((this.window.totalTokens / this.window.budget) * 100),
};
}
reset(budget?: number): void {
this.window = {
entries: [],
totalTokens: 0,
budget: budget ?? this.window.budget,
};
}
}
import { randomUUID } from "node:crypto";
function cryptoUid(): string {
return randomUUID();
}