feat: Add RTK (Rust Token Killer) integration for token optimization
- Add RTK utility module (src/utils/rtk.js) - Integrate RTK into BashTool for all bash commands - Integrate RTK into GitTool for git operations - Initialize RTK on bot startup - Support 60+ command types (git, npm, cargo, pytest, docker, etc.) - Track and report token savings per command - Graceful fallback when RTK is not available Expected savings: 60-90% token reduction for supported commands
This commit is contained in:
admin
@@ -6,6 +6,7 @@ import { WebSocketServer } from 'ws';
|
||||
import { spawn } from 'child_process';
|
||||
import fs from 'fs-extra';
|
||||
import path from 'path';
|
||||
import { getRTK } from "../utils/rtk.js";
|
||||
|
||||
export async function initBot(config, api, tools, skills) {
|
||||
const env = checkEnv();
|
||||
@@ -24,6 +25,7 @@ export async function initBot(config, api, tools, skills) {
|
||||
|
||||
// WebSocket for real-time updates
|
||||
const httpServer = createServer(app);
|
||||
// Initialize RTK integration\n const rtk = getRTK();\n await rtk.init();
|
||||
const wss = new WebSocketServer({ server: httpServer });
|
||||
|
||||
// Store active connections
|
||||
|
||||
@@ -1,14 +1,35 @@
|
||||
import { logger } from '../utils/logger.js';
|
||||
import { getRTK } from '../utils/rtk.js';
|
||||
import { spawn } from 'child_process';
|
||||
|
||||
export class BashTool {
|
||||
constructor() {
|
||||
this.name = 'bash';
|
||||
this.description = 'Execute shell commands';
|
||||
this.description = 'Execute shell commands with RTK token optimization';
|
||||
}
|
||||
|
||||
async execute(command, options = {}) {
|
||||
const { timeout = 300000, cwd = process.cwd() } = options;
|
||||
|
||||
// Check if RTK is available and command is supported
|
||||
const rtk = getRTK();
|
||||
const commandName = command.split(' ')[0];
|
||||
|
||||
if (rtk.enabled && rtk.isCommandSupported(commandName)) {
|
||||
const rtkResult = await rtk.optimizeCommand(commandName, command.split(' ').slice(1));
|
||||
|
||||
if (rtkResult.success) {
|
||||
return {
|
||||
success: true,
|
||||
stdout: rtkResult.stdout,
|
||||
stderr: rtkResult.stderr,
|
||||
optimized: true,
|
||||
savings: rtkResult.savings,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to original command if RTK is not available or failed
|
||||
logger.info(`🚀 Executing: ${command.substring(0, 100)}...`);
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
@@ -35,6 +56,8 @@ export class BashTool {
|
||||
success: true,
|
||||
stdout: stdout.trim(),
|
||||
stderr: stderr.trim(),
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
});
|
||||
} else {
|
||||
reject({
|
||||
@@ -42,6 +65,8 @@ export class BashTool {
|
||||
stdout: stdout.trim(),
|
||||
stderr: stderr.trim(),
|
||||
code,
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
});
|
||||
}
|
||||
});
|
||||
@@ -50,6 +75,8 @@ export class BashTool {
|
||||
reject({
|
||||
success: false,
|
||||
error: error.message,
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,20 +1,38 @@
|
||||
import { logger } from '../utils/logger.js';
|
||||
import { execa } from 'execa';
|
||||
import { getRTK } from '../utils/rtk.js';
|
||||
|
||||
export class GitTool {
|
||||
constructor() {
|
||||
this.name = 'git';
|
||||
this.description = 'Git operations';
|
||||
this.description = 'Git operations with RTK token optimization';
|
||||
}
|
||||
|
||||
async status() {
|
||||
try {
|
||||
const rtk = getRTK();
|
||||
|
||||
if (rtk.enabled) {
|
||||
const rtkResult = await rtk.optimizeCommand('git', ['status', '--short']);
|
||||
if (rtkResult.success) {
|
||||
return {
|
||||
success: true,
|
||||
status: rtkResult.stdout.trim() || 'clean',
|
||||
optimized: true,
|
||||
savings: rtkResult.savings,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to original command
|
||||
const { stdout } = await execa('git', ['status', '--short'], {
|
||||
cwd: process.cwd(),
|
||||
});
|
||||
return {
|
||||
success: true,
|
||||
status: stdout.trim() || 'clean',
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
@@ -26,14 +44,30 @@ export class GitTool {
|
||||
|
||||
async log(options = {}) {
|
||||
const { lines = 10 } = options;
|
||||
const rtk = getRTK();
|
||||
|
||||
try {
|
||||
if (rtk.enabled) {
|
||||
const rtkResult = await rtk.optimizeCommand('git', ['log', '--oneline', `-${lines}`]);
|
||||
if (rtkResult.success) {
|
||||
return {
|
||||
success: true,
|
||||
commits: rtkResult.stdout.trim().split('\n'),
|
||||
optimized: true,
|
||||
savings: rtkResult.savings,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to original command
|
||||
const { stdout } = await execa('git', ['log', '--oneline', `-${lines}`], {
|
||||
cwd: process.cwd(),
|
||||
});
|
||||
return {
|
||||
success: true,
|
||||
commits: stdout.trim().split('\n'),
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
@@ -44,13 +78,30 @@ export class GitTool {
|
||||
}
|
||||
|
||||
async branch() {
|
||||
const rtk = getRTK();
|
||||
|
||||
try {
|
||||
if (rtk.enabled) {
|
||||
const rtkResult = await rtk.optimizeCommand('git', ['branch', '--show-current']);
|
||||
if (rtkResult.success) {
|
||||
return {
|
||||
success: true,
|
||||
branch: rtkResult.stdout.trim(),
|
||||
optimized: true,
|
||||
savings: rtkResult.savings,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to original command
|
||||
const { stdout } = await execa('git', ['branch', '--show-current'], {
|
||||
cwd: process.cwd(),
|
||||
});
|
||||
return {
|
||||
success: true,
|
||||
branch: stdout.trim(),
|
||||
optimized: false,
|
||||
savings: 0,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
|
||||
170
src/utils/rtk.js
Normal file
170
src/utils/rtk.js
Normal file
@@ -0,0 +1,170 @@
|
||||
import { logger } from './logger.js';
|
||||
|
||||
/**
|
||||
* RTK (Rust Token Killer) Integration
|
||||
* Minimizes LLM token consumption by filtering and compressing command outputs
|
||||
*/
|
||||
export class RTKIntegration {
|
||||
constructor() {
|
||||
this.rtkPath = process.env.RTK_PATH || '/home/uroma2/.local/bin/rtk';
|
||||
this.enabled = false;
|
||||
this.version = null;
|
||||
}
|
||||
|
||||
async init() {
|
||||
try {
|
||||
// Check if RTK is installed
|
||||
const { execa } = await import('execa');
|
||||
const result = await execa(this.rtkPath, ['--version']);
|
||||
this.version = result.stdout.trim();
|
||||
this.enabled = true;
|
||||
logger.info(`✓ RTK integration initialized (v${this.version})`);
|
||||
return true;
|
||||
} catch (error) {
|
||||
logger.warn(`⚠ RTK not found at ${this.rtkPath}`);
|
||||
logger.warn(` Install with: curl -fsSL https://raw.githubusercontent.com/rtk-ai/rtk/master/install.sh | sh`);
|
||||
this.enabled = false;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a command is supported by RTK
|
||||
*/
|
||||
isCommandSupported(cmd) {
|
||||
const supportedCommands = [
|
||||
// Git commands
|
||||
'git', 'gh', 'gt',
|
||||
// Cargo commands
|
||||
'cargo',
|
||||
// Node/NPM commands
|
||||
'npm', 'pnpm', 'npx',
|
||||
// Python commands
|
||||
'ruff', 'pytest', 'pip', 'mypy',
|
||||
// Docker commands
|
||||
'docker', 'kubectl', 'aws',
|
||||
// Go commands
|
||||
'go', 'golangci-lint',
|
||||
// Ruby commands
|
||||
'rspec', 'rubocop', 'rake',
|
||||
// Dotnet commands
|
||||
'dotnet',
|
||||
// Testing frameworks
|
||||
'playwright', 'vitest', 'jest',
|
||||
];
|
||||
|
||||
return supportedCommands.some(supported => cmd.startsWith(supported));
|
||||
}
|
||||
|
||||
/**
|
||||
* Optimize a command with RTK
|
||||
* @param {string} command - The command to optimize
|
||||
* @param {string[]} args - Command arguments
|
||||
* @returns {Promise<{success: boolean, stdout?: string, stderr?: string, savings?: number}>}
|
||||
*/
|
||||
async optimizeCommand(command, args = []) {
|
||||
if (!this.enabled) {
|
||||
return { success: false, savings: 0 };
|
||||
}
|
||||
|
||||
if (!this.isCommandSupported(command)) {
|
||||
return { success: false, savings: 0 };
|
||||
}
|
||||
|
||||
try {
|
||||
const { execa } = await import('execa');
|
||||
const cmdString = `${command} ${args.join(' ')}`;
|
||||
|
||||
logger.info(`🚀 Executing with RTK: ${cmdString.substring(0, 100)}...`);
|
||||
|
||||
const result = await execa(this.rtkPath, [command, ...args], {
|
||||
timeout: 300000,
|
||||
cwd: process.cwd(),
|
||||
});
|
||||
|
||||
const savings = this.parseSavings(result.stdout);
|
||||
|
||||
logger.info(`✓ RTK optimized: ${savings > 0 ? `-${savings}% savings` : 'no savings'}`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
stdout: result.stdout.trim(),
|
||||
stderr: result.stderr.trim(),
|
||||
savings,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error(`✗ RTK execution failed: ${error.message}`);
|
||||
return {
|
||||
success: false,
|
||||
stderr: error.message,
|
||||
savings: 0,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse savings percentage from RTK output
|
||||
*/
|
||||
parseSavings(output) {
|
||||
const match = output.match(/-\s*(\d+)%\s+savings?/i);
|
||||
return match ? parseInt(match[1], 10) : 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get RTK tracking statistics
|
||||
*/
|
||||
async getTrackingStats() {
|
||||
if (!this.enabled) {
|
||||
return null;
|
||||
}
|
||||
|
||||
try {
|
||||
const { execa } = await import('execa');
|
||||
const result = await execa(this.rtkPath, ['gain']);
|
||||
|
||||
return {
|
||||
totalCommands: result.stdout.match(/Commands run: (\d+)/)?.[1] || '0',
|
||||
totalSavings: result.stdout.match(/Total savings: (\d+%)/)?.[1] || '0%',
|
||||
history: result.stdout,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error(`✗ Failed to get RTK tracking stats: ${error.message}`);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* List supported commands
|
||||
*/
|
||||
async listSupportedCommands() {
|
||||
if (!this.enabled) {
|
||||
return [];
|
||||
}
|
||||
|
||||
try {
|
||||
const { execa } = await import('execa');
|
||||
const result = await execa(this.rtkPath, ['--help']);
|
||||
|
||||
// Parse help output to extract supported commands
|
||||
const lines = result.stdout.split('\n');
|
||||
const supportedCommands = lines
|
||||
.filter(line => line.match(/^\s+[a-z]+/))
|
||||
.map(line => line.trim().split(/\s+/)[0]);
|
||||
|
||||
return supportedCommands;
|
||||
} catch (error) {
|
||||
logger.error(`✗ Failed to list supported commands: ${error.message}`);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Singleton instance
|
||||
let rtkInstance = null;
|
||||
|
||||
export function getRTK() {
|
||||
if (!rtkInstance) {
|
||||
rtkInstance = new RTKIntegration();
|
||||
}
|
||||
return rtkInstance;
|
||||
}
|
||||
Reference in New Issue
Block a user