fix: working atlas provider

Author: gagik

snippets/ai/auth/atlas-service.ts snippets/ai/atlas/atlas-service.tssnippets/ai/auth/atlas-service.ts renamed to snippets/ai/atlas/atlas-service.ts

@@ -1,4 +1,4 @@
-import { log } from '../logger';
+import { log } from '../../logger';
 import { AuthService } from './auth-service';
 import type { AtlasServiceConfig } from './util';
 import { throwIfAborted } from './util';

snippets/ai/auth/auth-service.ts snippets/ai/atlas/auth-service.tssnippets/ai/auth/auth-service.ts renamed to snippets/ai/atlas/auth-service.ts

@@ -6,7 +6,7 @@ import { EventEmitter } from 'events';
 import { createMongoDBOIDCPlugin } from '@mongodb-js/oidc-plugin';
 import { oidcServerRequestHandler } from '@mongodb-js/devtools-connect';
 import { Server } from 'http';
-import { log } from '../logger';
+import { log } from '../../logger';
 
 const redirectRequestHandler = oidcServerRequestHandler.bind(null, {
   productName: 'mongosh AI Suite',

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

@@ -1,6 +1,7 @@
 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

snippets/ai/decorators.ts

@@ -1,12 +1,9 @@
-import { AIQuery } from "./auth/atlas-ai-service";
+import process from "process";
 
 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;
@@ -38,13 +35,3 @@ export function createLoadingAnimation({signal, message = 'Loading'}: {signal: A
   };
   }
 
-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/helpers.ts

@@ -1,49 +1,29 @@
-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) => {
-    output('Opening authentication page in your default browser...');
-    await open(url);
-  },
-});
-
-const atlasService = new AtlasService(authService, {
-  ...config['atlas'],
-});
-
-const aiService = new AtlasAiService({
-  atlasService,
-  apiURLPreset: 'admin-api',
-});
-
-const mongoshCommandBuilder = new MongoshCommandBuilder();
+import { AiProvider } from './providers/ai-provider';
+import { getAtlasAiProvider } from './providers/atlas/atlas-ai-provider';
 
 class AI {
-  constructor(private readonly context: any, private readonly aiService: AtlasAiService) {
+  constructor(private readonly cliContext: any, private readonly ai: AiProvider) {
     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);
+        this.wrapFunction(methodName, method.bind(this));
       }
     }
+    const instanceState = this.cliContext.db._mongo._instanceState;
+    instanceState.registerPlugin(this);
+
+    this.wrapFunction(undefined, this.help.bind(this));
   }
 
-  wrapFunction(name: string, fn: Function) {
+  private wrapFunction(name: string | undefined, fn: Function) {
     const wrapperFn = (...args: string[]) => {
       return Object.assign(fn(...args), {
         [Symbol.for('@@mongosh.syntheticPromise')]: true,
@@ -52,49 +32,31 @@ class AI {
     wrapperFn.isDirectShellCommand = true;
     wrapperFn.returnsPromise = true;
 
-    const instanceState = this.context.db._mongo._instanceState;
+    const instanceState = this.cliContext.db._mongo._instanceState;
 
-    (instanceState as any).shellApi[`ai.${name}`] = (instanceState as any).context[`ai.${name}`] = wrapperFn;
-    instanceState.registerPlugin(this);
+    instanceState.shellApi[name ? `ai.${name}` : 'ai'] = instanceState.context[name ? `ai.${name}` : 'ai'] = wrapperFn;
   }
 
   @aiCommand
+  @withLoadingAnimation('Generating query...')
   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,
-          requestId: 'test',
-        },
-        {
-          connectionOptions: {
-            connectionString: 'mongodb://localhost:27017',
-          },
-          id: '1234',
-        },
-      );
-      loadingAnimation.stop();
+    return await this.ai.query(code);
+  }
 
-      const query = mongoshCommandBuilder.createMongoShellQuery(result.content);
-      setInput(query);
+  @aiCommand
+  @withLoadingAnimation('Thinking...')
+  async ask(code: string) {
+    return await this.ai.ask(code);
   }
 
   @aiCommand
-  @withLoadingAnimation('Generating help...')
   async help(...args: string[]) {
-    await new Promise(resolve => setTimeout(resolve, 1000));
+    this.ai.help();
   }
 }
 
-
 module.exports = (globalThis: any) => {
-  globalThis.ai = new AI(globalThis, aiService);
+  globalThis.ai = new AI(globalThis, getAtlasAiProvider());
 };