skunkworks - mongosh AI snippets

Author: gagik

snippets/ai/decorators.ts

@@ -14,10 +14,10 @@ export function aiCommand({
       if (requiresPrompt === false && args.length > 0) {
         throw new Error('This command does not accept any arguments');
       } else if (requiresPrompt && args.length === 0) {
-        throw new Error('Please specify a prompt to run');
+        throw new Error('Please specify arguments to run');
       }
       // Combine all arguments into a single string
-      const combinedString = args.join(' ');
+      const combinedString = args.join(' ').trim();
       // Call the original function with the combined string
       return value.call(this, combinedString);
     } as unknown as T;  // Cast the wrapped function to match the original type

snippets/ai/index.ts

@@ -4,7 +4,7 @@ import { getAtlasAiProvider } from './providers/atlas/atlas-ai-provider';
 import { getDocsAiProvider } from './providers/docs/docs-ai-provider';
 import { getAiSdkProvider, models } from './providers/generic/ai-sdk-provider';
 import { Config, ConfigSchema } from './config';
-import { CliContext, wrapAllFunctions, formatHelpCommands } from './helpers';
+import { CliContext, wrapAllFunctions } from './helpers';
 import chalk from 'chalk';
 
 class AI {
@@ -29,22 +29,29 @@ class AI {
     this.config = new Config(this.replConfig);
 
     // Set up provider change listener
-    this.config.on('change', (event) => {
+    this.config.on('change', async (event) => {
       switch (event.key) {
         case 'provider':
           this.ai = this.getProvider(event.value as ConfigSchema['provider']);
           break;
         case 'model':
-          if (Object.keys(models).includes(event.value as string)) {
+            if (!Object.keys(models).includes(this.config.get('provider') as keyof typeof models)) {
+              if (event.value == 'default') {
+                return;
+              }
+              await this.config.set('model', 'default');
+              throw new Error(`${this.config.get('provider')} does not support custom models`);
+            }
+          try {
             this.ai = getAiSdkProvider(
               models[this.config.get('provider') as keyof typeof models](
                 event.value as string,
               ),
               this.cliContext,
               this.config,
             );
-          } else {
-            throw new Error(`Invalid model: ${event.value}`);
+          } catch (error) {
+            throw new Error(`Invalid model, please ensure your name is correct: ${error}`);
           }
           break;
         default:

snippets/ai/package.json

@@ -1,5 +1,5 @@
 {
-  "name": "@gagik/snippet-ai",
+  "name": "@gagik.co/snippet-ai",
   "snippetName": "ai",
   "version": "0.0.1",
   "description": "Provides a ai command suite for mongosh to ask for MongoDB query expressions in natural language.",

snippets/ai/providers/ai-provider.ts

@@ -70,8 +70,12 @@ export abstract class AiProvider {
 
   /** @internal */
   setInput(text: string) {
-    let actualText = text.replace('\n\n', '\n');
-    process.stdin.unshift(actualText, 'utf-8');
+    const trimmedText = text.trim();
+    if (/[\n\r]/.test(trimmedText)) {
+      // If the text includes a newline or carriage return, we should enter editor mode first
+      process.stdin.unshift('.editor\n', 'utf-8');
+    }
+    process.stdin.unshift(trimmedText, 'utf-8');
   }
   /** @internal */
   respond(text: string) {
@@ -283,13 +287,14 @@ export abstract class AiProvider {
     },
   ): string {
     if (expectedOutput === 'mongoshCommand') {
+      // Often the models will return a command in a markdown code block, which we don't want
       return text
         .replace(/```\w*/g, '')
         .replace(/```/g, '')
         .trim();
     }
 
-    return text;
+    return chalk.blue.bold('Answer: ') + text;
   }
 
   abstract getResponse(

snippets/ai/providers/docs/docs-ai-provider.ts

@@ -20,7 +20,7 @@ export class DocsAiProvider extends AiProvider {
   ): Promise<string> {
     try {
     // Initialize conversation if not exists
-    if (!this.conversation) {
+    if (!this.docsConversation) {
       const conv = await this.aiService.createConversation({
         signal,
       });
@@ -37,7 +37,7 @@ export class DocsAiProvider extends AiProvider {
     let formattedResponse = response.content;
     if (expectedOutput === 'text') {
       // Format and display the response
-      formattedResponse = chalk.blue.bold('Answer:\n') + response.content;
+      formattedResponse = chalk.blue.bold('Answer: ') + response.content;
 
       // Add references if they exist
       if (response.references && response.references.length > 0) {