combine and simplify

Author: gagik

snippets/ai/config.ts

@@ -1,29 +1,35 @@
 import { AiProvider } from './providers/ai-provider';
 import { EventEmitter } from 'events';
 import { z } from 'zod';
-
+import chalk from 'chalk';
+import { inspect } from 'util';
 
 const configSchema = z.object({
   provider: z.enum(['docs', 'openai', 'mistral', 'atlas', 'ollama']),
   model: z.string(),
+  includeSampleDocs: z.boolean(),
+  defaultCollection: z.string().optional(),
 });
 
 const configKeys = Object.keys(configSchema.shape) as Array<keyof ConfigSchema>;
 
-
 export type ConfigSchema = z.infer<typeof configSchema>;
 type ConfigKeys = keyof ConfigSchema;
 
 const defaults: Record<ConfigKeys, any> = {
   provider: process.env.MONGOSH_AI_PROVIDER ?? 'docs',
   model: process.env.MONGOSH_AI_MODEL ?? 'default',
+  includeSampleDocs: process.env.MONGOSH_AI_INCLUDE_SAMPLE_DOCS ?? true,
+  defaultCollection: process.env.MONGOSH_AI_DEFAULT_COLLECTION,
 };
 
 export class Config extends EventEmitter<{
-  change: [{
-    key: ConfigKeys;
-    value: ConfigSchema[ConfigKeys];
-  }];
+  change: [
+    {
+      key: ConfigKeys;
+      value: ConfigSchema[ConfigKeys];
+    },
+  ];
 }> {
   private configMap: Record<string, any> = {};
 
@@ -68,6 +74,16 @@ export class Config extends EventEmitter<{
   }
 
   [Symbol.for('nodejs.util.inspect.custom')]() {
-    return this.configMap;
+    const lines = Object.entries(configSchema.shape).map(([key, schema]) => {
+      let type: string | undefined = undefined;
+      if (schema._def.typeName === 'ZodEnum') {
+        type = `${schema._def.values.join(' | ')}`;
+      }
+      const i = (value: any) => inspect(value, {colors: true});
+
+      return `  ${i(key)}: ${chalk.white(i(this.configMap[key]))},${type ? chalk.gray(` // ${type}`) : ''}`;
+    });
+
+    return `{\n${lines.join('\n')}\n}`;
   }
 }

snippets/ai/decorators.ts

