@aj-archipelago/cortex 1.0.4 → 1.0.6

package/tests/openai_api.test.js

@@ -0,0 +1,147 @@
+// openai_api.test.js
+
+import test from 'ava';
+import got from 'got';
+import axios from 'axios';
+import serverFactory from '../index.js';
+
+const API_BASE = 'http://localhost:4000/v1';
+
+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();
+  }
+});
+
+test('GET /models', async (t) => {
+  const response = await got(`${API_BASE}/models`, { responseType: 'json' });
+  t.is(response.statusCode, 200);
+  t.is(response.body.object, 'list');
+  t.true(Array.isArray(response.body.data));
+});
+
+test('POST /completions', async (t) => {
+  const response = await got.post(`${API_BASE}/completions`, {
+    json: {
+      model: 'gpt-3.5-turbo',
+      prompt: 'Word to your motha!',
+      stream: false,
+    },
+    responseType: 'json',
+  });
+
+  t.is(response.statusCode, 200);
+  t.is(response.body.object, 'text_completion');
+  t.true(Array.isArray(response.body.choices));
+});
+
+
+test('POST /chat/completions', async (t) => {
+  const response = await got.post(`${API_BASE}/chat/completions`, {
+    json: {
+      model: 'gpt-3.5-turbo',
+      messages: [{ role: 'user', content: 'Hello!' }],
+      stream: false,
+    },
+    responseType: 'json',
+  });
+
+  t.is(response.statusCode, 200);
+  t.is(response.body.object, 'chat.completion');
+  t.true(Array.isArray(response.body.choices));
+});
+
+async function connectToSSEEndpoint(url, endpoint, payload, t, customAssertions) {
+    return new Promise(async (resolve, reject) => {
+        try {
+            const instance = axios.create({
+                baseURL: url,
+                responseType: 'stream',
+            });
+        
+            const response = await instance.post(endpoint, payload);
+            const responseData = response.data;
+        
+            const incomingMessage = Array.isArray(responseData) && responseData.length > 0 ? responseData[0] : responseData;
+        
+            let eventCount = 0;
+        
+            incomingMessage.on('data', data => {
+                const events = data.toString().split('\n');
+        
+                events.forEach(event => {
+                    eventCount++;
+        
+                    if (event.trim() === '') return;
+
+                    if (event.trim() === 'data: [DONE]') {
+                        t.truthy(eventCount > 1);
+                        resolve();
+                        return;
+                    }
+        
+                    const message = event.replace(/^data: /, '');
+                    const messageJson = JSON.parse(message);
+                    
+                    customAssertions(t, messageJson);
+                });
+            });  
+        
+        } catch (error) {
+            console.error('Error connecting to SSE endpoint:', error);
+            reject(error);
+        }
+    });
+}
+
+test('POST SSE: /v1/completions should send a series of events and a [DONE] event', async (t) => {
+    const payload = {
+        model: 'gpt-3.5-turbo',
+        prompt: 'Word to your motha!',
+        stream: true,
+    };
+    
+    const url = 'http://localhost:4000/v1';
+    
+    const completionsAssertions = (t, messageJson) => {
+        t.truthy(messageJson.id);
+        t.is(messageJson.object, 'text_completion');
+        t.truthy(messageJson.choices[0].finish_reason === null || messageJson.choices[0].finish_reason === 'stop');
+    };
+    
+    await connectToSSEEndpoint(url, '/completions', payload, t, completionsAssertions);
+});
+
+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',
+        messages: [
+        {
+            role: 'user',
+            content: 'Hello!',
+        },
+        ],
+        stream: true,
+    };
+    
+    const url = 'http://localhost:4000/v1';
+    
+    const chatCompletionsAssertions = (t, messageJson) => {
+        t.truthy(messageJson.id);
+        t.is(messageJson.object, 'chat.completion.chunk');
+        t.truthy(messageJson.choices[0].delta);
+        t.truthy(messageJson.choices[0].finish_reason === null || messageJson.choices[0].finish_reason === 'stop');
+    };
+    
+    await connectToSSEEndpoint(url, '/chat/completions', payload, t, chatCompletionsAssertions);
+});  
+  

package/tests/palmChatPlugin.test.js

