chore: streaming

Author: gagik

snippets/ai/src/ai.ts

@@ -61,6 +61,11 @@ module.exports = async (globalThis: CliContext) => {
       await this.command(prompt);
     }
 
+    @aiCommand()
+    async find(prompt: string) {
+      await this.ai.aggregate(prompt);
+    }
+
     @aiCommand()
     async ask(prompt: string) {
       await this.ai.processResponse(prompt, {

snippets/ai/src/config.ts

@@ -18,7 +18,7 @@ type ConfigKeys = keyof ConfigSchema;
 const defaults: Record<ConfigKeys, ConfigSchema[ConfigKeys]> = {
   provider: process.env.MONGOSH_AI_PROVIDER ?? 'docs',
   model: process.env.MONGOSH_AI_MODEL ?? 'default',
-  includeSampleDocs: process.env.MONGOSH_AI_INCLUDE_SAMPLE_DOCS ?? true,
+  includeSampleDocs: process.env.MONGOSH_AI_INCLUDE_SAMPLE_DOCS ?? false,
   defaultCollection: process.env.MONGOSH_AI_DEFAULT_COLLECTION,
 };
 

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

@@ -1,4 +1,4 @@
-import { generateText, streamText, LanguageModel, ModelMessage } from 'ai';
+import { streamText, LanguageModel, ModelMessage } from 'ai';
 import type { Config } from '../config.js';
 import {
   type LoadingAnimation as LoadingAnimationType,
@@ -251,13 +251,18 @@ export class AiProvider {
     messages: ModelMessage[],
     { systemPrompt, signal }: Omit<GetResponseOptions, 'expectedOutput'>,
   ): Promise<string> {
-    const { text } = await generateText({
+    const result = streamText({
       model: this.model,
 
       messages: messages,
       system: systemPrompt,
       abortSignal: signal ?? AbortSignal.timeout(30_000),
     });
+
+    let text = '';
+    for await (const delta of result.textStream) {
+      text += delta;
+    }
     return text;
   }