Skip to content

Commit

Permalink
ai/core: add toTextStreamResponse to streamText result (#1318)
Browse files Browse the repository at this point in the history
  • Loading branch information
lgrammel authored Apr 10, 2024
1 parent 87d3db5 commit 5cd29bd
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 0 deletions.
5 changes: 5 additions & 0 deletions .changeset/pretty-poets-switch.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'ai': patch
---

ai/core: add toTextStreamResponse() method to streamText result
39 changes: 39 additions & 0 deletions packages/core/core/generate-text/stream-text.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -231,3 +231,42 @@ describe('result.fullStream', () => {
);
});
});

describe('result.toTextStreamResponse', () => {
it('should create a Response with a text stream', async () => {
const result = await experimental_streamText({
model: new MockLanguageModelV1({
doStream: async ({ prompt, mode }) => {
return {
stream: convertArrayToReadableStream([
{ type: 'text-delta', textDelta: 'Hello' },
{ type: 'text-delta', textDelta: ', ' },
{ type: 'text-delta', textDelta: 'world!' },
]),
rawCall: { rawPrompt: 'prompt', rawSettings: {} },
};
},
}),
prompt: 'test-input',
});

const response = result.toTextStreamResponse();

assert.strictEqual(response.status, 200);
assert.strictEqual(
response.headers.get('Content-Type'),
'text/plain; charset=utf-8',
);

// Read the chunks into an array
const reader = response.body!.getReader();
const chunks = [];
while (true) {
const { value, done } = await reader.read();
if (done) break;
chunks.push(new TextDecoder().decode(value));
}

assert.deepStrictEqual(chunks, ['Hello', ', ', 'world!']);
});
});
26 changes: 26 additions & 0 deletions packages/core/core/generate-text/stream-text.ts
Original file line number Diff line number Diff line change
Expand Up @@ -219,4 +219,30 @@ Stream callbacks that will be called when the stream emits events.
.pipeThrough(createCallbacksTransformer(callbacks))
.pipeThrough(createStreamDataTransformer());
}

/**
Creates a simple text stream response.
Each text delta is encoded as UTF-8 and sent as a separate chunk.
Non-text-delta events are ignored.
*/
toTextStreamResponse(init?: ResponseInit): Response {
const encoder = new TextEncoder();
return new Response(
this.textStream.pipeThrough(
new TransformStream({
transform(chunk, controller) {
controller.enqueue(encoder.encode(chunk));
},
}),
),
{
...init,
status: 200,
headers: {
'Content-Type': 'text/plain; charset=utf-8',
...init?.headers,
},
},
);
}
}

0 comments on commit 5cd29bd

Please sign in to comment.