@mongoosejs/studio 0.2.6 → 0.2.8

package/backend/actions/ChatMessage/executeScript.js

@@ -11,7 +11,7 @@ const ExecuteScriptParams = new Archetype({
   },
   chatMessageId: {
     $type: mongoose.Types.ObjectId,
-    $required: true,
+    $required: true
   },
   script: {
     $type: 'string'

package/backend/actions/Model/createChatMessage.js

@@ -0,0 +1,54 @@
+'use strict';
+
+const Archetype = require('archetype');
+const authorize = require('../../authorize');
+const callLLM = require('../../integrations/callLLM');
+const getModelDescriptions = require('../../helpers/getModelDescriptions');
+
+const CreateChatMessageParams = new Archetype({
+  model: {
+    $type: 'string',
+    $required: true
+  },
+  content: {
+    $type: 'string',
+    $required: true
+  },
+  documentData: {
+    $type: 'string'
+  },
+  roles: {
+    $type: ['string']
+  }
+}).compile('CreateChatMessageParams');
+
+module.exports = ({ db, options }) => async function createChatMessage(params) {
+  const { model, content, documentData, roles } = new CreateChatMessageParams(params);
+
+  await authorize('Model.createChatMessage', roles);
+
+  const Model = db.models[model];
+  if (Model == null) {
+    throw new Error(`Model ${model} not found`);
+  }
+
+  const modelDescriptions = getModelDescriptions({ models: { [Model.modelName]: Model } });
+  const context = [
+    modelDescriptions,
+    'Current draft document:\n' + (documentData || '')
+  ].join('\n\n');
+  const system = systemPrompt + '\n\n' + context + (options?.context ? '\n\n' + options.context : '');
+
+  const llmMessages = [{ role: 'user', content: [{ type: 'text', text: content }] }];
+  const res = await callLLM(llmMessages, system, options);
+
+  return { text: res.text };
+};
+
+const systemPrompt = `
+  You are a helpful assistant that drafts MongoDB documents for the user.
+
+  Use the model description and the current draft document to refine the user's intent.
+
+  Return only the updated document body as a JavaScript object literal. Do not use Markdown or code fences.
+`.trim();

package/backend/actions/Model/index.js

@@ -1,6 +1,7 @@
 'use strict';
 
 exports.addField = require('./addField');
+exports.createChatMessage = require('./createChatMessage');
 exports.createDocument = require('./createDocument');
 exports.deleteDocument = require('./deleteDocument');
 exports.deleteDocuments = require('./deleteDocuments');
@@ -12,5 +13,6 @@ exports.getDocumentsStream = require('./getDocumentsStream');
 exports.getCollectionInfo = require('./getCollectionInfo');
 exports.getIndexes = require('./getIndexes');
 exports.listModels = require('./listModels');
+exports.streamChatMessage = require('./streamChatMessage');
 exports.updateDocument = require('./updateDocument');
 exports.updateDocuments = require('./updateDocuments');

package/backend/actions/Model/streamChatMessage.js

@@ -0,0 +1,58 @@
+'use strict';
+
+const Archetype = require('archetype');
+const authorize = require('../../authorize');
+const streamLLM = require('../../integrations/streamLLM');
+const getModelDescriptions = require('../../helpers/getModelDescriptions');
+
+const StreamChatMessageParams = new Archetype({
+  model: {
+    $type: 'string',
+    $required: true
+  },
+  content: {
+    $type: 'string',
+    $required: true
+  },
+  documentData: {
+    $type: 'string'
+  },
+  roles: {
+    $type: ['string']
+  }
+}).compile('StreamChatMessageParams');
+
+module.exports = ({ db, options }) => async function* streamChatMessage(params) {
+  const { model, content, documentData, roles } = new StreamChatMessageParams(params);
+
+  await authorize('Model.streamChatMessage', roles);
+
+  const Model = db.models[model];
+  if (Model == null) {
+    throw new Error(`Model ${model} not found`);
+  }
+
+  const modelDescriptions = getModelDescriptions({ models: { [Model.modelName]: Model } });
+  const context = [
+    modelDescriptions,
+    'Current draft document:\n' + (documentData || '')
+  ].join('\n\n');
+  const system = systemPrompt + '\n\n' + context + (options?.context ? '\n\n' + options.context : '');
+
+  const llmMessages = [{ role: 'user', content: [{ type: 'text', text: content }] }];
+  const textStream = streamLLM(llmMessages, system, options);
+
+  for await (const textPart of textStream) {
+    yield { textPart };
+  }
+
+  return {};
+};
+
+const systemPrompt = `
+  You are a helpful assistant that drafts MongoDB documents for the user.
+
+  Use the model description and the current draft document to refine the user's intent.
+
+  Return only the updated document body as a JavaScript object literal. Do not use Markdown or code fences.
+`.trim();

package/backend/authorize.js

@@ -23,6 +23,7 @@ const actionsToRequiredRoles = {
   'Model.getDocumentsStream': ['owner', 'admin', 'member', 'readonly'],
   'Model.getIndexes': ['owner', 'admin', 'member', 'readonly'],
   'Model.listModels': ['owner', 'admin', 'member', 'readonly'],
+  'Model.streamChatMessage': ['owner', 'admin', 'member', 'readonly'],
   'Model.updateDocuments': ['owner', 'admin', 'member']
 };
 

package/backend/helpers/getModelDescriptions.js

@@ -17,8 +17,16 @@ const formatSchemaTypeInstance = schemaType => {
   return schemaType.instance;
 };
 
+const formatEnum = schemaType => {
+  if (!schemaType.options?.enum) {
+    return '';
+  }
+  return ` (enum: ${JSON.stringify(schemaType.options.enum)})`;
+};
+
 const formatSchemaPath = (path, schemaType) => `- ${path}: ${formatSchemaTypeInstance(schemaType)}` +
   formatRef(schemaType) +
+  formatEnum(schemaType) +
   (schemaType.schema ? formatNestedSchema(schemaType) : '');
 
 const listModelPaths = Model => [