@aj-archipelago/cortex 1.3.21 → 1.3.23

package/server/plugins/ollamaCompletionPlugin.js

@@ -0,0 +1,147 @@
+import ModelPlugin from './modelPlugin.js';
+import logger from '../../lib/logger.js';
+import { Transform } from 'stream';
+
+class OllamaCompletionPlugin extends ModelPlugin {
+
+  getRequestParameters(text, parameters, prompt) {
+    return {
+      data: {
+        model: parameters.ollamaModel,
+        prompt: text,
+        stream: parameters.stream
+      },
+      params: {}
+    };
+  }
+
+  logRequestData(data, responseData, prompt) {
+    const { stream, prompt: promptText, model } = data;
+
+    if (promptText) {
+      logger.info(`[ollama completion request sent to model ${model}]`);
+      const { length, units } = this.getLength(promptText);
+      const preview = this.shortenContent(promptText);
+      logger.verbose(`prompt ${units}: ${length}, content: "${preview}"`);
+      logger.info(`[completion request contained ${length} ${units}]`);
+    }
+
+    if (stream) {
+      logger.info(`[response received as an SSE stream]`);
+    } else if (responseData) {
+      const responseText = this.parseResponse(responseData);
+      const { length, units } = this.getLength(responseText);
+      logger.info(`[response received containing ${length} ${units}]`);
+      logger.verbose(`${this.shortenContent(responseText)}`);
+    }
+
+    prompt &&
+      prompt.debugInfo &&
+      (prompt.debugInfo += `\n${JSON.stringify(data)}`);
+  }
+
+  parseResponse(data) {
+    // If data is not a string (e.g. streaming), return as is
+    if (typeof data !== 'string') {
+      return data;
+    }
+
+    // Split into lines and filter empty ones
+    const lines = data.split('\n').filter(line => line.trim());
+    
+    let fullResponse = '';
+    
+    for (const line of lines) {
+      try {
+        const jsonObj = JSON.parse(line);
+        
+        if (jsonObj.response) {
+          // Unescape special sequences
+          const content = jsonObj.response
+            .replace(/\\n/g, '\n')
+            .replace(/\\"/g, '"')
+            .replace(/\\\\/g, '\\')
+            .replace(/\\u003c/g, '<')
+            .replace(/\\u003e/g, '>');
+            
+          fullResponse += content;
+        }
+      } catch (err) {
+        // If we can't parse the line as JSON, just skip it
+        continue;
+      }
+    }
+  
+    return fullResponse;
+  }
+
+  processStreamEvent(event, requestProgress) {
+    try {
+      const data = JSON.parse(event.data);
+      
+      // Handle the streaming response
+      if (data.response) {
+        // Unescape special sequences in the content
+        const content = data.response
+          .replace(/\\n/g, '\n')
+          .replace(/\\"/g, '"')
+          .replace(/\\\\/g, '\\')
+          .replace(/\\u003c/g, '<')
+          .replace(/\\u003e/g, '>');
+          
+        requestProgress.data = JSON.stringify(content);
+      }
+
+      // Check if this is the final message
+      if (data.done) {
+        requestProgress.data = '[DONE]';
+        requestProgress.progress = 1;
+      }
+
+      return requestProgress;
+    } catch (err) {
+      // If we can't parse the event data, return the progress as is
+      return requestProgress;
+    }
+  }
+
+  async execute(text, parameters, prompt, cortexRequest) {
+    const requestParameters = this.getRequestParameters(text, parameters, prompt);
+    cortexRequest.data = { ...(cortexRequest.data || {}), ...requestParameters.data };
+    cortexRequest.params = { ...(cortexRequest.params || {}), ...requestParameters.params };
+    
+    // For Ollama streaming, transform NDJSON to SSE format
+    if (parameters.stream) {
+      const response = await this.executeRequest(cortexRequest);
+      
+      // Create a transform stream that converts NDJSON to SSE format
+      const transformer = new Transform({
+        decodeStrings: false, // Keep as string
+        transform(chunk, encoding, callback) {
+          try {
+            const lines = chunk.toString().split('\n');
+            for (const line of lines) {
+              if (line.trim()) {
+                // Format as SSE data
+                this.push(`data: ${line}\n\n`);
+              }
+            }
+            callback();
+          } catch (err) {
+            callback(err);
+          }
+        }
+      });
+
+      // Pipe the response through our transformer
+      response.pipe(transformer);
+      
+      // Return the transformed stream
+      return transformer;
+    }
+    
+    return this.executeRequest(cortexRequest);
+  }
+}
+
+export default OllamaCompletionPlugin;

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 = '';
@@ -28,6 +44,13 @@ const processRestRequest = async (server, req, pathway, name, parameterMap = {})
             return Boolean(value);
         } else if (type === 'Int') {
             return parseInt(value, 10);
+        } else if (type === '[MultiMessage]' && Array.isArray(value)) {
+            return value.map(msg => ({
+                ...msg,
+                content: Array.isArray(msg.content) ? 
+                    JSON.stringify(msg.content) : 
+                    msg.content
+            }));
         } else {
             return value;
         }
@@ -58,8 +81,16 @@ const processRestRequest = async (server, req, pathway, name, parameterMap = {})
             `;
 
     const result = await server.executeOperation({ query, variables });
-    const resultText = result?.body?.singleResult?.data?.[name]?.result || result?.body?.singleResult?.errors?.[0]?.message || "";
 
+    // if we're streaming and there are errors, we return a standard error code
+    if (Boolean(req.body.stream)) {
+        if (result?.body?.singleResult?.errors) {
+            return `[ERROR] ${result.body.singleResult.errors[0].message.split(';')[0]}`;
+        }
+    }
+    
+    // otherwise errors can just be returned as a string
+    const resultText = result?.body?.singleResult?.data?.[name]?.result || result?.body?.singleResult?.errors?.[0]?.message || "";
     return resultText;
 };
 
@@ -86,7 +117,6 @@ const processIncomingStream = (requestId, res, jsonResponse, pathway) => {
 
         // If we haven't sent the stop message yet, do it now
         if (jsonResponse.choices?.[0]?.finish_reason !== "stop") {
-    
             let jsonEndStream = JSON.parse(JSON.stringify(jsonResponse));
     
             if (jsonEndStream.object === 'text_completion') {
@@ -116,7 +146,6 @@ const processIncomingStream = (requestId, res, jsonResponse, pathway) => {
     }
 
     const fillJsonResponse = (jsonResponse, inputText, _finishReason) => {
-
         jsonResponse.choices[0].finish_reason = null;
         if (jsonResponse.object === 'text_completion') {
             jsonResponse.choices[0].text = inputText;
@@ -129,6 +158,14 @@ const processIncomingStream = (requestId, res, jsonResponse, pathway) => {
 
     startStream(res);
 
+    // If the requestId is an error message, we can't continue
+    if (requestId.startsWith('[ERROR]')) {
+        fillJsonResponse(jsonResponse, requestId, "stop");
+        sendStreamData(jsonResponse);
+        finishStream(res, jsonResponse);
+        return;
+    }
+
     let subscription;
 
     subscription = pubsub.subscribe('REQUEST_PROGRESS', (data) => {
@@ -261,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({
@@ -297,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 {
@@ -309,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({
@@ -364,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;
@@ -376,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);
                 })