@aj-archipelago/cortex 1.1.0-beta.0 → 1.1.0

package/config.js

@@ -174,6 +174,12 @@ var config = convict({
         sensitive: true,
         env: 'STORAGE_CONNECTION_STRING'
     },
+    redisEncryptionKey: {
+        format: String,
+        default: null,
+        env: 'REDIS_ENCRYPTION_KEY',
+        sensitive: true
+    },
     dalleImageApiUrl: {
         format: String,
         default: 'null',

package/lib/crypto.js

@@ -0,0 +1,46 @@
+// This file is used to encrypt and decrypt data using the crypto library
+import logger from './logger.js';
+import crypto from 'crypto';
+
+// Encryption function
+function encrypt(text, key) {
+    if (!key) { return text; }
+    try {
+        key = tryBufferKey(key);
+        let iv = crypto.randomBytes(16);
+        let cipher = crypto.createCipheriv('aes-256-cbc', key, iv);
+        let encrypted = cipher.update(text, 'utf8', 'hex');
+        encrypted += cipher.final('hex');
+        return iv.toString('hex') + ':' + encrypted;
+    } catch (error) {
+        logger.error(`Encryption failed: ${error.message}`);
+        return null;
+    }
+}
+
+// Decryption function
+function decrypt(message, key) {
+    if (!key) { return message; }
+    try {
+        key = tryBufferKey(key);
+        let parts = message.split(':');
+        let iv = Buffer.from(parts.shift(), 'hex');
+        let encrypted = parts.join(':');
+        let decipher = crypto.createDecipheriv('aes-256-cbc', key, iv);
+        let decrypted = decipher.update(encrypted, 'hex', 'utf8');
+        decrypted += decipher.final('utf8');
+        return decrypted;
+    } catch (error) {
+        logger.error(`Decryption failed: ${error.message}`);
+        return null;
+    }
+}
+
+function tryBufferKey(key) {
+    if (key.length === 64) {
+        return Buffer.from(key, 'hex');
+    }   
+    return key;
+}
+
+export { encrypt, decrypt };

package/lib/keyValueStorageClient.js

@@ -1,15 +1,32 @@
 import Keyv from 'keyv';
 import { config } from '../config.js';
+import { encrypt, decrypt } from './crypto.js';
+import logger from './logger.js';
 
 const storageConnectionString = config.get('storageConnectionString');
 const cortexId = config.get('cortexId');
+const redisEncryptionKey = config.get('redisEncryptionKey');
 
 // Create a keyv client to store data
 const keyValueStorageClient = new Keyv(storageConnectionString, {
     ssl: true,
     abortConnect: false,
-    serialize: JSON.stringify,
-    deserialize: JSON.parse,
+    serialize: (data) => redisEncryptionKey ? encrypt(JSON.stringify(data), redisEncryptionKey) : JSON.stringify(data),
+    deserialize: (data) => {
+        try {
+            // Try to parse the data normally
+            return JSON.parse(data);
+        } catch (error) {
+            // If it fails, the data may be encrypted so attempt to decrypt it if we have a key
+            try {
+                return JSON.parse(decrypt(data, redisEncryptionKey));
+            } catch (decryptError) {
+                // If decryption also fails, log an error and return an empty object
+                logger.error(`Failed to parse or decrypt stored key value data: ${decryptError}`);
+                return {};
+            }
+        }
+    },
     namespace: `${cortexId}-cortex-context`
 });
 

package/lib/logger.js

@@ -1,22 +1,6 @@
 // logger.js
 import winston from 'winston';
 
-const format = winston.format.combine(
-    //winston.format.timestamp(),
-    winston.format.colorize({ all: true }),
-    winston.format.simple()
-);
-
-const transports = [
-    new winston.transports.Console({ format })
-];
-
-const logger = winston.createLogger({
-    level: process.env.NODE_ENV === 'production' ? 'info' : 'debug',
-    format: format,
-    transports: transports
-});
-
 winston.addColors({
     debug: 'green',
     verbose: 'blue',
@@ -26,4 +10,39 @@ winston.addColors({
     error: 'red'
 });
 
+const debugFormat = winston.format.combine(
+    winston.format.colorize({ all: true }),
+    winston.format.cli()
+);
+
+const prodFormat = winston.format.combine(
+    winston.format.simple()
+);
+
+const transports = process.env.NODE_ENV === 'production' ?
+    new winston.transports.Console({ level: 'info', format: prodFormat }) :
+    new winston.transports.Console({ level: 'debug', format: debugFormat });
+
+const logger = winston.createLogger({ transports });
+
+// Function to obscure sensitive URL parameters
+export const obscureUrlParams = url => {
+    try {
+        const urlObject = new URL(url);
+        urlObject.searchParams.forEach((value, name) => {
+            if (/token|key|password|secret|auth|apikey|access|passwd|credential/i.test(name)) {
+                urlObject.searchParams.set(name, '******');
+            }
+        });
+        return urlObject.toString();
+    } catch (e) {
+        if (e instanceof TypeError) {
+            logger.error('Error obscuring URL parameters - invalid URL.');
+            return url;
+        } else {
+            throw e;
+        }
+    }
+};
+
 export default logger;

package/lib/redisSubscription.js

@@ -3,90 +3,110 @@ import { config } from '../config.js';
 import pubsub from '../server/pubsub.js';
 import { requestState } from '../server/requestState.js';
 import logger from '../lib/logger.js';
+import { encrypt, decrypt } from '../lib/crypto.js';
 
 const connectionString = config.get('storageConnectionString');
-const channels = ['requestProgress', 'requestProgressSubscriptions'];
-let client;
+const redisEncryptionKey = config.get('redisEncryptionKey');
+const requestProgressChannel = 'requestProgress';
+const requestProgressSubscriptionsChannel = 'requestProgressSubscriptions';
+
+let subscriptionClient;
+let publisherClient;
 
 if (connectionString) {
-    logger.info(`Using Redis subscription for channel(s) ${channels.join(', ')}`);
+    logger.info(`Using Redis subscription for channel(s) ${requestProgressChannel}, ${requestProgressSubscriptionsChannel}`);
+    try {
+        subscriptionClient = connectionString && new Redis(connectionString);
+    } catch (error) {
+        logger.error(`Redis connection error: ${error}`);
+    }
+
+    logger.info(`Using Redis publish for channel(s) ${requestProgressChannel}, ${requestProgressSubscriptionsChannel}`);
     try {
-        client = connectionString && new Redis(connectionString);
+        publisherClient = connectionString && new Redis(connectionString);
     } catch (error) {
-        logger.error(`Redis connection error: ${JSON.stringify(error)}`);
+        logger.error(`Redis connection error: ${error}`);
+    }
+
+    if (redisEncryptionKey) {
+        logger.info('Using encryption for Redis');
+    } else {
+        logger.warn('REDIS_ENCRYPTION_KEY not set. Data stored in Redis will not be encrypted.');
     }
 
-    if (client) {
+    if (subscriptionClient) {
 
-        client.on('error', (error) => {
-            logger.error(`Redis client error: ${JSON.stringify(error)}`);
+        subscriptionClient.on('error', (error) => {
+            logger.error(`Redis subscriptionClient error: ${error}`);
         });
 
-        client.on('connect', () => {
-            client.subscribe('requestProgress', (error) => {
-                if (error) {
-                    logger.error(`Error subscribing to redis channel requestProgress: ${JSON.stringify(error)}`);
-                } else {
-                    logger.info(`Subscribed to channel requestProgress`);
-                }
-            });
-            client.subscribe('requestProgressSubscriptions', (error) => {
-                if (error) {
-                    logger.error(`Error subscribing to redis channel requestProgressSubscriptions: ${JSON.stringify(error)}`);
-                } else {
-                    logger.info(`Subscribed to channel requestProgressSubscriptions`);
-                }
+        subscriptionClient.on('connect', () => {
+            const channels = [requestProgressChannel, requestProgressSubscriptionsChannel];
+
+            channels.forEach(channel => {
+                subscriptionClient.subscribe(channel, (error) => {
+                    if (error) {
+                        logger.error(`Error subscribing to redis channel ${channel}: ${error}`);
+                    } else {
+                        logger.info(`Subscribed to channel ${channel}`);
+                    }
+                });
             });
         });
 
-        client.on('message', (channel, message) => {
-            if (channel === 'requestProgress') {
-                logger.debug(`Received message from ${channel}: ${message}`);
-                let parsedMessage;
+        subscriptionClient.on('message', (channel, message) => {
+            logger.debug(`Received message from ${channel}: ${message}`);
+            
+            let decryptedMessage = message;
 
+            if (channel === requestProgressChannel && redisEncryptionKey) {
                 try {
-                    parsedMessage = JSON.parse(message);
+                    decryptedMessage = decrypt(message, redisEncryptionKey);
                 } catch (error) {
-                    parsedMessage = message;
+                    logger.error(`Error decrypting message: ${error}`);
                 }
+            }
 
-                pubsubHandleMessage(parsedMessage);
-            } else {
-                if (channel === 'requestProgressSubscriptions') {
-                    logger.debug(`Received message from ${channel}: ${message}`);
-                    let parsedMessage;
-
-                    try {
-                        parsedMessage = JSON.parse(message);
-                    } catch (error) {
-                        parsedMessage = message;
-                    }
+            let parsedMessage = decryptedMessage;
+            try {
+                parsedMessage = JSON.parse(decryptedMessage);
+            } catch (error) {
+                logger.error(`Error parsing message: ${error}`);
+            }
 
+            switch(channel) {
+                case requestProgressChannel:
+                    pubsubHandleMessage(parsedMessage);
+                    break;
+                case requestProgressSubscriptionsChannel:
                     handleSubscription(parsedMessage);
-                }
+                    break;
+                default:
+                    logger.error(`Unsupported channel: ${channel}`);
+                    break;
             }
         });
-    }
-}
-
-
-let publisherClient;
-
-if (connectionString) {
-    logger.info(`Using Redis publish for channel(s) ${channels.join(', ')}`);
-    publisherClient = Redis.createClient(connectionString);
+    }   
 } else {
-    logger.info(`Using pubsub publish for channel ${channels[0]}`);
+    // No Redis connection, use pubsub for communication
+    logger.info(`Using pubsub publish for channel ${requestProgressChannel}`);
 }
 
 async function publishRequestProgress(data) {
     if (publisherClient) {
         try {
-            const message = JSON.stringify(data);
-            logger.debug(`Publishing message ${message} to channel ${channels[0]}`);
-            await publisherClient.publish(channels[0], message);
+            let message = JSON.stringify(data);
+            if (redisEncryptionKey) {
+                try {
+                    message = encrypt(message, redisEncryptionKey);
+                } catch (error) {
+                    logger.error(`Error encrypting message: ${error}`);
+                }
+            }
+            logger.debug(`Publishing message ${message} to channel ${requestProgressChannel}`);
+            await publisherClient.publish(requestProgressChannel, message);
         } catch (error) {
-            logger.error(`Error publishing message: ${JSON.stringify(error)}`);
+            logger.error(`Error publishing message: ${error}`);
         }
     } else {
         pubsubHandleMessage(data);
@@ -97,10 +117,10 @@ async function publishRequestProgressSubscription(data) {
     if (publisherClient) {
         try {
             const message = JSON.stringify(data);
-            logger.debug(`Publishing message ${message} to channel ${channels[1]}`);
-            await publisherClient.publish(channels[1], message);
+            logger.debug(`Publishing message ${message} to channel ${requestProgressSubscriptionsChannel}`);
+            await publisherClient.publish(requestProgressSubscriptionsChannel, message);
         } catch (error) {
-            logger.error(`Error publishing message: ${JSON.stringify(error)}`);
+            logger.error(`Error publishing message: ${error}`);
         }
     } else {
         handleSubscription(data);
@@ -113,7 +133,7 @@ function pubsubHandleMessage(data){
     try {
         pubsub.publish('REQUEST_PROGRESS', { requestProgress: data });
     } catch (error) {
-        logger.error(`Error publishing data to pubsub: ${JSON.stringify(error)}`);
+        logger.error(`Error publishing data to pubsub: ${error}`);
     }
 }
 
@@ -129,7 +149,6 @@ function handleSubscription(data){
     }
 }
 
-
 export {
-    client as subscriptionClient, publishRequestProgress, publishRequestProgressSubscription
+    subscriptionClient, publishRequestProgress, publishRequestProgressSubscription
 };

package/lib/request.js

@@ -241,7 +241,7 @@ const request = async (params, model, requestId, pathway) => {
         const response = await postRequest(params, model, requestId, pathway);
         const { error, data, cached } = response;
         if (cached) {
-            console.info(`<<< [${requestId}] served with cached response.`);
+            logger.info(`<<< [${requestId}] served with cached response.`);
         }
         if (error && error.length > 0) {
             const lastError = error[error.length - 1];

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aj-archipelago/cortex",
-  "version": "1.1.0-beta.0",
+  "version": "1.1.0",
   "description": "Cortex is a GraphQL API for AI. It provides a simple, extensible interface for using AI services from OpenAI, Azure and others.",
   "private": false,
   "repository": {

package/server/parser.js

@@ -1,3 +1,5 @@
+import logger from '../lib/logger.js';
+
 //simply trim and parse with given regex
 const regexParser = (text, regex) => {
     return text.trim().split(regex).map(s => s.trim()).filter(s => s.length);
@@ -25,7 +27,7 @@ const parseNumberedObjectList = (text, format) => {
             }
             result.push(obj);
         } catch (e) {
-            console.warn(`Failed to parse value in parseNumberedObjectList, value: ${value}, fields: ${fields}`);
+            logger.warn(`Failed to parse value in parseNumberedObjectList, value: ${value}, fields: ${fields}`);
         }
     }
 

package/server/pathwayResolver.js

@@ -229,7 +229,7 @@ class PathwayResolver {
     // Add a warning and log it
     logWarning(warning) {
         this.warnings.push(warning);
-        console.warn(warning);
+        logger.warn(warning);
     }
 
     // Here we choose how to handle long input - either summarize or chunk

package/server/plugins/modelPlugin.js

@@ -4,7 +4,7 @@ import HandleBars from '../../lib/handleBars.js';
 import { request } from '../../lib/request.js';
 import { encode } from 'gpt-3-encoder';
 import { getFirstNToken } from '../chunker.js';
-import logger from '../../lib/logger.js';
+import logger, { obscureUrlParams } from '../../lib/logger.js';
 
 const DEFAULT_MAX_TOKENS = 4096;
 const DEFAULT_MAX_RETURN_TOKENS = 256;
@@ -227,7 +227,7 @@ class ModelPlugin {
         const header = '>'.repeat(logMessage.length);
         logger.info(`${header}`);
         logger.info(`${logMessage}`);
-        logger.info(`>>> Making API request to ${url}`);
+        logger.info(`>>> Making API request to ${obscureUrlParams(url)}`);
     }
 
     logAIRequestFinished() {

package/server/plugins/openAiImagePlugin.js

@@ -24,7 +24,7 @@ class OpenAIImagePlugin extends ModelPlugin {
             requestDurationEstimator.startRequest(requestId);
             id = (await this.executeRequest(url, data, {}, { ...this.model.headers }, {}, requestId, pathway))?.id;
         } catch (error) {
-            const errMsg = `Error generating image: ${error?.message || JSON.stringify(error)}`;
+            const errMsg = `Error generating image: ${error?.message || error}`;
             logger.error(errMsg);
             return errMsg;
         }

package/server/plugins/openAiWhisperPlugin.js

@@ -224,7 +224,7 @@ class OpenAIWhisperPlugin extends ModelPlugin {
             // result = await Promise.all(mediaSplit.chunks.map(processChunk));
 
         } catch (error) {
-            const errMsg = `Transcribe error: ${error?.message || JSON.stringify(error)}`;
+            const errMsg = `Transcribe error: ${error?.message || error}`;
             logger.error(errMsg);
             return errMsg;
         }