Support newer replicate versions, fix onResponse callback when no onC…

…ompletion is passed (#631) Co-authored-by: Max Leiter <max.leiter@vercel.com>
vercel · Oct 11, 2023 · c5d1857 · c5d1857
1 parent 385bdcd
commit c5d1857
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 7 deletions.
diff --git a/.changeset/bright-fishes-cover.md b/.changeset/bright-fishes-cover.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix: return complete response in onFinish when onCompletion isn't passed
diff --git a/.changeset/friendly-elephants-return.md b/.changeset/friendly-elephants-return.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+replicate-stream: fix types for replicate@0.20.0+
diff --git a/examples/next-replicate/package.json b/examples/next-replicate/package.json
@@ -13,7 +13,7 @@
     "next": "13.4.12",
     "react": "18.2.0",
     "react-dom": "^18.2.0",
-    "replicate": "^0.16.0"
+    "replicate": "^0.20.1"
   },
   "devDependencies": {
     "@types/node": "^17.0.12",

diff --git a/packages/core/streams/ai-stream.ts b/packages/core/streams/ai-stream.ts
@@ -123,8 +123,8 @@ export function createCallbacksTransformer(
     async transform(message, controller): Promise<void> {
       controller.enqueue(textEncoder.encode(message));
 
+      aggregatedResponse += message;
       if (callbacks.onToken) await callbacks.onToken(message);
-      if (callbacks.onCompletion) aggregatedResponse += message;
     },
 
     async flush(): Promise<void> {

diff --git a/packages/core/streams/replicate-stream.ts b/packages/core/streams/replicate-stream.ts
@@ -17,7 +17,7 @@ interface Prediction {
   webhook?: string;
   webhook_events_filter?: ('start' | 'output' | 'logs' | 'completed')[];
   created_at: string;
-  updated_at: string;
+  updated_at?: string;
   completed_at?: string;
   urls: {
     get: string;

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml