@mongoosejs/studio 0.0.74 → 0.0.76

package/backend/actions/ChatMessage/executeScript.js

@@ -0,0 +1,74 @@
+'use strict';
+
+const Archetype = require('archetype');
+const mongoose = require('mongoose');
+const vm = require('vm');
+
+const ExecuteScriptParams = new Archetype({
+  userId: {
+    $type: mongoose.Types.ObjectId
+  },
+  chatMessageId: {
+    $type: mongoose.Types.ObjectId
+  },
+  script: {
+    $type: String
+  }
+}).compile('ExecuteScriptParams');
+
+module.exports = ({ db }) => async function executeScript(params) {
+  const { userId, chatMessageId, script } = new ExecuteScriptParams(params);
+  const ChatMessage = db.model('__Studio_ChatMessage');
+
+  const chatMessage = await ChatMessage.findById(chatMessageId);
+  if (!chatMessage) {
+    throw new Error('Chat message not found');
+  }
+
+  // Create a sandbox with the db object
+  const logs = [];
+  const sandbox = { db, console: {} };
+
+  // Capture console logs
+  sandbox.console.log = function() {
+    const args = Array.from(arguments);
+    logs.push(args.map(arg => typeof arg === 'object' ? JSON.stringify(arg) : arg).join(' '));
+  };
+
+  const context = vm.createContext(sandbox);
+
+  let output;
+  let error;
+
+  try {
+    // Execute the script in the sandbox
+    output = await vm.runInContext(wrappedScript(script), context);
+
+    chatMessage.script = script;
+    chatMessage.executionResult = { output, logs: logs.join('\n'), error: null };
+    await chatMessage.save();
+
+    return { chatMessage };
+  } catch (err) {
+    error = err.message;
+
+    // Update the chat message with the error
+    await ChatMessage.updateOne(
+      { _id: chatMessageId },
+      {
+        script,
+        executionResult: {
+          output: null,
+          logs: logs.join('\n'),
+          error
+        }
+      }
+    );
+
+    throw new Error(`Script execution failed: ${error}`);
+  }
+};
+
+const wrappedScript = script => `(async () => {
+  ${script}
+})()`;

package/backend/actions/ChatMessage/index.js

@@ -0,0 +1,3 @@
+'use strict';
+
+exports.executeScript = require('./executeScript');

package/backend/actions/ChatThread/createChatMessage.js

