Initial commit: kage-research project files
This commit is contained in:
shokollm
355
level3b.ts
Normal file
355
level3b.ts
Normal file
@@ -0,0 +1,355 @@
|
||||
/**
|
||||
* Level 3b: Context Management
|
||||
*
|
||||
* Features:
|
||||
* 1. Context pruning - Remove old messages when too long
|
||||
* 2. Context compression - Summarize old messages
|
||||
* 3. Token estimation
|
||||
* 4. Configurable limits
|
||||
*/
|
||||
|
||||
import { Agent, type AgentTool, type AgentMessage, type AgentEvent } from "@mariozechner/pi-agent-core";
|
||||
import { registerBuiltInApiProviders } from "@mariozechner/pi-ai";
|
||||
import type { Model } from "@mariozechner/pi-ai";
|
||||
import * as fs from "fs";
|
||||
import * as path from "path";
|
||||
import { exec } from "child_process";
|
||||
|
||||
// ============== CONFIG ==============
|
||||
const OPENROUTER_API_KEY = process.env.OPENROUTER_API_KEY || "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";
|
||||
process.env.OPENROUTER_API_KEY = OPENROUTER_API_KEY;
|
||||
|
||||
const model: Model<"openai-responses"> = {
|
||||
id: "stepfun/step-3.5-flash:free",
|
||||
name: "Step-3.5 Flash (Free)",
|
||||
api: "openai-responses",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
output: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
|
||||
// ============== CONTEXT MANAGER ==============
|
||||
interface ContextConfig {
|
||||
maxTokens?: number;
|
||||
pruneThreshold?: number; // When to start pruning
|
||||
keepRecent?: number; // How many recent messages to always keep
|
||||
compressionEnabled?: boolean;
|
||||
}
|
||||
|
||||
interface MessageWithTokens extends AgentMessage {
|
||||
_tokens?: number;
|
||||
}
|
||||
|
||||
class ContextManager {
|
||||
private maxTokens: number;
|
||||
private pruneThreshold: number;
|
||||
private keepRecent: number;
|
||||
private compressionEnabled: boolean;
|
||||
|
||||
// Stats
|
||||
private pruneCount = 0;
|
||||
private compressCount = 0;
|
||||
|
||||
constructor(config: ContextConfig = {}) {
|
||||
this.maxTokens = config.maxTokens || 100000; // Default 100k
|
||||
this.pruneThreshold = config.pruneThreshold || 80000; // Start pruning at 80k
|
||||
this.keepRecent = config.keepRecent || 10; // Keep last 10 messages
|
||||
this.compressionEnabled = config.compressionEnabled || false;
|
||||
}
|
||||
|
||||
// Estimate tokens (rough approximation: 1 token ≈ 4 characters)
|
||||
estimateTokens(message: AgentMessage): number {
|
||||
const msg = message as any;
|
||||
let text = "";
|
||||
|
||||
if (typeof msg.content === "string") {
|
||||
text = msg.content;
|
||||
} else if (Array.isArray(msg.content)) {
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") {
|
||||
text += block.text || "";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rough estimate: 1 token ≈ 4 characters
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
|
||||
// Calculate total tokens in messages
|
||||
calculateTotalTokens(messages: AgentMessage[]): number {
|
||||
return messages.reduce((sum, msg) => sum + this.estimateTokens(msg), 0);
|
||||
}
|
||||
|
||||
// Prune old messages
|
||||
prune(messages: AgentMessage[]): AgentMessage[] {
|
||||
const total = this.calculateTotalTokens(messages);
|
||||
|
||||
if (total < this.pruneThreshold) {
|
||||
return messages; // No pruning needed
|
||||
}
|
||||
|
||||
console.log(`✂️ Pruning context: ${total} tokens > ${this.pruneThreshold} threshold`);
|
||||
|
||||
// Keep system prompt (first message) if it's a system message
|
||||
let result: AgentMessage[] = [];
|
||||
if (messages.length > 0 && (messages[0] as any).role === "system") {
|
||||
result.push(messages[0]);
|
||||
}
|
||||
|
||||
// Keep recent messages
|
||||
const recent = messages.slice(-this.keepRecent);
|
||||
result = result.concat(recent);
|
||||
|
||||
// Add summary placeholder if we removed middle messages
|
||||
const removed = messages.length - result.length;
|
||||
if (removed > 1) {
|
||||
const summaryMsg: AgentMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: `[Context: ${removed} older messages removed for brevity]` }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
result.splice(1, 0, summaryMsg); // Insert after system prompt
|
||||
}
|
||||
|
||||
const newTotal = this.calculateTotalTokens(result);
|
||||
this.pruneCount++;
|
||||
|
||||
console.log(`✂️ Pruned: ${messages.length} → ${result.length} messages`);
|
||||
console.log(`✂️ Tokens: ${total} → ${newTotal}`);
|
||||
console.log(`✂️ (Total prunes: ${this.pruneCount})`);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// Compress messages (placeholder - would need LLM for real compression)
|
||||
compress(messages: AgentMessage[]): AgentMessage[] {
|
||||
// This is a simplified version - real compression would use an LLM
|
||||
console.log(`📦 Compression requested (${messages.length} messages)`);
|
||||
|
||||
// For now, just prune
|
||||
this.compressCount++;
|
||||
return this.prune(messages);
|
||||
}
|
||||
|
||||
// Transform context - call this before sending to LLM
|
||||
transform(messages: AgentMessage[]): AgentMessage[] {
|
||||
const total = this.calculateTotalTokens(messages);
|
||||
|
||||
if (total > this.maxTokens) {
|
||||
console.log(`⚠️ Context overflow: ${total} > ${this.maxTokens}, forcing prune`);
|
||||
return this.prune(messages);
|
||||
}
|
||||
|
||||
if (total > this.pruneThreshold && this.compressionEnabled) {
|
||||
return this.compress(messages);
|
||||
}
|
||||
|
||||
if (total > this.pruneThreshold) {
|
||||
return this.prune(messages);
|
||||
}
|
||||
|
||||
return messages;
|
||||
}
|
||||
|
||||
getStats() {
|
||||
return {
|
||||
maxTokens: this.maxTokens,
|
||||
pruneThreshold: this.pruneThreshold,
|
||||
keepRecent: this.keepRecent,
|
||||
compressionEnabled: this.compressionEnabled,
|
||||
pruneCount: this.pruneCount,
|
||||
compressCount: this.compressCount,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// ============== TOOLS ==============
|
||||
function createTools(cwd: string = process.cwd()): AgentTool[] {
|
||||
return [
|
||||
{
|
||||
name: "bash",
|
||||
label: "Run Command",
|
||||
description: "Run a shell command",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: {
|
||||
command: { type: "string", description: "Command to run" },
|
||||
},
|
||||
required: ["command"],
|
||||
} as const,
|
||||
execute: async (toolCallId: string, params: { command: string }) => {
|
||||
return new Promise((resolve) => {
|
||||
exec(params.command, { cwd }, (error, stdout, stderr) => {
|
||||
if (error) {
|
||||
resolve({
|
||||
content: [{ type: "text", text: stderr || error.message }],
|
||||
details: { command: params.command, exitCode: error.code },
|
||||
isError: true,
|
||||
});
|
||||
} else {
|
||||
resolve({
|
||||
content: [{ type: "text", text: stdout }],
|
||||
details: { command: params.command, exitCode: 0 },
|
||||
});
|
||||
}
|
||||
});
|
||||
});
|
||||
},
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
// ============== SHADOW WITH CONTEXT ==============
|
||||
class ShadowWithContext {
|
||||
private agent: Agent;
|
||||
private contextManager: ContextManager;
|
||||
public id: string;
|
||||
public messageCount = 0;
|
||||
|
||||
constructor(id: string, worktreePath: string, contextConfig?: ContextConfig) {
|
||||
this.id = id;
|
||||
this.contextManager = new ContextManager(contextConfig);
|
||||
|
||||
this.agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful coding assistant. Be concise.",
|
||||
model: model,
|
||||
tools: createTools(worktreePath) as any,
|
||||
messages: [],
|
||||
},
|
||||
convertToLlm: (messages: AgentMessage[]) => {
|
||||
// Transform context before sending to LLM
|
||||
const transformed = this.contextManager.transform(messages);
|
||||
|
||||
return transformed
|
||||
.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult")
|
||||
.map((m) => ({
|
||||
role: m.role,
|
||||
content: m.content,
|
||||
}));
|
||||
},
|
||||
});
|
||||
|
||||
this.agent.subscribe((event) => {
|
||||
if (event.type === "message_end") {
|
||||
this.messageCount++;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async run(message: string): Promise<void> {
|
||||
const msg: AgentMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: message }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
await this.agent.prompt(msg);
|
||||
}
|
||||
|
||||
getContextStats() {
|
||||
return {
|
||||
messageCount: this.messageCount,
|
||||
contextManager: this.contextManager.getStats(),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// ============== TEST ==============
|
||||
async function testContextPruning() {
|
||||
console.log("\n" + "=".repeat(50));
|
||||
console.log("TEST: Context Pruning");
|
||||
console.log("=".repeat(50));
|
||||
|
||||
// Create shadow with aggressive pruning (for testing)
|
||||
const shadow = new ShadowWithContext("test-1", "/tmp", {
|
||||
maxTokens: 5000,
|
||||
pruneThreshold: 2000,
|
||||
keepRecent: 3,
|
||||
compressionEnabled: false,
|
||||
});
|
||||
|
||||
console.log("Context config:", shadow.getContextStats().contextManager);
|
||||
|
||||
// Simulate many messages to trigger pruning
|
||||
const longText = "This is a test message with some content. ".repeat(50);
|
||||
|
||||
console.log("\n📝 Adding messages to trigger pruning...\n");
|
||||
|
||||
for (let i = 0; i < 15; i++) {
|
||||
const msg: AgentMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: `Message ${i}: ${longText}` }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
// Manually trigger context transform
|
||||
const messages = Array(15).fill(null).map((_, j) => ({
|
||||
role: j % 2 === 0 ? "user" as const : "assistant" as const,
|
||||
content: [{ type: "text" as const, text: `Message ${j}: ${longText}` }],
|
||||
timestamp: Date.now(),
|
||||
}));
|
||||
|
||||
const transformed = (shadow as any).contextManager.transform(messages);
|
||||
|
||||
if (transformed.length < messages.length) {
|
||||
console.log(`📊 After message ${i}: ${messages.length} → ${transformed.length} messages`);
|
||||
}
|
||||
}
|
||||
|
||||
console.log("\n📊 Final stats:", shadow.getContextStats());
|
||||
}
|
||||
|
||||
async function testActualAgent() {
|
||||
console.log("\n" + "=".repeat(50));
|
||||
console.log("TEST: Actual Agent with Context Management");
|
||||
console.log("=".repeat(50));
|
||||
|
||||
// Create with normal settings
|
||||
const shadow = new ShadowWithContext("test-2", "/tmp", {
|
||||
maxTokens: 50000,
|
||||
pruneThreshold: 30000,
|
||||
keepRecent: 10,
|
||||
});
|
||||
|
||||
console.log("\n🚀 Running agent with context management...\n");
|
||||
|
||||
// Run multiple turns to build up context
|
||||
await shadow.run("Say hello and run 'echo Hello 1'");
|
||||
console.log("📊 After turn 1:", shadow.getContextStats());
|
||||
|
||||
await shadow.run("Say hi and run 'echo Hello 2'");
|
||||
console.log("📊 After turn 2:", shadow.getContextStats());
|
||||
|
||||
await shadow.run("Run 'echo Hello 3'");
|
||||
console.log("📊 After turn 3:", shadow.getContextStats());
|
||||
|
||||
await shadow.run("Run 'echo Hello 4'");
|
||||
console.log("📊 After turn 4:", shadow.getContextStats());
|
||||
|
||||
await shadow.run("Run 'echo Hello 5'");
|
||||
console.log("📊 After turn 5:", shadow.getContextStats());
|
||||
|
||||
console.log("\n✅ Agent test complete!");
|
||||
console.log("📊 Final stats:", shadow.getContextStats());
|
||||
}
|
||||
|
||||
// ============== MAIN ==============
|
||||
async function main() {
|
||||
console.log("🧪 Level 3b: Context Management\n");
|
||||
|
||||
registerBuiltInApiProviders();
|
||||
|
||||
await testContextPruning();
|
||||
await testActualAgent();
|
||||
|
||||
console.log("\n✅ All tests complete!");
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
Reference in New Issue
Block a user