Skip to content

Commit

Permalink
Prevent signal parameter from being passed alongside input when calli…
Browse files Browse the repository at this point in the history
…ng `replicate.stream` (#254)

* Prevent signal parameter from being passed alongside input when calling replicate.stream

* Add test coverage for passing signal to replicate.stream
  • Loading branch information
mattt committed May 3, 2024
1 parent bb5ddaf commit 469f8da
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 3 deletions.
5 changes: 2 additions & 3 deletions index.js
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,7 @@ class Replicate {
* @yields {ServerSentEvent} Each streamed event from the prediction
*/
async *stream(ref, options) {
const { wait, ...data } = options;
const { wait, signal, ...data } = options;

const identifier = ModelVersionIdentifier.parse(ref);

Expand All @@ -296,11 +296,10 @@ class Replicate {
}

if (prediction.urls && prediction.urls.stream) {
const { signal } = options;
const stream = createReadableStream({
url: prediction.urls.stream,
fetch: this.fetch,
options: { signal },
...(signal ? { options: { signal } } : {}),
});

yield* streamAsyncIterator(stream);
Expand Down
2 changes: 2 additions & 0 deletions integration/cloudflare-worker/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,14 @@ export default {
const replicate = new Replicate({ auth: env.REPLICATE_API_TOKEN });

try {
const controller = new AbortController();
const output = replicate.stream(
"replicate/canary:30e22229542eb3f79d4f945dacb58d32001b02cc313ae6f54eef27904edf3272",
{
input: {
text: "Colin CloudFlare",
},
signal: controller.signal,
}
);
const stream = new ReadableStream({
Expand Down

0 comments on commit 469f8da

Please sign in to comment.