capsule-ai-cli
Version:
The AI Model Orchestrator - Intelligent multi-model workflows with device-locked licensing
399 lines • 16.4 kB
JavaScript
import { v4 as uuidv4 } from 'uuid';
import { configManager } from '../core/config.js';
import { stateService } from './state.js';
import { summarizerService } from './summarizer.js';
import { openRouterModelsService } from './openrouter-models.js';
import chalk from 'chalk';
export class ContextManager {
currentContext;
contexts = new Map();
constructor() {
this.currentContext = this.createNewContext();
this.loadSavedContexts();
}
createNewContext() {
const id = uuidv4();
const context = {
id,
messages: [],
metadata: {
created: new Date(),
lastModified: new Date(),
totalTokens: 0,
totalCost: 0,
modelsUsed: []
}
};
this.contexts.set(id, context);
return context;
}
getCurrentContext() {
return this.currentContext;
}
setCurrentContext(id) {
const context = this.contexts.get(id);
if (context) {
this.currentContext = context;
this.cleanOrphanedToolResults();
this.recalculateTokens();
}
else {
throw new Error(`Context ${id} not found`);
}
}
getContextById(id) {
return this.contexts.get(id);
}
addMessage(message) {
if (!message.metadata) {
message.metadata = {
timestamp: new Date(),
model: stateService.getModel(),
provider: stateService.getProvider()
};
}
this.currentContext.messages.push(message);
this.currentContext.metadata.lastModified = new Date();
if (message.metadata.tokens) {
this.currentContext.metadata.totalTokens += message.metadata.tokens;
}
else {
const contentLength = typeof message.content === 'string'
? message.content.length
: JSON.stringify(message.content).length;
message.metadata.tokens = Math.ceil(contentLength / 4);
this.currentContext.metadata.totalTokens += message.metadata.tokens;
}
if (message.metadata.cost) {
this.currentContext.metadata.totalCost += message.metadata.cost;
}
if (message.metadata.model && !this.currentContext.metadata.modelsUsed.includes(message.metadata.model)) {
this.currentContext.metadata.modelsUsed.push(message.metadata.model);
}
this.checkAndAutoCompact();
this.saveCurrentContext();
}
clearContext() {
const newContext = this.createNewContext();
this.currentContext = newContext;
}
getMessagesForModel(model, options) {
const limit = this.getTokenLimit(model);
if (options?.progressive) {
const recentLimit = Math.min(limit * 0.3, 30000);
return this.truncateContext(recentLimit);
}
return this.truncateContext(limit);
}
truncateContext(maxTokens) {
const messages = [...this.currentContext.messages];
let totalTokens = 0;
const truncatedMessages = [];
const systemMessages = messages.filter(m => m.role === 'system');
systemMessages.forEach(msg => {
totalTokens += msg.metadata?.tokens || Math.ceil(msg.content.length / 4);
truncatedMessages.push(msg);
});
for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i];
if (msg.role === 'system')
continue;
const msgTokens = msg.metadata?.tokens || Math.ceil(msg.content.length / 4);
if (totalTokens + msgTokens <= maxTokens * 0.9) {
truncatedMessages.unshift(msg);
totalTokens += msgTokens;
}
else {
break;
}
}
if (truncatedMessages.length < messages.length) {
truncatedMessages.unshift({
role: 'system',
content: `[Context truncated. Showing ${truncatedMessages.length} of ${messages.length} messages]`,
metadata: { timestamp: new Date() }
});
}
return truncatedMessages;
}
getTokenLimit(model) {
const contextLength = openRouterModelsService.getModelContextLength(model);
return contextLength || 128000;
}
estimateTokens(messages) {
return messages.reduce((total, msg) => {
return total + (msg.metadata?.tokens || Math.ceil(msg.content.length / 4));
}, 0);
}
getContextStats() {
const messages = this.currentContext.messages;
const timestamps = messages
.map(m => m.metadata?.timestamp)
.filter(t => t !== undefined);
return {
messageCount: messages.length,
tokenCount: this.currentContext.metadata.totalTokens,
totalCost: this.currentContext.metadata.totalCost,
modelsUsed: this.currentContext.metadata.modelsUsed,
oldestMessage: timestamps.length > 0 ? new Date(Math.min(...timestamps.map(t => t.getTime()))) : new Date(),
newestMessage: timestamps.length > 0 ? new Date(Math.max(...timestamps.map(t => t.getTime()))) : new Date()
};
}
switchContext(id) {
const context = this.contexts.get(id);
if (context) {
this.saveCurrentContext();
this.currentContext = context;
this.cleanOrphanedToolResults();
}
else {
throw new Error(`Context ${id} not found`);
}
}
listContexts() {
return Array.from(this.contexts.values()).map(ctx => ({
id: ctx.id,
created: ctx.metadata.created,
messageCount: ctx.messages.length
}));
}
deleteContext(id) {
if (id === this.currentContext.id) {
return false;
}
const deleted = this.contexts.delete(id);
if (deleted) {
const saved = this.getSavedContexts();
delete saved[id];
configManager.setConfig('contexts', saved);
}
return deleted;
}
saveCurrentContext() {
const contexts = this.getSavedContexts();
contexts[this.currentContext.id] = this.currentContext;
configManager.setConfig('contexts', contexts);
}
loadSavedContexts() {
const saved = this.getSavedContexts();
Object.entries(saved).forEach(([id, context]) => {
context.metadata.created = new Date(context.metadata.created);
context.metadata.lastModified = new Date(context.metadata.lastModified);
context.messages.forEach((msg) => {
if (msg.metadata?.timestamp) {
msg.metadata.timestamp = new Date(msg.metadata.timestamp);
}
if (typeof msg.content !== 'string' && msg.content) {
msg.content = JSON.stringify(msg.content);
}
});
this.contexts.set(id, context);
});
}
getSavedContexts() {
return configManager.getConfig().contexts || {};
}
exportContext(format = 'markdown') {
if (format === 'json') {
return JSON.stringify(this.currentContext, null, 2);
}
let output = `# Conversation Context\n\n`;
output += `**Created:** ${this.currentContext.metadata.created.toLocaleString()}\n`;
output += `**Messages:** ${this.currentContext.messages.length}\n`;
output += `**Total Tokens:** ${this.currentContext.metadata.totalTokens}\n`;
output += `**Total Cost:** $${this.currentContext.metadata.totalCost.toFixed(4)}\n`;
output += `**Models Used:** ${this.currentContext.metadata.modelsUsed.join(', ')}\n\n`;
output += `---\n\n`;
this.currentContext.messages.forEach(msg => {
const role = msg.role.charAt(0).toUpperCase() + msg.role.slice(1);
const model = msg.metadata?.model ? ` (${msg.metadata.model})` : '';
output += `### ${role}${model}\n\n`;
output += `${msg.content}\n\n`;
});
return output;
}
recalculateTokens() {
let totalTokens = 0;
this.currentContext.messages.forEach(msg => {
const content = typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
const estimatedTokens = Math.ceil(content.length / 4);
totalTokens += estimatedTokens;
if (!msg.metadata) {
msg.metadata = {};
}
msg.metadata.tokens = estimatedTokens;
});
this.currentContext.metadata.totalTokens = totalTokens;
}
cleanOrphanedToolResults() {
const messages = this.currentContext.messages;
const toolCallIds = new Set();
messages.forEach(msg => {
if (msg.role === 'assistant' && msg.tool_calls) {
for (const toolCall of msg.tool_calls) {
toolCallIds.add(toolCall.id);
}
}
});
const cleanedMessages = messages.filter(msg => {
if (msg.role === 'tool_result' && msg.tool_call_id) {
return toolCallIds.has(msg.tool_call_id);
}
return true;
});
if (cleanedMessages.length < messages.length) {
console.log(chalk.yellow(`\n⚠️ Cleaned ${messages.length - cleanedMessages.length} orphaned tool results from conversation history`));
this.currentContext.messages = cleanedMessages;
}
}
getContextSummary() {
const stats = this.getContextStats();
const currentModel = stateService.getModel();
const limit = this.getTokenLimit(currentModel);
const percentage = Math.round((stats.tokenCount / limit) * 100);
const warningIndicator = percentage > 80 ? ' ⚠️' : '';
return `${stats.tokenCount.toLocaleString()}/${(limit / 1000).toFixed(0)}k tokens (${percentage}%)${warningIndicator}`;
}
async checkAndAutoCompact() {
const currentModel = stateService.getModel();
const limit = this.getTokenLimit(currentModel);
const currentTokens = this.currentContext.metadata.totalTokens;
const usage = currentTokens / limit;
if (usage > 0.8 && this.currentContext.messages.length > 10) {
console.log(chalk.yellow(`\n⚠️ Context approaching limit (${Math.round(usage * 100)}%). Auto-compacting...`));
await this.autoCompact();
}
}
async autoCompact() {
const messages = this.currentContext.messages;
const systemMessages = messages.filter(m => m.role === 'system' && !m.metadata?.compacted);
const recentMessages = messages.slice(-10);
const messagesToSummarize = messages.slice(systemMessages.length, messages.length - 10).filter(m => !m.metadata?.compacted);
if (messagesToSummarize.length < 4)
return;
const summary = await this.createAISummary(messagesToSummarize);
const compactedMessages = [
...systemMessages,
{
role: 'system',
content: `[Previous conversation summary (${messagesToSummarize.length} messages)]\n${summary}`,
metadata: {
timestamp: new Date(),
tokens: Math.ceil(summary.length / 4),
compacted: true
}
},
...recentMessages
];
let newTokenCount = 0;
compactedMessages.forEach(msg => {
const contentLength = typeof msg.content === 'string'
? msg.content.length
: JSON.stringify(msg.content).length;
newTokenCount += msg.metadata?.tokens || Math.ceil(contentLength / 4);
});
const savedTokens = this.currentContext.metadata.totalTokens - newTokenCount;
console.log(chalk.green(`✓ Compacted ${messagesToSummarize.length} messages. Saved ${savedTokens.toLocaleString()} tokens.`));
this.currentContext.messages = compactedMessages;
this.currentContext.metadata.totalTokens = newTokenCount;
this.saveCurrentContext();
}
async createAISummary(messages) {
try {
return await summarizerService.summarizeMessages(messages);
}
catch (error) {
console.log(chalk.yellow('\nFailed to create AI summary, falling back to simple summary'));
return this.createSimpleSummary(messages);
}
}
createSimpleSummary(messages) {
const summary = [];
const exchanges = [];
let currentExchange = {};
messages.forEach(msg => {
if (msg.role === 'user') {
if (currentExchange.user) {
exchanges.push(currentExchange);
currentExchange = {};
}
const content = typeof msg.content === 'string' ? msg.content : '[multimodal content]';
currentExchange.user = content;
}
else if (msg.role === 'assistant' && currentExchange.user) {
const content = typeof msg.content === 'string' ? msg.content : '[assistant response with tools]';
currentExchange.assistant = content;
exchanges.push(currentExchange);
currentExchange = {};
}
});
if (currentExchange.user) {
exchanges.push(currentExchange);
}
exchanges.forEach((exchange) => {
const userPreview = exchange.user.substring(0, 80).replace(/\n/g, ' ');
const assistantPreview = exchange.assistant?.substring(0, 120).replace(/\n/g, ' ') || '[no response]';
summary.push(`• User: ${userPreview}${exchange.user.length > 80 ? '...' : ''}\n` +
` Assistant: ${assistantPreview}${exchange.assistant?.length > 120 ? '...' : ''}`);
});
return summary.join('\n\n');
}
async compactManuallyAI() {
const messages = this.currentContext.messages;
const currentTokens = this.currentContext.metadata.totalTokens;
if (messages.length <= 10) {
return {
success: false,
message: 'Not enough messages to compact (need more than 10)'
};
}
const hasRecentCompact = messages.some(m => m.metadata?.compacted &&
m.metadata?.timestamp &&
Date.now() - m.metadata.timestamp.getTime() < 300000);
if (hasRecentCompact) {
return {
success: false,
message: 'Already compacted recently. Wait a few minutes before compacting again.'
};
}
const beforeCount = messages.length;
await this.autoCompact();
const afterCount = this.currentContext.messages.length;
const newTokens = this.currentContext.metadata.totalTokens;
return {
success: true,
messagesCompacted: beforeCount - afterCount,
tokensSaved: currentTokens - newTokens
};
}
compactManually() {
const messages = this.currentContext.messages;
const currentTokens = this.currentContext.metadata.totalTokens;
if (messages.length <= 10) {
return {
success: false,
message: 'Not enough messages to compact (need more than 10)'
};
}
const hasRecentCompact = messages.some(m => m.metadata?.compacted &&
m.metadata?.timestamp &&
Date.now() - m.metadata.timestamp.getTime() < 300000);
if (hasRecentCompact) {
return {
success: false,
message: 'Already compacted recently. Wait a few minutes before compacting again.'
};
}
const beforeCount = messages.length;
this.autoCompact();
const afterCount = this.currentContext.messages.length;
const newTokens = this.currentContext.metadata.totalTokens;
return {
success: true,
messagesCompacted: beforeCount - afterCount,
tokensSaved: currentTokens - newTokens
};
}
}
export const contextManager = new ContextManager();
//# sourceMappingURL=context.js.map