1- import { EmptyAiProvider , type AiProvider } from './providers/ai-provider.js' ;
2- import { Config , ConfigSchema } from './config.js' ;
3- import { wrapAllFunctions , type CliContext } from './helpers.js' ;
4- import { getAiSdkProvider , models , type Models } from './providers/generic/ai-sdk-provider.js' ;
5- import { getDocsAiProvider } from './providers/docs/docs-ai-provider.js' ;
1+ import { AiProvider } from './providers/ai-provider.js' ;
2+ import { Config } from './config.js' ;
3+ import {
4+ formatHelpCommands ,
5+ wrapAllFunctions ,
6+ type CliContext ,
7+ } from './helpers.js' ;
8+ import { models } from './providers/models.js' ;
69import { aiCommand } from './decorators.js' ;
710import chalk from 'chalk' ;
811
9- module . exports = ( async ( globalThis : CliContext ) => {
10- class AI {
11- private readonly replConfig : {
12- set : ( key : string , value : unknown ) => Promise < void > ;
13- get : < T > ( key : string ) => Promise < T > ;
14- } ;
15-
16- private ai : AiProvider ;
17- public config : Config ;
18-
19- constructor ( private readonly cliContext : CliContext ) {
20- const instanceState = this . cliContext . db . _mongo . _instanceState ;
21-
22- this . replConfig = {
23- set : ( key , value ) =>
24- instanceState . evaluationListener . setConfig ( `snippet_ai_${ key } , value ) ,
25- get : ( key ) =>
26- instanceState . evaluationListener . getConfig ( `snippet_ai_${ key } ) ,
27- } ;
28-
29- this . config = new Config ( this . replConfig ) ;
30-
31- // Set up provider change listener
32- // eslint-disable-next-line @typescript-eslint/no-misused-promises
33- this . config . on ( 'change' , async ( event ) => {
34- switch ( event . key ) {
35- case 'provider' :
36- this . ai = this . getProvider ( event . value as ConfigSchema [ 'provider' ] ) ;
37- break ;
38- case 'model' :
39- if ( ! Object . keys ( models ) . includes ( this . config . get ( 'provider' ) as Models ) ) {
40- if ( event . value === 'default' ) {
41- return ;
42- }
43- await this . config . set ( 'model' , 'default' ) ;
44- throw new Error ( `${ this . config . get ( 'provider' ) } ) ;
45- }
46- try {
47- this . ai = getAiSdkProvider (
48- models [ this . config . get ( 'provider' ) as keyof typeof models ] (
49- event . value === 'default' ? undefined : event . value as string ,
50- ) ,
51- this . cliContext ,
52- this . config ,
53- ) ;
54- } catch ( error ) {
55- throw new Error ( `Invalid model, please ensure your name is correct: ${ error as string } ) ;
56- }
57- break ;
58- default :
59- break ;
60- }
61- } ) ;
62-
63- this . ai = this . getProvider (
64- process . env . MONGOSH_AI_PROVIDER as ConfigSchema [ 'provider' ] | undefined ,
65- ) ;
66- wrapAllFunctions ( this . cliContext , this ) ;
12+ module . exports = async ( globalThis : CliContext ) => {
13+ class AI {
14+ private readonly ai : AiProvider ;
15+ public config : Config ;
6716
68- }
69-
70- async setup ( ) {
71- await this . config . setup ( ) ;
72-
73- this . ai = this . getProvider ( this . config . get ( 'provider' ) ) ;
74- }
75-
76- private getProvider (
77- provider : ConfigSchema [ 'provider' ] | undefined ,
78- ) : AiProvider {
79- switch ( provider ) {
80- case 'docs' :
81- return getDocsAiProvider ( this . cliContext , this . config ) ;
82- case 'openai' :
83- case 'mistral' :
84- case 'ollama' : {
85- const model = this . config . get ( 'model' ) ;
86- return getAiSdkProvider (
87- models [ provider ] ( model === 'default' ? undefined : model ) ,
88- this . cliContext ,
89- this . config ,
90- ) ;
91- }
92- default :
93- return new EmptyAiProvider ( this . cliContext , this . config ) ;
17+ constructor ( { ai, config } : { ai : AiProvider ; config : Config } ) {
18+ this . ai = ai ;
19+ this . config = config ;
9420 }
95- }
96-
97- @aiCommand ( )
98- async shell ( prompt : string ) {
99- await this . ai . shell ( prompt ) ;
100- }
10121
102- @aiCommand ( )
103- async general ( prompt : string ) {
104- await this . ai . general ( prompt ) ;
105- }
22+ static async create ( cliContext : CliContext ) : Promise < AI > {
23+ const instanceState = cliContext . db . _mongo . _instanceState ;
24+
25+ const replConfig = {
26+ set : ( key : string , value : unknown ) =>
27+ instanceState . evaluationListener . setConfig (
28+ `snippet_ai_${ key } ,
29+ value ,
30+ ) ,
31+ get : < T > ( key : string ) : Promise < T > =>
32+ instanceState . evaluationListener . getConfig ( `snippet_ai_${ key } ) ,
33+ } ;
34+
35+ const config = await Config . create ( replConfig ) ;
36+
37+ const provider = new AiProvider (
38+ cliContext ,
39+ config ,
40+ models [ config . get ( 'provider' ) as keyof typeof models ] (
41+ config . get ( 'model' ) === 'default' ? undefined : config . get ( 'model' ) ,
42+ ) ,
43+ ) ;
44+
45+ const ai = new AI ( {
46+ ai : provider ,
47+ config,
48+ } ) ;
49+
50+ wrapAllFunctions ( cliContext , ai ) ;
51+ return ai ;
52+ }
10653
107- @aiCommand ( )
108- async data ( prompt : string ) {
109- await this . ai . data ( prompt ) ;
110- }
54+ @aiCommand ( )
55+ async shell ( prompt : string ) {
56+ await this . ai . shell ( prompt ) ;
57+ }
11158
112- @aiCommand ( )
113- async query ( prompt : string ) {
114- await this . ai . query ( prompt ) ;
115- }
59+ @aiCommand ( )
60+ async query ( prompt : string ) {
61+ await this . ai . query ( prompt ) ;
62+ }
11663
117- @aiCommand ( )
118- async ask ( prompt : string ) {
119- await this . ai . ask ( prompt ) ;
120- }
64+ @aiCommand ( )
65+ async aggregate ( prompt : string ) {
66+ await this . ai . aggregate ( prompt ) ;
67+ }
12168
122- @aiCommand ( )
123- async aggregate ( prompt : string ) {
124- await this . ai . aggregate ( prompt ) ;
125- }
69+ @aiCommand ( )
70+ async ask ( prompt : string ) {
71+ await this . ai . processResponse ( prompt , {
72+ systemPrompt :
73+ 'You are a MongoDB and mongosh expert. Give brief answers without any formatting.' ,
74+ expectedOutput : 'response' ,
75+ } ) ;
76+ }
12677
127- @aiCommand ( { requiresPrompt : false } )
128- help ( ) {
129- this . ai . help ( {
130- provider : this . config . get ( 'provider' ) ,
131- model : this . config . get ( 'model' ) ,
132- } ) ;
133- }
78+ @aiCommand ( { requiresPrompt : false } )
79+ help ( ) {
80+ const commands = [
81+ {
82+ cmd : 'ai.ask' ,
83+ desc : 'ask MongoDB questions' ,
84+ example : 'ai.ask how do I run queries in mongosh?' ,
85+ } ,
86+ {
87+ cmd : 'ai.query' ,
88+ desc : 'generate a MongoDB query' ,
89+ example : 'ai.query find documents where name = "Ada"' ,
90+ } ,
91+ {
92+ cmd : 'ai.aggregate' ,
93+ desc : 'generate a MongoDB aggregation' ,
94+ example : 'ai.aggregate find documents where name = "Ada"' ,
95+ } ,
96+ {
97+ cmd : 'ai.collection' ,
98+ desc : 'set the active collection' ,
99+ example : 'ai.collection("users")' ,
100+ } ,
101+ {
102+ cmd : 'ai.shell' ,
103+ desc : 'generate administrative mongosh commands' ,
104+ example : 'ai.shell insert a new sample document' ,
105+ } ,
106+ {
107+ cmd : 'ai.config' ,
108+ desc : 'configure the AI commands' ,
109+ example : 'ai.config.set("provider", "ollama")' ,
110+ } ,
111+ ] ;
112+
113+ this . ai . respond (
114+ formatHelpCommands ( commands , {
115+ provider : this . config . get ( 'provider' ) ,
116+ model : this . config . get ( 'model' ) ,
117+ collection : this . ai . activeCollection ,
118+ } ) ,
119+ ) ;
120+ }
134121
135- @aiCommand ( )
136- clear ( ) {
137- this . ai . clear ( ) ;
138- }
122+ @aiCommand ( )
123+ clear ( ) {
124+ this . ai . clear ( ) ;
125+ }
139126
140- @aiCommand ( )
141- collection ( name : string ) {
142- this . ai . collection ( name ) ;
143- }
127+ @aiCommand ( )
128+ collection ( name : string ) {
129+ this . ai . collection ( name ) ;
130+ }
144131
145- @aiCommand ( )
146- async provider ( provider : string ) {
147- await this . config . set ( 'provider' , provider ) ;
148- this . ai . respond ( `Switched to ${ chalk . blue ( provider ) } ) ;
149- }
132+ @aiCommand ( )
133+ async provider ( provider : string ) {
134+ await this . config . set ( 'provider' , provider ) ;
135+ this . ai . respond ( `Switched to ${ chalk . blue ( provider ) } ) ;
136+ }
150137
151- @aiCommand ( )
152- async model ( model : string ) {
153- await this . config . set ( 'model' , model ) ;
154- this . ai . respond ( `Switched to ${ chalk . blue ( model ) } ) ;
155- }
138+ @aiCommand ( )
139+ async model ( model : string ) {
140+ await this . config . set ( 'model' , model ) ;
141+ this . ai . respond ( `Switched to ${ chalk . blue ( model ) } ) ;
142+ }
156143
157- [ Symbol . for ( 'nodejs.util.inspect.custom' ) ] ( ) {
158- this . help ( ) ;
159- return '' ;
144+ [ Symbol . for ( 'nodejs.util.inspect.custom' ) ] ( ) {
145+ this . help ( ) ;
146+ return '' ;
147+ }
160148 }
161- }
162-
163-
164- ( globalThis as unknown as CliContext ) . ai = new AI ( globalThis as unknown as CliContext ) ;
165- await ( globalThis as unknown as {
166- ai : AI ;
167- } ) . ai . setup ( ) ;
168149
169- } ) ;
150+ ( globalThis as unknown as CliContext ) . ai = await AI . create (
151+ globalThis as unknown as CliContext ,
152+ ) ;
153+ } ;
0 commit comments