chore: basic structure

Author: gagik

snippets/ai/auth/atlas-ai-service.ts

@@ -35,7 +35,7 @@ type AIFeatureEnablement = {
   };
 };
 
-type AIQuery = {
+export type AIQuery = {
   content: {
     query: Record<
       'filter' | 'project' | 'collation' | 'sort' | 'skip' | 'limit',

snippets/ai/auth/atlas-service.ts

@@ -1,3 +1,4 @@
+import { log } from '../logger';
 import { AuthService } from './auth-service';
 import type { AtlasServiceConfig } from './util';
 import { throwIfAborted } from './util';
@@ -62,13 +63,13 @@ export class AtlasService {
       throw new Error('Network traffic is not allowed');
     }
     throwIfAborted(init?.signal as AbortSignal);
-    const authHeaders = this.authService.getAuthHeaders();
+    const authHeaders = await this.authService.getAuthHeaders();
     const finalHeaders = {
       ...authHeaders,
       ...this.options?.defaultHeaders,
       ...init?.headers,
     };
-    console.log('AtlasService: Making a fetch', {
+    log.debug('AtlasService: Making a fetch', {
       url,
       headers: finalHeaders,
       method: init?.method || 'GET',
@@ -82,7 +83,7 @@ export class AtlasService {
       res.headers.forEach((value, key) => {
         responseHeadersObj[key] = value;
       });
-      console.log('AtlasService: Received API response', {
+      log.debug('AtlasService: Received API response', {
         url,
         status: res.status,
         statusText: res.statusText,
@@ -91,7 +92,7 @@ export class AtlasService {
       await throwIfNotOk(res);
       return res;
     } catch (err) {
-      console.log('AtlasService: Fetch errored', {
+      log.error('AtlasService: Fetch errored', {
         url,
         error: err instanceof Error ? err.message : err,
         stack: err instanceof Error ? err.stack : undefined,
@@ -104,7 +105,7 @@ export class AtlasService {
     init?: RequestInit,
   ): Promise<Response> {
     const authHeaders = await this.authService.getAuthHeaders();
-    console.log('AtlasService: Authenticated fetch', { url, authHeaders });
+    log.debug('AtlasService: Authenticated fetch', { url, authHeaders });
     const res = await this.fetch(url, {
       ...init,
       headers: {

snippets/ai/auth/auth-service.ts

@@ -6,27 +6,11 @@ import { EventEmitter } from 'events';
 import { createMongoDBOIDCPlugin } from '@mongodb-js/oidc-plugin';
 import { oidcServerRequestHandler } from '@mongodb-js/devtools-connect';
 import { Server } from 'http';
-
-class OIDCLogger extends EventEmitter {
-  debug(...args: unknown[]) {
-    console.debug(...args);
-    this.emit('debug', ...args);
-  }
-
-  info(...args: unknown[]) {
-    console.info(...args);
-    this.emit('info', ...args);
-  }
-
-  error(...args: unknown[]) {
-    console.error(...args);
-    this.emit('error', ...args);
-  }
-}
+import { log } from '../logger';
 
 const redirectRequestHandler = oidcServerRequestHandler.bind(null, {
-  productName: 'Compass',
-  productDocsLink: 'https://www.mongodb.com/docs/compass',
+  productName: 'mongosh AI Suite',
+  productDocsLink: 'https://www.mongodb.com/docs/mongosh',
 });
 
 export class AuthService extends AtlasAuthService {
@@ -51,7 +35,7 @@ export class AuthService extends AtlasAuthService {
     const { openBrowser, ...atlasConfig } = config;
     this.config = atlasConfig;
     this.openBrowser = openBrowser;
-    console.log('Initializing AuthService with config:', {
+    log.debug('Initializing AuthService with config:', {
       ...this.config,
     });
 
@@ -73,7 +57,7 @@ export class AuthService extends AtlasAuthService {
         await this.openBrowser(url);
       },
       allowedFlows: ['auth-code', 'device-auth'],
-      logger: new OIDCLogger(),
+      logger: log,
     });
   }
 
@@ -148,7 +132,7 @@ export class AuthService extends AtlasAuthService {
         )
       ).accessToken;
 
-      console.log('token', this.token);
+      log.debug('token', this.token);
       return {
         primaryEmail: 'test@example.com',
         sub: 'test',

snippets/ai/decorators.ts

@@ -0,0 +1,37 @@
+import { output } from "./helpers";
+import { createLoadingAnimation } from "./helpers";
+
+export function aiCommand<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[]) {
+      // 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);
+  }
+
+export function withLoadingAnimation(message: string = 'Loading...') {
+  return function<T extends Function>(
+    value: T,
+    context: ClassMethodDecoratorContext
+  ): T {
+    const wrappedFunction = async function(this: any, ...args: any[]) {
+      const signal = AbortSignal.timeout(10000);
+      const loadingAnimation = createLoadingAnimation({signal, message});
+      loadingAnimation.start();
+
+      try {
+        const result = await value.call(this, ...args);
+        return result;
+      } finally {
+        loadingAnimation.stop();
+      }
+    } as unknown as T;
+    return wrappedFunction;
+  };
+}

snippets/ai/helpers.ts

@@ -0,0 +1,50 @@
+import { AIQuery } from "./auth/atlas-ai-service";
+
+export function output(text: string) {
+  process.stdout.write(`${text}`);
+}
+
+export function setInput(text: string) {
+  process.stdin.unshift(text);
+}
+export function createLoadingAnimation({signal, message = 'Loading'}: {signal: AbortSignal, message?: string}): {
+  start: (message?: string) => void;
+  stop: () => void;
+} {
+  const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
+  let i = 0;
+  let interval: NodeJS.Timeout | null = null;
+  
+  return {
+    start() {
+      interval = setInterval(() => {
+        const frame = frames[i = ++i % frames.length];
+        process.stdout.write(`\r${frame} ${message}`);
+      }, 80);
+
+      signal.addEventListener('abort', () => {
+        if (interval) {
+          clearInterval(interval);
+          process.stdout.write('\r\x1b[K'); // Clear the line
+        }
+      }, { once: true });
+    },
+    stop() {
+      if (interval) {
+        clearInterval(interval);
+        process.stdout.write('\r\x1b[K'); // Clear the line
+      }
+    }
+  };
+  }
+
+export class MongoshCommandBuilder {
+  createMongoShellQuery(params: AIQuery['content']): string {
+    const {filter, project, collation, sort, skip, limit} = params.query;
+    
+  return `db.collection.find(
+    ${filter},
+    ${project ? `{ projection: ${project} }` : '{}'}
+  )${collation ? `.collation(${collation})` : ''}${sort ? `.sort(${sort})` : ''}${skip ? `.skip(${skip})` : ''}${limit ? `.limit(${limit})` : ''}`
+  };
+}

snippets/ai/index.ts

@@ -2,12 +2,14 @@ import { AuthService } from './auth/auth-service';
 import { AtlasService } from './auth/atlas-service';
 import { AtlasAiService } from './auth/atlas-ai-service';
 import { config } from './auth/util';
+import { createLoadingAnimation, MongoshCommandBuilder, output, setInput } from './helpers';
 import open from 'open';
+import { aiCommand, withLoadingAnimation } from './decorators';
 
 const authService = new AuthService({
   ...config['atlas'],
   openBrowser: async (url: string) => {
-    console.log('\nOpening authentication page in your default browser...');
+    output('Opening authentication page in your default browser...');
     await open(url);
   },
 });
@@ -21,18 +23,53 @@ const aiService = new AtlasAiService({
   apiURLPreset: 'admin-api',
 });
 
-module.exports = (globalThis: any) => {
-  globalThis.ai = {
-    login: async () => {
-      await authService.signIn();
-    },
-    explain: async (code: string) => {
+const mongoshCommandBuilder = new MongoshCommandBuilder();
+
+class AI {
+  constructor(private readonly context: any, private readonly aiService: AtlasAiService) {
+    const methods = Object.getOwnPropertyNames(Object.getPrototypeOf(this))
+      .filter(name => {
+        const descriptor = Object.getOwnPropertyDescriptor(Object.getPrototypeOf(this), name);
+        return descriptor && typeof descriptor.value === 'function' && name !== 'constructor';
+      });
+    console.log('Class methods:', methods);
+    
+    // for all methods, wrap them with the wrapFunction method
+    for (const methodName of methods) {
+      const method = (this as any)[methodName];
+      if (typeof method === 'function' && method.isDirectShellCommand) {
+        this.wrapFunction(methodName, method);
+      }
+    }
+  }
+
+  wrapFunction(name: string, fn: Function) {
+    const wrapperFn = (...args: string[]) => {
+      return Object.assign(fn(...args), {
+        [Symbol.for('@@mongosh.syntheticPromise')]: true,
+      });
+    };
+    wrapperFn.isDirectShellCommand = true;
+    wrapperFn.returnsPromise = true;
+
+    const instanceState = this.context.db._mongo._instanceState;
+
+    (instanceState as any).shellApi[`ai.${name}`] = (instanceState as any).context[`ai.${name}`] = wrapperFn;
+    instanceState.registerPlugin(this);
+  }
+
+  @aiCommand
+  async query(code: string) {
+    const signal = AbortSignal.timeout(10000);
+      const loadingAnimation = createLoadingAnimation({signal, message: 'Generating query...'});
+      loadingAnimation.start();
+
       const result = await aiService.getQueryFromUserInput(
         {
           userInput: code,
           databaseName: 'test',
           collectionName: 'test',
-          signal: AbortSignal.timeout(10000),
+          signal,
           requestId: 'test',
         },
         {
@@ -42,10 +79,22 @@ module.exports = (globalThis: any) => {
           id: '1234',
         },
       );
-      return JSON.stringify(result.content.query);
-    },
-    ask: async (question: string) => {
-      return 'test';
-    },
-  };
+      loadingAnimation.stop();
+
+      const query = mongoshCommandBuilder.createMongoShellQuery(result.content);
+      setInput(query);
+  }
+
+  @aiCommand
+  @withLoadingAnimation('Generating help...')
+  async help(...args: string[]) {
+    await new Promise(resolve => setTimeout(resolve, 1000));
+  }
+}
+
+
+module.exports = (globalThis: any) => {
+  globalThis.ai = new AI(globalThis, aiService);
 };
+
+

snippets/ai/logger.ts

@@ -0,0 +1,28 @@
+import EventEmitter from "events";
+
+const IS_DEBUG = process.env.DEBUG === 'true';
+
+class Logger extends EventEmitter {
+    debug(...args: unknown[]) {
+      if (IS_DEBUG) {
+        console.debug(...args);
+        this.emit('debug', ...args);
+      }
+    }
+  
+    info(...args: unknown[]) {
+      if (IS_DEBUG) {
+        console.info(...args);
+        this.emit('info', ...args);
+      }
+    }
+  
+    error(...args: unknown[]) {
+      if (IS_DEBUG) {
+        console.error(...args);
+        this.emit('error', ...args);
+      }
+    }
+  }
+  
+  export const log = new Logger();