Skip to content

Commit 496c573

Browse files
gagikgagik
committed
full docs chatbot suite
1 parent f9d222f commit 496c573

File tree

7 files changed

+207
-143
lines changed

7 files changed

+207
-143
lines changed

‎snippets/ai/decorators.ts‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,3 @@
1-
import { output } from "./helpers";
2-
import { createLoadingAnimation } from "./helpers";
3-
4-
51
export function aiCommand<T extends Function>(
62
value: T,
73
// eslint-disable-next-line @typescript-eslint/no-unused-vars

‎snippets/ai/helpers.ts‎

Lines changed: 41 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,49 @@
1-
import process from "process";
2-
1+
import chalk from 'chalk';
2+
import process from 'process';
33

44
export function output(text: string) {
55
process.stdout.write(`${text}`);
66
}
77

8-
export function createLoadingAnimation({message = 'Loading'}: {message?: string}): {
9-
start: (signal: AbortSignal) => void;
10-
stop: () => void;
11-
} {
12-
const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
13-
let i = 0;
14-
let interval: NodeJS.Timeout | null = null;
15-
16-
return {
17-
start(
18-
signal: AbortSignal,
19-
) {
20-
interval = setInterval(() => {
21-
const frame = frames[i = ++i % frames.length];
22-
process.stdout.write(`\r${frame} ${message}`);
23-
}, 80);
24-
25-
signal.addEventListener('abort', () => {
26-
if (interval) {
27-
clearInterval(interval);
28-
process.stdout.write('\r\x1b[K'); // Clear the line
29-
}
30-
}, { once: true });
31-
},
32-
stop() {
33-
if (interval) {
34-
clearInterval(interval);
35-
process.stdout.write('\r\x1b[K'); // Clear the line
36-
}
8+
export class LoadingAnimation {
9+
private frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
10+
private i = 0;
11+
private interval: NodeJS.Timeout | null = null;
12+
private abortListener: (() => void) | null = null;
13+
private message: string;
14+
private signal: AbortSignal | null = null;
15+
16+
constructor({ message = 'Loading' }: { message?: string }) {
17+
this.message = message;
18+
}
19+
20+
start(signal: AbortSignal): void {
21+
if (this.interval) {
22+
return;
3723
}
38-
};
24+
25+
this.interval = setInterval(() => {
26+
const frame = this.frames[(this.i = ++this.i % this.frames.length)];
27+
process.stdout.write(chalk.blue(`\r${frame} ${this.message}`));
28+
}, 80);
29+
30+
this.abortListener = () => {
31+
this.stop();
32+
};
33+
34+
signal.addEventListener('abort', this.abortListener, { once: true });
3935
}
4036

37+
stop(): void {
38+
if (this.signal && this.abortListener) {
39+
this.signal.removeEventListener('abort', this.abortListener);
40+
this.abortListener = null;
41+
}
42+
43+
if (this.interval) {
44+
clearInterval(this.interval);
45+
process.stdout.write('\r\x1b[K'); // Clear the line
46+
this.interval = null;
47+
}
48+
}
49+
}

‎snippets/ai/index.ts‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,24 @@ import { AiProvider } from './providers/ai-provider';
33
import { getDocsAiProvider } from './providers/docs/docs-ai-provider';
44

55
class AI {
6-
constructor(private readonly cliContext: any, private readonly ai: AiProvider) {
7-
const methods = Object.getOwnPropertyNames(Object.getPrototypeOf(this))
8-
.filter(name => {
9-
const descriptor = Object.getOwnPropertyDescriptor(Object.getPrototypeOf(this), name);
10-
return descriptor && typeof descriptor.value === 'function' && name !== 'constructor';
11-
});
12-
6+
constructor(
7+
private readonly cliContext: any,
8+
private readonly ai: AiProvider,
9+
) {
10+
const methods = Object.getOwnPropertyNames(
11+
Object.getPrototypeOf(this),
12+
).filter((name) => {
13+
const descriptor = Object.getOwnPropertyDescriptor(
14+
Object.getPrototypeOf(this),
15+
name,
16+
);
17+
return (
18+
descriptor &&
19+
typeof descriptor.value === 'function' &&
20+
name !== 'constructor'
21+
);
22+
});
23+
1324
// for all methods, wrap them with the wrapFunction method
1425
for (const methodName of methods) {
1526
const method = (this as any)[methodName];
@@ -34,7 +45,9 @@ class AI {
3445

3546
const instanceState = this.cliContext.db._mongo._instanceState;
3647

37-
instanceState.shellApi[name ? `ai.${name}` : 'ai'] = instanceState.context[name ? `ai.${name}` : 'ai'] = wrapperFn;
48+
instanceState.shellApi[name ? `ai.${name}` : 'ai'] = instanceState.context[
49+
name ? `ai.${name}` : 'ai'
50+
] = wrapperFn;
3851
}
3952

4053
@aiCommand
@@ -47,14 +60,17 @@ class AI {
4760
return await this.ai.ask(code);
4861
}
4962

63+
@aiCommand
64+
async aggregate(code: string) {
65+
return await this.ai.aggregate(code);
66+
}
67+
5068
@aiCommand
5169
async help(...args: string[]) {
5270
this.ai.help();
5371
}
5472
}
5573

5674
module.exports = (globalThis: any) => {
57-
globalThis.ai = new AI(globalThis, getDocsAiProvider());
75+
globalThis.ai = new AI(globalThis, getDocsAiProvider(globalThis));
5876
};
59-
60-

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
1-
import process from "process";
2-
import { createLoadingAnimation } from "../helpers";
1+
import process from 'process';
2+
import { LoadingAnimation } from '../helpers';
3+
4+
export type CliContext = any;
35

46
export abstract class AiProvider {
5-
thinking: { start: (signal: AbortSignal) => void; stop: () => void };
7+
thinking: LoadingAnimation;
68

7-
constructor() {
8-
this.thinking = createLoadingAnimation({
9+
constructor(private readonly cliContext: CliContext) {
10+
this.thinking = new LoadingAnimation({
911
message: 'Thinking...',
1012
});
1113
}
@@ -17,7 +19,7 @@ export abstract class AiProvider {
1719
connectionString: 'mongodb://localhost:27017',
1820
},
1921
id: '1234',
20-
}
22+
};
2123
}
2224

2325
/** @internal */
@@ -26,12 +28,11 @@ export abstract class AiProvider {
2628
collectionName: string;
2729
} {
2830
return {
29-
databaseName: 'test',
30-
collectionName: 'test',
31-
}
31+
databaseName: this.cliContext.db._name,
32+
collectionName: '',
33+
};
3234
}
3335

34-
3536
/** @internal */
3637
setInput(text: string) {
3738
process.stdin.unshift(text);
@@ -51,6 +52,7 @@ export abstract class AiProvider {
5152
`);
5253
}
5354

55+
abstract aggregate(prompt: string): Promise<void>;
5456
abstract query(prompt: string): Promise<void>;
5557
abstract ask(prompt: string): Promise<void>;
5658
}

‎snippets/ai/providers/atlas/atlas-ai-provider.ts‎

Lines changed: 39 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,30 @@
11
import { aiCommand } from '../../decorators';
22
import { output } from '../../helpers';
3-
import { AiProvider } from '../ai-provider';
3+
import { AiProvider, CliContext } from '../ai-provider';
44
import { AtlasAiService, AIAggregation, AIQuery } from './atlas-ai-service';
55
import { AtlasService } from './atlas-service';
66
import { AuthService } from './auth-service';
77
import { config } from './util';
88
import open from 'open';
99

1010
export class AtlasAiProvider extends AiProvider {
11-
constructor(private readonly aiService: AtlasAiService) {
12-
super();
11+
constructor(
12+
private readonly aiService: AtlasAiService,
13+
cliContext: CliContext,
14+
) {
15+
super(cliContext);
1316
}
14-
17+
1518
async query(prompt: string): Promise<void> {
16-
const query = await this.aiService.getQueryFromUserInput({
17-
userInput: prompt,
18-
signal: AbortSignal.timeout(10000),
19-
requestId: crypto.randomUUID(),
20-
...this.getDatabaseContext(),
21-
}, this.getConnectionInfo());
19+
const query = await this.aiService.getQueryFromUserInput(
20+
{
21+
userInput: prompt,
22+
signal: AbortSignal.timeout(10000),
23+
requestId: crypto.randomUUID(),
24+
...this.getDatabaseContext(),
25+
},
26+
this.getConnectionInfo(),
27+
);
2228

2329
this.setInput(this.createMongoShellQuery(query.content));
2430
}
@@ -28,48 +34,53 @@ export class AtlasAiProvider extends AiProvider {
2834
}
2935

3036
async aggregate(prompt: string): Promise<void> {
31-
const aggregation = await this.aiService.getAggregationFromUserInput({
32-
userInput: prompt,
33-
signal: AbortSignal.timeout(10000),
34-
requestId: crypto.randomUUID(),
35-
...this.getDatabaseContext(),
36-
}, this.getConnectionInfo());
37+
const aggregation = await this.aiService.getAggregationFromUserInput(
38+
{
39+
userInput: prompt,
40+
signal: AbortSignal.timeout(10000),
41+
requestId: crypto.randomUUID(),
42+
...this.getDatabaseContext(),
43+
},
44+
this.getConnectionInfo(),
45+
);
3746

3847
this.respond(this.createMongoShellAggregation(aggregation.content));
3948
}
4049

4150
private createMongoShellQuery(params: AIQuery['content']): string {
42-
const {filter, project, collation, sort, skip, limit} = params.query;
43-
44-
return `db.collection.find(
51+
const { filter, project, collation, sort, skip, limit } = params.query;
52+
53+
return `db.collection.find(
4554
${filter},
4655
${project ? `{ projection: ${project} }` : '{}'}
47-
)${collation ? `.collation(${collation})` : ''}${sort ? `.sort(${sort})` : ''}${skip ? `.skip(${skip})` : ''}${limit ? `.limit(${limit})` : ''}`
48-
};
56+
)${collation ? `.collation(${collation})` : ''}${sort ? `.sort(${sort})` : ''}${skip ? `.skip(${skip})` : ''}${limit ? `.limit(${limit})` : ''}`;
57+
}
4958

50-
private createMongoShellAggregation(params: AIAggregation['content']): string {
51-
const {aggregation} = params;
59+
private createMongoShellAggregation(
60+
params: AIAggregation['content'],
61+
): string {
62+
const { aggregation } = params;
5263
return `db.collection.aggregate(${aggregation?.pipeline})`;
5364
}
5465
}
5566

56-
export function getAtlasAiProvider(): AtlasAiProvider {
67+
export function getAtlasAiProvider(cliContext: CliContext): AtlasAiProvider {
5768
const authService = new AuthService({
5869
...config['atlas'],
5970
openBrowser: async (url: string) => {
6071
output('Opening authentication page in your default browser...');
6172
await open(url);
6273
},
6374
});
64-
75+
6576
const atlasService = new AtlasService(authService, {
6677
...config['atlas'],
6778
});
68-
79+
6980
const aiService = new AtlasAiService({
7081
atlasService,
7182
apiURLPreset: 'admin-api',
7283
});
7384

74-
return new AtlasAiProvider(aiService);
75-
}
85+
return new AtlasAiProvider(aiService, cliContext);
86+
}

0 commit comments

Comments
 (0)