chore: cleanup commands

Author: gagik

snippets/ai/README.md

@@ -1,4 +1,4 @@
-# AI Command Snippets
+# AI Snippets for mongosh
 
 > [!CAUTION]
 > This is an experimental, early-stage snippet that is not meant for production use. 
@@ -9,11 +9,9 @@ This snippet adds a suite of commands accessible with the `ai` command. This inc
 |---------|-------------|---------|
 | `ai.ask` | Ask questions about MongoDB | `ai.ask how do I run queries in mongosh?` |
 | `ai.data` | Generate data-related mongosh commands | `ai.data insert some sample user info` |
-| `ai.query` | Generate a MongoDB query | `ai.query find documents where name = "Ada"` |
-| `ai.aggregate` | Generate a MongoDB aggregation | `ai.aggregate find documents where name = "Ada"` |
-| `ai.collection` | Set the active collection | `ai.collection("users")` |
-| `ai.shell` | Generate general mongosh commands | `ai.shell get sharding info` |
-| `ai.general` | Ask general questions to your model | `ai.general what is the meaning of life?`
+| `ai.query` | Generate a MongoDB query or aggregation | `ai.query find documents where name = "Ada"` |
+| `ai.collection` | Set the active collection | `ai.collection users` |
+| `ai.command` | Generate general mongosh commands _alias:_ `ai.cmd` | `ai.command get sharding info` |
 | `ai.config` | Configure the AI commands | `ai.config.set("provider", "ollama")` |
 
 This currently supports 5 different AI providers: `docs, openai | mistral | atlas | ollama` and any model they support. For cloud providers, you can specify the API key with `MONGOSH_AI_API_KEY`.

snippets/ai/src/ai.ts

