feat: add chat transport and AI chat helpers to @trigger.dev/sdk
Two new subpath exports: @trigger.dev/sdk/chat (frontend, browser-safe): - TriggerChatTransport — ChatTransport implementation for useChat - createChatTransport() — factory function - TriggerChatTransportOptions type @trigger.dev/sdk/ai (backend, adds to existing ai.tool/ai.currentToolOptions): - chatTask() — pre-typed task wrapper with auto-pipe - pipeChat() — pipe StreamTextResult to realtime stream - CHAT_STREAM_KEY constant - ChatTaskPayload type - ChatTaskOptions type - PipeChatOptions type Co-authored-by: Eric Allam <eric@trigger.dev>
This commit is contained in:
@@ -0,0 +1,132 @@
|
||||
import { task as createTask } from "@trigger.dev/sdk";
|
||||
import type { Task } from "@trigger.dev/core/v3";
|
||||
import type { ChatTaskPayload } from "./types.js";
|
||||
import { pipeChat } from "./pipeChat.js";
|
||||
|
||||
/**
|
||||
* Options for defining a chat task.
|
||||
*
|
||||
* This is a simplified version of the standard task options with the payload
|
||||
* pre-typed as `ChatTaskPayload`.
|
||||
*/
|
||||
export type ChatTaskOptions<TIdentifier extends string> = {
|
||||
/** Unique identifier for the task */
|
||||
id: TIdentifier;
|
||||
|
||||
/** Optional description of the task */
|
||||
description?: string;
|
||||
|
||||
/** Retry configuration */
|
||||
retry?: {
|
||||
maxAttempts?: number;
|
||||
factor?: number;
|
||||
minTimeoutInMs?: number;
|
||||
maxTimeoutInMs?: number;
|
||||
randomize?: boolean;
|
||||
};
|
||||
|
||||
/** Queue configuration */
|
||||
queue?: {
|
||||
name?: string;
|
||||
concurrencyLimit?: number;
|
||||
};
|
||||
|
||||
/** Machine preset for the task */
|
||||
machine?: {
|
||||
preset?: string;
|
||||
};
|
||||
|
||||
/** Maximum duration in seconds */
|
||||
maxDuration?: number;
|
||||
|
||||
/**
|
||||
* The main run function for the chat task.
|
||||
*
|
||||
* Receives a `ChatTaskPayload` with the conversation messages, chat session ID,
|
||||
* and trigger type.
|
||||
*
|
||||
* **Auto-piping:** If this function returns a value that has a `.toUIMessageStream()` method
|
||||
* (like a `StreamTextResult` from `streamText()`), the stream will automatically be piped
|
||||
* to the frontend via the chat realtime stream. If you need to pipe from deeper in your
|
||||
* code, use `pipeChat()` instead and don't return the result.
|
||||
*/
|
||||
run: (payload: ChatTaskPayload) => Promise<unknown>;
|
||||
};
|
||||
|
||||
/**
|
||||
* An object that has a `toUIMessageStream()` method, like the result of `streamText()`.
|
||||
*/
|
||||
type UIMessageStreamable = {
|
||||
toUIMessageStream: (...args: any[]) => AsyncIterable<unknown> | ReadableStream<unknown>;
|
||||
};
|
||||
|
||||
function isUIMessageStreamable(value: unknown): value is UIMessageStreamable {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
"toUIMessageStream" in value &&
|
||||
typeof (value as any).toUIMessageStream === "function"
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a Trigger.dev task pre-configured for AI SDK chat.
|
||||
*
|
||||
* This is a convenience wrapper around `task()` from `@trigger.dev/sdk` that:
|
||||
* - **Pre-types the payload** as `ChatTaskPayload` — no manual typing needed
|
||||
* - **Auto-pipes the stream** if the `run` function returns a `StreamTextResult`
|
||||
*
|
||||
* Requires `@trigger.dev/sdk` to be installed (it's a peer dependency).
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { chatTask } from "@trigger.dev/ai";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
* import { openai } from "@ai-sdk/openai";
|
||||
*
|
||||
* // Simple: return streamText result — auto-piped to the frontend
|
||||
* export const myChatTask = chatTask({
|
||||
* id: "my-chat-task",
|
||||
* run: async ({ messages }) => {
|
||||
* return streamText({
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(messages),
|
||||
* });
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { chatTask, pipeChat } from "@trigger.dev/ai";
|
||||
*
|
||||
* // Complex: use pipeChat() from deep inside your agent code
|
||||
* export const myAgentTask = chatTask({
|
||||
* id: "my-agent-task",
|
||||
* run: async ({ messages }) => {
|
||||
* await runComplexAgentLoop(messages);
|
||||
* // pipeChat() called internally by the agent loop
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export function chatTask<TIdentifier extends string>(
|
||||
options: ChatTaskOptions<TIdentifier>
|
||||
): Task<TIdentifier, ChatTaskPayload, unknown> {
|
||||
const { run: userRun, ...restOptions } = options;
|
||||
|
||||
return createTask<TIdentifier, ChatTaskPayload, unknown>({
|
||||
...restOptions,
|
||||
run: async (payload: ChatTaskPayload) => {
|
||||
const result = await userRun(payload);
|
||||
|
||||
// If the run function returned a StreamTextResult or similar,
|
||||
// automatically pipe it to the chat stream
|
||||
if (isUIMessageStreamable(result)) {
|
||||
await pipeChat(result);
|
||||
}
|
||||
|
||||
return result;
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
import { realtimeStreams } from "@trigger.dev/core/v3";
|
||||
|
||||
/**
|
||||
* The default stream key used for chat transport communication.
|
||||
*
|
||||
* Both `TriggerChatTransport` (frontend) and `pipeChat` (backend) use this key
|
||||
* by default to ensure they communicate over the same stream.
|
||||
*/
|
||||
export const CHAT_STREAM_KEY = "chat";
|
||||
|
||||
/**
|
||||
* Options for `pipeChat`.
|
||||
*/
|
||||
export type PipeChatOptions = {
|
||||
/**
|
||||
* Override the stream key to pipe to.
|
||||
* Must match the `streamKey` option on `TriggerChatTransport`.
|
||||
*
|
||||
* @default "chat"
|
||||
*/
|
||||
streamKey?: string;
|
||||
|
||||
/**
|
||||
* An AbortSignal to cancel the stream.
|
||||
*/
|
||||
signal?: AbortSignal;
|
||||
|
||||
/**
|
||||
* The target run ID to pipe the stream to.
|
||||
* @default "self" (current run)
|
||||
*/
|
||||
target?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* An object that has a `toUIMessageStream()` method, like the result of `streamText()` from the AI SDK.
|
||||
*/
|
||||
type UIMessageStreamable = {
|
||||
toUIMessageStream: (...args: any[]) => AsyncIterable<unknown> | ReadableStream<unknown>;
|
||||
};
|
||||
|
||||
function isUIMessageStreamable(value: unknown): value is UIMessageStreamable {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
"toUIMessageStream" in value &&
|
||||
typeof (value as any).toUIMessageStream === "function"
|
||||
);
|
||||
}
|
||||
|
||||
function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
Symbol.asyncIterator in value
|
||||
);
|
||||
}
|
||||
|
||||
function isReadableStream(value: unknown): value is ReadableStream<unknown> {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
typeof (value as any).getReader === "function"
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Pipes a chat stream to the realtime stream, making it available to the
|
||||
* `TriggerChatTransport` on the frontend.
|
||||
*
|
||||
* Accepts any of:
|
||||
* - A `StreamTextResult` from the AI SDK (has `.toUIMessageStream()`)
|
||||
* - An `AsyncIterable` of `UIMessageChunk`s
|
||||
* - A `ReadableStream` of `UIMessageChunk`s
|
||||
*
|
||||
* This must be called from inside a Trigger.dev task's `run` function.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { task } from "@trigger.dev/sdk";
|
||||
* import { pipeChat, type ChatTaskPayload } from "@trigger.dev/ai";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
*
|
||||
* export const myChatTask = task({
|
||||
* id: "my-chat-task",
|
||||
* run: async (payload: ChatTaskPayload) => {
|
||||
* const result = streamText({
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(payload.messages),
|
||||
* });
|
||||
*
|
||||
* await pipeChat(result);
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* // Deep inside your agent library — pipeChat works from anywhere inside a task
|
||||
* async function runAgentLoop(messages: CoreMessage[]) {
|
||||
* const result = streamText({ model, messages });
|
||||
* await pipeChat(result);
|
||||
* }
|
||||
* ```
|
||||
*
|
||||
* @param source - A StreamTextResult, AsyncIterable, or ReadableStream of UIMessageChunks
|
||||
* @param options - Optional configuration
|
||||
* @returns A promise that resolves when the stream has been fully piped
|
||||
*/
|
||||
export async function pipeChat(
|
||||
source: UIMessageStreamable | AsyncIterable<unknown> | ReadableStream<unknown>,
|
||||
options?: PipeChatOptions
|
||||
): Promise<void> {
|
||||
const streamKey = options?.streamKey ?? CHAT_STREAM_KEY;
|
||||
|
||||
// Resolve the source to an AsyncIterable or ReadableStream
|
||||
let stream: AsyncIterable<unknown> | ReadableStream<unknown>;
|
||||
|
||||
if (isUIMessageStreamable(source)) {
|
||||
stream = source.toUIMessageStream();
|
||||
} else if (isAsyncIterable(source) || isReadableStream(source)) {
|
||||
stream = source;
|
||||
} else {
|
||||
throw new Error(
|
||||
"pipeChat: source must be a StreamTextResult (with .toUIMessageStream()), " +
|
||||
"an AsyncIterable, or a ReadableStream"
|
||||
);
|
||||
}
|
||||
|
||||
// Pipe to the realtime stream
|
||||
const instance = realtimeStreams.pipe(streamKey, stream, {
|
||||
signal: options?.signal,
|
||||
target: options?.target,
|
||||
});
|
||||
|
||||
await instance.wait();
|
||||
}
|
||||
+11
-11
@@ -8,7 +8,7 @@ export type TriggerChatTransportOptions = {
|
||||
* The Trigger.dev task ID to trigger for chat completions.
|
||||
* This task will receive the chat messages as its payload.
|
||||
*/
|
||||
taskId: string;
|
||||
task: string;
|
||||
|
||||
/**
|
||||
* An access token for authenticating with the Trigger.dev API.
|
||||
@@ -36,8 +36,8 @@ export type TriggerChatTransportOptions = {
|
||||
|
||||
/**
|
||||
* The stream key where the task pipes UIMessageChunk data.
|
||||
* Your task must pipe the AI SDK stream to this same key using
|
||||
* `streams.pipe(streamKey, result.toUIMessageStream())`.
|
||||
* When using `chatTask()` or `pipeChat()`, this is handled automatically.
|
||||
* Only set this if you're using a custom stream key.
|
||||
*
|
||||
* @default "chat"
|
||||
*/
|
||||
@@ -59,15 +59,16 @@ export type TriggerChatTransportOptions = {
|
||||
};
|
||||
|
||||
/**
|
||||
* The payload shape that TriggerChatTransport sends to the triggered task.
|
||||
* The payload shape that the transport sends to the triggered task.
|
||||
*
|
||||
* Use this type to type your task's `run` function payload:
|
||||
* When using `chatTask()`, the payload is automatically typed — you don't need
|
||||
* to import this type. When using `task()` directly, use this type to annotate
|
||||
* your payload:
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { task, streams } from "@trigger.dev/sdk";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
* import type { ChatTaskPayload } from "@trigger.dev/ai";
|
||||
* import { task } from "@trigger.dev/sdk";
|
||||
* import { pipeChat, type ChatTaskPayload } from "@trigger.dev/ai";
|
||||
*
|
||||
* export const myChatTask = task({
|
||||
* id: "my-chat-task",
|
||||
@@ -76,9 +77,7 @@ export type TriggerChatTransportOptions = {
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(payload.messages),
|
||||
* });
|
||||
*
|
||||
* const { waitUntilComplete } = streams.pipe("chat", result.toUIMessageStream());
|
||||
* await waitUntilComplete();
|
||||
* await pipeChat(result);
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
@@ -110,6 +109,7 @@ export type ChatTaskPayload<TMessage extends UIMessage = UIMessage> = {
|
||||
|
||||
/**
|
||||
* Internal state for tracking active chat sessions, used for stream reconnection.
|
||||
* @internal
|
||||
*/
|
||||
export type ChatSessionState = {
|
||||
runId: string;
|
||||
|
||||
@@ -24,7 +24,8 @@
|
||||
"./package.json": "./package.json",
|
||||
".": "./src/v3/index.ts",
|
||||
"./v3": "./src/v3/index.ts",
|
||||
"./ai": "./src/v3/ai.ts"
|
||||
"./ai": "./src/v3/ai.ts",
|
||||
"./chat": "./src/v3/chat.ts"
|
||||
},
|
||||
"sourceDialects": [
|
||||
"@triggerdotdev/source"
|
||||
@@ -37,6 +38,9 @@
|
||||
],
|
||||
"ai": [
|
||||
"dist/commonjs/v3/ai.d.ts"
|
||||
],
|
||||
"chat": [
|
||||
"dist/commonjs/v3/chat.d.ts"
|
||||
]
|
||||
}
|
||||
},
|
||||
@@ -123,6 +127,17 @@
|
||||
"types": "./dist/commonjs/v3/ai.d.ts",
|
||||
"default": "./dist/commonjs/v3/ai.js"
|
||||
}
|
||||
},
|
||||
"./chat": {
|
||||
"import": {
|
||||
"@triggerdotdev/source": "./src/v3/chat.ts",
|
||||
"types": "./dist/esm/v3/chat.d.ts",
|
||||
"default": "./dist/esm/v3/chat.js"
|
||||
},
|
||||
"require": {
|
||||
"types": "./dist/commonjs/v3/chat.d.ts",
|
||||
"default": "./dist/commonjs/v3/chat.js"
|
||||
}
|
||||
}
|
||||
},
|
||||
"main": "./dist/commonjs/v3/index.js",
|
||||
|
||||
@@ -3,11 +3,16 @@ import {
|
||||
isSchemaZodEsque,
|
||||
Task,
|
||||
type inferSchemaIn,
|
||||
type PipeStreamOptions,
|
||||
type TaskOptions,
|
||||
type TaskSchema,
|
||||
type TaskWithSchema,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import type { UIMessage } from "ai";
|
||||
import { dynamicTool, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai";
|
||||
import { metadata } from "./metadata.js";
|
||||
import { streams } from "./streams.js";
|
||||
import { createTask } from "./shared.js";
|
||||
|
||||
const METADATA_KEY = "tool.execute.options";
|
||||
|
||||
@@ -116,3 +121,240 @@ export const ai = {
|
||||
tool: toolFromTask,
|
||||
currentToolOptions: getToolOptionsFromMetadata,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Chat transport helpers — backend side
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* The default stream key used for chat transport communication.
|
||||
* Both `TriggerChatTransport` (frontend) and `pipeChat`/`chatTask` (backend)
|
||||
* use this key by default.
|
||||
*/
|
||||
export const CHAT_STREAM_KEY = "chat";
|
||||
|
||||
/**
|
||||
* The payload shape that the chat transport sends to the triggered task.
|
||||
*
|
||||
* When using `chatTask()`, the payload is automatically typed — you don't need
|
||||
* to import this type. Use this type only if you're using `task()` directly
|
||||
* with `pipeChat()`.
|
||||
*/
|
||||
export type ChatTaskPayload<TMessage extends UIMessage = UIMessage> = {
|
||||
/** The conversation messages */
|
||||
messages: TMessage[];
|
||||
|
||||
/** The unique identifier for the chat session */
|
||||
chatId: string;
|
||||
|
||||
/**
|
||||
* The trigger type:
|
||||
* - `"submit-message"`: A new user message
|
||||
* - `"regenerate-message"`: Regenerate the last assistant response
|
||||
*/
|
||||
trigger: "submit-message" | "regenerate-message";
|
||||
|
||||
/** The ID of the message to regenerate (only for `"regenerate-message"`) */
|
||||
messageId?: string;
|
||||
|
||||
/** Custom metadata from the frontend */
|
||||
metadata?: unknown;
|
||||
};
|
||||
|
||||
/**
|
||||
* Options for `pipeChat`.
|
||||
*/
|
||||
export type PipeChatOptions = {
|
||||
/**
|
||||
* Override the stream key. Must match the `streamKey` on `TriggerChatTransport`.
|
||||
* @default "chat"
|
||||
*/
|
||||
streamKey?: string;
|
||||
|
||||
/** An AbortSignal to cancel the stream. */
|
||||
signal?: AbortSignal;
|
||||
|
||||
/**
|
||||
* The target run ID to pipe to.
|
||||
* @default "self" (current run)
|
||||
*/
|
||||
target?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* An object with a `toUIMessageStream()` method (e.g. `StreamTextResult` from `streamText()`).
|
||||
*/
|
||||
type UIMessageStreamable = {
|
||||
toUIMessageStream: (...args: any[]) => AsyncIterable<unknown> | ReadableStream<unknown>;
|
||||
};
|
||||
|
||||
function isUIMessageStreamable(value: unknown): value is UIMessageStreamable {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
"toUIMessageStream" in value &&
|
||||
typeof (value as any).toUIMessageStream === "function"
|
||||
);
|
||||
}
|
||||
|
||||
function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
|
||||
return typeof value === "object" && value !== null && Symbol.asyncIterator in value;
|
||||
}
|
||||
|
||||
function isReadableStream(value: unknown): value is ReadableStream<unknown> {
|
||||
return typeof value === "object" && value !== null && typeof (value as any).getReader === "function";
|
||||
}
|
||||
|
||||
/**
|
||||
* Pipes a chat stream to the realtime stream, making it available to the
|
||||
* `TriggerChatTransport` on the frontend.
|
||||
*
|
||||
* Accepts:
|
||||
* - A `StreamTextResult` from `streamText()` (has `.toUIMessageStream()`)
|
||||
* - An `AsyncIterable` of `UIMessageChunk`s
|
||||
* - A `ReadableStream` of `UIMessageChunk`s
|
||||
*
|
||||
* Must be called from inside a Trigger.dev task's `run` function.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { task } from "@trigger.dev/sdk";
|
||||
* import { pipeChat, type ChatTaskPayload } from "@trigger.dev/sdk/ai";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
*
|
||||
* export const myChatTask = task({
|
||||
* id: "my-chat-task",
|
||||
* run: async (payload: ChatTaskPayload) => {
|
||||
* const result = streamText({
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(payload.messages),
|
||||
* });
|
||||
*
|
||||
* await pipeChat(result);
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* // Works from anywhere inside a task — even deep in your agent code
|
||||
* async function runAgentLoop(messages: CoreMessage[]) {
|
||||
* const result = streamText({ model, messages });
|
||||
* await pipeChat(result);
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
export async function pipeChat(
|
||||
source: UIMessageStreamable | AsyncIterable<unknown> | ReadableStream<unknown>,
|
||||
options?: PipeChatOptions
|
||||
): Promise<void> {
|
||||
const streamKey = options?.streamKey ?? CHAT_STREAM_KEY;
|
||||
|
||||
let stream: AsyncIterable<unknown> | ReadableStream<unknown>;
|
||||
|
||||
if (isUIMessageStreamable(source)) {
|
||||
stream = source.toUIMessageStream();
|
||||
} else if (isAsyncIterable(source) || isReadableStream(source)) {
|
||||
stream = source;
|
||||
} else {
|
||||
throw new Error(
|
||||
"pipeChat: source must be a StreamTextResult (with .toUIMessageStream()), " +
|
||||
"an AsyncIterable, or a ReadableStream"
|
||||
);
|
||||
}
|
||||
|
||||
const pipeOptions: PipeStreamOptions = {};
|
||||
if (options?.signal) {
|
||||
pipeOptions.signal = options.signal;
|
||||
}
|
||||
if (options?.target) {
|
||||
pipeOptions.target = options.target;
|
||||
}
|
||||
|
||||
const { waitUntilComplete } = streams.pipe(streamKey, stream, pipeOptions);
|
||||
await waitUntilComplete();
|
||||
}
|
||||
|
||||
/**
|
||||
* Options for defining a chat task.
|
||||
*
|
||||
* Extends the standard `TaskOptions` but pre-types the payload as `ChatTaskPayload`
|
||||
* and overrides `run` to accept `ChatTaskPayload` directly.
|
||||
*
|
||||
* **Auto-piping:** If the `run` function returns a value with `.toUIMessageStream()`
|
||||
* (like a `StreamTextResult`), the stream is automatically piped to the frontend.
|
||||
* For complex flows, use `pipeChat()` manually from anywhere in your code.
|
||||
*/
|
||||
export type ChatTaskOptions<TIdentifier extends string> = Omit<
|
||||
TaskOptions<TIdentifier, ChatTaskPayload, unknown>,
|
||||
"run"
|
||||
> & {
|
||||
/**
|
||||
* The run function for the chat task.
|
||||
*
|
||||
* Receives a `ChatTaskPayload` with the conversation messages, chat session ID,
|
||||
* and trigger type.
|
||||
*
|
||||
* **Auto-piping:** If this function returns a value with `.toUIMessageStream()`,
|
||||
* the stream is automatically piped to the frontend.
|
||||
*/
|
||||
run: (payload: ChatTaskPayload) => Promise<unknown>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Creates a Trigger.dev task pre-configured for AI SDK chat.
|
||||
*
|
||||
* - **Pre-types the payload** as `ChatTaskPayload` — no manual typing needed
|
||||
* - **Auto-pipes the stream** if `run` returns a `StreamTextResult`
|
||||
* - For complex flows, use `pipeChat()` from anywhere inside your task code
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { chatTask } from "@trigger.dev/sdk/ai";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
* import { openai } from "@ai-sdk/openai";
|
||||
*
|
||||
* // Simple: return streamText result — auto-piped to the frontend
|
||||
* export const myChatTask = chatTask({
|
||||
* id: "my-chat-task",
|
||||
* run: async ({ messages }) => {
|
||||
* return streamText({
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(messages),
|
||||
* });
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { chatTask, pipeChat } from "@trigger.dev/sdk/ai";
|
||||
*
|
||||
* // Complex: pipeChat() from deep in your agent code
|
||||
* export const myAgentTask = chatTask({
|
||||
* id: "my-agent-task",
|
||||
* run: async ({ messages }) => {
|
||||
* await runComplexAgentLoop(messages);
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export function chatTask<TIdentifier extends string>(
|
||||
options: ChatTaskOptions<TIdentifier>
|
||||
): Task<TIdentifier, ChatTaskPayload, unknown> {
|
||||
const { run: userRun, ...restOptions } = options;
|
||||
|
||||
return createTask<TIdentifier, ChatTaskPayload, unknown>({
|
||||
...restOptions,
|
||||
run: async (payload: ChatTaskPayload) => {
|
||||
const result = await userRun(payload);
|
||||
|
||||
// Auto-pipe if the run function returned a StreamTextResult or similar
|
||||
if (isUIMessageStreamable(result)) {
|
||||
await pipeChat(result);
|
||||
}
|
||||
|
||||
return result;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
/**
|
||||
* @module @trigger.dev/sdk/chat
|
||||
*
|
||||
* Browser-safe module for AI SDK chat transport integration.
|
||||
* Use this on the frontend with the AI SDK's `useChat` hook.
|
||||
*
|
||||
* For backend helpers (`chatTask`, `pipeChat`), use `@trigger.dev/sdk/ai` instead.
|
||||
*
|
||||
* @example
|
||||
* ```tsx
|
||||
* import { useChat } from "@ai-sdk/react";
|
||||
* import { TriggerChatTransport } from "@trigger.dev/sdk/chat";
|
||||
*
|
||||
* function Chat({ accessToken }: { accessToken: string }) {
|
||||
* const { messages, sendMessage, status } = useChat({
|
||||
* transport: new TriggerChatTransport({
|
||||
* task: "my-chat-task",
|
||||
* accessToken,
|
||||
* }),
|
||||
* });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
|
||||
import type { ChatTransport, UIMessage, UIMessageChunk, ChatRequestOptions } from "ai";
|
||||
import { ApiClient, SSEStreamSubscription } from "@trigger.dev/core/v3";
|
||||
|
||||
const DEFAULT_STREAM_KEY = "chat";
|
||||
const DEFAULT_BASE_URL = "https://api.trigger.dev";
|
||||
const DEFAULT_STREAM_TIMEOUT_SECONDS = 120;
|
||||
|
||||
/**
|
||||
* Options for creating a TriggerChatTransport.
|
||||
*/
|
||||
export type TriggerChatTransportOptions = {
|
||||
/**
|
||||
* The Trigger.dev task ID to trigger for chat completions.
|
||||
* This task should be defined using `chatTask()` from `@trigger.dev/sdk/ai`,
|
||||
* or a regular `task()` that uses `pipeChat()`.
|
||||
*/
|
||||
task: string;
|
||||
|
||||
/**
|
||||
* An access token for authenticating with the Trigger.dev API.
|
||||
*
|
||||
* This must be a token with permission to trigger the task. You can use:
|
||||
* - A **trigger public token** created via `auth.createTriggerPublicToken(taskId)` (recommended for frontend use)
|
||||
* - A **secret API key** (for server-side use only — never expose in the browser)
|
||||
*
|
||||
* Can also be a function that returns a token string, useful for dynamic token refresh.
|
||||
*/
|
||||
accessToken: string | (() => string);
|
||||
|
||||
/**
|
||||
* Base URL for the Trigger.dev API.
|
||||
* @default "https://api.trigger.dev"
|
||||
*/
|
||||
baseURL?: string;
|
||||
|
||||
/**
|
||||
* The stream key where the task pipes UIMessageChunk data.
|
||||
* When using `chatTask()` or `pipeChat()`, this is handled automatically.
|
||||
* Only set this if you're using a custom stream key.
|
||||
*
|
||||
* @default "chat"
|
||||
*/
|
||||
streamKey?: string;
|
||||
|
||||
/**
|
||||
* Additional headers to include in API requests to Trigger.dev.
|
||||
*/
|
||||
headers?: Record<string, string>;
|
||||
|
||||
/**
|
||||
* The number of seconds to wait for the realtime stream to produce data
|
||||
* before timing out.
|
||||
*
|
||||
* @default 120
|
||||
*/
|
||||
streamTimeoutSeconds?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Internal state for tracking active chat sessions.
|
||||
* @internal
|
||||
*/
|
||||
type ChatSessionState = {
|
||||
runId: string;
|
||||
publicAccessToken: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* A custom AI SDK `ChatTransport` that runs chat completions as durable Trigger.dev tasks.
|
||||
*
|
||||
* When `sendMessages` is called, the transport:
|
||||
* 1. Triggers a Trigger.dev task with the chat messages as payload
|
||||
* 2. Subscribes to the task's realtime stream to receive `UIMessageChunk` data
|
||||
* 3. Returns a `ReadableStream<UIMessageChunk>` that the AI SDK processes natively
|
||||
*
|
||||
* @example
|
||||
* ```tsx
|
||||
* import { useChat } from "@ai-sdk/react";
|
||||
* import { TriggerChatTransport } from "@trigger.dev/sdk/chat";
|
||||
*
|
||||
* function Chat({ accessToken }: { accessToken: string }) {
|
||||
* const { messages, sendMessage, status } = useChat({
|
||||
* transport: new TriggerChatTransport({
|
||||
* task: "my-chat-task",
|
||||
* accessToken,
|
||||
* }),
|
||||
* });
|
||||
*
|
||||
* // ... render messages
|
||||
* }
|
||||
* ```
|
||||
*
|
||||
* On the backend, define the task using `chatTask` from `@trigger.dev/sdk/ai`:
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* import { chatTask } from "@trigger.dev/sdk/ai";
|
||||
* import { streamText, convertToModelMessages } from "ai";
|
||||
*
|
||||
* export const myChatTask = chatTask({
|
||||
* id: "my-chat-task",
|
||||
* run: async ({ messages }) => {
|
||||
* return streamText({
|
||||
* model: openai("gpt-4o"),
|
||||
* messages: convertToModelMessages(messages),
|
||||
* });
|
||||
* },
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
private readonly taskId: string;
|
||||
private readonly resolveAccessToken: () => string;
|
||||
private readonly baseURL: string;
|
||||
private readonly streamKey: string;
|
||||
private readonly extraHeaders: Record<string, string>;
|
||||
private readonly streamTimeoutSeconds: number;
|
||||
|
||||
private sessions: Map<string, ChatSessionState> = new Map();
|
||||
|
||||
constructor(options: TriggerChatTransportOptions) {
|
||||
this.taskId = options.task;
|
||||
this.resolveAccessToken =
|
||||
typeof options.accessToken === "function"
|
||||
? options.accessToken
|
||||
: () => options.accessToken as string;
|
||||
this.baseURL = options.baseURL ?? DEFAULT_BASE_URL;
|
||||
this.streamKey = options.streamKey ?? DEFAULT_STREAM_KEY;
|
||||
this.extraHeaders = options.headers ?? {};
|
||||
this.streamTimeoutSeconds = options.streamTimeoutSeconds ?? DEFAULT_STREAM_TIMEOUT_SECONDS;
|
||||
}
|
||||
|
||||
sendMessages = async (
|
||||
options: {
|
||||
trigger: "submit-message" | "regenerate-message";
|
||||
chatId: string;
|
||||
messageId: string | undefined;
|
||||
messages: UIMessage[];
|
||||
abortSignal: AbortSignal | undefined;
|
||||
} & ChatRequestOptions
|
||||
): Promise<ReadableStream<UIMessageChunk>> => {
|
||||
const { trigger, chatId, messageId, messages, abortSignal, body, metadata } = options;
|
||||
|
||||
const payload = {
|
||||
messages,
|
||||
chatId,
|
||||
trigger,
|
||||
messageId,
|
||||
metadata,
|
||||
...(body ?? {}),
|
||||
};
|
||||
|
||||
const currentToken = this.resolveAccessToken();
|
||||
const apiClient = new ApiClient(this.baseURL, currentToken);
|
||||
|
||||
const triggerResponse = await apiClient.triggerTask(this.taskId, {
|
||||
payload: JSON.stringify(payload),
|
||||
options: {
|
||||
payloadType: "application/json",
|
||||
},
|
||||
});
|
||||
|
||||
const runId = triggerResponse.id;
|
||||
const publicAccessToken =
|
||||
"publicAccessToken" in triggerResponse
|
||||
? (triggerResponse as { publicAccessToken?: string }).publicAccessToken
|
||||
: undefined;
|
||||
|
||||
this.sessions.set(chatId, {
|
||||
runId,
|
||||
publicAccessToken: publicAccessToken ?? currentToken,
|
||||
});
|
||||
|
||||
return this.subscribeToStream(runId, publicAccessToken ?? currentToken, abortSignal);
|
||||
};
|
||||
|
||||
reconnectToStream = async (
|
||||
options: {
|
||||
chatId: string;
|
||||
} & ChatRequestOptions
|
||||
): Promise<ReadableStream<UIMessageChunk> | null> => {
|
||||
const session = this.sessions.get(options.chatId);
|
||||
if (!session) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return this.subscribeToStream(session.runId, session.publicAccessToken, undefined);
|
||||
};
|
||||
|
||||
private subscribeToStream(
|
||||
runId: string,
|
||||
accessToken: string,
|
||||
abortSignal: AbortSignal | undefined
|
||||
): ReadableStream<UIMessageChunk> {
|
||||
const headers: Record<string, string> = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
...this.extraHeaders,
|
||||
};
|
||||
|
||||
const subscription = new SSEStreamSubscription(
|
||||
`${this.baseURL}/realtime/v1/streams/${runId}/${this.streamKey}`,
|
||||
{
|
||||
headers,
|
||||
signal: abortSignal,
|
||||
timeoutInSeconds: this.streamTimeoutSeconds,
|
||||
}
|
||||
);
|
||||
|
||||
return new ReadableStream<UIMessageChunk>({
|
||||
start: async (controller) => {
|
||||
try {
|
||||
const sseStream = await subscription.subscribe();
|
||||
const reader = sseStream.getReader();
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
|
||||
if (done) {
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
|
||||
if (abortSignal?.aborted) {
|
||||
reader.cancel();
|
||||
reader.releaseLock();
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
|
||||
controller.enqueue(value.chunk as UIMessageChunk);
|
||||
}
|
||||
} catch (readError) {
|
||||
reader.releaseLock();
|
||||
throw readError;
|
||||
}
|
||||
} catch (error) {
|
||||
if (error instanceof Error && error.name === "AbortError") {
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
|
||||
controller.error(error);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a new `TriggerChatTransport` instance.
|
||||
*
|
||||
* @example
|
||||
* ```tsx
|
||||
* import { useChat } from "@ai-sdk/react";
|
||||
* import { createChatTransport } from "@trigger.dev/sdk/chat";
|
||||
*
|
||||
* const transport = createChatTransport({
|
||||
* task: "my-chat-task",
|
||||
* accessToken: publicAccessToken,
|
||||
* });
|
||||
*
|
||||
* function Chat() {
|
||||
* const { messages, sendMessage } = useChat({ transport });
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
export function createChatTransport(options: TriggerChatTransportOptions): TriggerChatTransport {
|
||||
return new TriggerChatTransport(options);
|
||||
}
|
||||
Reference in New Issue
Block a user