mirror of
https://github.com/samanhappy/mcphub.git
synced 2026-01-01 04:08:52 -05:00
feat: implement MCP output compression mechanism
- Add compressionService.ts with AI-powered compression using OpenAI - Add compression configuration to SystemConfig type - Update SystemConfig database entity with compression field - Integrate compression with mcpService tool call responses - Add migration support for compression settings - Add comprehensive unit tests (23 new tests) - Support graceful fallback when compression is unavailable Co-authored-by: samanhappy <2755122+samanhappy@users.noreply.github.com>
This commit is contained in:
@@ -33,6 +33,9 @@ export class SystemConfig {
|
||||
@Column({ type: 'boolean', nullable: true })
|
||||
enableSessionRebuild?: boolean;
|
||||
|
||||
@Column({ type: 'simple-json', nullable: true })
|
||||
compression?: Record<string, any>;
|
||||
|
||||
@CreateDateColumn({ name: 'created_at', type: 'timestamp' })
|
||||
createdAt: Date;
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@ export class SystemConfigRepository {
|
||||
oauth: {},
|
||||
oauthServer: {},
|
||||
enableSessionRebuild: false,
|
||||
compression: {},
|
||||
});
|
||||
config = await this.repository.save(config);
|
||||
}
|
||||
|
||||
266
src/services/compressionService.ts
Normal file
266
src/services/compressionService.ts
Normal file
@@ -0,0 +1,266 @@
|
||||
import OpenAI from 'openai';
|
||||
import { getSmartRoutingConfig, SmartRoutingConfig } from '../utils/smartRouting.js';
|
||||
import { getSystemConfigDao } from '../dao/index.js';
|
||||
|
||||
/**
|
||||
* Compression configuration interface
|
||||
*/
|
||||
export interface CompressionConfig {
|
||||
enabled: boolean;
|
||||
model?: string;
|
||||
maxInputTokens?: number;
|
||||
targetReductionRatio?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Default compression configuration
|
||||
*/
|
||||
const DEFAULT_COMPRESSION_CONFIG: CompressionConfig = {
|
||||
enabled: false,
|
||||
model: 'gpt-4o-mini',
|
||||
maxInputTokens: 100000,
|
||||
targetReductionRatio: 0.5,
|
||||
};
|
||||
|
||||
/**
|
||||
* Get compression configuration from system settings
|
||||
*/
|
||||
export async function getCompressionConfig(): Promise<CompressionConfig> {
|
||||
try {
|
||||
const systemConfigDao = getSystemConfigDao();
|
||||
const systemConfig = await systemConfigDao.get();
|
||||
const compressionSettings = systemConfig?.compression || {};
|
||||
|
||||
return {
|
||||
enabled: compressionSettings.enabled ?? DEFAULT_COMPRESSION_CONFIG.enabled,
|
||||
model: compressionSettings.model ?? DEFAULT_COMPRESSION_CONFIG.model,
|
||||
maxInputTokens: compressionSettings.maxInputTokens ?? DEFAULT_COMPRESSION_CONFIG.maxInputTokens,
|
||||
targetReductionRatio:
|
||||
compressionSettings.targetReductionRatio ?? DEFAULT_COMPRESSION_CONFIG.targetReductionRatio,
|
||||
};
|
||||
} catch (error) {
|
||||
console.warn('Failed to get compression config, using defaults:', error);
|
||||
return DEFAULT_COMPRESSION_CONFIG;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if compression is available and enabled
|
||||
*/
|
||||
export async function isCompressionEnabled(): Promise<boolean> {
|
||||
const config = await getCompressionConfig();
|
||||
if (!config.enabled) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if we have OpenAI API key configured (via smart routing config)
|
||||
const smartRoutingConfig = await getSmartRoutingConfig();
|
||||
return !!smartRoutingConfig.openaiApiKey;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get OpenAI client for compression
|
||||
*/
|
||||
async function getOpenAIClient(smartRoutingConfig: SmartRoutingConfig): Promise<OpenAI | null> {
|
||||
if (!smartRoutingConfig.openaiApiKey) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return new OpenAI({
|
||||
apiKey: smartRoutingConfig.openaiApiKey,
|
||||
baseURL: smartRoutingConfig.openaiApiBaseUrl || 'https://api.openai.com/v1',
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Estimate token count for a string (rough approximation)
|
||||
* Uses ~4 characters per token as a rough estimate
|
||||
*/
|
||||
export function estimateTokenCount(text: string): number {
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if content should be compressed based on token count
|
||||
*/
|
||||
export function shouldCompress(content: string, maxInputTokens: number): boolean {
|
||||
const estimatedTokens = estimateTokenCount(content);
|
||||
// Only compress if content is larger than a reasonable threshold
|
||||
const compressionThreshold = Math.min(maxInputTokens * 0.1, 1000);
|
||||
return estimatedTokens > compressionThreshold;
|
||||
}
|
||||
|
||||
/**
|
||||
* Compress MCP tool output using AI
|
||||
*
|
||||
* @param content The MCP tool output content to compress
|
||||
* @param context Optional context about the tool that generated this output
|
||||
* @returns Compressed content or original content if compression fails/is disabled
|
||||
*/
|
||||
export async function compressOutput(
|
||||
content: string,
|
||||
context?: {
|
||||
toolName?: string;
|
||||
serverName?: string;
|
||||
},
|
||||
): Promise<{ compressed: string; originalLength: number; compressedLength: number; wasCompressed: boolean }> {
|
||||
const originalLength = content.length;
|
||||
|
||||
// Check if compression is enabled
|
||||
const compressionConfig = await getCompressionConfig();
|
||||
if (!compressionConfig.enabled) {
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
|
||||
// Check if content should be compressed
|
||||
if (!shouldCompress(content, compressionConfig.maxInputTokens || 100000)) {
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const smartRoutingConfig = await getSmartRoutingConfig();
|
||||
const openai = await getOpenAIClient(smartRoutingConfig);
|
||||
|
||||
if (!openai) {
|
||||
console.warn('Compression enabled but OpenAI API key not configured');
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
|
||||
const targetRatio = compressionConfig.targetReductionRatio || 0.5;
|
||||
const toolContext = context?.toolName ? `from tool "${context.toolName}"` : '';
|
||||
const serverContext = context?.serverName ? `on server "${context.serverName}"` : '';
|
||||
|
||||
const systemPrompt = `You are a data compression assistant. Your task is to compress MCP (Model Context Protocol) tool outputs while preserving all essential information.
|
||||
|
||||
Guidelines:
|
||||
- Remove redundant information, formatting, and verbose descriptions
|
||||
- Preserve all data values, identifiers, and critical information
|
||||
- Keep error messages and status information intact
|
||||
- Maintain structured data (JSON, arrays) in a compact but readable format
|
||||
- Target approximately ${Math.round(targetRatio * 100)}% reduction in size
|
||||
- If the content cannot be meaningfully compressed, return it as-is
|
||||
|
||||
The output is ${toolContext} ${serverContext}.`;
|
||||
|
||||
const userPrompt = `Compress the following MCP tool output while preserving all essential information:
|
||||
|
||||
${content}`;
|
||||
|
||||
const response = await openai.chat.completions.create({
|
||||
model: compressionConfig.model || 'gpt-4o-mini',
|
||||
messages: [
|
||||
{ role: 'system', content: systemPrompt },
|
||||
{ role: 'user', content: userPrompt },
|
||||
],
|
||||
temperature: 0.1,
|
||||
max_tokens: Math.ceil(estimateTokenCount(content) * targetRatio * 1.5),
|
||||
});
|
||||
|
||||
const compressedContent = response.choices[0]?.message?.content;
|
||||
|
||||
if (!compressedContent) {
|
||||
console.warn('Compression returned empty result, using original content');
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
|
||||
const compressedLength = compressedContent.length;
|
||||
|
||||
// Only use compressed version if it's actually smaller
|
||||
if (compressedLength >= originalLength) {
|
||||
console.log('Compression did not reduce size, using original content');
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
|
||||
const reductionPercent = ((originalLength - compressedLength) / originalLength * 100).toFixed(1);
|
||||
console.log(`Compressed output: ${originalLength} -> ${compressedLength} chars (${reductionPercent}% reduction)`);
|
||||
|
||||
return {
|
||||
compressed: compressedContent,
|
||||
originalLength,
|
||||
compressedLength,
|
||||
wasCompressed: true,
|
||||
};
|
||||
} catch (error) {
|
||||
console.error('Compression failed, using original content:', error);
|
||||
return {
|
||||
compressed: content,
|
||||
originalLength,
|
||||
compressedLength: originalLength,
|
||||
wasCompressed: false,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Compress tool call result content
|
||||
* This handles the MCP tool result format with content array
|
||||
*/
|
||||
export async function compressToolResult(
|
||||
result: any,
|
||||
context?: {
|
||||
toolName?: string;
|
||||
serverName?: string;
|
||||
},
|
||||
): Promise<any> {
|
||||
// Check if compression is enabled first
|
||||
const compressionEnabled = await isCompressionEnabled();
|
||||
if (!compressionEnabled) {
|
||||
return result;
|
||||
}
|
||||
|
||||
// Handle error results - don't compress error messages
|
||||
if (result?.isError) {
|
||||
return result;
|
||||
}
|
||||
|
||||
// Handle content array format
|
||||
if (!result?.content || !Array.isArray(result.content)) {
|
||||
return result;
|
||||
}
|
||||
|
||||
const compressedContent = await Promise.all(
|
||||
result.content.map(async (item: any) => {
|
||||
// Only compress text content
|
||||
if (item?.type !== 'text' || !item?.text) {
|
||||
return item;
|
||||
}
|
||||
|
||||
const compressionResult = await compressOutput(item.text, context);
|
||||
|
||||
return {
|
||||
...item,
|
||||
text: compressionResult.compressed,
|
||||
};
|
||||
}),
|
||||
);
|
||||
|
||||
return {
|
||||
...result,
|
||||
content: compressedContent,
|
||||
};
|
||||
}
|
||||
@@ -27,6 +27,7 @@ import { getDataService } from './services.js';
|
||||
import { getServerDao, getSystemConfigDao, ServerConfigWithName } from '../dao/index.js';
|
||||
import { initializeAllOAuthClients } from './oauthService.js';
|
||||
import { createOAuthProvider } from './mcpOAuthProvider.js';
|
||||
import { compressToolResult } from './compressionService.js';
|
||||
|
||||
const servers: { [sessionId: string]: Server } = {};
|
||||
|
||||
@@ -1260,7 +1261,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
const result = await openApiClient.callTool(cleanToolName, finalArgs, passthroughHeaders);
|
||||
|
||||
console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`);
|
||||
return {
|
||||
const openApiResult = {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
@@ -1268,6 +1269,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
},
|
||||
],
|
||||
};
|
||||
return compressToolResult(openApiResult, {
|
||||
toolName: cleanToolName,
|
||||
serverName: targetServerInfo.name,
|
||||
});
|
||||
}
|
||||
|
||||
// Call the tool on the target server (MCP servers)
|
||||
@@ -1297,7 +1302,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
);
|
||||
|
||||
console.log(`Tool invocation result: ${JSON.stringify(result)}`);
|
||||
return result;
|
||||
return compressToolResult(result, {
|
||||
toolName,
|
||||
serverName: targetServerInfo.name,
|
||||
});
|
||||
}
|
||||
|
||||
// Regular tool handling
|
||||
@@ -1356,7 +1364,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
);
|
||||
|
||||
console.log(`OpenAPI tool invocation result: ${JSON.stringify(result)}`);
|
||||
return {
|
||||
const openApiResult = {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
@@ -1364,6 +1372,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
},
|
||||
],
|
||||
};
|
||||
return compressToolResult(openApiResult, {
|
||||
toolName: cleanToolName,
|
||||
serverName: serverInfo.name,
|
||||
});
|
||||
}
|
||||
|
||||
// Handle MCP servers
|
||||
@@ -1374,6 +1386,7 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
|
||||
const separator = getNameSeparator();
|
||||
const prefix = `${serverInfo.name}${separator}`;
|
||||
const originalToolName = request.params.name;
|
||||
request.params.name = request.params.name.startsWith(prefix)
|
||||
? request.params.name.substring(prefix.length)
|
||||
: request.params.name;
|
||||
@@ -1383,7 +1396,10 @@ export const handleCallToolRequest = async (request: any, extra: any) => {
|
||||
serverInfo.options || {},
|
||||
);
|
||||
console.log(`Tool call result: ${JSON.stringify(result)}`);
|
||||
return result;
|
||||
return compressToolResult(result, {
|
||||
toolName: originalToolName,
|
||||
serverName: serverInfo.name,
|
||||
});
|
||||
} catch (error) {
|
||||
console.error(`Error handling CallToolRequest: ${error}`);
|
||||
return {
|
||||
|
||||
@@ -173,6 +173,12 @@ export interface SystemConfig {
|
||||
oauth?: OAuthProviderConfig; // OAuth provider configuration for upstream MCP servers
|
||||
oauthServer?: OAuthServerConfig; // OAuth authorization server configuration for MCPHub itself
|
||||
enableSessionRebuild?: boolean; // Controls whether server session rebuild is enabled
|
||||
compression?: {
|
||||
enabled?: boolean; // Enable/disable AI compression of MCP tool outputs
|
||||
model?: string; // AI model to use for compression (default: 'gpt-4o-mini')
|
||||
maxInputTokens?: number; // Maximum input tokens for compression (default: 100000)
|
||||
targetReductionRatio?: number; // Target reduction ratio, 0.0-1.0 (default: 0.5)
|
||||
};
|
||||
}
|
||||
|
||||
export interface UserConfig {
|
||||
|
||||
@@ -117,6 +117,7 @@ export async function migrateToDatabase(): Promise<boolean> {
|
||||
oauth: settings.systemConfig.oauth || {},
|
||||
oauthServer: settings.systemConfig.oauthServer || {},
|
||||
enableSessionRebuild: settings.systemConfig.enableSessionRebuild,
|
||||
compression: settings.systemConfig.compression || {},
|
||||
};
|
||||
await systemConfigRepo.update(systemConfig);
|
||||
console.log(' - System configuration updated');
|
||||
|
||||
428
tests/services/compressionService.test.ts
Normal file
428
tests/services/compressionService.test.ts
Normal file
@@ -0,0 +1,428 @@
|
||||
// Mock the DAO module before imports
|
||||
jest.mock('../../src/dao/index.js', () => ({
|
||||
getSystemConfigDao: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock smart routing config
|
||||
jest.mock('../../src/utils/smartRouting.js', () => ({
|
||||
getSmartRoutingConfig: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock OpenAI
|
||||
jest.mock('openai', () => {
|
||||
return {
|
||||
__esModule: true,
|
||||
default: jest.fn().mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: jest.fn(),
|
||||
},
|
||||
},
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
import {
|
||||
getCompressionConfig,
|
||||
isCompressionEnabled,
|
||||
estimateTokenCount,
|
||||
shouldCompress,
|
||||
compressOutput,
|
||||
compressToolResult,
|
||||
} from '../../src/services/compressionService.js';
|
||||
import { getSystemConfigDao } from '../../src/dao/index.js';
|
||||
import { getSmartRoutingConfig } from '../../src/utils/smartRouting.js';
|
||||
import OpenAI from 'openai';
|
||||
|
||||
describe('CompressionService', () => {
|
||||
const mockSystemConfigDao = {
|
||||
get: jest.fn(),
|
||||
getSection: jest.fn(),
|
||||
update: jest.fn(),
|
||||
updateSection: jest.fn(),
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
(getSystemConfigDao as jest.Mock).mockReturnValue(mockSystemConfigDao);
|
||||
});
|
||||
|
||||
describe('getCompressionConfig', () => {
|
||||
it('should return default config when no config is set', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({});
|
||||
|
||||
const config = await getCompressionConfig();
|
||||
|
||||
expect(config).toEqual({
|
||||
enabled: false,
|
||||
model: 'gpt-4o-mini',
|
||||
maxInputTokens: 100000,
|
||||
targetReductionRatio: 0.5,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return configured values when set', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: {
|
||||
enabled: true,
|
||||
model: 'gpt-4o',
|
||||
maxInputTokens: 50000,
|
||||
targetReductionRatio: 0.3,
|
||||
},
|
||||
});
|
||||
|
||||
const config = await getCompressionConfig();
|
||||
|
||||
expect(config).toEqual({
|
||||
enabled: true,
|
||||
model: 'gpt-4o',
|
||||
maxInputTokens: 50000,
|
||||
targetReductionRatio: 0.3,
|
||||
});
|
||||
});
|
||||
|
||||
it('should use defaults for missing values', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: {
|
||||
enabled: true,
|
||||
},
|
||||
});
|
||||
|
||||
const config = await getCompressionConfig();
|
||||
|
||||
expect(config).toEqual({
|
||||
enabled: true,
|
||||
model: 'gpt-4o-mini',
|
||||
maxInputTokens: 100000,
|
||||
targetReductionRatio: 0.5,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return defaults on error', async () => {
|
||||
mockSystemConfigDao.get.mockRejectedValue(new Error('Test error'));
|
||||
|
||||
const config = await getCompressionConfig();
|
||||
|
||||
expect(config).toEqual({
|
||||
enabled: false,
|
||||
model: 'gpt-4o-mini',
|
||||
maxInputTokens: 100000,
|
||||
targetReductionRatio: 0.5,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('isCompressionEnabled', () => {
|
||||
it('should return false when compression is disabled', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: false },
|
||||
});
|
||||
|
||||
const enabled = await isCompressionEnabled();
|
||||
|
||||
expect(enabled).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false when enabled but no API key', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: '',
|
||||
});
|
||||
|
||||
const enabled = await isCompressionEnabled();
|
||||
|
||||
expect(enabled).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true when enabled and API key is set', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
});
|
||||
|
||||
const enabled = await isCompressionEnabled();
|
||||
|
||||
expect(enabled).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('estimateTokenCount', () => {
|
||||
it('should estimate tokens for short text', () => {
|
||||
const text = 'Hello world';
|
||||
const tokens = estimateTokenCount(text);
|
||||
|
||||
// 11 characters / 4 = 2.75, ceil = 3
|
||||
expect(tokens).toBe(3);
|
||||
});
|
||||
|
||||
it('should estimate tokens for longer text', () => {
|
||||
const text = 'This is a longer piece of text that should have more tokens';
|
||||
const tokens = estimateTokenCount(text);
|
||||
|
||||
// 59 characters / 4 = 14.75, ceil = 15
|
||||
expect(tokens).toBe(15);
|
||||
});
|
||||
|
||||
it('should handle empty string', () => {
|
||||
const tokens = estimateTokenCount('');
|
||||
|
||||
expect(tokens).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('shouldCompress', () => {
|
||||
it('should return false for small content', () => {
|
||||
const content = 'Small content';
|
||||
const result = shouldCompress(content, 100000);
|
||||
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for large content', () => {
|
||||
// Create content larger than the threshold (1000 tokens = ~4000 chars)
|
||||
const content = 'x'.repeat(5000);
|
||||
const result = shouldCompress(content, 100000);
|
||||
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
|
||||
it('should use 10% of maxInputTokens as threshold', () => {
|
||||
// With maxInputTokens = 1000, threshold = 100 tokens = ~400 chars
|
||||
const smallContent = 'x'.repeat(300);
|
||||
const largeContent = 'x'.repeat(500);
|
||||
|
||||
expect(shouldCompress(smallContent, 1000)).toBe(false);
|
||||
expect(shouldCompress(largeContent, 1000)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('compressOutput', () => {
|
||||
it('should return original content when compression is disabled', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: false },
|
||||
});
|
||||
|
||||
const content = 'Test content';
|
||||
const result = await compressOutput(content);
|
||||
|
||||
expect(result).toEqual({
|
||||
compressed: content,
|
||||
originalLength: content.length,
|
||||
compressedLength: content.length,
|
||||
wasCompressed: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return original content when content is too small', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true, maxInputTokens: 100000 },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
});
|
||||
|
||||
const content = 'Small content';
|
||||
const result = await compressOutput(content);
|
||||
|
||||
expect(result.wasCompressed).toBe(false);
|
||||
expect(result.compressed).toBe(content);
|
||||
});
|
||||
|
||||
it('should return original content when no API key is configured', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: '',
|
||||
});
|
||||
|
||||
const content = 'x'.repeat(5000);
|
||||
const result = await compressOutput(content);
|
||||
|
||||
expect(result.wasCompressed).toBe(false);
|
||||
expect(result.compressed).toBe(content);
|
||||
});
|
||||
|
||||
it('should compress content when enabled and content is large', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
openaiApiBaseUrl: 'https://api.openai.com/v1',
|
||||
});
|
||||
|
||||
const originalContent = 'x'.repeat(5000);
|
||||
const compressedContent = 'y'.repeat(2000);
|
||||
|
||||
// Mock OpenAI response
|
||||
const mockCreate = jest.fn().mockResolvedValue({
|
||||
choices: [{ message: { content: compressedContent } }],
|
||||
});
|
||||
|
||||
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
const result = await compressOutput(originalContent, {
|
||||
toolName: 'test-tool',
|
||||
serverName: 'test-server',
|
||||
});
|
||||
|
||||
expect(result.wasCompressed).toBe(true);
|
||||
expect(result.compressed).toBe(compressedContent);
|
||||
expect(result.originalLength).toBe(originalContent.length);
|
||||
expect(result.compressedLength).toBe(compressedContent.length);
|
||||
});
|
||||
|
||||
it('should return original content when compressed is larger', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
openaiApiBaseUrl: 'https://api.openai.com/v1',
|
||||
});
|
||||
|
||||
const originalContent = 'x'.repeat(5000);
|
||||
const largerContent = 'y'.repeat(6000);
|
||||
|
||||
const mockCreate = jest.fn().mockResolvedValue({
|
||||
choices: [{ message: { content: largerContent } }],
|
||||
});
|
||||
|
||||
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
const result = await compressOutput(originalContent);
|
||||
|
||||
expect(result.wasCompressed).toBe(false);
|
||||
expect(result.compressed).toBe(originalContent);
|
||||
});
|
||||
|
||||
it('should return original content on API error', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true, model: 'gpt-4o-mini', maxInputTokens: 100000 },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
openaiApiBaseUrl: 'https://api.openai.com/v1',
|
||||
});
|
||||
|
||||
const mockCreate = jest.fn().mockRejectedValue(new Error('API error'));
|
||||
|
||||
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
const content = 'x'.repeat(5000);
|
||||
const result = await compressOutput(content);
|
||||
|
||||
expect(result.wasCompressed).toBe(false);
|
||||
expect(result.compressed).toBe(content);
|
||||
});
|
||||
});
|
||||
|
||||
describe('compressToolResult', () => {
|
||||
it('should return original result when compression is disabled', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: false },
|
||||
});
|
||||
|
||||
const result = {
|
||||
content: [{ type: 'text', text: 'Test output' }],
|
||||
};
|
||||
|
||||
const compressed = await compressToolResult(result);
|
||||
|
||||
expect(compressed).toEqual(result);
|
||||
});
|
||||
|
||||
it('should not compress error results', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
});
|
||||
|
||||
const result = {
|
||||
content: [{ type: 'text', text: 'Error message' }],
|
||||
isError: true,
|
||||
};
|
||||
|
||||
const compressed = await compressToolResult(result);
|
||||
|
||||
expect(compressed).toEqual(result);
|
||||
});
|
||||
|
||||
it('should handle results without content array', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
});
|
||||
|
||||
const result = { someOtherField: 'value' };
|
||||
|
||||
const compressed = await compressToolResult(result);
|
||||
|
||||
expect(compressed).toEqual(result);
|
||||
});
|
||||
|
||||
it('should only compress text content items', async () => {
|
||||
mockSystemConfigDao.get.mockResolvedValue({
|
||||
compression: { enabled: true, maxInputTokens: 100000 },
|
||||
});
|
||||
(getSmartRoutingConfig as jest.Mock).mockResolvedValue({
|
||||
openaiApiKey: 'test-api-key',
|
||||
openaiApiBaseUrl: 'https://api.openai.com/v1',
|
||||
});
|
||||
|
||||
const largeText = 'x'.repeat(5000);
|
||||
const compressedText = 'y'.repeat(2000);
|
||||
|
||||
const mockCreate = jest.fn().mockResolvedValue({
|
||||
choices: [{ message: { content: compressedText } }],
|
||||
});
|
||||
|
||||
(OpenAI as unknown as jest.Mock).mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
const result = {
|
||||
content: [
|
||||
{ type: 'text', text: largeText },
|
||||
{ type: 'image', data: 'base64data' },
|
||||
],
|
||||
};
|
||||
|
||||
const compressed = await compressToolResult(result);
|
||||
|
||||
expect(compressed.content[0].text).toBe(compressedText);
|
||||
expect(compressed.content[1]).toEqual({ type: 'image', data: 'base64data' });
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user