@@ -37,7 +37,7 @@ module.exports = async (globalThis: CliContext) => {
       const provider = new AiProvider(
         cliContext,
         config,
-        models[config.get('provider') as keyof typeof models](
+        models[config.get('provider')](
           config.get('model') === 'default' ? undefined : config.get('model'),
         ),
       );
@@ -52,18 +52,13 @@ module.exports = async (globalThis: CliContext) => {
     }
 
     @aiCommand()
-    async shell(prompt: string) {
+    async command(prompt: string) {
       await this.ai.shell(prompt);
     }
 
     @aiCommand()
-    async query(prompt: string) {
-      await this.ai.query(prompt);
-    }
-
-    @aiCommand()
-    async aggregate(prompt: string) {
-      await this.ai.aggregate(prompt);
+    async cmd(prompt: string) {
+      await this.command(prompt);
     }
 
     @aiCommand()
@@ -84,24 +79,19 @@ module.exports = async (globalThis: CliContext) => {
           example: 'ai.ask how do I run queries in mongosh?',
         },
         {
-          cmd: 'ai.query',
-          desc: 'generate a MongoDB query',
-          example: 'ai.query find documents where name = "Ada"',
-        },
-        {
-          cmd: 'ai.aggregate',
-          desc: 'generate a MongoDB aggregation',
-          example: 'ai.aggregate find documents where name = "Ada"',
+          cmd: 'ai.find',
+          desc: 'generate a MongoDB query or aggregation',
+          example: 'ai.find documents where name = "Ada"',
         },
         {
           cmd: 'ai.collection',
           desc: 'set the active collection',
-          example: 'ai.collection("users")',
+          example: 'ai.collection users',
         },
         {
-          cmd: 'ai.shell',
-          desc: 'generate administrative mongosh commands',
-          example: 'ai.shell insert a new sample document',
+          cmd: 'ai.command',
+          desc: `Generate mongosh commands`,
+          example: 'ai.command insert a new sample document | alias: ai.cmd',
         },
         {
           cmd: 'ai.config',

snippets/ai/src/helpers.ts

@@ -1,4 +1,4 @@
-import chalk from "chalk";
+import chalk from 'chalk';
 
 export function output(text: string) {
   process.stdout.write(`${text}`);
@@ -92,12 +92,17 @@ export function wrapFunction(
   ] = wrapperFn;
 }
 
-export function wrapAllFunctions(cliContext: CliContext, currentInstance: unknown) {
+export function wrapAllFunctions(
+  cliContext: CliContext,
+  currentInstance: unknown,
+) {
   const instanceState = cliContext.db._mongo._instanceState;
   const instance = currentInstance as {
-    [key: string]: (...args: unknown[]) => Record<string, unknown> | {
-      isDirectShellCommand: boolean;
-    }
+    [key: string]: (...args: unknown[]) =>
+      | Record<string, unknown>
+      | {
+          isDirectShellCommand: boolean;
+        };
   };
 
   const methods = Object.getOwnPropertyNames(
@@ -114,12 +119,14 @@ export function wrapAllFunctions(cliContext: CliContext, currentInstance: unknow
     );
   });
 
-
-
   // for all methods, wrap them with the wrapFunction method
   for (const methodName of methods) {
     const method = instance[methodName];
-    if (typeof method === 'function' && (method as unknown as { isDirectShellCommand: boolean }).isDirectShellCommand) {
+    if (
+      typeof method === 'function' &&
+      (method as unknown as { isDirectShellCommand: boolean })
+        .isDirectShellCommand
+    ) {
       wrapFunction(cliContext, instance, methodName, method.bind(instance));
     }
   }
@@ -146,13 +153,14 @@ export function formatHelpCommands(
   const formattedCommands = commands
     .map((c) => {
       const padding = ' '.repeat(maxCmdLength - c.cmd.length);
-      const base = `  ${chalk.yellow(c.cmd)}${padding} ${chalk.white(c.desc)}`;
+      const base = `  ${chalk.cyan(c.cmd)}${padding} ${chalk.white(c.desc)}`;
       return c.example ? `${base} ${chalk.gray(`| ${c.example}`)}` : base;
     })
     .join('\n');
 
-  return `${chalk.blue.bold('AI command suite for mongosh')}
-${chalk.gray(`Collection: ${chalk.white.bold(collection ?? 'not set')}. Set it with ai.collection("collection_name")`)}\n
+  return `${chalk.cyan.bold('mongosh AI snippet')}
+${chalk.yellow.bold('Note: This snippet is experimental and not meant for production use.')}
+${chalk.gray(`Collection: ${chalk.white.bold(collection ?? 'not set')}. Set it with ${chalk.white.bold('ai.collection("collection_name")')}`)}\n
 ${formattedCommands}\n
 ${chalk.gray(`Using ${chalk.white.bold(provider)} as provider and its ${chalk.white.bold(model)} model`)}
   `.trim();

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

@@ -229,20 +229,6 @@ export class AiProvider {
     });
   }
 
-  async query(prompt: string): Promise<void> {
-    const signal = AbortSignal.timeout(30_000);
-    await this.ensureCollectionName(prompt);
-
-    await this.processResponse(prompt, {
-      systemPrompt: await this.getSystemPrompt(
-        "You generate the exact valid mongosh command that matches the user's request.",
-        { includeSampleDocs: true },
-      ),
-      signal,
-      expectedOutput: 'command',
-    });
-  }
-
   private formatResponse({
     response,
     expectedOutput,

snippets/ai/src/providers/models.ts

@@ -8,9 +8,7 @@ export const models = {
       baseURL: 'https://knowledge.mongodb.com/api/v1',
       apiKey: '',
       headers: {
-        // TODO: Change to actual origin
         'X-Request-Origin': 'mongodb-compass',
-        origin: 'mongodb-compass',
         'user-agent':
           'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) MongoDBCompass/1.47.0 Chrome/138.0.7204.251 Electron/37.6.0 Safari/537.36',
       },
@@ -31,12 +29,12 @@ export const models = {
   },
   mistral(model = 'mistral-small-latest') {
     return createMistral({
-      apiKey: process.env.MONGOSH_AI_API_KEY,
+      apiKey: process.env.MONGOSH_AI_MISTRAL_API_KEY,
     }).languageModel(model);
   },
   openai(model = 'gpt-4.1-mini') {
     return createOpenAI({
-      apiKey: process.env.MONGOSH_AI_API_KEY,
+      apiKey: process.env.MONGOSH_AI_OPENAI_API_KEY,
     }).languageModel(model);
   },
 };