@aj-archipelago/cortex 1.0.5 → 1.0.7

package/{graphql → server}/graphql.js

@@ -1,24 +1,29 @@
-import { createServer } from 'http';
-import {
-    ApolloServerPluginDrainHttpServer,
-    ApolloServerPluginLandingPageLocalDefault,
-} from 'apollo-server-core';
+// graphql.js
+// Setup the Apollo server and Express middleware
+
+import { ApolloServerPluginDrainHttpServer } from '@apollo/server/plugin/drainHttpServer';
+import { ApolloServerPluginLandingPageLocalDefault } from '@apollo/server/plugin/landingPage/default';
+import { ApolloServer } from '@apollo/server';
+import { expressMiddleware } from '@apollo/server/express4';
 import { makeExecutableSchema } from '@graphql-tools/schema';
 import { WebSocketServer } from 'ws';
 import { useServer } from 'graphql-ws/lib/use/ws';
 import express from 'express';
-import { ApolloServer } from 'apollo-server-express';
+import http from 'http';
 import Keyv from 'keyv';
+import cors from 'cors';
 import { KeyvAdapter } from '@apollo/utils.keyvadapter';
-import responseCachePlugin from 'apollo-server-plugin-response-cache';
+import responseCachePlugin from '@apollo/server-plugin-response-cache';
 import subscriptions from './subscriptions.js';
 import { buildLimiters } from '../lib/request.js';
 import { cancelRequestResolver } from './resolver.js';
 import { buildPathways, buildModels } from '../config.js';
 import { requestState } from './requestState.js';
+import { buildRestEndpoints } from './rest.js';
 