@@ -0,0 +1,128 @@
+'use strict';
+
+const Archetype = require('archetype');
+const getModelDescriptions = require('../../helpers/getModelDescriptions');
+const mongoose = require('mongoose');
+
+const CreateChatMessageParams = new Archetype({
+  chatThreadId: {
+    $type: mongoose.Types.ObjectId
+  },
+  userId: {
+    $type: mongoose.Types.ObjectId
+  },
+  content: {
+    $type: String
+  }
+}).compile('CreateChatMessageParams');
+
+const systemPrompt = `
+You are a data querying assistant who writes scripts for users accessing MongoDB data using Node.js and Mongoose.
+
+Keep scripts concise. Avoid unnecessary comments, error handling, and temporary variables.
+
+Do not write any imports or require() statements, that will cause the script to break.
+
+Assume the user has pre-defined schemas and models. Do not define any new schemas or models for the user.
+
+Use async/await where possible. Assume top-level await is allowed.
+
+Think carefully about the user's input and identify the models referred to by the user's query.
+
+Format output as Markdown, including code fences for any scripts the user requested.
+
+Add a brief text description of what the script does.
+
+If the user's query is best answered with a chart, return a Chart.js 4 configuration as \`return { $chart: chartJSConfig };\`. Disable ChartJS animation by default unless user asks for it.
+
+Example output:
+
+The following script counts the number of users which are not deleted.
+
+\`\`\`javascript
+const users = await db.model('User').find({ isDeleted: false });
+return { numUsers: users.length };
+\`\`\`
+
+-----------
+
+Here is a description of the user's models. Assume these are the only models available in the system unless explicitly instructed otherwise by the user.
+`.trim();
+
+module.exports = ({ db }) => async function createChatMessage(params) {
+  const { chatThreadId, userId, content, script } = new CreateChatMessageParams(params);
+  const ChatThread = db.model('__Studio_ChatThread');
+  const ChatMessage = db.model('__Studio_ChatMessage');
+
+  // Check that the user owns the thread
+  const chatThread = await ChatThread.findOne({ _id: chatThreadId });
+  if (!chatThread) {
+    throw new Error('Chat thread not found');
+  }
+  if (userId != null && chatThread.userId.toString() !== userId.toString()) {
+    throw new Error('Not authorized');
+  }
+
+  const messages = await ChatMessage.find({ chatThreadId }).sort({ createdAt: 1 });
+  const llmMessages = messages.map(m => ({
+    role: m.role,
+    content: m.content
+  }));
+  llmMessages.push({ role: 'user', content });
+
+  if (chatThread.title == null) {
+    getChatCompletion([
+      { role: 'system', content: 'Summarize the following chat thread in 6 words or less, as a helpful thread title' },
+      ...llmMessages
+    ]).then(res => {
+      const title = res.choices?.[0]?.message?.content;
+      chatThread.title = title;
+      return chatThread.save();
+    }).catch(() => {});
+  }
+
+  llmMessages.unshift({
+    role: 'system',
+    content: systemPrompt + getModelDescriptions(db)
+  });
+
+  // Create the chat message and get OpenAI response in parallel
+  const chatMessages = await Promise.all([
+    ChatMessage.create({
+      chatThreadId,
+      role: 'user',
+      content,
+      script,
+      executionResult: null
+    }),
+    getChatCompletion(llmMessages).then(res => {
+      const content = res.choices?.[0]?.message?.content || '';
+      console.log('Content', content, res);
+      return ChatMessage.create({
+        chatThreadId,
+        role: 'assistant',
+        content
+      });
+    })
+  ]);
+
+  return { chatMessages };
+};
+
+async function getChatCompletion(messages, options = {}) {
+  const response = await fetch('https://api.openai.com/v1/chat/completions', {
+    method: 'POST',
+    headers: {
+      Authorization: `Bearer ${process.env.OPENAI_API_KEY}`,
+      'Content-Type': 'application/json'
+    },
+    body: JSON.stringify({
+      model: 'gpt-4o',
+      max_tokens: 2500,
+      ...options,
+      messages
+    })
+  });
+
+  return await response.json();
+};

package/backend/actions/ChatThread/createChatThread.js

@@ -0,0 +1,19 @@
+'use strict';
+
+const Archetype = require('archetype');
+const mongoose = require('mongoose');
+
+const CreateChatThreadParams = new Archetype({
+  userId: {
+    $type: mongoose.Types.ObjectId
+  }
+}).compile('CreateChatThreadParams');
+
+module.exports = ({ db }) => async function createChatThread(params) {
+  const { userId } = new CreateChatThreadParams(params);
+  const ChatThread = db.model('__Studio_ChatThread');
+
+  const chatThread = await ChatThread.create({ userId });
+
+  return { chatThread };
+};

package/backend/actions/ChatThread/getChatThread.js

@@ -0,0 +1,36 @@
+'use strict';
+
+const Archetype = require('archetype');
+const mongoose = require('mongoose');
+
+const GetChatThreadParams = new Archetype({
+  chatThreadId: {
+    $type: mongoose.Types.ObjectId
+  },
+  userId: {
+    $type: mongoose.Types.ObjectId
+  }
+}).compile('GetChatThreadParams');
+
+module.exports = ({ db }) => async function getChatThread(params) {
+  const { chatThreadId, userId } = new GetChatThreadParams(params);
+  const ChatThread = db.model('__Studio_ChatThread');
+  const ChatMessage = db.model('__Studio_ChatMessage');
+
+  const chatThread = await ChatThread.findById(chatThreadId);
+
+  if (!chatThread) {
+    throw new Error('Chat thread not found');
+  }
+  if (userId && chatThread.userId.toString() !== userId.toString()) {
+    throw new Error('Not authorized');
+  }
+
+  const chatMessages = await ChatMessage.find({ chatThreadId })
+    .sort({ createdAt: -1 });
+
+  return {
+    chatThread,
+    chatMessages: chatMessages.reverse() // Return in chronological order
+  };
+};

package/backend/actions/ChatThread/index.js

@@ -0,0 +1,6 @@
+'use strict';
+
+exports.createChatMessage = require('./createChatMessage');
+exports.createChatThread = require('./createChatThread');
+exports.getChatThread = require('./getChatThread');
+exports.listChatThreads = require('./listChatThreads');

package/backend/actions/ChatThread/listChatThreads.js