@@ -0,0 +1,256 @@
+// test_palmChatPlugin.js
+import test from 'ava';
+import PalmChatPlugin from '../server/plugins/palmChatPlugin.js';
+import { mockConfig } from './mocks.js';
+
+test.beforeEach((t) => {
+  const pathway = 'testPathway';
+  const palmChatPlugin = new PalmChatPlugin(mockConfig, pathway);
+  t.context = { palmChatPlugin };
+});
+
+test('convertMessagesToPalm', (t) => {
+  const { palmChatPlugin } = t.context;
+  const messages = [
+    { role: 'system', content: 'System Message' },
+    { role: 'user', content: 'User Message' },
+    { role: 'user', content: 'User Message 2'},
+  ];
+
+  const expectedResult = {
+    modifiedMessages: [
+      { author: 'user', content: 'User Message\nUser Message 2' },
+    ],
+    context: 'System Message',
+  };
+
+  t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+});
+
+test('convertMessagesToPalm - already PaLM format', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [
+      { author: 'user', content: 'User Message' },
+      { author: 'user', content: 'User Message 2'},
+    ];
+  
+    const expectedResult = {
+      modifiedMessages: [
+        { author: 'user', content: 'User Message\nUser Message 2' },
+      ],
+      context: '',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+
+  test('convertMessagesToPalm - empty string roles', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [
+      { role: '', content: 'Empty role message' },
+      { role: 'user', content: 'User Message' },
+    ];
+  
+    const expectedResult = {
+      modifiedMessages: [
+        { author: 'user', content: 'User Message' },
+      ],
+      context: '',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+  
+  test('convertMessagesToPalm - consecutive system messages', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [
+      { role: 'system', content: 'System Message 1' },
+      { role: 'system', content: 'System Message 2' },
+      { role: 'user', content: 'User Message' },
+    ];
+  
+    const expectedResult = {
+      modifiedMessages: [
+        { author: 'user', content: 'User Message' },
+      ],
+      context: 'System Message 1\nSystem Message 2',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+  
+  test('convertMessagesToPalm - multiple authors', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [
+      { role: 'system', content: 'System Message' },
+      { author: 'user1', content: 'User1 Message' },
+      { author: 'user1', content: 'User1 Message 2' },
+      { author: 'user2', content: 'User2 Message' },
+      { author: 'assistant', content: 'Assistant Message' },
+    ];
+  
+    const expectedResult = {
+      modifiedMessages: [
+        { author: 'user1', content: 'User1 Message\nUser1 Message 2' },
+        { author: 'user2', content: 'User2 Message' },
+        { author: 'assistant', content: 'Assistant Message' },
+      ],
+      context: 'System Message',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+  
+  test('convertMessagesToPalm - no messages', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [];
+  
+    const expectedResult = {
+      modifiedMessages: [],
+      context: '',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+  
+  test('convertMessagesToPalm - only system messages', (t) => {
+    const { palmChatPlugin } = t.context;
+    const messages = [
+      { role: 'system', content: 'System Message 1' },
+      { role: 'system', content: 'System Message 2' },
+    ];
+  
+    const expectedResult = {
+      modifiedMessages: [],
+      context: 'System Message 1\nSystem Message 2',
+    };
+  
+    t.deepEqual(palmChatPlugin.convertMessagesToPalm(messages), expectedResult);
+  });
+
+test('getCompiledContext', (t) => {
+  const { palmChatPlugin } = t.context;
+  const text = 'Hello';
+  const parameters = { name: 'John' };
+  const context = '{{text}} from {{name}}';
+
+  const expectedResult = 'Hello from John';
+
+  t.is(palmChatPlugin.getCompiledContext(text, parameters, context), expectedResult);
+});
+
+test('getCompiledExamples', (t) => {
+  const { palmChatPlugin } = t.context;
+  const text = 'Greetings';
+  const parameters = { name: 'Jane' };
+  const examples = [
+    {
+      input: { content: 'Input: {{text}} from {{name}}' },
+      output: { content: 'Output: {{text}} to {{name}}' },
+    },
+  ];
+
+  const expectedResult = [
+    {
+      input: { content: 'Input: Greetings from Jane' },
+      output: { content: 'Output: Greetings to Jane' },
+    },
+  ];
+
+  t.deepEqual(palmChatPlugin.getCompiledExamples(text, parameters, examples), expectedResult);
+});
+
+test('getRequestParameters', (t) => {
+  const { palmChatPlugin } = t.context;
+  const text = 'Hello';
+  const parameters = { stream: false, name: 'John'};
+  const messages = [
+    { role: 'system', content: 'System Message' },
+    { role: 'user', content: 'Hello' },
+    { role: 'assistant', content: 'What can I do for you?' },
+    { role: 'user', content: 'Be my assistant!' },
+  ];
+  const prompt = { context: '{{text}} from {{name}}', examples: [], messages };
+
+  const requestParameters = palmChatPlugin.getRequestParameters(text, parameters, prompt);
+  const requestMessages = requestParameters.instances[0].messages;
+
+  t.is(requestMessages[0].author, 'user');
+  t.is(requestMessages[0].content, 'Hello');
+});
+
+test('getSafetyAttributes', (t) => {
+  const { palmChatPlugin } = t.context;
+  const responseData = {
+    predictions: [
+      {
+        safetyAttributes: {
+          blocked: false,
+        },
+      },
+    ],
+  };
+
+  const expectedResult = {
+    blocked: false,
+  };
+
+  t.deepEqual(palmChatPlugin.getSafetyAttributes(responseData), expectedResult);
+});
+
+test('parseResponse', (t) => {
+  const { palmChatPlugin } = t.context;
+  const responseData = {
+    predictions: [
+      {
+        candidates: [
+          {
+            content: 'Hello, how can I help you today?',
+          },
+        ],
+      },
+    ],
+  };
+
+  const expectedResult = 'Hello, how can I help you today?';
+
+  t.is(palmChatPlugin.parseResponse(responseData), expectedResult);
+});
+
+test('logRequestData', (t) => {
+  const { palmChatPlugin } = t.context;
+  const data = {
+    instances: [
+      {
+        messages: [
+          { author: 'user', content: 'Hello' },
+          { author: 'assistant', content: 'How can I help you?' },
+        ],
+      },
+    ],
+  };
+  const responseData = {
+    predictions: [
+      {
+        candidates: [
+          {
+            content: 'Hello, how can I help you today?',
+          },
+        ],
+      },
+    ],
+  };
+  const prompt = { debugInfo: '' };
+
+  const consoleLog = console.log;
+  let logOutput = '';
+  console.log = (msg) => (logOutput += msg + '\n');
+
+  palmChatPlugin.logRequestData(data, responseData, prompt);
+
+  console.log = consoleLog;
+
+  t.true(logOutput.includes('Message 1:'));
+  t.true(logOutput.includes('Message 2:'));
+  t.true(logOutput.includes('> Hello, how can I help you today?'));
+});

package/tests/palmCompletionPlugin.test.js

@@ -0,0 +1,87 @@
+// palmCompletionPlugin.test.js
+
+import test from 'ava';
+import PalmCompletionPlugin from '../server/plugins/palmCompletionPlugin.js';
+import { mockConfig } from './mocks.js';
+
+test.beforeEach((t) => {
+  const pathway = 'testPathway';
+  const palmCompletionPlugin = new PalmCompletionPlugin(mockConfig, pathway);
+  t.context = { palmCompletionPlugin };
+});
+
+test('getRequestParameters', (t) => {
+  const { palmCompletionPlugin } = t.context;
+  const text = 'Hello';
+  const parameters = { stream: false, name: 'John' };
+  const prompt = {prompt:'{{text}} from {{name}}'};
+
+  const requestParameters = palmCompletionPlugin.getRequestParameters(text, parameters, prompt);
+  const requestPrompt = requestParameters.instances[0].prompt;
+
+  t.is(requestPrompt, 'Hello from John');
+});
+
+test('parseResponse', (t) => {
+  const { palmCompletionPlugin } = t.context;
+  const responseData = {
+    predictions: [
+      {
+        content: 'Hello, how can I help you today?',
+      },
+    ],
+  };
+
+  const expectedResult = 'Hello, how can I help you today?';
+
+  t.is(palmCompletionPlugin.parseResponse(responseData), expectedResult);
+});
+
+test('getSafetyAttributes', (t) => {
+  const { palmCompletionPlugin } = t.context;
+  const responseData = {
+    predictions: [
+      {
+        safetyAttributes: {
+          blocked: false,
+        },
+      },
+    ],
+  };
+
+  const expectedResult = {
+    blocked: false,
+  };
+
+  t.deepEqual(palmCompletionPlugin.getSafetyAttributes(responseData), expectedResult);
+});
+
+test('logRequestData', (t) => {
+  const { palmCompletionPlugin } = t.context;
+  const data = {
+    instances: [
+      {
+        prompt: 'Hello, how can I help you?',
+      },
+    ],
+  };
+  const responseData = {
+    predictions: [
+      {
+        content: 'Hello, how can I help you today?',
+      },
+    ],
+  };
+  const prompt = { debugInfo: '' };
+
+  const consoleLog = console.log;
+  let logOutput = '';
+  console.log = (msg) => (logOutput += msg + '\n');
+
+  palmCompletionPlugin.logRequestData(data, responseData, prompt);
+
+  console.log = consoleLog;
+
+  t.true(logOutput.includes('Hello, how can I help you?'));
+  t.true(logOutput.includes('> Hello, how can I help you today?'));
+});

package/tests/pathwayResolver.test.js

@@ -1,5 +1,5 @@
 import test from 'ava';
-import { PathwayResolver } from '../graphql/pathwayResolver.js';
+import { PathwayResolver } from '../server/pathwayResolver.js';
 import sinon from 'sinon';
 import { mockConfig, mockPathwayString } from './mocks.js';
 

package/tests/server.js

@@ -0,0 +1,23 @@
+import 'dotenv/config'
+import { ApolloServer } from 'apollo-server';
+import { config } from '../config.js';
+import typeDefsresolversFactory from '../index.js';
+
+let typeDefs;
+let resolvers;
+
+const initTypeDefsResolvers = async () => {
+    const result = await typeDefsresolversFactory();
+    typeDefs = result.typeDefs;
+    resolvers = result.resolvers;
+};
+
+export const startTestServer = async () => {
+    await initTypeDefsResolvers();
+
+    return new ApolloServer({
+        typeDefs,
+        resolvers,
+        context: () => ({ config, requestState: {} }),
+    });
+};

package/tests/truncateMessages.test.js

@@ -1,6 +1,6 @@
 // ModelPlugin.test.js
 import test from 'ava';
-import ModelPlugin from '../graphql/plugins/modelPlugin.js';
+import ModelPlugin from '../server/plugins/modelPlugin.js';
 import { encode } from 'gpt-3-encoder';
 import { mockConfig, mockPathwayString } from './mocks.js';
 

package/graphql/plugins/openAiChatPlugin.js

@@ -1,46 +0,0 @@
-// OpenAIChatPlugin.js
-import ModelPlugin from './modelPlugin.js';
-
-class OpenAIChatPlugin extends ModelPlugin {
-    constructor(config, pathway) {
-        super(config, pathway);
-    }
-
-    // Set up parameters specific to the OpenAI Chat API
-    getRequestParameters(text, parameters, prompt) {
-        const { modelPromptText, modelPromptMessages, tokenLength } = this.getCompiledPrompt(text, parameters, prompt);
-        const { stream } = parameters;
-    
-        // Define the model's max token length
-        const modelTargetTokenLength = this.getModelMaxTokenLength() * this.getPromptTokenRatio();
-    
-        let requestMessages = modelPromptMessages || [{ "role": "user", "content": modelPromptText }];
-    
-        // Check if the token length exceeds the model's max token length
-        if (tokenLength > modelTargetTokenLength) {
-            // Remove older messages until the token length is within the model's limit
-            requestMessages = this.truncateMessagesToTargetLength(requestMessages, modelTargetTokenLength);
-        }
-    
-        const requestParameters = {
-        messages: requestMessages,
-        temperature: this.temperature ?? 0.7,
-        stream
-        };
-    
-        return requestParameters;
-    }
-
-    // Execute the request to the OpenAI Chat API
-    async execute(text, parameters, prompt) {
-        const url = this.requestUrl(text);
-        const requestParameters = this.getRequestParameters(text, parameters, prompt);
-
-        const data = { ...(this.model.params || {}), ...requestParameters };
-        const params = {};
-        const headers = this.model.headers || {};
-        return this.executeRequest(url, data, params, headers, prompt);
-    }
-}
-
-export default OpenAIChatPlugin;