@aj-archipelago/cortex 1.3.22 → 1.3.24

package/server/rest.js

@@ -6,6 +6,22 @@ import { requestState } from './requestState.js';
 import { v4 as uuidv4 } from 'uuid';
 import logger from '../lib/logger.js';
 import { getSingleTokenChunks } from './chunker.js';
+import axios from 'axios';
+
+const getOllamaModels = async (ollamaUrl) => {
+    try {
+        const response = await axios.get(`${ollamaUrl}/api/tags`);
+        return response.data.models.map(model => ({
+            id: `ollama-${model.name}`,
+            object: 'model',
+            owned_by: 'ollama',
+            permission: ''
+        }));
+    } catch (error) {
+        logger.error(`Error fetching Ollama models: ${error.message}`);
+        return [];
+    }
+};
 
 const chunkTextIntoTokens = (() => {
     let partialToken = '';
@@ -282,7 +298,14 @@ function buildRestEndpoints(pathways, app, server, config) {
         // Create OpenAI compatible endpoints
         app.post('/v1/completions', async (req, res) => {
             const modelName = req.body.model || 'gpt-3.5-turbo';
-            const pathwayName = openAICompletionModels[modelName] || openAICompletionModels['*'];
+            let pathwayName;
+
+            if (modelName.startsWith('ollama-')) {
+                pathwayName = 'sys_ollama_completion';
+                req.body.ollamaModel = modelName.replace('ollama-', '');
+            } else {
+                pathwayName = openAICompletionModels[modelName] || openAICompletionModels['*'];
+            }
 
             if (!pathwayName) {
                 res.status(404).json({
@@ -318,7 +341,6 @@ function buildRestEndpoints(pathways, app, server, config) {
             if (Boolean(req.body.stream)) {
                 jsonResponse.id = `cmpl-${resultText}`;
                 jsonResponse.choices[0].finish_reason = null;
-                //jsonResponse.object = "text_completion.chunk";
 
                 processIncomingStream(resultText, res, jsonResponse, pathway);
             } else {
@@ -330,7 +352,14 @@ function buildRestEndpoints(pathways, app, server, config) {
         
         app.post('/v1/chat/completions', async (req, res) => {
             const modelName = req.body.model || 'gpt-3.5-turbo';
-            const pathwayName = openAIChatModels[modelName] || openAIChatModels['*'];
+            let pathwayName;
+
+            if (modelName.startsWith('ollama-')) {
+                pathwayName = 'sys_ollama_chat';
+                req.body.ollamaModel = modelName.replace('ollama-', '');
+            } else {
+                pathwayName = openAIChatModels[modelName] || openAIChatModels['*'];
+            }
 
             if (!pathwayName) {
                 res.status(404).json({
@@ -385,8 +414,11 @@ function buildRestEndpoints(pathways, app, server, config) {
         app.get('/v1/models', async (req, res) => {
             const openAIModels = { ...openAIChatModels, ...openAICompletionModels };
             const defaultModelId = 'gpt-3.5-turbo';
+            let models = [];
 
-            const models = Object.entries(openAIModels)
+            // Get standard OpenAI-compatible models, filtering out our internal pathway models
+            models = Object.entries(openAIModels)
+                .filter(([modelId]) => !['ollama-chat', 'ollama-completion'].includes(modelId))
                 .map(([modelId]) => {
                     if (modelId.includes('*')) {
                         modelId = defaultModelId;
@@ -397,7 +429,16 @@ function buildRestEndpoints(pathways, app, server, config) {
                         owned_by: 'openai',
                         permission: '',
                     };
-                })
+                });
+
+            // Get Ollama models if configured
+            if (config.get('ollamaUrl')) {
+                const ollamaModels = await getOllamaModels(config.get('ollamaUrl'));
+                models = [...models, ...ollamaModels];
+            }
+
+            // Filter out duplicates and sort
+            models = models
                 .filter((model, index, self) => {
                     return index === self.findIndex((m) => m.id === model.id);
                 })

package/tests/multimodal_conversion.test.js

@@ -271,12 +271,18 @@ test('Pathological cases', async (t) => {
 
     t.is(geminiSystem15.parts[0].text, 'You are a helpful assistant.');
     t.is(geminiSystem15.parts[1].text, 'You are also very knowledgeable.');
+
     t.is(geminiMessages15.length, 3);
+    // First user message combines "Hello" and "Another greeting"
     t.is(geminiMessages15[0].role, 'user');
     t.is(geminiMessages15[0].parts[0].text, 'Hello');
     t.is(geminiMessages15[0].parts[1].text, 'Another greeting');
+    
+    // Assistant message "Hi there!"
     t.is(geminiMessages15[1].role, 'assistant');
     t.is(geminiMessages15[1].parts[0].text, 'Hi there!');
+     
+    // Final user message combines "How are you?", image content, and "Another question"
     t.is(geminiMessages15[2].role, 'user');
     t.is(geminiMessages15[2].parts[0].text, 'How are you?');
     t.is(geminiMessages15[2].parts[1].text, 'What\'s this?');
@@ -310,6 +316,79 @@ test('Empty message array', async (t) => {
     t.is(geminiMessages15.length, 0);
 });
 
+// Test simple string array content
+test('Simple string array content', async (t) => {
+    const { gemini15 } = createPlugins();
+    
+    const messages = [
+        { role: 'user', content: "Initial message" },
+        { role: 'assistant', content: [
+            "\"Searchin' for my lost shaker of salt...\"\n",
+        ]},
+        { role: 'user', content: [
+            "Here's another simple string in an array",
+        ]}
+    ];
+
+    const { modifiedMessages } = gemini15.convertMessagesToGemini(messages);
+
+    t.is(modifiedMessages.length, 3);
+    t.is(modifiedMessages[0].role, 'user');
+    t.is(modifiedMessages[0].parts.length, 1);
+    t.is(modifiedMessages[0].parts[0].text, "Initial message");
+    t.is(modifiedMessages[1].role, 'assistant');
+    t.is(modifiedMessages[1].parts.length, 1);
+    t.is(modifiedMessages[1].parts[0].text, "\"Searchin' for my lost shaker of salt...\"\n");
+    t.is(modifiedMessages[2].role, 'user');
+    t.is(modifiedMessages[2].parts.length, 1);
+    t.is(modifiedMessages[2].parts[0].text, "Here's another simple string in an array");
+});
+
+// Test string-encoded multimodal content
+test('String-encoded multimodal content', async (t) => {
+    const { gemini15 } = createPlugins();
+    
+    const messages = [
+        { role: 'user', content: [
+            JSON.stringify({
+                type: 'text',
+                text: 'What is in this image?'
+            }),
+            JSON.stringify({
+                type: 'image_url',
+                image_url: { url: 'gs://my-bucket/image.jpg' }
+            })
+        ]},
+        { role: 'assistant', content: [
+            JSON.stringify({
+                type: 'text',
+                text: 'I see a cat.'
+            })
+        ]},
+        { role: 'user', content: [
+            JSON.stringify({
+                type: 'text',
+                text: 'Is it a big cat?'
+            })
+        ]}
+    ];
+
+    const { modifiedMessages } = gemini15.convertMessagesToGemini(messages);
+
+    t.is(modifiedMessages.length, 3);
+    t.is(modifiedMessages[0].role, 'user');
+    t.is(modifiedMessages[0].parts.length, 2);
+    t.is(modifiedMessages[0].parts[0].text, 'What is in this image?');
+    t.true('fileData' in modifiedMessages[0].parts[1]);
+    t.is(modifiedMessages[0].parts[1].fileData.fileUri, 'gs://my-bucket/image.jpg');
+    t.is(modifiedMessages[1].role, 'assistant');
+    t.is(modifiedMessages[1].parts.length, 1);
+    t.is(modifiedMessages[1].parts[0].text, 'I see a cat.');
+    t.is(modifiedMessages[2].role, 'user');
+    t.is(modifiedMessages[2].parts.length, 1);
+    t.is(modifiedMessages[2].parts[0].text, 'Is it a big cat?');
+});
+
 // Test messages with only system messages
 test('Only system messages', async (t) => {
     const { openai, claude, gemini, gemini15 } = createPlugins();
@@ -417,3 +496,93 @@ test('Gemini 1.5 image URL edge cases', t => {
     // Verify we only have one part (the text)
     t.is(modifiedMessages[0].parts.length, 1, 'Should only have the text part');
 });
+
+// Test multiple images in single message for Claude
+test('Multiple images in single Claude message', async (t) => {
+    const { claude } = createPlugins();
+    
+    const multiImageMessage = [
+        { role: 'user', content: [
+            { type: 'text', text: 'Compare these images:' },
+            { type: 'image_url', image_url: { url: sampleBase64Image } },
+            { type: 'text', text: 'with this one:' },
+            { type: 'image_url', image_url: { url: sampleBase64Image } },
+            { type: 'image_url', gcs: 'gs://cortex-bucket/image.jpg' }
+        ]}
+    ];
+
+    const { modifiedMessages } = await claude.convertMessagesToClaudeVertex(multiImageMessage);
+
+    t.is(modifiedMessages.length, 1);
+    t.is(modifiedMessages[0].role, 'user');
+    t.is(modifiedMessages[0].content.length, 4);
+    t.is(modifiedMessages[0].content[0].text, 'Compare these images:');
+    t.true(modifiedMessages[0].content[1].source.type === 'base64');
+    t.is(modifiedMessages[0].content[2].text, 'with this one:');
+    t.true(modifiedMessages[0].content[3].source.type === 'base64');
+});
+
+// Test conversation history with mixed image types
+test('Conversation history with mixed image types', async (t) => {
+    const { claude, gemini15 } = createPlugins();
+    
+    const conversationHistory = [
+        { role: 'system', content: 'You are a visual analysis assistant.' },
+        { role: 'user', content: [
+            { type: 'text', text: 'What\'s in this image?' },
+            { type: 'image_url', image_url: { url: sampleBase64Image } }
+        ]},
+        { role: 'assistant', content: 'I see a landscape.' },
+        { role: 'user', content: [
+            { type: 'text', text: 'Compare it with this:' },
+            { type: 'image_url', gcs: 'gs://cortex-bucket/image2.jpg' }
+        ]},
+        { role: 'assistant', content: 'The second image shows a different scene.' },
+        { role: 'user', content: 'Which one do you prefer?' }
+    ];
+
+    // Test Claude conversion
+    const { system: claudeSystem, modifiedMessages: claudeMessages } = await claude.convertMessagesToClaudeVertex(conversationHistory);
+
+    t.is(claudeSystem, 'You are a visual analysis assistant.');
+    t.is(claudeMessages.length, 5);
+    t.is(claudeMessages[1].content[0].text, 'I see a landscape.');
+    t.is(claudeMessages[3].content[0].text, 'The second image shows a different scene.');
+    t.is(claudeMessages[4].content[0].text, 'Which one do you prefer?');
+
+    // Test Gemini 1.5 conversion
+    const { system: geminiSystem15, modifiedMessages: geminiMessages15 } = gemini15.convertMessagesToGemini(conversationHistory);
+
+    t.is(geminiSystem15.parts[0].text, 'You are a visual analysis assistant.');
+    t.is(geminiMessages15.length, 5);
+    t.true('inlineData' in geminiMessages15[0].parts[1]);
+    t.is(geminiMessages15[1].parts[0].text, 'I see a landscape.');
+    t.true('fileData' in geminiMessages15[2].parts[1]);
+    t.is(geminiMessages15[2].parts[1].fileData.fileUri, 'gs://cortex-bucket/image2.jpg');
+    t.is(geminiMessages15[3].parts[0].text, 'The second image shows a different scene.');
+    t.is(geminiMessages15[4].parts[0].text, 'Which one do you prefer?');
+});
+
+// Test handling of large images
+test('Large image handling', async (t) => {
+    const { claude, gemini15 } = createPlugins();
+    
+    // Create a large base64 string (>10MB)
+    const largeSampleImage = 'data:image/jpeg;base64,' + 'A'.repeat(10 * 1024 * 1024);
+    
+    const largeImageMessage = [
+        { role: 'user', content: [
+            { type: 'text', text: 'Check this large image:' },
+            { type: 'image_url', image_url: { url: largeSampleImage } }
+        ]}
+    ];
+
+    // Both Claude and Gemini should handle or reject oversized images gracefully
+    const { modifiedMessages: claudeMessages } = await claude.convertMessagesToClaudeVertex(largeImageMessage);
+    const { modifiedMessages: geminiMessages } = gemini15.convertMessagesToGemini(largeImageMessage);
+
+    // Verify both models handle the oversized image appropriately
+    // (The exact behavior - rejection vs. compression - should match the model's specifications)
+    t.is(claudeMessages[0].content[0].text, 'Check this large image:');
+    t.is(geminiMessages[0].parts[0].text, 'Check this large image:');
+});

package/tests/openai_api.test.js

@@ -48,7 +48,7 @@ test('POST /completions', async (t) => {
 test('POST /chat/completions', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-4o',
       messages: [{ role: 'user', content: 'Hello!' }],
       stream: false,
     },
@@ -63,7 +63,7 @@ test('POST /chat/completions', async (t) => {
 test('POST /chat/completions with multimodal content', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'claude-3.5-sonnet',
+      model: 'gpt-4o',
       messages: [{
         role: 'user',
         content: [
@@ -153,7 +153,7 @@ test('POST SSE: /v1/completions should send a series of events and a [DONE] even
 
 test('POST SSE: /v1/chat/completions should send a series of events and a [DONE] event', async (t) => {
     const payload = {
-        model: 'gpt-3.5-turbo',
+        model: 'gpt-4o',
         messages: [
         {
             role: 'user',
@@ -177,7 +177,7 @@ test('POST SSE: /v1/chat/completions should send a series of events and a [DONE]
 
 test('POST SSE: /v1/chat/completions with multimodal content should send a series of events and a [DONE] event', async (t) => {
     const payload = {
-        model: 'claude-3.5-sonnet',
+        model: 'gpt-4o',
         messages: [{
           role: 'user',
           content: [
@@ -213,7 +213,7 @@ test('POST SSE: /v1/chat/completions with multimodal content should send a serie
 test('POST /chat/completions should handle multimodal content for non-multimodal model', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-4o',
       messages: [{
         role: 'user',
         content: [
@@ -242,7 +242,7 @@ test('POST /chat/completions should handle multimodal content for non-multimodal
 
 test('POST SSE: /v1/chat/completions should handle streaming multimodal content for non-multimodal model', async (t) => {
   const payload = {
-    model: 'gpt-3.5-turbo',
+    model: 'gpt-4o',
     messages: [{
       role: 'user',
       content: [
@@ -282,7 +282,7 @@ test('POST SSE: /v1/chat/completions should handle streaming multimodal content
 test('POST /chat/completions should handle malformed multimodal content', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'claude-3.5-sonnet',
+      model: 'gpt-4o',
       messages: [{
         role: 'user',
         content: [
@@ -310,7 +310,7 @@ test('POST /chat/completions should handle malformed multimodal content', async
 test('POST /chat/completions should handle invalid image data', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'claude-3.5-sonnet',
+      model: 'gpt-4o',
       messages: [{
         role: 'user',
         content: [
@@ -361,7 +361,7 @@ test('POST /completions should handle model parameters', async (t) => {
 test('POST /chat/completions should handle function calling', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-4o',
       messages: [{ role: 'user', content: 'What is the weather in Boston?' }],
       functions: [{
         name: 'get_weather',
@@ -401,7 +401,7 @@ test('POST /chat/completions should handle function calling', async (t) => {
 test('POST /chat/completions should validate response format', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-4o',
       messages: [{ role: 'user', content: 'Hello!' }],
       stream: false,
     },
@@ -426,7 +426,7 @@ test('POST /chat/completions should validate response format', async (t) => {
 test('POST /chat/completions should handle system messages', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-4o',
       messages: [
         { role: 'system', content: 'You are a helpful assistant.' },
         { role: 'user', content: 'Hello!' }
@@ -443,10 +443,29 @@ test('POST /chat/completions should handle system messages', async (t) => {
 });
 
 test('POST /chat/completions should handle errors gracefully', async (t) => {
+  const error = await t.throwsAsync(
+    () => got.post(`${API_BASE}/chat/completions`, {
+      json: {
+        // Missing required model field
+        messages: [{ role: 'user', content: 'Hello!' }],
+      },
+      responseType: 'json',
+    })
+  );
+  
+  t.is(error.response.statusCode, 404);
+});
+
+test('POST /chat/completions should handle token limits', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      // Missing required model field
-      messages: [{ role: 'user', content: 'Hello!' }],
+      model: 'gpt-4o',
+      messages: [{ 
+        role: 'user', 
+        content: 'Hello!'.repeat(5000) // Very long message
+      }],
+      max_tokens: 100,
+      stream: false,
     },
     responseType: 'json',
   });
@@ -455,17 +474,16 @@ test('POST /chat/completions should handle errors gracefully', async (t) => {
   t.is(response.body.object, 'chat.completion');
   t.true(Array.isArray(response.body.choices));
   t.truthy(response.body.choices[0].message.content);
-});
+});  
 
-test('POST /chat/completions should handle token limits', async (t) => {
+test('POST /chat/completions should return complete responses from gpt-4o', async (t) => {
   const response = await got.post(`${API_BASE}/chat/completions`, {
     json: {
-      model: 'gpt-3.5-turbo',
-      messages: [{ 
-        role: 'user', 
-        content: 'Hello!'.repeat(5000) // Very long message
-      }],
-      max_tokens: 100,
+      model: 'gpt-4o',
+      messages: [
+        { role: 'system', content: 'You are a helpful assistant. Always end your response with the exact string "END_MARKER_XYZ".' },
+        { role: 'user', content: 'Say hello and explain why complete responses matter.' }
+      ],
       stream: false,
     },
     responseType: 'json',
@@ -474,6 +492,8 @@ test('POST /chat/completions should handle token limits', async (t) => {
   t.is(response.statusCode, 200);
   t.is(response.body.object, 'chat.completion');
   t.true(Array.isArray(response.body.choices));
-  t.truthy(response.body.choices[0].message.content);
-});  
+  console.log('GPT-4o Response:', JSON.stringify(response.body.choices[0].message.content));
+  const content = response.body.choices[0].message.content;
+  t.regex(content, /END_MARKER_XYZ$/);
+}); 
   

package/tests/streaming.test.js

@@ -0,0 +1,197 @@
+import test from 'ava';
+import serverFactory from '../index.js';
+import { PathwayResolver } from '../server/pathwayResolver.js';
+import OpenAIChatPlugin from '../server/plugins/openAiChatPlugin.js';
+import GeminiChatPlugin from '../server/plugins/geminiChatPlugin.js';
+import Gemini15ChatPlugin from '../server/plugins/gemini15ChatPlugin.js';
+import Claude3VertexPlugin from '../server/plugins/claude3VertexPlugin.js';
+import { config } from '../config.js';
+
+let testServer;
+
+test.before(async () => {
+  process.env.CORTEX_ENABLE_REST = 'true';
+  const { server, startServer } = await serverFactory();
+  startServer && await startServer();
+  testServer = server;
+});
+
+test.after.always('cleanup', async () => {
+  if (testServer) {
+    await testServer.stop();
+  }
+});
+
+// Helper function to create a PathwayResolver with a specific plugin
+function createResolverWithPlugin(pluginClass, modelName = 'test-model') {
+  // Map plugin classes to their corresponding model types
+  const pluginToModelType = {
+    OpenAIChatPlugin: 'OPENAI-VISION',
+    GeminiChatPlugin: 'GEMINI-VISION',
+    Gemini15ChatPlugin: 'GEMINI-1.5-VISION',
+    Claude3VertexPlugin: 'CLAUDE-3-VERTEX'
+  };
+
+  const modelType = pluginToModelType[pluginClass.name];
+  if (!modelType) {
+    throw new Error(`Unknown plugin class: ${pluginClass.name}`);
+  }
+
+  const pathway = {
+    name: 'test-pathway',
+    model: modelName,
+    prompt: 'test prompt'
+  };
+  
+  const model = {
+    name: modelName,
+    type: modelType
+  };
+  
+  const resolver = new PathwayResolver({ 
+    config,
+    pathway,
+    args: {},
+    endpoints: { [modelName]: model }
+  });
+  
+  resolver.modelExecutor.plugin = new pluginClass(pathway, model);
+  return resolver;
+}
+
+// Test OpenAI Chat Plugin Streaming
+test('OpenAI Chat Plugin - processStreamEvent handles content chunks correctly', async t => {
+  const resolver = createResolverWithPlugin(OpenAIChatPlugin);
+  const plugin = resolver.modelExecutor.plugin;
+  
+  // Test regular content chunk
+  const contentEvent = {
+    data: JSON.stringify({
+      id: 'test-id',
+      choices: [{
+        delta: { content: 'test content' },
+        finish_reason: null
+      }]
+    })
+  };
+  
+  let progress = plugin.processStreamEvent(contentEvent, {});
+  t.is(progress.data, contentEvent.data);
+  t.falsy(progress.progress);
+  
+  // Test stream end
+  const endEvent = {
+    data: JSON.stringify({
+      id: 'test-id',
+      choices: [{
+        delta: {},
+        finish_reason: 'stop'
+      }]
+    })
+  };
+  
+  progress = plugin.processStreamEvent(endEvent, {});
+  t.is(progress.progress, 1);
+});
+
+// Test Gemini Chat Plugin Streaming
+test('Gemini Chat Plugin - processStreamEvent handles content chunks correctly', async t => {
+  const resolver = createResolverWithPlugin(GeminiChatPlugin);
+  const plugin = resolver.modelExecutor.plugin;
+  
+  // Test regular content chunk
+  const contentEvent = {
+    data: JSON.stringify({
+      candidates: [{
+        content: {
+          parts: [{ text: 'test content' }]
+        },
+        finishReason: null
+      }]
+    })
+  };
+  
+  let progress = plugin.processStreamEvent(contentEvent, {});
+  t.truthy(progress.data, 'Should have data');
+  const parsedData = JSON.parse(progress.data);
+  t.truthy(parsedData.candidates, 'Should have candidates array');
+  t.truthy(parsedData.candidates[0].content, 'Should have content object');
+  t.truthy(parsedData.candidates[0].content.parts, 'Should have parts array');
+  t.is(parsedData.candidates[0].content.parts[0].text, 'test content', 'Content should match');
+  t.falsy(progress.progress);
+  
+  // Test stream end with STOP
+  const endEvent = {
+    data: JSON.stringify({
+      candidates: [{
+        content: {
+          parts: [{ text: '' }]
+        },
+        finishReason: 'STOP'
+      }]
+    })
+  };
+  
+  progress = plugin.processStreamEvent(endEvent, {});
+  t.is(progress.progress, 1);
+});
+
+// Test Gemini 15 Chat Plugin Streaming
+test('Gemini 15 Chat Plugin - processStreamEvent handles safety blocks', async t => {
+  const resolver = createResolverWithPlugin(Gemini15ChatPlugin);
+  const plugin = resolver.modelExecutor.plugin;
+  
+  // Test safety block
+  const safetyEvent = {
+    data: JSON.stringify({
+      candidates: [{
+        safetyRatings: [{ blocked: true }]
+      }]
+    })
+  };
+  
+  const progress = plugin.processStreamEvent(safetyEvent, {});
+  t.true(progress.data.includes('Response blocked'));
+  t.is(progress.progress, 1);
+});
+
+// Test Claude 3 Vertex Plugin Streaming
+test('Claude 3 Vertex Plugin - processStreamEvent handles message types', async t => {
+  const resolver = createResolverWithPlugin(Claude3VertexPlugin);
+  const plugin = resolver.modelExecutor.plugin;
+  
+  // Test message start
+  const startEvent = {
+    data: JSON.stringify({
+      type: 'message_start',
+      message: { id: 'test-id' }
+    })
+  };
+  
+  let progress = plugin.processStreamEvent(startEvent, {});
+  t.true(JSON.parse(progress.data).choices[0].delta.role === 'assistant');
+  
+  // Test content block
+  const contentEvent = {
+    data: JSON.stringify({
+      type: 'content_block_delta',
+      delta: {
+        type: 'text_delta',
+        text: 'test content'
+      }
+    })
+  };
+  
+  progress = plugin.processStreamEvent(contentEvent, {});
+  t.true(JSON.parse(progress.data).choices[0].delta.content === 'test content');
+  
+  // Test message stop
+  const stopEvent = {
+    data: JSON.stringify({
+      type: 'message_stop'
+    })
+  };
+  
+  progress = plugin.processStreamEvent(stopEvent, {});
+  t.is(progress.progress, 1);
+});