@@ -1,13 +1,26 @@
-export function aiCommand<T extends Function>(
+export interface AiCommandOptions {
+  requiresPrompt?: boolean;
+}
+
+export function aiCommand({
+  requiresPrompt = true,
+}: AiCommandOptions = {}) {
+  return function decorator<T extends Function>(
     value: T,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     context: ClassMethodDecoratorContext
   ): T & { isDirectShellCommand: true } {
     const wrappedFunction = function(this: any, ...args: any[]) {
+      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');
+      }
       // Combine all arguments into a single string
       const combinedString = args.join(' ');
       // 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
     return Object.assign(wrappedFunction, { isDirectShellCommand: true } as const);
   }
+}

snippets/ai/helpers.ts

@@ -68,7 +68,7 @@ export function wrapFunction(
   cliContext: CliContext,
   instance: any,
   name: string | undefined,
-  fn: Function
+  fn: Function,
 ) {
   const wrapperFn = (...args: string[]) => {
     return Object.assign(fn(...args), {
@@ -88,11 +88,11 @@ export function wrapFunction(
 export function wrapAllFunctions(cliContext: CliContext, instance: any) {
   const instanceState = cliContext.db._mongo._instanceState;
   const methods = Object.getOwnPropertyNames(
-    Object.getPrototypeOf(instance)
+    Object.getPrototypeOf(instance),
   ).filter((name) => {
     const descriptor = Object.getOwnPropertyDescriptor(
       Object.getPrototypeOf(instance),
-      name
+      name,
     );
     return (
       descriptor &&
@@ -119,16 +119,26 @@ interface HelpCommand {
   example?: string;
 }
 
-export function formatHelpCommands(commands: HelpCommand[], provider: string, model: string): string {
-  const maxCmdLength = Math.max(...commands.map(c => c.cmd.length));
-  const formattedCommands = commands.map(c => {
-    const padding = ' '.repeat(maxCmdLength - c.cmd.length);
-    const base = `  ${chalk.yellow(c.cmd)}${padding} ${chalk.white(c.desc)}`;
-    return c.example ? `${base} ${chalk.gray(`| ${c.example}`)}` : base;
-  }).join('\n');
+export function formatHelpCommands(
+  commands: HelpCommand[],
+  {
+    provider,
+    model,
+    collection,
+  }: { provider: string; model: string; collection?: string },
+): string {
+  const maxCmdLength = Math.max(...commands.map((c) => c.cmd.length));
+  const formattedCommands = commands
+    .map((c) => {
+      const padding = ' '.repeat(maxCmdLength - c.cmd.length);
+      const base = `  ${chalk.yellow(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(`Using ${chalk.white.bold(provider)} as provider and its ${chalk.white.bold(model)} model`)}\n
-${formattedCommands}
+${chalk.gray(`Collection: ${chalk.white.bold(collection ?? 'not set')}. Set it with 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/index.ts

@@ -2,12 +2,10 @@ import { aiCommand } from './decorators';
 import { AiProvider, EmptyAiProvider } from './providers/ai-provider';
 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 { getAiSdkProvider, models } from './providers/generic/ai-sdk-provider';
 import { Config, ConfigSchema } from './config';
 import { CliContext, wrapAllFunctions, formatHelpCommands } from './helpers';
+import chalk from 'chalk';
 
 class AI {
   private readonly replConfig: {
@@ -37,12 +35,13 @@ class AI {
           this.ai = this.getProvider(event.value as ConfigSchema['provider']);
           break;
         case 'model':
-          if (Object.keys(models).includes(event.value)) {
+          if (Object.keys(models).includes(event.value as string)) {
             this.ai = getAiSdkProvider(
               models[this.config.get('provider') as keyof typeof models](
-                event.value,
+                event.value as string,
               ),
               this.cliContext,
+              this.config,
             );
           } else {
             throw new Error(`Invalid model: ${event.value}`);
@@ -53,7 +52,9 @@ class AI {
       }
     });
 
-    this.ai = this.getProvider(process.env.MONGOSH_AI_PROVIDER as ConfigSchema['provider'] | undefined);
+    this.ai = this.getProvider(
+      process.env.MONGOSH_AI_PROVIDER as ConfigSchema['provider'] | undefined,
+    );
     wrapAllFunctions(this.cliContext, this);
 
     this.setupConfig();
@@ -65,66 +66,85 @@ class AI {
     this.ai = this.getProvider(this.config.get('provider'));
   }
 
-  private getProvider(provider: ConfigSchema['provider'] | undefined): AiProvider {
+  private getProvider(
+    provider: ConfigSchema['provider'] | undefined,
+  ): AiProvider {
     switch (provider) {
       case 'docs':
-        return getDocsAiProvider(this.cliContext);
+        return getDocsAiProvider(this.cliContext, this.config);
       case 'atlas':
-        return getAtlasAiProvider(this.cliContext);
+        return getAtlasAiProvider(this.cliContext, this.config);
       case 'openai':
       case 'mistral':
       case 'ollama':
         const model = this.config.get('model');
         return getAiSdkProvider(
           models[provider](model === 'default' ? undefined : model),
           this.cliContext,
+          this.config,
         );
       default:
-        return new EmptyAiProvider(this.cliContext);
+        return new EmptyAiProvider(this.cliContext, this.config);
     }
   }
 
-  @aiCommand
-  async command(prompt: string) {
-    await this.ai.command(prompt);
+  @aiCommand()
+  async shell(prompt: string) {
+    await this.ai.shell(prompt);
+  }
+
+  @aiCommand()
+  async data(prompt: string) {
+    await this.ai.data(prompt);
   }
 
-  @aiCommand
+  @aiCommand()
   async query(prompt: string) {
     await this.ai.query(prompt);
   }
 
-  @aiCommand
+  @aiCommand()
   async ask(prompt: string) {
     await this.ai.ask(prompt);
   }
 
-  @aiCommand
+  @aiCommand()
   async aggregate(prompt: string) {
     await this.ai.aggregate(prompt);
   }
 
-  @aiCommand
-  async help(...args: string[]) {
-    const commands = [
-      { cmd: 'ai.ask', desc: 'ask questions', example: 'ai.ask how do I run queries in mongosh?' },
-      { cmd: 'ai.command', desc: 'generate any mongosh command', example: 'ai.command create a new database' },
-      { 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.config', desc: 'configure the AI commands', example: 'ai.config.set("provider", "ollama")' }
-    ];
-
-    this.ai.respond(
-      formatHelpCommands(
-        commands,
-        this.config.get('provider'),
-        this.config.get('model')
-      )
-    );
+  @aiCommand({requiresPrompt: false})
+  async help() {
+    this.ai.help({
+      provider: this.config.get('provider'),
+      model: this.config.get('model'),
+    });
+  }
+
+  @aiCommand()
+  async clear() {
+    this.ai.clear();
+  }
+
+  @aiCommand()
+  async collection(name: string) {
+    await this.ai.collection(name);
+  }
+
+  @aiCommand()
+  async provider(provider: string) {
+    this.config.set('provider', provider);
+    this.ai.respond(`Switched to ${chalk.blue(provider)} provider`);
+  }
+
+  @aiCommand()
+  async model(model: string) {
+    this.config.set('model', model);
+    this.ai.respond(`Switched to ${chalk.blue(model)} model`);
   }
 
   [Symbol.for('nodejs.util.inspect.custom')]() {
-    this.ai.help();
+    this.help();
     return '';
   }
 }