refactor(js/ai/chat): make sendStream sync (#1732)

Author: pavelgj

js/ai/src/chat.ts

@@ -15,6 +15,7 @@
  */
 
 import { StreamingCallback, z } from '@genkit-ai/core';
+import { Channel } from '@genkit-ai/core/async';
 import {
   ATTR_PREFIX,
   SPAN_TYPE_ATTR,
@@ -30,7 +31,6 @@ import {
   MessageData,
   Part,
   generate,
-  generateStream,
 } from './index.js';
 import {
   BaseGenerateOptions,
@@ -154,23 +154,12 @@ export class Chat {
           },
         },
         async (metadata) => {
-          let resolvedOptions: ChatGenerateOptions<O, CustomOptions>;
+          let resolvedOptions = resolveSendOptions(options);
           let streamingCallback:
             | StreamingCallback<GenerateResponseChunk>
             | undefined = undefined;
 
-          // string
-          if (typeof options === 'string') {
-            resolvedOptions = {
-              prompt: options,
-            } as ChatGenerateOptions<O, CustomOptions>;
-          } else if (Array.isArray(options)) {
-            // Part[]
-            resolvedOptions = {
-              prompt: options,
-            } as ChatGenerateOptions<O, CustomOptions>;
-          } else {
-            resolvedOptions = options as ChatGenerateOptions<O, CustomOptions>;
+          if (resolvedOptions.onChunk || resolvedOptions.streamingCallback) {
             streamingCallback =
               resolvedOptions.onChunk ?? resolvedOptions.streamingCallback;
           }
@@ -204,66 +193,23 @@ export class Chat {
     CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
   >(
     options: string | Part[] | GenerateStreamOptions<O, CustomOptions>
-  ): Promise<GenerateStreamResponse<z.infer<O>>> {
-    return runWithSession(this.session.registry, this.session, () =>
-      runInNewSpan(
-        this.session.registry,
-        {
-          metadata: { name: 'send' },
-          labels: {
-            [SPAN_TYPE_ATTR]: 'helper',
-            [SESSION_ID_ATTR]: this.session.id,
-            [THREAD_NAME_ATTR]: this.threadName,
-          },
-        },
-        async (metadata) => {
-          let resolvedOptions;
-
-          // string
-          if (typeof options === 'string') {
-            resolvedOptions = {
-              prompt: options,
-            } as GenerateStreamOptions<O, CustomOptions>;
-          } else if (Array.isArray(options)) {
-            // Part[]
-            resolvedOptions = {
-              prompt: options,
-            } as GenerateStreamOptions<O, CustomOptions>;
-          } else {
-            resolvedOptions = options as GenerateStreamOptions<
-              O,
-              CustomOptions
-            >;
-          }
-          metadata.input = resolvedOptions;
-
-          const { response, stream } = await generateStream(
-            this.session.registry,
-            {
-              ...(await this.requestBase),
-              messages: this.messages,
-              ...resolvedOptions,
-            }
-          );
+  ): GenerateStreamResponse<z.infer<O>> {
+    let channel = new Channel<GenerateResponseChunk>();
+    let resolvedOptions = resolveSendOptions(options);
 
-          return {
-            response: response.finally(async () => {
-              const resolvedResponse = await response;
-              this.requestBase = Promise.resolve({
-                ...(await this.requestBase),
-                // these things may get changed by tools calling within generate.
-                tools: resolvedResponse?.request?.tools,
-                toolChoice: resolvedResponse?.request?.toolChoice,
-                config: resolvedResponse?.request?.config,
-              });
-              this.updateMessages(resolvedResponse.messages);
-              metadata.output = JSON.stringify(resolvedResponse);
-            }),
-            stream,
-          };
-        }
-      )
+    const sent = this.send({
+      ...resolvedOptions,
+      onChunk: (chunk) => channel.send(chunk),
+    });
+    sent.then(
+      () => channel.close(),
+      (err) => channel.error(err)
     );
+
+    return {
+      response: sent,
+      stream: channel,
+    };
   }
 
   get messages(): MessageData[] {
@@ -287,3 +233,27 @@ function getPreamble(msgs?: MessageData[]) {
 function stripPreamble(msgs?: MessageData[]) {
   return msgs?.filter((m) => !m.metadata?.preamble);
 }
+
+function resolveSendOptions<
+  O extends z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny,
+>(
+  options: string | Part[] | ChatGenerateOptions<O, CustomOptions>
+): ChatGenerateOptions<O, CustomOptions> {
+  let resolvedOptions: ChatGenerateOptions<O, CustomOptions>;
+
+  // string
+  if (typeof options === 'string') {
+    resolvedOptions = {
+      prompt: options,
+    } as ChatGenerateOptions<O, CustomOptions>;
+  } else if (Array.isArray(options)) {
+    // Part[]
+    resolvedOptions = {
+      prompt: options,
+    } as ChatGenerateOptions<O, CustomOptions>;
+  } else {
+    resolvedOptions = options as ChatGenerateOptions<O, CustomOptions>;
+  }
+  return resolvedOptions;
+}

js/genkit/tests/chat_test.ts

@@ -66,7 +66,7 @@ describe('chat', () => {
 
   it('maintains history in the session with streaming', async () => {
     const chat = ai.chat();
-    let { response, stream } = await chat.sendStream('hi');
+    let { response, stream } = chat.sendStream('hi');
 
     let chunks: string[] = [];
     for await (const chunk of stream) {
@@ -75,7 +75,7 @@ describe('chat', () => {
     assert.strictEqual((await response).text, 'Echo: hi; config: {}');
     assert.deepStrictEqual(chunks, ['3', '2', '1']);
 
-    ({ response, stream } = await chat.sendStream('bye'));
+    ({ response, stream } = chat.sendStream('bye'));
 
     chunks = [];
     for await (const chunk of stream) {

js/genkit/tests/session_test.ts

@@ -170,7 +170,7 @@ describe('session', () => {
     const session = ai.createSession();
     const chat = session.chat();
 
-    let { response, stream } = await chat.sendStream('hi');
+    let { response, stream } = chat.sendStream('hi');
 
     let chunks: string[] = [];
     for await (const chunk of stream) {
@@ -179,7 +179,7 @@ describe('session', () => {
     assert.strictEqual((await response).text, 'Echo: hi; config: {}');
     assert.deepStrictEqual(chunks, ['3', '2', '1']);
 
-    ({ response, stream } = await chat.sendStream('bye'));
+    ({ response, stream } = chat.sendStream('bye'));
 
     chunks = [];
     for await (const chunk of stream) {