@aj-archipelago/cortex 1.1.7 → 1.1.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aj-archipelago/cortex",
-  "version": "1.1.7",
+  "version": "1.1.8",
   "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": {
@@ -35,6 +35,7 @@
     "@datastructures-js/deque": "^1.0.4",
     "@graphql-tools/schema": "^9.0.12",
     "@keyv/redis": "^2.5.4",
+    "@langchain/openai": "^0.0.24",
     "axios": "^1.3.4",
     "axios-cache-interceptor": "^1.0.1",
     "bottleneck": "^2.19.5",
@@ -42,6 +43,7 @@
     "compromise": "^14.8.1",
     "compromise-paragraphs": "^0.1.0",
     "convict": "^6.2.3",
+    "eventsource-parser": "^1.1.2",
     "express": "^4.18.2",
     "form-data": "^4.0.0",
     "google-auth-library": "^8.8.0",

package/pathways/test_langchain.mjs

@@ -2,12 +2,7 @@
 // LangChain Cortex integration test
 
 // Import required modules
-import { OpenAI } from "langchain/llms";
-//import { PromptTemplate } from "langchain/prompts";
-//import { LLMChain, ConversationChain } from "langchain/chains";
-import { initializeAgentExecutor } from "langchain/agents";
-import { SerpAPI, Calculator } from "langchain/tools";
-//import { BufferMemory } from "langchain/memory";
+import { ChatOpenAI } from "@langchain/openai";
 
 export default {
 
@@ -15,89 +10,22 @@ export default {
     resolver: async (parent, args, contextValue, _info) => {
 
         const { config } = contextValue;
-        const env = config.getEnv();
 
         // example of reading from a predefined config variable
         const openAIApiKey = config.get('openaiApiKey');
-        // example of reading straight from environment
-        const serpApiKey = env.SERPAPI_API_KEY;
 
-        const model = new OpenAI({ openAIApiKey: openAIApiKey, temperature: 0 });
-        const tools = [new SerpAPI( serpApiKey ), new Calculator()];
-
-        const executor = await initializeAgentExecutor(
-            tools,
-            model,
-            "zero-shot-react-description"
-            );
+        const model = new ChatOpenAI({ openAIApiKey: openAIApiKey, temperature: 0 });
 
         console.log(`====================`);
-        console.log("Loaded langchain agent.");
+        console.log("Loaded langchain.");
         const input = args.text;
         console.log(`Executing with input "${input}"...`);
-        const result = await executor.call({ input });
-        console.log(`Got output ${result.output}`);
-        console.log(`====================`);
-
-        return result?.output;
-    },
-
-    /*
-    // Agent test case
-    resolver: async (parent, args, contextValue, info) => {
-
-        const { config } = contextValue;
-        const openAIApiKey = config.get('openaiApiKey');
-        const serpApiKey = config.get('serpApiKey');
-
-        const model = new OpenAI({ openAIApiKey: openAIApiKey, temperature: 0 });
-        const tools = [new SerpAPI( serpApiKey ), new Calculator()];
-
-        const executor = await initializeAgentExecutor(
-            tools,
-            model,
-            "zero-shot-react-description"
-            );
-
-        console.log(`====================`);
-        console.log("Loaded langchain agent.");
-        const input = args.text;
-        console.log(`Executing with input "${input}"...`);
-        const result = await executor.call({ input });
-        console.log(`Got output ${result.output}`);
-        console.log(`====================`);
-
-        return result?.output;
-    },
-    */
-    // Simplest test case
-    /*
-    resolver: async (parent, args, contextValue, info) => {
-     
-        const { config } = contextValue;
-        const openAIApiKey = config.get('openaiApiKey');
-
-        const model = new OpenAI({ openAIApiKey: openAIApiKey, temperature: 0.9 });
-
-        const template = "What is a good name for a company that makes {product}?";
-
-        const prompt = new PromptTemplate({
-            template: template,
-            inputVariables: ["product"],
-        });
-
-        const chain = new LLMChain({ llm: model, prompt: prompt });  
-
+        const result = await model.invoke(input);
+        console.log(`Got output "${result.content}"`);
         console.log(`====================`);
-        console.log(`Calling langchain with prompt: ${prompt?.template}`);
-        console.log(`Input text: ${args.text}`);
-        const res = await chain.call({ product: args.text });
-        console.log(`Result: ${res?.text}`);
-        console.log(`====================`);     
 
-        return res?.text?.trim();
+        return result?.content;
     },
-    */
 };
 
 

package/server/modelExecutor.js

@@ -20,6 +20,7 @@ import OpenAIVisionPlugin from './plugins/openAiVisionPlugin.js';
 import GeminiChatPlugin from './plugins/geminiChatPlugin.js';
 import GeminiVisionPlugin from './plugins/geminiVisionPlugin.js';
 import AzureBingPlugin from './plugins/azureBingPlugin.js';
+import Claude3VertexPlugin from './plugins/claude3VertexPlugin.js';
 
 class ModelExecutor {
     constructor(pathway, model) {
@@ -84,6 +85,9 @@ class ModelExecutor {
             case 'AZURE-BING':
                 plugin = new AzureBingPlugin(pathway, model);
                 break;
+            case 'CLAUDE-3-VERTEX':
+                plugin = new Claude3VertexPlugin(pathway, model);
+                break;
             default:
                 throw new Error(`Unsupported model type: ${model.type}`);
         }

package/server/pathwayResolver.js

@@ -1,6 +1,5 @@
 import { ModelExecutor } from './modelExecutor.js';
 import { modelEndpoints } from '../lib/requestExecutor.js';
-// eslint-disable-next-line import/no-extraneous-dependencies
 import { v4 as uuidv4 } from 'uuid';
 import { encode } from '../lib/encodeCache.js';
 import { getFirstNToken, getLastNToken, getSemanticChunks } from './chunker.js';
@@ -11,6 +10,8 @@ import { requestState } from './requestState.js';
 import { callPathway } from '../lib/pathwayTools.js';
 import { publishRequestProgress } from '../lib/redisSubscription.js';
 import logger from '../lib/logger.js';
+// eslint-disable-next-line import/no-extraneous-dependencies
+import { createParser } from 'eventsource-parser';
 
 const modelTypesExcludedFromProgressUpdates = ['OPENAI-DALLE2', 'OPENAI-DALLE3'];
 
@@ -69,136 +70,112 @@ class PathwayResolver {
         this.pathwayPrompt = pathway.prompt;
     }
 
-    // This code handles async and streaming responses.  In either case, we use
-    // 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.
+    // This code handles async and streaming responses for either long-running
+    // tasks or streaming model responses
     async asyncResolve(args) {
-        const MAX_RETRY_COUNT = 3;
-        let attempt = 0;
         let streamErrorOccurred = false;
+        let responseData = null;
 
-        while (attempt < MAX_RETRY_COUNT) {
-            const responseData = await this.executePathway(args);
+        try {
+            responseData = await this.executePathway(args);
+        }
+        catch (error) {
+            if (!args.async) {
+                publishRequestProgress({
+                    requestId: this.requestId,
+                    progress: 1,
+                    data: '[DONE]',
+                });
+            }
+            return;
+        }
 
-            if (args.async || typeof responseData === 'string') {
-                const { completedCount, totalCount } = requestState[this.requestId];
-                requestState[this.requestId].data = responseData;
-                
-                // if model type is OPENAI-IMAGE
-                if (!modelTypesExcludedFromProgressUpdates.includes(this.model.type)) {
-                    await publishRequestProgress({
-                            requestId: this.requestId,
-                            progress: completedCount / totalCount,
-                            data: JSON.stringify(responseData),
-                    });
-                }
-            } else {
-                try {
-                    const incomingMessage = responseData;
-
-                    let messageBuffer = '';
-                    let streamEnded = false;
-
-                    const processStreamSSE = (data) => {
-                        try {
-                            //logger.info(`\n\nReceived stream data for requestId ${this.requestId}: ${data.toString()}`);
-                            let events = data.toString().split('\n');
-                            
-                            //events = "data: {\"id\":\"chatcmpl-20bf1895-2fa7-4ef9-abfe-4d142aba5817\",\"object\":\"chat.completion.chunk\",\"created\":1689303423723,\"model\":\"gpt-4\",\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":{\"error\":{\"message\":\"The server had an error while processing your request. Sorry about that!\",\"type\":\"server_error\",\"param\":null,\"code\":null}}},\"finish_reason\":null}]}\n\n".split("\n");
-
-                            for (let event of events) {
-                                if (streamErrorOccurred) break;
-                                
-                                // skip empty events
-                                if (!(event.trim() === '')) {
-                                    //logger.info(`Processing stream event for requestId ${this.requestId}: ${event}`);
-                                    messageBuffer += event.replace(/^data: /, '');
-
-                                    const requestProgress = {
-                                        requestId: this.requestId,
-                                        data: messageBuffer,
-                                    }
-
-                                    // check for end of stream or in-stream errors
-                                    if (messageBuffer.trim() === '[DONE]') {
-                                        requestProgress.progress = 1;
-                                    } else {
-                                        let parsedMessage;
-                                        try {
-                                            parsedMessage = JSON.parse(messageBuffer);
-                                            messageBuffer = '';
-                                        } catch (error) {
-                                            // incomplete stream message, try to buffer more data
-                                            return;
-                                        }
-
-                                        // error can be in different places in the message
-                                        const streamError = parsedMessage?.error || parsedMessage?.choices?.[0]?.delta?.content?.error || parsedMessage?.choices?.[0]?.text?.error;
-                                        if (streamError) {
-                                            streamErrorOccurred = true;
-                                            logger.error(`Stream error: ${streamError.message}`);
-                                            incomingMessage.off('data', processStreamSSE);
-                                            return;
-                                        }
-
-                                        // finish reason can be in different places in the message
-                                        const finishReason = parsedMessage?.choices?.[0]?.finish_reason || parsedMessage?.candidates?.[0]?.finishReason;
-                                        if (finishReason?.toLowerCase() === 'stop') {
-                                            requestProgress.progress = 1;
-                                        } else {
-                                            if (finishReason?.toLowerCase() === 'safety') {
-                                                const safetyRatings = JSON.stringify(parsedMessage?.candidates?.[0]?.safetyRatings) || '';
-                                                logger.warn(`Request ${this.requestId} was blocked by the safety filter. ${safetyRatings}`);
-                                                requestProgress.data = `\n\nResponse blocked by safety filter: ${safetyRatings}`;
-                                                requestProgress.progress = 1;
-                                            }
-                                        }
-                                    }
-
-                                    try {
-                                        if (!streamEnded) {
-                                            //logger.info(`Publishing stream message to requestId ${this.requestId}: ${message}`);
-                                            publishRequestProgress(requestProgress);
-                                            streamEnded = requestProgress.progress === 1;
-                                        }
-                                    } catch (error) {
-                                        logger.error(`Could not publish the stream message: "${messageBuffer}", ${error}`);
-                                    }
-                                }
-                            }
-                        } catch (error) {
-                            logger.error(`Could not process stream data: ${error}`);
-                        }
+        // If the response is a string, it's a regular long running response
+        if (args.async || typeof responseData === 'string') {
+            const { completedCount, totalCount } = requestState[this.requestId];
+            requestState[this.requestId].data = responseData;
+            
+            // some models don't support progress updates
+            if (!modelTypesExcludedFromProgressUpdates.includes(this.model.type)) {
+                await publishRequestProgress({
+                        requestId: this.requestId,
+                        progress: completedCount / totalCount,
+                        data: JSON.stringify(responseData),
+                });
+            }
+        // If the response is an object, it's a streaming response
+        } else {
+            try {
+                const incomingMessage = responseData;
+                let streamEnded = false;
+
+                const onParse = (event) => {
+                    let requestProgress = {
+                        requestId: this.requestId
                     };
 
-                    if (incomingMessage) {
-                        await new Promise((resolve, reject) => {
-                            incomingMessage.on('data', processStreamSSE);
-                            incomingMessage.on('end', resolve);
-                            incomingMessage.on('error', reject);
-                        });
+                    logger.debug(`Received event: ${event.type}`);
+
+                    if (event.type === 'event') {
+                        logger.debug('Received event!')
+                        logger.debug(`id: ${event.id || '<none>'}`)
+                        logger.debug(`name: ${event.name || '<none>'}`)
+                        logger.debug(`data: ${event.data}`)
+                    } else if (event.type === 'reconnect-interval') {
+                        logger.debug(`We should set reconnect interval to ${event.value} milliseconds`)
+                    }
+
+                    try {
+                        requestProgress = this.modelExecutor.plugin.processStreamEvent(event, requestProgress);
+                    } catch (error) {
+                        streamErrorOccurred = true;
+                        logger.error(`Stream error: ${error.message}`);
+                        incomingMessage.off('data', processStream);
+                        return;
+                    }
+
+                    try {
+                        if (!streamEnded && requestProgress.data) {
+                            //logger.info(`Publishing stream message to requestId ${this.requestId}: ${message}`);
+                            publishRequestProgress(requestProgress);
+                            streamEnded = requestProgress.progress === 1;
+                        }
+                    } catch (error) {
+                        logger.error(`Could not publish the stream message: "${event.data}", ${error}`);
                     }
 
-                } catch (error) {
-                    logger.error(`Could not subscribe to stream: ${error}`);
                 }
+                
+                const sseParser = createParser(onParse);
+
+                const processStream = (data) => {
+                    //logger.warn(`RECEIVED DATA: ${JSON.stringify(data.toString())}`);
+                    sseParser.feed(data.toString());
+                }
+
+                if (incomingMessage) {
+                    await new Promise((resolve, reject) => {
+                        incomingMessage.on('data', processStream);
+                        incomingMessage.on('end', resolve);
+                        incomingMessage.on('error', reject);
+                    });
+                }
+
+            } catch (error) {
+                logger.error(`Could not subscribe to stream: ${error}`);
             }
 
             if (streamErrorOccurred) {
-                attempt++;
-                logger.error(`Stream attempt ${attempt} failed. Retrying...`);
-                streamErrorOccurred = false; // Reset the flag for the next attempt
+                logger.error(`Stream read failed. Finishing stream...`);
+                publishRequestProgress({
+                    requestId: this.requestId,
+                    progress: 1,
+                    data: '[DONE]',
+                });
             } else {
                 return;
             }
         }
-        // if all retries failed, publish the stream end message
-        publishRequestProgress({
-                requestId: this.requestId,
-                progress: 1,
-                data: '[DONE]',
-        });
     }
 
     async resolve(args) {

package/server/plugins/azureBingPlugin.js

@@ -1,5 +1,6 @@
 import ModelPlugin from './modelPlugin.js';
 import logger from '../../lib/logger.js';
+import { config } from '../../config.js';
 
 class AzureBingPlugin extends ModelPlugin {
     constructor(pathway, model) {
@@ -18,6 +19,9 @@ class AzureBingPlugin extends ModelPlugin {
     }
 
     async execute(text, parameters, prompt, cortexRequest) {
+        if(!config.getEnv()["AZURE_BING_KEY"]){
+            throw new Error("AZURE_BING_KEY is not set in the environment variables!");
+        }
         const requestParameters = this.getRequestParameters(text, parameters, prompt);
 
         cortexRequest.data = requestParameters.data;

package/server/plugins/claude3VertexPlugin.js

@@ -0,0 +1,126 @@
+import OpenAIVisionPlugin from './openAiVisionPlugin.js';
+
+class Claude3VertexPlugin extends OpenAIVisionPlugin {
+
+    parseResponse(data) 
+    {
+        if (!data) {
+            return data;
+        }
+
+        const { content } = data;
+
+        // if the response is an array, return the text property of the first item
+        // if the type property is 'text'
+        if (content && Array.isArray(content) && content[0].type === 'text') {
+            return content[0].text;
+        } else {
+            return data;
+        }
+    }
+
+    // This code converts messages to the format required by the Claude Vertex API
+    convertMessagesToClaudeVertex(messages) {
+        let modifiedMessages = [];
+        let system = '';
+        let lastAuthor = '';
+
+        // Claude needs system messages in a separate field
+        const systemMessages = messages.filter(message => message.role === 'system');
+        if (systemMessages.length > 0) {
+            system = systemMessages.map(message => message.content).join('\n');
+            modifiedMessages = messages.filter(message => message.role !== 'system');
+        } else {
+            modifiedMessages = messages;
+        }
+
+        // remove any empty messages
+        modifiedMessages = modifiedMessages.filter(message => message.content);
+
+        // combine any consecutive messages from the same author
+        var combinedMessages = [];
+
+        modifiedMessages.forEach((message) => {
+        if (message.role === lastAuthor) {
+            combinedMessages[combinedMessages.length - 1].content += '\n' + message.content;
+        } else {
+            combinedMessages.push(message);
+            lastAuthor = message.role;
+        }
+        });
+
+        modifiedMessages = combinedMessages;
+
+        // Claude vertex requires an even number of messages
+        if (modifiedMessages.length % 2 === 0) {
+            modifiedMessages = modifiedMessages.slice(1);
+        }
+
+        return {
+            system,
+            modifiedMessages,
+        };
+    }
+
+    getRequestParameters(text, parameters, prompt, cortexRequest) {
+        const requestParameters = super.getRequestParameters(text, parameters, prompt, cortexRequest);
+        const { system, modifiedMessages } = this.convertMessagesToClaudeVertex(requestParameters.messages);
+        requestParameters.system = system;
+        requestParameters.messages = modifiedMessages;  
+        requestParameters.max_tokens = this.getModelMaxReturnTokens();
+        requestParameters.anthropic_version = 'vertex-2023-10-16';
+        return requestParameters;
+    }
+
+    async execute(text, parameters, prompt, cortexRequest) {
+        const requestParameters = this.getRequestParameters(text, parameters, prompt, cortexRequest);
+        const { stream } = parameters;
+
+        cortexRequest.data = { ...(cortexRequest.data || {}), ...requestParameters };
+        cortexRequest.params = {}; // query params
+        cortexRequest.stream = stream;
+        cortexRequest.url = cortexRequest.stream ? `${cortexRequest.url}:streamRawPredict` : `${cortexRequest.url}:rawPredict`;
+
+        const gcpAuthTokenHelper = this.config.get('gcpAuthTokenHelper');
+        const authToken = await gcpAuthTokenHelper.getAccessToken();
+        cortexRequest.headers.Authorization = `Bearer ${authToken}`;
+
+        return this.executeRequest(cortexRequest);
+    }
+
+    processStreamEvent(event, requestProgress) {
+        const eventData = JSON.parse(event.data);
+        switch (eventData.type) {
+            case 'message_start':
+                requestProgress.data = JSON.stringify(eventData.message);
+                break;
+            case 'content_block_start':
+                break;
+            case 'ping':
+                break;
+            case 'content_block_delta':
+                if (eventData.delta.type === 'text_delta') {
+                    requestProgress.data = JSON.stringify(eventData.delta.text);
+                }
+                break;
+            case 'content_block_stop':
+                break;
+            case 'message_delta':
+                break;
+            case 'message_stop':
+                requestProgress.data = '[DONE]';
+                requestProgress.progress = 1;
+                break;
+            case 'error':
+                requestProgress.data = `\n\n*** ${eventData.error.message || eventData.error} ***`;
+                requestProgress.progress = 1;
+                break;
+        }
+
+        return requestProgress;
+
+    }
+
+}
+
+export default Claude3VertexPlugin;

package/server/plugins/geminiChatPlugin.js

@@ -5,8 +5,18 @@ import logger from '../../lib/logger.js';
 const mergeResults = (data) => {
     let output = '';
     let safetyRatings = [];
+    const RESPONSE_BLOCKED = 'The response was blocked because the input or response potentially violates policies. Try rephrasing the prompt or adjusting the parameter settings.';
 
     for (let chunk of data) {
+        const { promptfeedback } = chunk;
+        if (promptfeedback) {
+            const { blockReason } = promptfeedback;
+            if (blockReason) {
+                logger.warn(`Response blocked due to prompt feedback: ${blockReason}`);
+                return {mergedResult: RESPONSE_BLOCKED, safetyRatings: safetyRatings};
+            }
+        }
+
         const { candidates } = chunk;
         if (!candidates || !candidates.length) {
             continue;
@@ -15,7 +25,8 @@ const mergeResults = (data) => {
         // If it was blocked, return the blocked message
         if (candidates[0].safetyRatings.some(rating => rating.blocked)) {
             safetyRatings = candidates[0].safetyRatings;
-            return {mergedResult: 'The response was blocked because the input or response potentially violates policies. Try rephrasing the prompt or adjusting the parameter settings.', safetyRatings: safetyRatings};
+            logger.warn(`Response blocked due to safety ratings: ${JSON.stringify(safetyRatings, null, 2)}`);
+            return {mergedResult: RESPONSE_BLOCKED, safetyRatings: safetyRatings};
         }
 
         // Append the content of the first part of the first candidate to the output

package/server/plugins/modelPlugin.js

@@ -236,8 +236,11 @@ class ModelPlugin {
 
     getLength(data) {
         const isProd = config.get('env') === 'production';
-        const length = isProd ? data.length : encode(data).length;
-        const units = isProd ? 'characters' : 'tokens';
+        let length = 0;
+        let units = isProd ? 'characters' : 'tokens';
+        if (data) {
+           length = isProd ? data.length : encode(data).length;
+        }
         return {length, units};
     }
 
@@ -288,6 +291,42 @@ class ModelPlugin {
         }
     }
 
+    processStreamEvent(event, requestProgress) {
+        // check for end of stream or in-stream errors
+        if (event.data.trim() === '[DONE]') {
+            requestProgress.progress = 1;
+        } else {
+            let parsedMessage;
+            try {
+                parsedMessage = JSON.parse(event.data);
+                requestProgress.data = event.data;
+            } catch (error) {
+                throw new Error(`Could not parse stream data: ${error}`);
+            }
+
+            // error can be in different places in the message
+            const streamError = parsedMessage?.error || parsedMessage?.choices?.[0]?.delta?.content?.error || parsedMessage?.choices?.[0]?.text?.error;
+            if (streamError) {
+                throw new Error(streamError);
+            }
+
+            // finish reason can be in different places in the message
+            const finishReason = parsedMessage?.choices?.[0]?.finish_reason || parsedMessage?.candidates?.[0]?.finishReason;
+            if (finishReason?.toLowerCase() === 'stop') {
+                requestProgress.progress = 1;
+            } else {
+                if (finishReason?.toLowerCase() === 'safety') {
+                    const safetyRatings = JSON.stringify(parsedMessage?.candidates?.[0]?.safetyRatings) || '';
+                    logger.warn(`Request ${this.requestId} was blocked by the safety filter. ${safetyRatings}`);
+                    requestProgress.data = `\n\nResponse blocked by safety filter: ${safetyRatings}`;
+                    requestProgress.progress = 1;
+                }
+            }
+        }
+        return requestProgress;
+    }
+
+
 }
 
 export default ModelPlugin;

package/server/plugins/openAiChatPlugin.js

@@ -87,6 +87,7 @@ class OpenAIChatPlugin extends ModelPlugin {
 
     // Parse the response from the OpenAI Chat API
     parseResponse(data) {
+        if(!data) return "";
         const { choices } = data;
         if (!choices || !choices.length) {
             return data;

package/server/plugins/openAiWhisperPlugin.js

@@ -100,7 +100,13 @@ function alignSubtitles(subtitles, format) {
     const result = [];
 
     function preprocessStr(str) {
-        return str.trim().replace(/(\n\n)(?!\n)/g, '\n\n\n');
+        try{
+            if(!str) return '';
+            return str.trim().replace(/(\n\n)(?!\n)/g, '\n\n\n');
+        }catch(e){
+            logger.error(`An error occurred in content text preprocessing: ${e}`);
+            return '';
+        }
     }
 
     function shiftSubtitles(subtitle, shiftOffset) {

package/server/plugins/palmChatPlugin.js

@@ -14,6 +14,9 @@ class PalmChatPlugin extends ModelPlugin {
         let modifiedMessages = [];
         let lastAuthor = '';
     
+        // remove any empty messages
+        messages = messages.filter(message => message.content);
+
         messages.forEach(message => {
             const { role, author, content } = message;
     
@@ -153,7 +156,7 @@ class PalmChatPlugin extends ModelPlugin {
     parseResponse(data) {
         const { predictions } = data;
         if (!predictions || !predictions.length) {
-            return null;
+            return data;
         }
     
         // Get the candidates array from the first prediction

package/server/rest.js

@@ -148,6 +148,10 @@ const processIncomingStream = (requestId, res, jsonResponse) => {
                 } else if (messageJson.candidates) {
                     const { content, finishReason } = messageJson.candidates[0];
                     fillJsonResponse(jsonResponse, content.parts[0].text, finishReason);
+                } else if (messageJson.content) {
+                    const text = messageJson.content?.[0]?.text || '';
+                    const finishReason = messageJson.stop_reason;
+                    fillJsonResponse(jsonResponse, text, finishReason);
                 } else {
                     fillJsonResponse(jsonResponse, messageJson, null);
                 }