Add ability to stream results after useTaskTrigger

ericallam · ericallam · commit e7ca7d360f0a · 2024-11-14T12:28:11.000Z
diff --git a/packages/react-hooks/src/hooks/useRealtime.ts b/packages/react-hooks/src/hooks/useRealtime.ts
@@ -9,7 +9,7 @@ import { useApiClient, UseApiClientOptions } from "./useApiClient.js";
 export type UseRealtimeRunOptions = UseApiClientOptions & {
   id?: string;
   enabled?: boolean;
-  throttleInMs?: number;
+  experimental_throttleInMs?: number;
 };
 
 export type UseRealtimeRunInstance<TTask extends AnyTask = AnyTask> = {
@@ -37,7 +37,7 @@ export type UseRealtimeRunInstance<TTask extends AnyTask = AnyTask> = {
  * ```
  */
 export function useRealtimeRun<TTask extends AnyTask>(
-  runId: string,
+  runId?: string,
   options?: UseRealtimeRunOptions
 ): UseRealtimeRunInstance<TTask> {
   const hookId = useId();
@@ -71,13 +71,17 @@ export function useRealtimeRun<TTask extends AnyTask>(
 
   const triggerRequest = useCallback(async () => {
     try {
+      if (!runId) {
+        return;
+      }
+
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
       await processRealtimeRun(
         runId,
         apiClient,
-        throttle(mutateRun, options?.throttleInMs),
+        throttle(mutateRun, options?.experimental_throttleInMs),
         abortControllerRef
       );
     } catch (err) {
@@ -96,6 +100,10 @@ export function useRealtimeRun<TTask extends AnyTask>(
       return;
     }
 
+    if (!runId) {
+      return;
+    }
+
     triggerRequest().finally(() => {});
 
     return () => {
@@ -130,7 +138,7 @@ export function useRealtimeRunWithStreams<
   TTask extends AnyTask = AnyTask,
   TStreams extends Record<string, any> = Record<string, any>,
 >(
-  runId: string,
+  runId?: string,
   options?: UseRealtimeRunOptions
 ): UseRealtimeRunWithStreamsInstance<TTask, TStreams> {
   const hookId = useId();
@@ -181,14 +189,18 @@ export function useRealtimeRunWithStreams<
 
   const triggerRequest = useCallback(async () => {
     try {
+      if (!runId) {
+        return;
+      }
+
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
       await processRealtimeRunWithStreams(
         runId,
         apiClient,
-        throttle(mutateRun, options?.throttleInMs),
-        throttle(mutateStreams, options?.throttleInMs),
+        throttle(mutateRun, options?.experimental_throttleInMs),
+        throttle(mutateStreams, options?.experimental_throttleInMs),
         streamsRef,
         abortControllerRef
       );
@@ -208,6 +220,10 @@ export function useRealtimeRunWithStreams<
       return;
     }
 
+    if (!runId) {
+      return;
+    }
+
     triggerRequest().finally(() => {});
 
     return () => {
@@ -272,7 +288,7 @@ export function useRealtimeRunsWithTag<TTask extends AnyTask>(
       await processRealtimeRunsWithTag(
         tag,
         apiClient,
-        throttle(mutateRuns, options?.throttleInMs),
+        throttle(mutateRuns, options?.experimental_throttleInMs),
         runsRef,
         abortControllerRef
       );
@@ -359,7 +375,7 @@ export function useRealtimeBatch<TTask extends AnyTask>(
       await processRealtimeBatch(
         batchId,
         apiClient,
-        throttle(mutateRuns, options?.throttleInMs),
+        throttle(mutateRuns, options?.experimental_throttleInMs),
         runsRef,
         abortControllerRef
       );
diff --git a/packages/react-hooks/src/hooks/useTaskTrigger.ts b/packages/react-hooks/src/hooks/useTaskTrigger.ts
@@ -11,18 +11,27 @@ import {
   TaskRunOptions,
 } from "@trigger.dev/core/v3";
 import useSWRMutation from "swr/mutation";
-import { useApiClient } from "./useApiClient.js";
+import { useApiClient, UseApiClientOptions } from "./useApiClient.js";
+import {
+  useRealtimeRun,
+  UseRealtimeRunInstance,
+  useRealtimeRunWithStreams,
+  UseRealtimeRunWithStreamsInstance,
+} from "./useRealtime.js";
 
 export interface TriggerInstance<TTask extends AnyTask> {
   submit: (payload: TaskPayload<TTask>) => void;
   isLoading: boolean;
   handle?: RunHandleFromTypes<InferRunTypes<TTask>>;
 }
 
+export type UseTaskTriggerOptions = UseApiClientOptions;
+
 export function useTaskTrigger<TTask extends AnyTask>(
-  id: TaskIdentifier<TTask>
+  id: TaskIdentifier<TTask>,
+  options?: UseTaskTriggerOptions
 ): TriggerInstance<TTask> {
-  const apiClient = useApiClient();
+  const apiClient = useApiClient(options);
 
   async function triggerTask(
     id: string,
@@ -62,3 +71,58 @@ export function useTaskTrigger<TTask extends AnyTask>(
     handle: mutation.data as RunHandleFromTypes<InferRunTypes<TTask>>,
   };
 }
+
+export type UseRealtimeTaskTriggerOptions = UseTaskTriggerOptions & {
+  enabled?: boolean;
+  experimental_throttleInMs?: number;
+};
+
+export type RealtimeTriggerInstanceWithStreams<
+  TTask extends AnyTask,
+  TStreams extends Record<string, any> = Record<string, any>,
+> = UseRealtimeRunWithStreamsInstance<TTask, TStreams> & {
+  submit: (payload: TaskPayload<TTask>) => void;
+  isLoading: boolean;
+  handle?: RunHandleFromTypes<InferRunTypes<TTask>>;
+};
+
+export function useRealtimeTaskTriggerWithStreams<
+  TTask extends AnyTask,
+  TStreams extends Record<string, any> = Record<string, any>,
+>(
+  id: TaskIdentifier<TTask>,
+  options?: UseRealtimeTaskTriggerOptions
+): RealtimeTriggerInstanceWithStreams<TTask, TStreams> {
+  const triggerInstance = useTaskTrigger<TTask>(id, options);
+  const realtimeInstance = useRealtimeRunWithStreams<TTask, TStreams>(triggerInstance.handle?.id, {
+    ...options,
+    accessToken: triggerInstance.handle?.publicAccessToken ?? options?.accessToken,
+  });
+
+  return {
+    ...realtimeInstance,
+    ...triggerInstance,
+  };
+}
+
+export type RealtimeTriggerInstance<TTask extends AnyTask> = UseRealtimeRunInstance<TTask> & {
+  submit: (payload: TaskPayload<TTask>) => void;
+  isLoading: boolean;
+  handle?: RunHandleFromTypes<InferRunTypes<TTask>>;
+};
+
+export function useRealtimeTaskTrigger<TTask extends AnyTask>(
+  id: TaskIdentifier<TTask>,
+  options?: UseRealtimeTaskTriggerOptions
+): RealtimeTriggerInstance<TTask> {
+  const triggerInstance = useTaskTrigger<TTask>(id, options);
+  const realtimeInstance = useRealtimeRun<TTask>(triggerInstance.handle?.id, {
+    ...options,
+    accessToken: triggerInstance.handle?.publicAccessToken ?? options?.accessToken,
+  });
+
+  return {
+    ...realtimeInstance,
+    ...triggerInstance,
+  };
+}
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/references/nextjs-realtime/package.json b/references/nextjs-realtime/package.json
@@ -12,6 +12,7 @@
   "dependencies": {
     "@ai-sdk/openai": "^0.0.72",
     "@fal-ai/serverless-client": "^0.15.0",
+    "@radix-ui/react-dialog": "^1.0.3",
     "@radix-ui/react-icons": "^1.3.0",
     "@radix-ui/react-scroll-area": "^1.2.0",
     "@radix-ui/react-slot": "^1.1.0",
diff --git a/references/nextjs-realtime/src/app/ai/[id]/ClientAiDetails.tsx b/references/nextjs-realtime/src/app/ai/[id]/ClientAiDetails.tsx
@@ -8,7 +8,6 @@ function AiRunDetailsWrapper({ runId, accessToken }: { runId: string; accessToke
   const { run, streams, error } = useRealtimeRunWithStreams<typeof openaiStreaming, STREAMS>(
     runId,
     {
-      throttleInMs: 500,
       accessToken,
       baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
     }
diff --git a/references/nextjs-realtime/src/app/page.tsx b/references/nextjs-realtime/src/app/page.tsx
@@ -1,6 +1,7 @@
 import RunButton from "@/components/RunButton";
 import BatchRunButton from "@/components/BatchRunButton";
 import TriggerButton from "@/components/TriggerButton";
+import TriggerButtonWithStreaming from "@/components/TriggerButtonWithStreaming";
 import { ImageUploadDropzone } from "@/components/ImageUploadButton";
 import { auth } from "@trigger.dev/sdk/v3";
 
@@ -24,7 +25,8 @@ export default async function Home() {
       <div className="flex items-center space-x-4 justify-center w-full">
         <RunButton />
         <BatchRunButton />
-        <TriggerButton publicAccessToken={publicAccessToken} />
+        <TriggerButton accessToken={publicAccessToken} />
+        <TriggerButtonWithStreaming accessToken={publicAccessToken} />
       </div>
     </main>
   );
diff --git a/references/nextjs-realtime/src/components/TriggerButton.tsx b/references/nextjs-realtime/src/components/TriggerButton.tsx
@@ -2,12 +2,15 @@
 
 import { Button } from "@/components/ui/button";
 import { type openaiStreaming } from "@/trigger/ai";
-import { TriggerAuthContext, useTaskTrigger } from "@trigger.dev/react-hooks";
+import { useTaskTrigger } from "@trigger.dev/react-hooks";
 import { useRouter } from "next/navigation";
 import { useEffect } from "react";
 
-function TriggerButton() {
-  const { submit, handle, isLoading } = useTaskTrigger<typeof openaiStreaming>("openai-streaming");
+export default function TriggerButton({ accessToken }: { accessToken: string }) {
+  const { submit, handle, isLoading } = useTaskTrigger<typeof openaiStreaming>("openai-streaming", {
+    accessToken,
+    baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
+  });
   const router = useRouter();
 
   useEffect(() => {
@@ -33,17 +36,3 @@ function TriggerButton() {
     </Button>
   );
 }
-
-export default function TriggerButtonClientWrapper({
-  publicAccessToken,
-}: {
-  publicAccessToken: string;
-}) {
-  return (
-    <TriggerAuthContext.Provider
-      value={{ accessToken: publicAccessToken, baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL }}
-    >
-      <TriggerButton />
-    </TriggerAuthContext.Provider>
-  );
-}
diff --git a/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx b/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx
diff --git a/references/nextjs-realtime/src/components/ui/dialog.tsx b/references/nextjs-realtime/src/components/ui/dialog.tsx

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,6 @@ function AiRunDetailsWrapper({ runId, accessToken }: { runId: string; accessToke`
`8`	`8`	`const { run, streams, error } = useRealtimeRunWithStreams<typeof openaiStreaming, STREAMS>(`
`9`	`9`	`runId,`
`10`	`10`	`{`
`11`		`- throttleInMs: 500,`
`12`	`11`	`accessToken,`
`13`	`12`	`baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,`
`14`	`13`	`}`