+// Utility functions
+// Server plugins
 const getPlugins = (config) => {
-    // server plugins
     const plugins = [
         ApolloServerPluginLandingPageLocalDefault({ embed: true }), // For local development.   
     ];
@@ -39,41 +44,8 @@ const getPlugins = (config) => {
     return { plugins, cache };
 }
 
-const buildRestEndpoints = (pathways, app, server, config) => {
-  for (const [name, pathway] of Object.entries(pathways)) {
-    // Only expose endpoints for enabled pathways that explicitly want to expose a REST endpoint
-    if (pathway.disabled || !config.get('enableRestEndpoints')) continue;
-
-    const fieldVariableDefs = pathway.typeDef(pathway).restDefinition || [];
-
-    app.post(`/rest/${name}`, async (req, res) => {
-      const variables = fieldVariableDefs.reduce((acc, variableDef) => {
-        if (Object.prototype.hasOwnProperty.call(req.body, variableDef.name)) {
-          acc[variableDef.name] = req.body[variableDef.name];
-        }
-        return acc;
-      }, {});
-
-      const variableParams = fieldVariableDefs.map(({ name, type }) => `$${name}: ${type}`).join(', ');
-      const queryArgs = fieldVariableDefs.map(({ name }) => `${name}: $${name}`).join(', ');
-
-      const query = `
-                query ${name}(${variableParams}) {
-                    ${name}(${queryArgs}) {
-                        contextId
-                        previousResult
-                        result
-                    }
-                }
-            `;
-
-      const result = await server.executeOperation({ query, variables });
-      res.json(result.data[name]);
-    });
-  }
-};
 
-//typeDefs
+// Type Definitions for GraphQL
 const getTypedefs = (pathways) => {
 
     const defaultTypeDefs = `#graphql
@@ -111,6 +83,7 @@ const getTypedefs = (pathways) => {
     return typeDefs.join('\n');
 }
 
+// Resolvers for GraphQL
 const getResolvers = (config, pathways) => {
     const resolverFunctions = {};
     for (const [name, pathway] of Object.entries(pathways)) {
@@ -118,6 +91,7 @@ const getResolvers = (config, pathways) => {
         resolverFunctions[name] = (parent, args, contextValue, info) => {
             // add shared state to contextValue
             contextValue.pathway = pathway;
+            contextValue.config = config;
             return pathway.rootResolver(parent, args, contextValue, info);
         }
     }
@@ -131,7 +105,7 @@ const getResolvers = (config, pathways) => {
     return resolvers;
 }
 
-//graphql api build factory method
+// Build the server including the GraphQL schema and REST endpoints
 const build = async (config) => {
     // First perform config build
     await buildPathways(config);
@@ -150,9 +124,9 @@ const build = async (config) => {
 
     const { plugins, cache } = getPlugins(config);
 
-    const app = express()
+    const app = express();
 
-    const httpServer = createServer(app);
+    const httpServer = http.createServer(app);
 
     // Creating the WebSocket server
     const wsServer = new WebSocketServer({
@@ -182,35 +156,69 @@ const build = async (config) => {
                         },
                     };
                 },
-            }]),
-        context: ({ req, res }) => ({ req, res, config, requestState }),
+            }
+        ]),
     });
 
     // If CORTEX_API_KEY is set, we roll our own auth middleware - usually not used if you're being fronted by a proxy
     const cortexApiKey = config.get('cortexApiKey');
+    if (cortexApiKey) {
+        app.use((req, res, next) => {
+            let providedApiKey = req.headers['cortex-api-key'] || req.query['cortex-api-key'];
+            if (!providedApiKey) {
+                providedApiKey = req.headers['authorization'];
+                providedApiKey = providedApiKey?.startsWith('Bearer ') ? providedApiKey.slice(7) : providedApiKey;
+            }
+
+            if (cortexApiKey && cortexApiKey !== providedApiKey) {
+                if (req.baseUrl === '/graphql' || req.headers['content-type'] === 'application/graphql') {
+                    res.status(401)
+                    .set('WWW-Authenticate', 'Cortex-Api-Key')
+                    .set('X-Cortex-Api-Key-Info', 'Server requires Cortex API Key')
+                    .json({
+                            errors: [
+                                {
+                                    message: 'Unauthorized',
+                                    extensions: {
+                                        code: 'UNAUTHENTICATED',
+                                    },
+                                },
+                            ],
+                        });
+                } else {
+                    res.status(401)
+                    .set('WWW-Authenticate', 'Cortex-Api-Key')
+                    .set('X-Cortex-Api-Key-Info', 'Server requires Cortex API Key')
+                    .send('Unauthorized');
+                }
+            } else {
+                next();
+            }
+        });
+    };
 
-    app.use((req, res, next) => {
-        if (cortexApiKey && req.headers.cortexApiKey !== cortexApiKey && req.query.cortexApiKey !== cortexApiKey) {
-            res.status(401).send('Unauthorized');
-        } else {
-            next();
-        }
-    });
-    
-    // Use the JSON body parser middleware for REST endpoints
+    // Parse the body for REST endpoints
     app.use(express.json());
-        
-    // add the REST endpoints
-    buildRestEndpoints(pathways, app, server, config);
 
-    // if local start server
+    // Server Startup Function
     const startServer = async () => {
         await server.start();
-        server.applyMiddleware({ app });
+        app.use(
+            '/graphql',
+
+            cors(),
+
+            expressMiddleware(server, {
+                context: async ({ req, res }) => ({ req, res, config, requestState }),
+            }),
+        );
+            
+        // add the REST endpoints
+        buildRestEndpoints(pathways, app, server, config);
 
         // Now that our HTTP server is fully set up, we can listen to it.
         httpServer.listen(config.get('PORT'), () => {
-            console.log(`🚀 Server is now running at http://localhost:${config.get('PORT')}${server.graphqlPath}`);
+            console.log(`🚀 Server is now running at http://localhost:${config.get('PORT')}/graphql`);
         });
     };
 

package/{graphql → server}/pathwayPrompter.js

@@ -6,40 +6,37 @@ import OpenAIWhisperPlugin from './plugins/openAiWhisperPlugin.js';
 import LocalModelPlugin from './plugins/localModelPlugin.js';
 import PalmChatPlugin from './plugins/palmChatPlugin.js';
 import PalmCompletionPlugin from './plugins/palmCompletionPlugin.js';
+import PalmCodeCompletionPlugin from './plugins/palmCodeCompletionPlugin.js';
 
 class PathwayPrompter {
-    constructor({ config, pathway }) {
-
-        const modelName = pathway.model || config.get('defaultModelName');
-        const model = config.get('models')[modelName];
-
-        if (!model) {
-            throw new Error(`Model ${modelName} not found in config`);
-        }
-
+    constructor(config, pathway, modelName, model) {
+        
         let plugin;
 
         switch (model.type) {
             case 'OPENAI-CHAT':
-                plugin = new OpenAIChatPlugin(config, pathway);
+                plugin = new OpenAIChatPlugin(config, pathway, modelName, model);
                 break;
             case 'AZURE-TRANSLATE':
-                plugin = new AzureTranslatePlugin(config, pathway);
+                plugin = new AzureTranslatePlugin(config, pathway, modelName, model);
                 break;
             case 'OPENAI-COMPLETION':
-                plugin = new OpenAICompletionPlugin(config, pathway);
+                plugin = new OpenAICompletionPlugin(config, pathway, modelName, model);
                 break;
-            case 'OPENAI_WHISPER':
-                plugin = new OpenAIWhisperPlugin(config, pathway);
+            case 'OPENAI-WHISPER':
+                plugin = new OpenAIWhisperPlugin(config, pathway, modelName, model);
                 break;
             case 'LOCAL-CPP-MODEL':
-                plugin = new LocalModelPlugin(config, pathway);
+                plugin = new LocalModelPlugin(config, pathway, modelName, model);
                 break;
             case 'PALM-CHAT':
-                plugin = new PalmChatPlugin(config, pathway);
+                plugin = new PalmChatPlugin(config, pathway, modelName, model);
                 break;
             case 'PALM-COMPLETION':
-                plugin = new PalmCompletionPlugin(config, pathway);
+                plugin = new PalmCompletionPlugin(config, pathway, modelName, model);
+                break;
+            case 'PALM-CODE-COMPLETION':
+                plugin = new PalmCodeCompletionPlugin(config, pathway, modelName, model);
                 break;
             default:
                 throw new Error(`Unsupported model type: ${model.type}`);

package/{graphql → server}/pathwayResolver.js

@@ -20,9 +20,31 @@ class PathwayResolver {
         this.warnings = [];
         this.requestId = uuidv4();
         this.responseParser = new PathwayResponseParser(pathway);
-        this.pathwayPrompter = new PathwayPrompter({ config, pathway });
+        this.modelName = [
+            pathway.model,
+            args?.model,
+            pathway.inputParameters?.model,
+            config.get('defaultModelName')
+            ].find(modelName => modelName && config.get('models').hasOwnProperty(modelName));
+        this.model = config.get('models')[this.modelName];
+
+        if (!this.model) {
+            throw new Error(`Model ${this.modelName} not found in config`);
+        }
+
+        const specifiedModelName = pathway.model || args?.model || pathway.inputParameters?.model;
+
+        if (this.modelName !== (specifiedModelName)) {
+            if (specifiedModelName) {
+                this.logWarning(`Specified model ${specifiedModelName} not found in config, using ${this.modelName} instead.`);
+            } else {
+                this.logWarning(`No model specified in the pathway, using ${this.modelName}.`);
+            }
+        }
+
         this.previousResult = '';
         this.prompts = [];
+        this.pathwayPrompter = new PathwayPrompter(this.config, this.pathway, this.modelName, this.model);
 
         Object.defineProperty(this, 'pathwayPrompt', {
             get() {
@@ -42,66 +64,61 @@ class PathwayResolver {
     }
 
     async asyncResolve(args) {
-        // Wait with a sleep promise for the race condition to resolve
-        // const results = await Promise.all([this.promptAndParse(args), await new Promise(resolve => setTimeout(resolve, 250))]);
-        const data = await this.promptAndParse(args);
-        // Process the results for async
-        if(args.async || typeof data === 'string') { // if async flag set or processed async and got string response
+        const responseData = await this.promptAndParse(args);
+
+        // Either we're dealing with an async request or a stream
+        if(args.async || typeof responseData === 'string') {
             const { completedCount, totalCount } = requestState[this.requestId];
-            requestState[this.requestId].data = data;
+            requestState[this.requestId].data = responseData;
             pubsub.publish('REQUEST_PROGRESS', {
                 requestProgress: {
                     requestId: this.requestId,
                     progress: completedCount / totalCount,
-                    data: JSON.stringify(data),
+                    data: JSON.stringify(responseData),
                 }
             });
-        } else { //stream
-            for (const handle of data) {
-                handle.on('data', data => {
-                    console.log(data.toString());
-                    const lines = data.toString().split('\n').filter(line => line.trim() !== '');
-                    for (const line of lines) {
-                        const message = line.replace(/^data: /, '');
-                        if (message === '[DONE]') {
-                            // Send stream finished message
-                            pubsub.publish('REQUEST_PROGRESS', {
-                                requestProgress: {
-                                    requestId: this.requestId,
-                                    data: null,
-                                    progress: 1,
-                                }
-                            });
-                            return; // Stream finished
+        } else { // stream
+            try {
+                const incomingMessage = Array.isArray(responseData) && responseData.length > 0 ? responseData[0] : responseData;
+                incomingMessage.on('data', data => {
+                    const events = data.toString().split('\n');
+            
+                    events.forEach(event => {
+                        if (event.trim() === '') return; // Skip empty lines
+            
+                        const message = event.replace(/^data: /, '');
+                        
+                        //console.log(`====================================`);
+                        //console.log(`STREAM EVENT: ${event}`);
+                        //console.log(`MESSAGE: ${message}`);
+            
+                        const requestProgress = {
+                            requestId: this.requestId,
+                            data: message,
                         }
+            
+                        if (message.trim() === '[DONE]') {
+                            requestProgress.progress = 1;
+                        }
+            
                         try {
-                            const parsed = JSON.parse(message);
-                            const result = this.pathwayPrompter.plugin.parseResponse(parsed)
-
                             pubsub.publish('REQUEST_PROGRESS', {
-                                requestProgress: {
-                                    requestId: this.requestId,
-                                    data: JSON.stringify(result)
-                                }
+                                requestProgress: requestProgress
                             });
                         } catch (error) {
                             console.error('Could not JSON parse stream message', message, error);
                         }
-                    }
+                    });
                 });
-
-                // data.on('end', () => {
-                //     console.log("stream done");
-                // });
+            } catch (error) {
+                console.error('Could not subscribe to stream', error);
             }
-            
         }
     }
 
     async resolve(args) {
+        // Either we're dealing with an async request, stream, or regular request
         if (args.async || args.stream) {
-            // Asyncronously process the request
-            // this.asyncResolve(args);
             if (!requestState[this.requestId]) {
                 requestState[this.requestId] = {}
             }
@@ -161,7 +178,7 @@ class PathwayResolver {
         }
 
         // chunk the text and return the chunks with newline separators
-        return getSemanticChunks(text, chunkTokenLength);
+        return getSemanticChunks(text, chunkTokenLength, this.pathway.inputFormat);
     }
 
     truncate(str, n) {
@@ -292,7 +309,7 @@ class PathwayResolver {
         let result = '';
 
         // If this text is empty, skip applying the prompt as it will likely be a nonsensical result
-        if (!/^\s*$/.test(text)) {
+        if (!/^\s*$/.test(text) || parameters?.file) {
             result = await this.pathwayPrompter.execute(text, { ...parameters, ...this.savedContext }, prompt, this);
         } else {
             result = text;

package/{graphql → server}/plugins/azureTranslatePlugin.js

@@ -2,8 +2,8 @@
 import ModelPlugin from './modelPlugin.js';
 
 class AzureTranslatePlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
     
     // Set up parameters specific to the Azure Translate API

package/{graphql → server}/plugins/localModelPlugin.js

@@ -4,8 +4,8 @@ import { execFileSync } from 'child_process';
 import { encode } from 'gpt-3-encoder';
 
 class LocalModelPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
 
     // if the input starts with a chatML response, just return that

package/{graphql → server}/plugins/modelPlugin.js

@@ -6,19 +6,13 @@ import { encode } from 'gpt-3-encoder';
 import { getFirstNToken } from '../chunker.js';
 
 const DEFAULT_MAX_TOKENS = 4096;
+const DEFAULT_MAX_RETURN_TOKENS = 256;
 const DEFAULT_PROMPT_TOKEN_RATIO = 0.5;
 
 class ModelPlugin {
-    constructor(config, pathway) {
-        // If the pathway specifies a model, use that, otherwise use the default
-        this.modelName = pathway.model || config.get('defaultModelName');
-        // Get the model from the config
-        this.model = config.get('models')[this.modelName];
-        // If the model doesn't exist, throw an exception
-        if (!this.model) {
-            throw new Error(`Model ${this.modelName} not found in config`);
-        }
-
+    constructor(config, pathway, modelName, model) {
+        this.modelName = modelName;
+        this.model = model;
         this.config = config;
         this.environmentVariables = config.getEnv();
         this.temperature = pathway.temperature;
@@ -143,6 +137,10 @@ class ModelPlugin {
         return (this.promptParameters.maxTokenLength ?? this.model.maxTokenLength ?? DEFAULT_MAX_TOKENS);
     }
 
+    getModelMaxReturnTokens() {
+        return (this.promptParameters.maxReturnTokens ?? this.model.maxReturnTokens ?? DEFAULT_MAX_RETURN_TOKENS);
+    }
+
     getPromptTokenRatio() {
         // TODO: Is this the right order of precedence? inputParameters should maybe be second?
         return this.promptParameters.inputParameters?.tokenRatio ?? this.promptParameters.tokenRatio ?? DEFAULT_PROMPT_TOKEN_RATIO;

package/{graphql → server}/plugins/openAiChatPlugin.js

@@ -3,8 +3,8 @@ import ModelPlugin from './modelPlugin.js';
 import { encode } from 'gpt-3-encoder';
 
 class OpenAIChatPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
 
     // convert to OpenAI messages array format if necessary
@@ -90,7 +90,7 @@ class OpenAIChatPlugin extends ModelPlugin {
     parseResponse(data) {
         const { choices } = data;
         if (!choices || !choices.length) {
-            return null;
+            return data;
         }
 
         // if we got a choices array back with more than one choice, return the whole array
@@ -108,8 +108,9 @@ class OpenAIChatPlugin extends ModelPlugin {
         const separator = `\n=== ${this.pathwayName}.${this.requestCount++} ===\n`;
         console.log(separator);
     
-        if (data && data.messages && data.messages.length > 1) {
-            data.messages.forEach((message, index) => {
+        const { stream, messages } = data;
+        if (messages && messages.length > 1) {
+            messages.forEach((message, index) => {
                 const words = message.content.split(" ");
                 const tokenCount = encode(message.content).length;
                 const preview = words.length < 41 ? message.content : words.slice(0, 20).join(" ") + " ... " + words.slice(-20).join(" ");
@@ -117,11 +118,15 @@ class OpenAIChatPlugin extends ModelPlugin {
                 console.log(`\x1b[36mMessage ${index + 1}: Role: ${message.role}, Tokens: ${tokenCount}, Content: "${preview}"\x1b[0m`);
             });
         } else {
-            console.log(`\x1b[36m${data.messages[0].content}\x1b[0m`);
+            console.log(`\x1b[36m${messages[0].content}\x1b[0m`);
         }
     
-        console.log(`\x1b[34m> ${this.parseResponse(responseData)}\x1b[0m`);
-    
+        if (stream) {
+            console.log(`\x1b[34m> Response is streaming...\x1b[0m`);
+        } else {
+            console.log(`\x1b[34m> ${this.parseResponse(responseData)}\x1b[0m`);
+        }
+
         prompt && prompt.debugInfo && (prompt.debugInfo += `${separator}${JSON.stringify(data)}`);
     }
 }

package/{graphql → server}/plugins/openAiCompletionPlugin.js

@@ -15,8 +15,8 @@ const truncatePromptIfNecessary = (text, textTokenCount, modelMaxTokenCount, tar
 }
 
 class OpenAICompletionPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
 
     // Set up parameters specific to the OpenAI Completion API
@@ -108,10 +108,16 @@ class OpenAICompletionPlugin extends ModelPlugin {
         const separator = `\n=== ${this.pathwayName}.${this.requestCount++} ===\n`;
         console.log(separator);
     
+        const stream = data.stream;
         const modelInput = data.prompt;
     
         console.log(`\x1b[36m${modelInput}\x1b[0m`);
-        console.log(`\x1b[34m> ${this.parseResponse(responseData)}\x1b[0m`);
+
+        if (stream) {
+            console.log(`\x1b[34m> Response is streaming...\x1b[0m`);
+        } else {
+            console.log(`\x1b[34m> ${this.parseResponse(responseData)}\x1b[0m`);
+        }
     
         prompt && prompt.debugInfo && (prompt.debugInfo += `${separator}${JSON.stringify(data)}`);
     }

package/{graphql → server}/plugins/openAiWhisperPlugin.js

@@ -19,6 +19,7 @@ const pipeline = promisify(stream.pipeline);
 
 
 const API_URL = config.get('whisperMediaApiUrl');
+const WHISPER_TS_API_URL  = config.get('whisperTSApiUrl');
 
 function alignSubtitles(subtitles) {
     const result = [];
@@ -74,14 +75,14 @@ const downloadFile = async (fileUrl) => {
             fs.unlink(localFilePath, () => {
                 reject(error);
             });
-            throw error;
+            //throw error;
         }
     });
 };
 
 class OpenAIWhisperPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
 
     async getMediaChunks(file, requestId) {
@@ -115,11 +116,28 @@ class OpenAIWhisperPlugin extends ModelPlugin {
 
     // Execute the request to the OpenAI Whisper API
     async execute(text, parameters, prompt, pathwayResolver) {
-        const { responseFormat } = parameters;
+        const { responseFormat, wordTimestamped } = parameters;
         const url = this.requestUrl(text);
         const params = {};
         const { modelPromptText } = this.getCompiledPrompt(text, parameters, prompt);
 
+        const processTS = async (uri) => {
+            if (wordTimestamped) {
+                if (!WHISPER_TS_API_URL) {
+                    throw new Error(`WHISPER_TS_API_URL not set for word timestamped processing`);
+                }
+
+                try {
+                    // const res = await axios.post(WHISPER_TS_API_URL, { params: { fileurl: uri } });
+                    const res = await this.executeRequest(WHISPER_TS_API_URL, {fileurl:uri},{},{});
+                    return res;
+                } catch (err) {
+                    console.log(`Error getting word timestamped data from api:`, err);
+                    throw err;
+                }
+            }
+        }
+
         const processChunk = async (chunk) => {
             try {
                 const { language, responseFormat } = parameters;
@@ -159,7 +177,6 @@ class OpenAIWhisperPlugin extends ModelPlugin {
 
         let chunks = []; // array of local file paths
         try {
-
             const uris = await this.getMediaChunks(file, requestId); // array of remote file uris
             if (!uris || !uris.length) {
                 throw new Error(`Error in getting chunks from media helper for file ${file}`);
@@ -169,7 +186,13 @@ class OpenAIWhisperPlugin extends ModelPlugin {
 
             // sequential download of chunks
             for (const uri of uris) {
-                chunks.push(await downloadFile(uri));
+                if (wordTimestamped) { // get word timestamped data 
+                    sendProgress(); // no download needed auto progress 
+                    const ts = await processTS(uri);
+                    result.push(ts);
+                } else {
+                    chunks.push(await downloadFile(uri));
+                }
                 sendProgress();
             }
 
@@ -210,7 +233,7 @@ class OpenAIWhisperPlugin extends ModelPlugin {
             }
         }
 
-        if (['srt','vtt'].includes(responseFormat)) { // align subtitles for formats
+        if (['srt','vtt'].includes(responseFormat) || wordTimestamped) { // align subtitles for formats
             return alignSubtitles(result);
         }
         return result.join(` `);

package/{graphql → server}/plugins/palmChatPlugin.js

@@ -4,8 +4,8 @@ import { encode } from 'gpt-3-encoder';
 import HandleBars from '../../lib/handleBars.js';
 
 class PalmChatPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
+    constructor(config, pathway, modelName, model) {
+        super(config, pathway, modelName, model);
     }
 
     // Convert to PaLM messages array format if necessary
@@ -92,10 +92,8 @@ class PalmChatPlugin extends ModelPlugin {
         const context = this.getCompiledContext(text, parameters, prompt.context || palmMessages.context || '');
         const examples = this.getCompiledExamples(text, parameters, prompt.examples || []);
         
-        // For PaLM right now, the max return tokens is 1024, regardless of the max context length
-        // I can't think of a time you'd want to constrain it to fewer at the moment.
-        const max_tokens = 1024//this.getModelMaxTokenLength() - tokenLength;
-    
+        const max_tokens = this.getModelMaxReturnTokens();
+        
         if (max_tokens < 0) {
             throw new Error(`Prompt is too long to successfully call the model at ${tokenLength} tokens.  The model will not be called.`);
         }