@aj-archipelago/cortex 1.1.2 → 1.1.3

package/lib/redisSubscription.js

@@ -46,7 +46,7 @@ if (connectionString) {
             channels.forEach(channel => {
                 subscriptionClient.subscribe(channel, (error) => {
                     if (error) {
-                        logger.error(`Error subscribing to redis channel ${channel}: ${error}`);
+                        logger.error(`Error subscribing to Redis channel ${channel}: ${error}`);
                     } else {
                         logger.info(`Subscribed to channel ${channel}`);
                     }
@@ -55,26 +55,22 @@ if (connectionString) {
         });
 
         subscriptionClient.on('message', (channel, message) => {
-            logger.debug(`Received message from ${channel}: ${message}`);
+            logger.debug(`Received message from Redis channel ${channel}: ${message}`);
             
-            let decryptedMessage;
-
-            if (channel === requestProgressChannel && redisEncryptionKey) {
-                try {
-                    decryptedMessage = decrypt(message, redisEncryptionKey);
-                } catch (error) {
-                    logger.error(`Error decrypting message: ${error}`);
-                }
-            }
-
-            decryptedMessage = decryptedMessage || message;
-
             let parsedMessage;
 
             try {
-                parsedMessage = JSON.parse(decryptedMessage);
+                parsedMessage = JSON.parse(message);
             } catch (error) {
-                logger.error(`Error parsing message: ${error}`);
+                if (channel === requestProgressChannel && redisEncryptionKey) {
+                    try {
+                        parsedMessage = JSON.parse(decrypt(message, redisEncryptionKey));
+                    } catch (error) {
+                        logger.error(`Error parsing or decrypting message: ${error}`);
+                    }
+                } else {
+                    logger.error(`Error parsing message: ${error}`);
+                }
             }
 
             switch(channel) {
@@ -95,8 +91,8 @@ if (connectionString) {
     logger.info(`Using pubsub publish for channel ${requestProgressChannel}`);
 }
 
-async function publishRequestProgress(data, useRedis = true) {
-    if (publisherClient && useRedis) {
+async function publishRequestProgress(data) {
+    if (publisherClient && requestState?.[data?.requestId]?.useRedis) {
         try {
             let message = JSON.stringify(data);
             if (redisEncryptionKey) {
@@ -106,10 +102,10 @@ async function publishRequestProgress(data, useRedis = true) {
                     logger.error(`Error encrypting message: ${error}`);
                 }
             }
-            logger.debug(`Publishing message ${message} to channel ${requestProgressChannel}`);
+            logger.debug(`Publishing request progress ${message} to Redis channel ${requestProgressChannel}`);
             await publisherClient.publish(requestProgressChannel, message);
         } catch (error) {
-            logger.error(`Error publishing message: ${error}`);
+            logger.error(`Error publishing request progress to Redis: ${error}`);
         }
     } else {
         pubsubHandleMessage(data);
@@ -119,11 +115,30 @@ async function publishRequestProgress(data, useRedis = true) {
 async function publishRequestProgressSubscription(data) {
     if (publisherClient) {
         try {
-            const message = JSON.stringify(data);
-            logger.debug(`Publishing message ${message} to channel ${requestProgressSubscriptionsChannel}`);
-            await publisherClient.publish(requestProgressSubscriptionsChannel, message);
+            const requestIds = data;
+            const idsToForward = [];
+            // If any of these requests belong to this instance, we can just start and handle them locally
+            for (const requestId of requestIds) {
+                if (requestState[requestId]) {
+                    if (!requestState[requestId].started) {
+                        requestState[requestId].started = true;
+                        requestState[requestId].useRedis = false;
+                        logger.info(`Starting local execution for registered async request: ${requestId}`);
+                        const { resolver, args } = requestState[requestId];
+                        resolver(args, false);
+                    }
+                } else {
+                    idsToForward.push(requestId);
+                }
+            }
+
+            if (idsToForward.length > 0) {
+                const message = JSON.stringify(idsToForward);
+                logger.debug(`Sending subscription request(s) to channel ${requestProgressSubscriptionsChannel} for remote execution: ${message}`);
+                await publisherClient.publish(requestProgressSubscriptionsChannel, message);
+            }
         } catch (error) {
-            logger.error(`Error publishing message: ${error}`);
+            logger.error(`Error handling subscription: ${error}`);
         }
     } else {
         handleSubscription(data);
@@ -132,11 +147,11 @@ async function publishRequestProgressSubscription(data) {
 
 function pubsubHandleMessage(data){
     const message = JSON.stringify(data);
-    logger.debug(`Publishing message to pubsub: ${message}`);
+    logger.debug(`Publishing request progress to local subscribers: ${message}`);
     try {
         pubsub.publish('REQUEST_PROGRESS', { requestProgress: data });
     } catch (error) {
-        logger.error(`Error publishing data to pubsub: ${error}`);
+        logger.error(`Error publishing request progress to local subscribers: ${error}`);
     }
 }
 
@@ -145,7 +160,8 @@ function handleSubscription(data){
     for (const requestId of requestIds) {
         if (requestState[requestId] && !requestState[requestId].started) {
             requestState[requestId].started = true;
-            logger.info(`Subscription starting async requestProgress, requestId: ${requestId}`);
+            requestState[requestId].useRedis = true;
+            logger.info(`Starting execution for registered async request: ${requestId}`);
             const { resolver, args } = requestState[requestId];
             resolver(args);
         }

package/lib/request.js

@@ -46,7 +46,6 @@ const buildLimiters = (config) => {
         if (connection) {
             limiterOptions.id = `${cortexId}-${name}-limiter`; // Unique id for each limiter
             limiterOptions.connection = connection;  // Shared Redis connection
-            limiterOptions.clearDatastore = true;    // Clear Redis datastore on startup
         }
 
         limiters[name] = new Bottleneck(limiterOptions);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aj-archipelago/cortex",
-  "version": "1.1.2",
+  "version": "1.1.3",
   "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/pathwayResolver.js

@@ -70,7 +70,7 @@ class PathwayResolver {
     // the graphql subscription to send progress updates to the client.  Most of 
     // the time the client will be an external client, but it could also be the
     // Cortex REST api code.
-    async asyncResolve(args, useRedis = true) {
+    async asyncResolve(args) {
         const MAX_RETRY_COUNT = 3;
         let attempt = 0;
         let streamErrorOccurred = false;
@@ -88,7 +88,7 @@ class PathwayResolver {
                             requestId: this.requestId,
                             progress: completedCount / totalCount,
                             data: JSON.stringify(responseData),
-                    }, useRedis);
+                    });
                 }
             } else {
                 try {
@@ -140,7 +140,7 @@ class PathwayResolver {
 
                                     try {
                                         //logger.info(`Publishing stream message to requestId ${this.requestId}: ${message}`);
-                                        publishRequestProgress(requestProgress, useRedis);
+                                        publishRequestProgress(requestProgress);
                                     } catch (error) {
                                         logger.error(`Could not publish the stream message: "${messageBuffer}", ${error}`);
                                     }

package/server/rest.js

@@ -167,11 +167,10 @@ const processIncomingStream = (requestId, res, jsonResponse) => {
     // Fire the resolver for the async requestProgress
     logger.info(`Rest Endpoint starting async requestProgress, requestId: ${requestId}`);
     const { resolver, args } = requestState[requestId];
-    // The false here means never use a Redis subscription channel
-    // to handle these streaming messages. This is because we are
-    // guaranteed in this case that the stream is going to the same
-    // client.
-    resolver(args, false);
+    requestState[requestId].useRedis = false;
+    requestState[requestId].started = true;
+
+    resolver(args);
 
     return subscription;
   

package/server/subscriptions.js

@@ -1,11 +1,13 @@
 import pubsub from './pubsub.js';
 import { withFilter } from 'graphql-subscriptions';
 import { publishRequestProgressSubscription } from '../lib/redisSubscription.js';
+import logger from '../lib/logger.js';
 
 const subscriptions = {
     requestProgress: {
         subscribe: withFilter(
             (_, args, __, _info) => {
+                logger.debug(`Client requested subscription for request ids: ${args.requestIds}`);
                 publishRequestProgressSubscription(args.requestIds);
                 return pubsub.asyncIterator(['REQUEST_PROGRESS'])
             },