feat: support message compaction
This commit is contained in:
Eric Allam
+1114
-43
File diff suppressed because it is too large
Load Diff
@@ -661,8 +661,18 @@ function define<TPart>(opts: RealtimeDefineStreamOptions): RealtimeDefinedStream
|
||||
read(runId, options) {
|
||||
return read(runId, opts.id, options);
|
||||
},
|
||||
append(value, options) {
|
||||
return append(opts.id, value as BodyInit, options);
|
||||
async append(value, options) {
|
||||
// Use a single-write writer so objects are serialized the same way
|
||||
// as stream.writer() — the raw append API sends BodyInit which
|
||||
// doesn't serialize objects correctly for SSE consumers.
|
||||
const { waitUntilComplete } = writer(opts.id, {
|
||||
...options,
|
||||
spanName: "streams.append()",
|
||||
execute: ({ write }) => {
|
||||
write(value);
|
||||
},
|
||||
});
|
||||
await waitUntilComplete();
|
||||
},
|
||||
writer(options) {
|
||||
return writer(opts.id, options);
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import type { UIMessage } from "ai";
|
||||
import { useChat } from "@ai-sdk/react";
|
||||
import type { TriggerChatTransport } from "@trigger.dev/sdk/chat";
|
||||
import type { CompactionChunkData } from "@trigger.dev/sdk/ai";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
import { MODEL_OPTIONS } from "@/lib/models";
|
||||
@@ -11,10 +12,10 @@ function ToolInvocation({ part }: { part: any }) {
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
const toolName =
|
||||
part.type === "dynamic-tool"
|
||||
? (part.toolName ?? "tool")
|
||||
? part.toolName ?? "tool"
|
||||
: part.type.startsWith("tool-")
|
||||
? part.type.slice(5)
|
||||
: "tool";
|
||||
? part.type.slice(5)
|
||||
: "tool";
|
||||
const state = part.state ?? "input-available";
|
||||
const args = part.input;
|
||||
const result = part.output;
|
||||
@@ -102,7 +103,9 @@ function ResearchProgress({ part }: { part: any }) {
|
||||
{data.completedUrls.length > 0 && (
|
||||
<div className="mt-1 space-y-0.5 text-blue-400">
|
||||
{data.completedUrls.map((url, i) => (
|
||||
<div key={i} className="truncate">✓ {url}</div>
|
||||
<div key={i} className="truncate">
|
||||
✓ {url}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
@@ -132,9 +135,7 @@ function DebugPanel({
|
||||
const [open, setOpen] = useState(false);
|
||||
|
||||
const runUrl =
|
||||
session?.runId && dashboardUrl
|
||||
? `${dashboardUrl}/runs/${session.runId}`
|
||||
: undefined;
|
||||
session?.runId && dashboardUrl ? `${dashboardUrl}/runs/${session.runId}` : undefined;
|
||||
|
||||
const latestTtfb = ttfbHistory.length > 0 ? ttfbHistory[ttfbHistory.length - 1]! : undefined;
|
||||
const avgTtfb =
|
||||
@@ -155,17 +156,17 @@ function DebugPanel({
|
||||
status === "streaming"
|
||||
? "bg-green-500"
|
||||
: session?.runId
|
||||
? "bg-yellow-500"
|
||||
: "bg-gray-300"
|
||||
? "bg-yellow-500"
|
||||
: "bg-gray-300"
|
||||
}`}
|
||||
/>
|
||||
<span>{status}</span>
|
||||
{latestTtfb && (
|
||||
<span className="font-mono text-blue-600">TTFB {latestTtfb.ttfbMs.toLocaleString()}ms</span>
|
||||
)}
|
||||
{session?.runId && (
|
||||
<span className="font-mono">{session.runId.slice(0, 16)}...</span>
|
||||
<span className="font-mono text-blue-600">
|
||||
TTFB {latestTtfb.ttfbMs.toLocaleString()}ms
|
||||
</span>
|
||||
)}
|
||||
{session?.runId && <span className="font-mono">{session.runId}</span>}
|
||||
<span className="ml-auto text-gray-400">{open ? "▲" : "▼"}</span>
|
||||
</button>
|
||||
|
||||
@@ -362,7 +363,9 @@ export function Chat({
|
||||
{/* Messages */}
|
||||
<div className="flex-1 space-y-4 overflow-y-auto p-4">
|
||||
{messages.length === 0 && (
|
||||
<p className="pt-20 text-center text-sm text-gray-400">Send a message to start chatting.</p>
|
||||
<p className="pt-20 text-center text-sm text-gray-400">
|
||||
Send a message to start chatting.
|
||||
</p>
|
||||
)}
|
||||
|
||||
{messages.map((message, messageIndex) => (
|
||||
@@ -373,9 +376,7 @@ export function Chat({
|
||||
<div className={`max-w-[80%] ${message.role === "user" ? "" : "w-full"}`}>
|
||||
<div
|
||||
className={`rounded-lg px-4 py-2 text-sm ${
|
||||
message.role === "user"
|
||||
? "bg-blue-600 text-white"
|
||||
: "bg-gray-100 text-gray-900"
|
||||
message.role === "user" ? "bg-blue-600 text-white" : "bg-gray-100 text-gray-900"
|
||||
}`}
|
||||
>
|
||||
{message.parts.map((part, i) => {
|
||||
@@ -386,8 +387,7 @@ export function Chat({
|
||||
key={i}
|
||||
animated
|
||||
isAnimating={
|
||||
status === "streaming" &&
|
||||
messageIndex === messages.length - 1
|
||||
status === "streaming" && messageIndex === messages.length - 1
|
||||
}
|
||||
>
|
||||
{part.text}
|
||||
@@ -414,13 +414,41 @@ export function Chat({
|
||||
return <ResearchProgress key={i} part={part} />;
|
||||
}
|
||||
|
||||
if (part.type === "data-compaction") {
|
||||
const data = (part as any).data as CompactionChunkData;
|
||||
return (
|
||||
<div
|
||||
key={i}
|
||||
className={`my-2 flex items-center gap-2 rounded-md border px-3 py-2 text-xs ${
|
||||
data.status === "compacting"
|
||||
? "border-blue-200 bg-blue-50 text-blue-700"
|
||||
: "border-amber-200 bg-amber-50 text-amber-700"
|
||||
}`}
|
||||
>
|
||||
<span>{data.status === "compacting" ? "⏳" : "✂️"}</span>
|
||||
<span>
|
||||
{data.status === "compacting"
|
||||
? `Compacting conversation${
|
||||
data.totalTokens
|
||||
? ` (${data.totalTokens.toLocaleString()} tokens)`
|
||||
: ""
|
||||
}...`
|
||||
: "Conversation compacted"}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
if (part.type.startsWith("tool-") || part.type === "dynamic-tool") {
|
||||
return <ToolInvocation key={i} part={part} />;
|
||||
}
|
||||
|
||||
if (part.type.startsWith("data-")) {
|
||||
return (
|
||||
<div key={i} className="my-1 rounded border border-gray-200 bg-gray-50 p-2 text-xs text-gray-500">
|
||||
<div
|
||||
key={i}
|
||||
className="my-1 rounded border border-gray-200 bg-gray-50 p-2 text-xs text-gray-500"
|
||||
>
|
||||
<span className="font-medium">{part.type}</span>
|
||||
<pre className="mt-1 overflow-x-auto whitespace-pre-wrap">
|
||||
{JSON.stringify((part as any).data, null, 2)}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { chat, ai, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
|
||||
import { logger, schemaTask, task, prompts } from "@trigger.dev/sdk";
|
||||
import { streamText, tool, dynamicTool, stepCountIs, generateId, createProviderRegistry } from "ai";
|
||||
import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
|
||||
import { logger, task, prompts } from "@trigger.dev/sdk";
|
||||
import { streamText, generateText, tool, dynamicTool, stepCountIs, generateId, createProviderRegistry } from "ai";
|
||||
import type { LanguageModel, Tool as AITool, UIMessage } from "ai";
|
||||
import { openai } from "@ai-sdk/openai";
|
||||
import { anthropic } from "@ai-sdk/anthropic";
|
||||
@@ -16,9 +16,24 @@ import TurndownService from "turndown";
|
||||
import { DEFAULT_MODEL, REASONING_MODELS } from "@/lib/models";
|
||||
|
||||
const turndown = new TurndownService();
|
||||
const COMPACT_AFTER_TOKENS = Number(process.env.COMPACT_AFTER_TOKENS) || 80_000;
|
||||
|
||||
const registry = createProviderRegistry({ openai, anthropic });
|
||||
|
||||
const compactionPrompt = prompts.define({
|
||||
id: "ai-chat-compaction",
|
||||
model: "openai:gpt-4o-mini",
|
||||
content: `You are a conversation compactor. You will receive a transcript of a multi-turn conversation between a user and an assistant.
|
||||
|
||||
Produce a concise summary that captures:
|
||||
- The topics discussed and questions asked
|
||||
- Any key facts, answers, or decisions reached
|
||||
- Important context needed to continue the conversation naturally
|
||||
|
||||
Write in third person (e.g. "The user asked about..." / "The assistant explained...").
|
||||
Keep it under 300 words. Do not include greetings or filler.`,
|
||||
});
|
||||
|
||||
const systemPrompt = prompts.define({
|
||||
id: "ai-chat-system",
|
||||
model: "openai:gpt-4o",
|
||||
@@ -157,59 +172,40 @@ const userToolDefs = chat.local<{
|
||||
}>({ id: "userToolDefs" });
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Subtask: deep research — fetches multiple URLs and streams progress
|
||||
// back to the parent chat via chat.stream using data-* chunks
|
||||
// Deep research — fetches multiple URLs and synthesizes the results.
|
||||
// Plain tool (not a subtask) to avoid parallel wait issues.
|
||||
// --------------------------------------------------------------------------
|
||||
export const deepResearch = schemaTask({
|
||||
id: "deep-research",
|
||||
const deepResearch = tool({
|
||||
description:
|
||||
"Research a topic by fetching multiple URLs and synthesizing the results. " +
|
||||
"Streams progress updates to the chat as it works.",
|
||||
schema: z.object({
|
||||
inputSchema: z.object({
|
||||
query: z.string().describe("The research query or topic"),
|
||||
urls: z.array(z.string().url()).describe("URLs to fetch and analyze"),
|
||||
}),
|
||||
run: async ({ query, urls }) => {
|
||||
// Access chat context from the parent chat.task — typed via typeof aiChat
|
||||
const { chatId, clientData } = ai.chatContextOrThrow<typeof aiChat>();
|
||||
console.log(`Deep research for chat ${chatId}, user ${clientData?.userId}`);
|
||||
|
||||
execute: async ({ query, urls }) => {
|
||||
const partId = generateId();
|
||||
const results: { url: string; status: number; snippet: string }[] = [];
|
||||
|
||||
// Stream progress using data-research-progress chunks.
|
||||
// Using the same id means each write updates the same part in the message.
|
||||
function streamProgress(progress: {
|
||||
status: "fetching" | "done";
|
||||
query: string;
|
||||
current: number;
|
||||
total: number;
|
||||
currentUrl?: string;
|
||||
completedUrls: string[];
|
||||
}) {
|
||||
return chat.stream.writer({
|
||||
target: "root",
|
||||
execute: ({ write }) => {
|
||||
write({
|
||||
type: "data-research-progress" as any,
|
||||
id: partId,
|
||||
data: progress,
|
||||
});
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
for (let i = 0; i < urls.length; i++) {
|
||||
const url = urls[i]!;
|
||||
|
||||
// Update progress — fetching
|
||||
const { waitUntilComplete } = streamProgress({
|
||||
status: "fetching",
|
||||
query,
|
||||
current: i + 1,
|
||||
total: urls.length,
|
||||
currentUrl: url,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
// Stream progress — runs in the chat.task process, so no target needed
|
||||
const { waitUntilComplete } = chat.stream.writer({
|
||||
execute: ({ write }) => {
|
||||
write({
|
||||
type: "data-research-progress" as any,
|
||||
id: partId,
|
||||
data: {
|
||||
status: "fetching" as const,
|
||||
query,
|
||||
current: i + 1,
|
||||
total: urls.length,
|
||||
currentUrl: url,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
await waitUntilComplete();
|
||||
|
||||
@@ -236,13 +232,21 @@ export const deepResearch = schemaTask({
|
||||
}
|
||||
}
|
||||
|
||||
// Final progress update — done
|
||||
const { waitUntilComplete: waitForDone } = streamProgress({
|
||||
status: "done",
|
||||
query,
|
||||
current: urls.length,
|
||||
total: urls.length,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
// Final progress — done
|
||||
const { waitUntilComplete: waitForDone } = chat.stream.writer({
|
||||
execute: ({ write }) => {
|
||||
write({
|
||||
type: "data-research-progress" as any,
|
||||
id: partId,
|
||||
data: {
|
||||
status: "done" as const,
|
||||
query,
|
||||
current: urls.length,
|
||||
total: urls.length,
|
||||
completedUrls: results.map((r) => r.url),
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
await waitForDone();
|
||||
|
||||
@@ -255,6 +259,46 @@ export const aiChat = chat.task({
|
||||
clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }),
|
||||
idleTimeoutInSeconds: 60,
|
||||
chatAccessTokenTTL: "2h",
|
||||
compaction: {
|
||||
shouldCompact: ({ totalTokens }) => (totalTokens ?? 0) > COMPACT_AFTER_TOKENS,
|
||||
summarize: async ({ messages }) => {
|
||||
const resolved = await compactionPrompt.resolve({});
|
||||
return generateText({
|
||||
model: registry.languageModel(resolved.model ?? "openai:gpt-4o-mini"),
|
||||
messages: [...messages, { role: "user" as const, content: resolved.text }],
|
||||
...resolved.toAISDKTelemetry(),
|
||||
}).then((r) => r.text);
|
||||
},
|
||||
compactUIMessages: ({ uiMessages, summary }) => {
|
||||
return [
|
||||
{
|
||||
id: generateId(),
|
||||
role: "assistant" as const,
|
||||
parts: [{ type: "text" as const, text: `[Conversation summary]\n\n${summary}` }],
|
||||
},
|
||||
...uiMessages.slice(-2),
|
||||
];
|
||||
},
|
||||
},
|
||||
prepareMessages: ({ messages, reason }) => {
|
||||
// Add Anthropic cache breaks to the last message for prompt caching.
|
||||
// Applied everywhere — run(), compaction rebuilds, compaction results.
|
||||
if (messages.length === 0) return messages;
|
||||
const last = messages[messages.length - 1]!;
|
||||
return [
|
||||
...messages.slice(0, -1),
|
||||
{
|
||||
...last,
|
||||
providerOptions: {
|
||||
...last.providerOptions,
|
||||
anthropic: {
|
||||
...(last.providerOptions?.anthropic as Record<string, unknown> | undefined),
|
||||
cacheControl: { type: "ephemeral" },
|
||||
},
|
||||
},
|
||||
},
|
||||
];
|
||||
},
|
||||
uiMessageStreamOptions: {
|
||||
sendReasoning: true,
|
||||
onError: (error) => {
|
||||
@@ -362,12 +406,21 @@ export const aiChat = chat.task({
|
||||
update: { runId, publicAccessToken: chatAccessToken },
|
||||
});
|
||||
},
|
||||
onCompacted: async ({ summary, totalTokens, messageCount, chatId, turn }) => {
|
||||
logger.info("Conversation compacted", {
|
||||
chatId,
|
||||
turn,
|
||||
totalTokens,
|
||||
messageCount,
|
||||
summaryLength: summary.length,
|
||||
});
|
||||
},
|
||||
onTurnStart: async ({ chatId, uiMessages }) => {
|
||||
// Persist messages so mid-stream refresh still shows the user message.
|
||||
// Deferred — runs in parallel with streaming, awaited before onTurnComplete.
|
||||
chat.defer(prisma.chat.update({ where: { id: chatId }, data: { messages: uiMessages as any } }));
|
||||
},
|
||||
onTurnComplete: async ({ chatId, uiMessages, runId, chatAccessToken, lastEventId, clientData, stopped }) => {
|
||||
onTurnComplete: async ({ chatId, uiMessages, runId, chatAccessToken, lastEventId }) => {
|
||||
// Persist final messages + assistant response + stream position
|
||||
await prisma.chat.update({
|
||||
where: { id: chatId },
|
||||
@@ -426,11 +479,11 @@ export const aiChat = chat.task({
|
||||
telemetry: clientData?.userId ? { userId: clientData.userId } : undefined,
|
||||
}),
|
||||
...(modelOverride ? { model: getModel(modelOverride) } : {}),
|
||||
messages,
|
||||
messages: messages,
|
||||
tools: {
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
deepResearch: ai.tool(deepResearch),
|
||||
deepResearch,
|
||||
...dynamicTools,
|
||||
},
|
||||
stopWhen: stepCountIs(10),
|
||||
@@ -553,11 +606,11 @@ export const aiChatRaw = task({
|
||||
const result = streamText({
|
||||
...chat.toStreamTextOptions({ registry }),
|
||||
...(modelOverride ? { model: getModel(modelOverride) } : {}),
|
||||
messages,
|
||||
messages: messages,
|
||||
tools: {
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
deepResearch: ai.tool(deepResearch),
|
||||
deepResearch,
|
||||
...dynamicTools,
|
||||
},
|
||||
stopWhen: stepCountIs(10),
|
||||
@@ -569,6 +622,33 @@ export const aiChatRaw = task({
|
||||
...(useReasoning ? { thinking: { type: "enabled", budgetTokens: 10000 } } : {}),
|
||||
},
|
||||
},
|
||||
// Low-level compaction using chat.compact() — gives full control
|
||||
// while chat.compact handles the decision tree + stream chunks
|
||||
prepareStep: async ({ messages: stepMessages, steps }) => {
|
||||
// Custom logic before/around compaction
|
||||
const lastStep = steps.at(-1);
|
||||
if (lastStep?.usage.totalTokens) {
|
||||
logger.info("Raw task: step usage", { totalTokens: lastStep.usage.totalTokens, turn });
|
||||
}
|
||||
|
||||
const result = await chat.compact(stepMessages, steps, {
|
||||
threshold: COMPACT_AFTER_TOKENS,
|
||||
summarize: async (msgs) => {
|
||||
const resolved = await compactionPrompt.resolve({});
|
||||
return generateText({
|
||||
model: registry.languageModel(resolved.model ?? "openai:gpt-4o-mini"),
|
||||
...resolved.toAISDKTelemetry(),
|
||||
messages: [...msgs, { role: "user" as const, content: resolved.text }],
|
||||
}).then((r) => r.text);
|
||||
},
|
||||
});
|
||||
|
||||
if (result.type === "compacted") {
|
||||
logger.info("Raw task: compacted", { summary: result.summary.slice(0, 100) });
|
||||
}
|
||||
|
||||
return result.type === "skipped" ? undefined : result;
|
||||
},
|
||||
});
|
||||
|
||||
let response: UIMessage | undefined;
|
||||
@@ -663,7 +743,7 @@ export const aiChatSession = task({
|
||||
tools: {
|
||||
inspectEnvironment,
|
||||
webFetch,
|
||||
deepResearch: ai.tool(deepResearch),
|
||||
deepResearch,
|
||||
},
|
||||
stopWhen: stepCountIs(10),
|
||||
abortSignal: turn.signal,
|
||||
@@ -674,6 +754,31 @@ export const aiChatSession = task({
|
||||
...(useReasoning ? { thinking: { type: "enabled", budgetTokens: 10000 } } : {}),
|
||||
},
|
||||
},
|
||||
// Low-level compaction — same pattern as raw task
|
||||
prepareStep: async ({ messages: stepMessages, steps }) => {
|
||||
const lastStep = steps.at(-1);
|
||||
if (lastStep?.usage.totalTokens) {
|
||||
logger.info("Session: step usage", { totalTokens: lastStep.usage.totalTokens, turn: turn.number });
|
||||
}
|
||||
|
||||
const result = await chat.compact(stepMessages, steps, {
|
||||
threshold: COMPACT_AFTER_TOKENS,
|
||||
summarize: async (msgs) => {
|
||||
const resolved = await compactionPrompt.resolve({});
|
||||
return generateText({
|
||||
model: registry.languageModel(resolved.model ?? "openai:gpt-4o-mini"),
|
||||
...resolved.toAISDKTelemetry(),
|
||||
messages: [...msgs, { role: "user" as const, content: resolved.text }],
|
||||
}).then((r) => r.text);
|
||||
},
|
||||
});
|
||||
|
||||
if (result.type === "compacted") {
|
||||
logger.info("Session: compacted", { summary: result.summary.slice(0, 100) });
|
||||
}
|
||||
|
||||
return result.type === "skipped" ? undefined : result;
|
||||
},
|
||||
});
|
||||
|
||||
await turn.complete(result);
|
||||
|
||||
@@ -4,7 +4,7 @@ import { prismaExtension } from "@trigger.dev/build/extensions/prisma";
|
||||
export default defineConfig({
|
||||
project: process.env.TRIGGER_PROJECT_REF!,
|
||||
dirs: ["./src/trigger"],
|
||||
maxDuration: 300,
|
||||
maxDuration: 3600,
|
||||
build: {
|
||||
extensions: [
|
||||
prismaExtension({
|
||||
|
||||
Reference in New Issue
Block a user