tryaii-mcp-server
Version:
TryAII MCP Server - 15+ AI models with comparison, cost tracking, and collective intelligence
558 lines • 24.8 kB
JavaScript
import express, { Router } from 'express';
import { authenticateApiKeyEnhanced, requirePermissionEnhanced, requireSufficientBalance, createSessionTracking, completeSessionTracking } from '../middleware/integratedAuth.js';
import { asyncHandler } from '../utils/asyncHandler.js';
import { ValidationError } from '../errors/CustomErrors.js';
import { sessionManager } from '../services/sessionManager.js';
import { balanceService } from '../services/BalanceService.js';
import { sessionService } from '../services/SessionService.js';
import { logger } from '../utils/logger.js';
import { calculateActualCost } from '../utils/costCalculator.js';
import fs from 'fs/promises';
import path from 'path';
import { fileURLToPath } from 'url';
const router = Router();
// Endpoint for internal service to save HTML reports
router.post('/save-report', express.json({ limit: '10mb' }), asyncHandler(async (req, res) => {
const { uuid, htmlContent } = req.body;
if (!uuid || !htmlContent) {
throw new ValidationError('Both uuid and htmlContent are required.');
}
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
const outputDir = path.resolve(__dirname, '../../../output');
const fileName = `brains-response-${uuid}.html`;
const filePath = path.join(outputDir, fileName);
const gatewayPublicUrl = process.env.GATEWAY_PUBLIC_URL || 'https://tryaii-mcp.onrender.com';
const publicUrl = `${gatewayPublicUrl}/output/${fileName}`;
try {
await fs.mkdir(outputDir, { recursive: true });
await fs.writeFile(filePath, htmlContent, 'utf-8');
logger.info('Successfully saved HTML report', { uuid, path: filePath, url: publicUrl });
res.status(200).json({
success: true,
message: 'Report saved successfully.',
url: publicUrl
});
}
catch (error) {
logger.error('Failed to save HTML report', { uuid, error: error instanceof Error ? error.message : String(error) });
throw new Error('Failed to save HTML report.');
}
}));
// Apply enhanced authentication to all OTHER API routes
router.use(authenticateApiKeyEnhanced);
router.use(completeSessionTracking());
/**
* Validation helpers
*/
const validateModelId = (modelId) => {
if (!modelId || typeof modelId !== 'string' || modelId.trim().length === 0) {
throw new ValidationError('modelId must be a non-empty string');
}
};
const validateMessage = (message) => {
if (!message || typeof message !== 'string' || message.trim().length === 0) {
throw new ValidationError('message must be a non-empty string');
}
if (message.length > 10000) {
throw new ValidationError('message must be less than 10,000 characters');
}
};
const validateModelIds = (modelIds) => {
if (!Array.isArray(modelIds)) {
throw new ValidationError('modelIds must be an array');
}
if (modelIds.length === 0) {
throw new ValidationError('modelIds array cannot be empty');
}
if (modelIds.length > 5) {
throw new ValidationError('modelIds array cannot contain more than 5 models');
}
modelIds.forEach((id, index) => {
if (!id || typeof id !== 'string' || id.trim().length === 0) {
throw new ValidationError(`modelIds[${index}] must be a non-empty string`);
}
});
};
/**
* POST /api/v2/chat - Enhanced chat endpoint with full integration
*/
router.post('/v2/chat', requirePermissionEnhanced('chat'), requireSufficientBalance(0.001), // Minimum balance check
createSessionTracking('chat'), asyncHandler(async (req, res) => {
const { modelId, message, enableWebSearch, temperature, maxTokens, conversationHistory } = req.body;
validateModelId(modelId);
validateMessage(message);
const sessionId = req.sessionId;
try {
// Get or create MCP session
const mcpSession = await sessionManager.getOrCreateSession(req.user.userId, req.user.apiKeyId);
logger.info('Processing enhanced chat request', {
userId: req.user.userId,
sessionId,
mcpSessionId: mcpSession.sessionId,
modelId,
requestId: req.sessionContext.requestId
});
// Record start time for latency calculation
const modelStartTime = Date.now();
// Call MCP service
const result = await sessionManager.queueRequest(mcpSession.sessionId, 'chat_with_model', {
modelId,
message,
enableWebSearch,
temperature,
maxTokens,
conversationHistory
});
const modelLatency = Date.now() - modelStartTime;
// 🆕 Extract actual cost and token data from AI provider response (including thinking tokens)
const actualCost = result?.cost || 0.01; // Use actual cost from AI provider
const actualInputTokens = result?.inputTokens || Math.ceil(message.length / 4);
const actualOutputTokens = result?.outputTokens || Math.ceil((result?.response?.length || 0) / 4);
const actualThinkingTokens = result?.thinkingTokens || 0; // 🆕 Reasoning tokens
// 🆕 Calculate cost breakdown (thinking tokens priced same as output tokens)
const costBreakdown = result?.costBreakdown || {
inputCost: (actualInputTokens / 1000) * 0.01,
outputCost: (actualOutputTokens / 1000) * 0.03,
thinkingCost: (actualThinkingTokens / 1000) * 0.03 // Same as output pricing
};
// Record model interaction in session with enhanced data
await sessionService.addModelInteraction(sessionId, 'chat', {
modelId,
provider: result?.provider || 'tryaii',
inputTokens: actualInputTokens,
outputTokens: actualOutputTokens,
thinkingTokens: actualThinkingTokens, // 🆕 Reasoning tokens
totalTokens: actualInputTokens + actualOutputTokens + actualThinkingTokens, // 🆕 Total tokens
cost: actualCost,
// costBreakdown, // 🆕 TODO: Add after updating interface
latency: modelLatency,
success: !!result?.response,
errorMessage: result?.error,
// 🆕 Enhanced metadata
reasoningUsed: actualThinkingTokens > 0,
webSearchUsed: result?.webSearchUsed || false,
// responsePreview field removed for storage optimization
});
// 🆕 Deduct balance atomically with enhanced metadata
const balanceResult = await balanceService.deductBalance(req.user.userId, actualCost, `Chat with ${modelId}`, {
...req.sessionContext.clientMetadata,
// 🆕 Enhanced token and cost metadata
tokenUsage: {
totalTokens: actualInputTokens + actualOutputTokens + actualThinkingTokens,
thinkingTokens: actualThinkingTokens,
reasoningModelsUsed: actualThinkingTokens > 0 ? 1 : 0
},
costBreakdown
});
// Update session with balance result
await sessionService.updateSessionBalance(sessionId, 'chat', balanceResult);
if (!balanceResult.success) {
logger.error('Balance deduction failed for chat', {
userId: req.user.userId,
sessionId,
error: balanceResult.errorMessage
});
}
const response = {
success: true,
data: result,
metadata: {
requestId: req.sessionContext.requestId,
sessionId,
cost: actualCost,
balanceAfter: balanceResult.balanceAfter || 0,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Enhanced chat request failed', {
userId: req.user.userId,
sessionId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* POST /api/v2/compare - Enhanced compare endpoint with full integration
*
* ⚠️ ARCHITECTURE ISSUE: Currently this creates a single queue entry that internally
* breaks down into separate model API calls. This can cause inconsistent timing
* where some models in a comparison finish quickly while others wait in queue.
*
* PROPOSED SOLUTION: Implement batch processing where:
* 1. Reserve processing slots for ALL models upfront when queued
* 2. Process all models in parallel with consistent start times
* 3. Fail the entire comparison if resource allocation fails
*
* TODO: Refactor to use batched resource allocation in sessionManager
*/
router.post('/v2/compare', requirePermissionEnhanced('compare'), requireSufficientBalance(0.05), // Higher minimum for multiple models
createSessionTracking('compare'), asyncHandler(async (req, res) => {
const { modelIds, message, enableWebSearch, temperature, maxTokens } = req.body;
validateModelIds(modelIds);
validateMessage(message);
const sessionId = req.sessionId;
try {
const mcpSession = await sessionManager.getOrCreateSession(req.user.userId, req.user.apiKeyId);
logger.info('Processing enhanced compare request', {
userId: req.user.userId,
sessionId,
mcpSessionId: mcpSession.sessionId,
modelIds,
requestId: req.sessionContext.requestId,
// 🔥 ADD: Log queue status to help identify timing issues
queueStatus: sessionManager.getSessionStatus(mcpSession.sessionId)
});
const modelStartTime = Date.now();
// 🔥 CURRENT ISSUE: This single queue entry will internally break into
// separate API calls that may not execute simultaneously
const result = await sessionManager.queueRequest(mcpSession.sessionId, 'compare_models', {
modelIds,
message,
enableWebSearch,
temperature,
maxTokens,
// 🔥 ADD: Flag to indicate this needs batch processing
requiresBatchProcessing: true,
modelCount: modelIds.length
});
const totalLatency = Date.now() - modelStartTime;
// Calculate actual total cost using model interactions instead of hardcoded values
let totalCost = 0;
if (result?.responses) {
// Use the centralized cost calculator
const costResult = calculateActualCost({ content: [{ text: JSON.stringify(result) }] }, 'compare_models');
totalCost = costResult.totalCost;
// Log actual vs estimated cost comparison
logger.info('Enhanced compare cost calculation', {
actualCost: totalCost,
modelCount: modelIds.length,
costBreakdown: costResult.breakdown,
source: 'centralizedCalculator'
});
for (const [modelId, response] of Object.entries(result.responses)) {
// Find the specific model cost from breakdown
const modelBreakdown = costResult.breakdown.find((b) => b.modelId === modelId);
const modelCost = modelBreakdown?.cost || (totalCost / modelIds.length); // Fallback to average
await sessionService.addModelInteraction(sessionId, 'compare', {
modelId,
provider: 'tryaii',
inputTokens: Math.ceil(message.length / 4),
outputTokens: Math.ceil((response?.response?.length || 0) / 4),
cost: modelCost,
latency: totalLatency / modelIds.length, // Average latency
fullResponse: response?.response || '',
success: !!response?.response,
errorMessage: response?.error,
// responsePreview field removed for storage optimization
});
}
}
else {
// Fallback if no response data
logger.warn('No response data for compare models - using fallback cost');
totalCost = Math.min(modelIds.length * 0.025, 0.25); // Conservative fallback
}
// Deduct total balance
const balanceResult = await balanceService.deductBalance(req.user.userId, totalCost, `Compare ${modelIds.length} models`, req.sessionContext.clientMetadata);
await sessionService.updateSessionBalance(sessionId, 'compare', balanceResult);
const response = {
success: true,
data: result,
metadata: {
requestId: req.sessionContext.requestId,
sessionId,
cost: totalCost,
balanceAfter: balanceResult.balanceAfter || 0,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Enhanced compare request failed', {
userId: req.user.userId,
sessionId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* POST /api/v2/brains - Enhanced brains endpoint with full integration
*/
router.post('/v2/brains', requirePermissionEnhanced('brains'), requireSufficientBalance(0.10), // Higher cost for 5 top models
createSessionTracking('brains'), asyncHandler(async (req, res) => {
const { question, enableWebSearch, temperature, maxTokens } = req.body;
if (!question || typeof question !== 'string' || question.trim().length === 0) {
throw new ValidationError('question must be a non-empty string');
}
const sessionId = req.sessionId;
try {
const mcpSession = await sessionManager.getOrCreateSession(req.user.userId, req.user.apiKeyId);
logger.info('Processing enhanced brains request', {
userId: req.user.userId,
sessionId,
mcpSessionId: mcpSession.sessionId,
questionLength: question.length,
requestId: req.sessionContext.requestId
});
const modelStartTime = Date.now();
const result = await sessionManager.queueRequest(mcpSession.sessionId, 'brains', {
question,
enableWebSearch,
temperature,
maxTokens
});
const totalLatency = Date.now() - modelStartTime;
// Calculate actual total cost using centralized cost calculator
let totalCost = 0;
// Use the centralized cost calculator for brains
const costResult = calculateActualCost({ content: [{ text: JSON.stringify(result) }] }, 'brains');
totalCost = costResult.totalCost;
// Log actual vs estimated cost comparison
logger.info('Enhanced brains cost calculation', {
actualCost: totalCost,
modelCount: costResult.modelCount,
successfulModels: costResult.successfulModels,
costBreakdown: costResult.breakdown,
source: 'centralizedCalculator'
});
// Record interactions for the brains models with actual costs
for (const breakdown of costResult.breakdown) {
await sessionService.addModelInteraction(sessionId, 'brains', {
modelId: breakdown.modelId,
provider: 'tryaii',
inputTokens: Math.ceil(question.length / 4),
outputTokens: 100, // Estimated, could be improved with actual data
cost: breakdown.cost,
latency: totalLatency / costResult.modelCount,
success: breakdown.success,
errorMessage: breakdown.success ? undefined : 'Model execution failed'
});
}
const balanceResult = await balanceService.deductBalance(req.user.userId, totalCost, 'Brains collective intelligence', req.sessionContext.clientMetadata);
await sessionService.updateSessionBalance(sessionId, 'brains', balanceResult);
const response = {
success: true,
data: result,
metadata: {
requestId: req.sessionContext.requestId,
sessionId,
cost: totalCost,
balanceAfter: balanceResult.balanceAfter || 0,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Enhanced brains request failed', {
userId: req.user.userId,
sessionId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* GET /api/v2/models - Enhanced models list endpoint (free)
*/
router.get('/v2/models', requirePermissionEnhanced('models'), createSessionTracking('models'), asyncHandler(async (req, res) => {
const { provider } = req.query;
const sessionId = req.sessionId;
try {
const mcpSession = await sessionManager.getOrCreateSession(req.user.userId, req.user.apiKeyId);
logger.info('Processing enhanced models list request', {
userId: req.user.userId,
sessionId,
provider,
requestId: req.sessionContext.requestId
});
const modelStartTime = Date.now();
const result = await sessionManager.queueRequest(mcpSession.sessionId, 'list_available_models', {
provider: provider
});
const latency = Date.now() - modelStartTime;
// Record model interaction (free endpoint)
await sessionService.addModelInteraction(sessionId, 'models', {
modelId: 'list_models',
provider: 'tryaii',
inputTokens: 0,
outputTokens: 0,
cost: 0,
latency,
success: !!result,
errorMessage: result?.error
});
// No balance deduction for models endpoint - it's free
await sessionService.updateSessionBalance(sessionId, 'models', {
success: true,
balanceAfter: await balanceService.getCurrentBalance(req.user.userId)
});
const response = {
success: true,
data: result,
metadata: {
requestId: req.sessionContext.requestId,
sessionId,
cost: 0,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Enhanced models request failed', {
userId: req.user.userId,
sessionId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* GET /api/v2/session/analytics - Get session analytics
*/
router.get('/v2/session/analytics', requirePermissionEnhanced('usage'), asyncHandler(async (req, res) => {
const { days = 30 } = req.query;
try {
const analytics = await sessionService.getSessionAnalytics(req.user.userId, parseInt(days) || 30);
const response = {
success: true,
data: analytics,
metadata: {
requestId: req.sessionContext.requestId,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Failed to get session analytics', {
userId: req.user.userId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* GET /api/v2/balance - Get current balance
*/
router.get('/v2/balance', requirePermissionEnhanced('usage'), asyncHandler(async (req, res) => {
try {
const balance = await balanceService.getCurrentBalance(req.user.userId);
const recentTransactions = await balanceService.getRecentTransactions(req.user.userId, 10);
const response = {
success: true,
data: {
currentBalance: balance,
currency: 'USD',
recentTransactions
},
metadata: {
requestId: req.sessionContext.requestId,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Failed to get balance', {
userId: req.user.userId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
/**
* POST /api/brains - Compatibility endpoint for backward compatibility
* Redirects to /api/v2/brains with parameter mapping
*/
router.post('/brains', requirePermissionEnhanced('brains'), requireSufficientBalance(0.001), createSessionTracking('brains'), asyncHandler(async (req, res) => {
// Map old parameters to new format
const { question, query, enableWebSearch, temperature, maxTokens } = req.body;
// Support both 'question' and 'query' parameters for backward compatibility
const finalQuestion = question || query;
if (!finalQuestion) {
throw new ValidationError('Either question or query parameter is required');
}
const sessionId = req.sessionId;
try {
// Get or create MCP session
const mcpSession = await sessionManager.getOrCreateSession(req.user.userId, req.user.apiKeyId);
logger.info('Processing compatibility brains request', {
userId: req.user.userId,
sessionId,
mcpSessionId: mcpSession.sessionId,
requestId: req.sessionContext.requestId
});
const modelStartTime = Date.now();
// Call MCP service with mapped parameters
const result = await sessionManager.queueRequest(mcpSession.sessionId, 'brains', {
question: finalQuestion,
enableWebSearch: enableWebSearch || false,
temperature: temperature || 0.7,
maxTokens: maxTokens || 12000
});
const totalLatency = Date.now() - modelStartTime;
// Use the centralized cost calculator for brains
const costResult = calculateActualCost({ content: [{ text: JSON.stringify(result) }] }, 'brains');
const totalCost = costResult.totalCost;
logger.info('Compatibility brains cost calculation', {
actualCost: totalCost,
modelCount: costResult.modelCount,
successfulModels: costResult.successfulModels,
costBreakdown: costResult.breakdown,
source: 'compatibilityEndpoint'
});
// Record interactions for the brains models with actual costs
for (const breakdown of costResult.breakdown) {
await sessionService.addModelInteraction(sessionId, 'brains', {
modelId: breakdown.modelId,
provider: 'tryaii',
inputTokens: Math.ceil(finalQuestion.length / 4),
outputTokens: 100, // Estimated, could be improved with actual data
cost: breakdown.cost,
latency: totalLatency / costResult.modelCount,
success: breakdown.success,
errorMessage: breakdown.success ? undefined : 'Model execution failed'
});
}
const balanceResult = await balanceService.deductBalance(req.user.userId, totalCost, 'Brains collective intelligence (compatibility endpoint)', req.sessionContext.clientMetadata);
await sessionService.updateSessionBalance(sessionId, 'brains', balanceResult);
const response = {
success: true,
data: result,
metadata: {
requestId: req.sessionContext.requestId,
sessionId,
cost: totalCost,
balanceAfter: balanceResult.balanceAfter || 0,
timestamp: new Date().toISOString(),
processingTime: Date.now() - req.sessionContext.startTime
}
};
res.json(response);
}
catch (error) {
logger.error('Compatibility brains request failed', {
userId: req.user.userId,
sessionId,
error: error instanceof Error ? error.message : String(error)
});
throw error;
}
}));
export default router;
//# sourceMappingURL=enhancedApi.js.map