@@ -0,0 +1,24 @@
+'use strict';
+
+const Archetype = require('archetype');
+const mongoose = require('mongoose');
+
+const ListChatThreadsParams = new Archetype({
+  userId: {
+    $type: mongoose.Types.ObjectId
+  }
+}).compile('ListChatThreadsParams');
+
+module.exports = ({ db }) => async function listChatThreads(params) {
+  // Just validate the params object, but no actual parameters needed
+  const { userId } = new ListChatThreadsParams(params);
+  const ChatThread = db.model('__Studio_ChatThread');
+
+  // Get all chat threads
+  const chatThreads = await ChatThread.find(userId ? { userId } : {})
+    .sort({ updatedAt: -1 }); // Sort by most recently updated
+
+  return {
+    chatThreads
+  };
+};

package/backend/actions/index.js

@@ -1,5 +1,7 @@
 'use strict';
 
+exports.ChatMessage = require('./ChatMessage');
+exports.ChatThread = require('./ChatThread');
 exports.Dashboard = require('./Dashboard');
 exports.Model = require('./Model');
 exports.Script = require('./Script');

package/backend/db/chatMessageSchema.js

@@ -0,0 +1,30 @@
+'use strict';
+
+const mongoose = require('mongoose');
+
+const chatMessageSchema = new mongoose.Schema({
+  chatThreadId: {
+    type: 'ObjectId',
+    required: true
+  },
+  role: {
+    type: String,
+    enum: ['user', 'assistant'],
+    required: true
+  },
+  content: {
+    type: String,
+    required: true
+  },
+  script: {
+    type: String,
+    required: false
+  },
+  executionResult: {
+    output: mongoose.Schema.Types.Mixed,
+    logs: String,
+    error: String
+  }
+}, { timestamps: true });
+
+module.exports = chatMessageSchema;

package/backend/db/chatThreadSchema.js

@@ -0,0 +1,15 @@
+'use strict';
+
+const mongoose = require('mongoose');
+
+const chatThreadSchema = new mongoose.Schema({
+  title: {
+    type: String
+  },
+  userId: {
+    type: mongoose.ObjectId,
+    ref: 'User'
+  }
+}, { timestamps: true });
+
+module.exports = chatThreadSchema;

package/backend/helpers/getModelDescriptions.js

@@ -0,0 +1,18 @@
+'use strict';
+
+const listModelPaths = Model => [
+  ...Object.entries(Model.schema.paths).map(
+    ([path, schemaType]) => `- ${path}: ${schemaType.instance}`
+      + (schemaType.options?.ref ? ' (ref: ' + schemaType.options.ref + ')' : '')
+  ),
+  ...Object.entries(Model.schema.virtuals).filter(([path, virtual]) => virtual.options?.ref).map(
+    ([path, virtual]) => `- ${path}: Virtual (ref: ${virtual.options.ref})`
+  )
+].join('\n');
+
+const getModelDescriptions = db => Object.values(db.models).filter(Model => !Model.modelName.startsWith('__Studio')).map(Model => `
+${Model.modelName} (collection: ${Model.collection.collectionName})
+${listModelPaths(Model)}
+`.trim()).join('\n\n');
+
+module.exports = getModelDescriptions;

package/backend/index.js

@@ -4,12 +4,19 @@ const Actions = require('./actions');
 const { applySpec } = require('extrovert');
 const mongoose = require('mongoose');
 
+const chatMessageSchema = require('./db/chatMessageSchema');
+const chatThreadSchema = require('./db/chatThreadSchema');
 const dashboardSchema = require('./db/dashboardSchema');
 
+const getModelDescriptions = require('./helpers/getModelDescriptions');
+
 module.exports = function backend(db) {
   db = db || mongoose.connection;
-  
+
   const Dashboard = db.model('__Studio_Dashboard', dashboardSchema, 'studio__dashboards');
+  const ChatMessage = db.model('__Studio_ChatMessage', chatMessageSchema, 'studio__chatMessages');
+  const ChatThread = db.model('__Studio_ChatThread', chatThreadSchema, 'studio__chatThreads');
+
   const actions = applySpec(Actions, { db });
   return actions;
-};
+};

package/express.js

@@ -63,6 +63,7 @@ module.exports = async function(apiUrl, conn, options) {
           if (!user || !roles) {
             throw new Error('Not authorized');
           }
+          req.params.initiatedById = user._id;
           req.params.roles = roles;
 
           next();