feat(cohere): support AsyncIterable for Cohere (#871)

Co-authored-by: Tomeu <tomeu@cohere.com>
Co-authored-by: Max Leiter <max.leiter@vercel.com>

Author: 3 people

.changeset/eighty-days-turn.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+cohere-stream: support AsyncIterable

packages/core/package.json

@@ -101,6 +101,7 @@
     "@vercel/ai-tsconfig": "workspace:*",
     "@vitejs/plugin-react": "4.2.0",
     "@vitejs/plugin-vue": "4.5.0",
+    "cohere-ai": "^7.6.2",
     "eslint": "^7.32.0",
     "eslint-config-vercel-ai": "workspace:*",
     "jsdom": "^23.0.0",

packages/core/streams/cohere-stream.test.ts

@@ -1,13 +1,19 @@
+import { CohereClient } from 'cohere-ai';
 import {
   CohereStream,
   StreamingTextResponse,
   experimental_StreamData,
 } from '.';
-import { cohereChunks } from '../tests/snapshots/cohere';
+import { cohereChatChunks, cohereChunks } from '../tests/snapshots/cohere';
 import { readAllChunks } from '../tests/utils/mock-client';
 import { DEFAULT_TEST_URL, createMockServer } from '../tests/utils/mock-server';
 
 const server = createMockServer([
+  {
+    url: 'https://api.cohere.ai/v1/chat',
+    chunks: cohereChatChunks,
+    formatChunk: chunk => `${JSON.stringify(chunk)}\n`,
+  },
   {
     url: DEFAULT_TEST_URL,
     chunks: cohereChunks,
@@ -28,9 +34,27 @@ describe('CohereStream', () => {
     server.close();
   });
 
+  it('should be able to parse Chat Streaming API and receive the streamed response', async () => {
+    const co = new CohereClient({
+      token: 'cohere-token',
+    });
+    const cohereResponse = await co.chatStream({
+      message: 'hi there!',
+    });
+
+    const stream = CohereStream(cohereResponse);
+    const response = new StreamingTextResponse(stream);
+    expect(await readAllChunks(response)).toEqual([
+      ' Hello',
+      ',',
+      ' world',
+      '.',
+      ' ',
+    ]);
+  });
+
   it('should be able to parse SSE and receive the streamed response', async () => {
     const stream = CohereStream(await fetch(DEFAULT_TEST_URL));
-
     const response = new StreamingTextResponse(stream);
 
     expect(await readAllChunks(response)).toEqual([

packages/core/streams/cohere-stream.ts

@@ -1,11 +1,25 @@
 import {
   type AIStreamCallbacksAndOptions,
   createCallbacksTransformer,
+  readableFromAsyncIterable,
 } from './ai-stream';
 import { createStreamDataTransformer } from './stream-data';
 
 const utf8Decoder = new TextDecoder('utf-8');
 
+// Full types
+// @see: https://github.com/cohere-ai/cohere-typescript/blob/c2eceb4a845098240ba0bc44e3787ccf75e268e8/src/api/types/StreamedChatResponse.ts
+interface StreamChunk {
+  text?: string;
+  eventType:
+    | 'stream-start'
+    | 'search-queries-generation'
+    | 'search-results'
+    | 'text-generation'
+    | 'citation-generation'
+    | 'stream-end';
+}
+
 async function processLines(
   lines: string[],
   controller: ReadableStreamDefaultController<string>,
@@ -63,13 +77,30 @@ function createParser(res: Response) {
   });
 }
 
+async function* streamable(stream: AsyncIterable<StreamChunk>) {
+  for await (const chunk of stream) {
+    if (chunk.eventType === 'text-generation') {
+      const text = chunk.text;
+      if (text) yield text;
+    }
+  }
+}
+
 export function CohereStream(
-  reader: Response,
+  reader: Response | AsyncIterable<StreamChunk>,
   callbacks?: AIStreamCallbacksAndOptions,
 ): ReadableStream {
-  return createParser(reader)
-    .pipeThrough(createCallbacksTransformer(callbacks))
-    .pipeThrough(
-      createStreamDataTransformer(callbacks?.experimental_streamData),
-    );
+  if (Symbol.asyncIterator in reader) {
+    return readableFromAsyncIterable(streamable(reader))
+      .pipeThrough(createCallbacksTransformer(callbacks))
+      .pipeThrough(
+        createStreamDataTransformer(callbacks?.experimental_streamData),
+      );
+  } else {
+    return createParser(reader)
+      .pipeThrough(createCallbacksTransformer(callbacks))
+      .pipeThrough(
+        createStreamDataTransformer(callbacks?.experimental_streamData),
+      );
+  }
 }

packages/core/tests/snapshots/cohere.ts

@@ -1,3 +1,27 @@
+export const cohereChatChunks = [
+  { text: ' Hello', is_finished: false, event_type: 'text-generation' },
+  { text: ',', is_finished: false, event_type: 'text-generation' },
+  { text: ' world', is_finished: false, event_type: 'text-generation' },
+  { text: '.', is_finished: false, event_type: 'text-generation' },
+  { text: ' ', is_finished: false, event_type: 'text-generation' },
+  {
+    is_finished: true,
+    event_type: 'text-generation',
+    finish_reason: 'COMPLETE',
+    response: {
+      id: 'a15fdf82-758e-4d35-a520-40fbb7308631',
+      generations: [
+        {
+          id: '9700d8ef-a8e5-4eb8-8288-3f0e0a1daed5',
+          text: ' Hello, world. ',
+          finish_reason: 'COMPLETE',
+        },
+      ],
+      prompt: 'Hello',
+    },
+  },
+];
+
 export const cohereChunks = [
   { text: ' Hello', is_finished: false },
   { text: ',', is_finished: false },