1371 lines
51 KiB
TypeScript
1371 lines
51 KiB
TypeScript
import { db } from "~~/server/lib/db";
|
|
import * as z from 'zod';
|
|
import { type MessageEntity } from '~/composables/useChat';
|
|
import { promises as fs } from 'fs';
|
|
import { isLoopFinished, type ModelMessage, rerank, type RerankingModel, streamText, type StreamTextTransform, type Tool, tool } from "ai";
|
|
import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
|
|
import { topicEvents } from "~~/server/utils/events";
|
|
import { nanoid } from "nanoid";
|
|
import { Monty, MontyRuntimeError, MontySyntaxError, MontyTypingError } from '@pydantic/monty';
|
|
import { type Model } from "~/composables/useModels";
|
|
import { eq } from "drizzle-orm";
|
|
import { buildFocusedMessageTree, buildMessageTree, marshallMessages } from "~~/utils/message";
|
|
import path from "path";
|
|
import { isRerankingProvider } from "~~/server/utils/ai-provider";
|
|
import { generateFileToken } from "~~/server/utils/file-token";
|
|
|
|
export default defineEventHandler(async (event) => {
|
|
await protectRoute(event);
|
|
|
|
const topicId = getRouterParam(event, 'topicId')!;
|
|
const userId = event.context.user!.id as string;
|
|
|
|
const result = await readValidatedBody(event, z.object({
|
|
parentMessageId: z.string().optional(),
|
|
modelId: z.string(),
|
|
rerank: z.object({
|
|
modelId: z.string(),
|
|
providerApiKey: z.string().optional(),
|
|
}).optional(),
|
|
args: z.record(z.string(), z.any()).optional(),
|
|
providerApiKey: z.string().optional(),
|
|
}).safeParse);
|
|
if (!result.success) {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: 'Bad Request',
|
|
message: result.error.issues.map(issue => issue.message).join(', '),
|
|
});
|
|
}
|
|
|
|
const { parentMessageId, modelId, rerank: rerankConfig, providerApiKey, args } = result.data;
|
|
|
|
const topic = await db.query.topics.findFirst({
|
|
where: {
|
|
id: topicId,
|
|
userId,
|
|
},
|
|
with: {
|
|
agent: true,
|
|
messages: {
|
|
orderBy: {
|
|
createdAt: 'asc',
|
|
},
|
|
with: {
|
|
parts: {
|
|
orderBy: {
|
|
createdAt: 'asc',
|
|
},
|
|
with: {
|
|
toolCall: true,
|
|
}
|
|
},
|
|
attachments: {
|
|
with: {
|
|
file: true,
|
|
}
|
|
},
|
|
generation: true,
|
|
}
|
|
},
|
|
},
|
|
});
|
|
if (!topic) {
|
|
throw createError({
|
|
statusCode: 404,
|
|
statusMessage: 'Not Found',
|
|
message: 'Topic not found',
|
|
});
|
|
}
|
|
|
|
const model = await db.query.models.findFirst({
|
|
where: {
|
|
id: modelId,
|
|
userId,
|
|
},
|
|
with: {
|
|
provider: true,
|
|
}
|
|
});
|
|
|
|
if (!model) {
|
|
throw createError({
|
|
statusCode: 404,
|
|
statusMessage: 'Not Found',
|
|
message: 'Model not found',
|
|
});
|
|
}
|
|
|
|
const providerDetails = await getProviderDetails(model.provider, providerApiKey, model);
|
|
if (!providerDetails.ok) {
|
|
switch (providerDetails.error) {
|
|
case GatewayFetchError.NoProviderApiKey: {
|
|
throw createError({
|
|
statusCode: 400,
|
|
message: `${model.provider.type} provider requires an API key`,
|
|
});
|
|
}
|
|
case GatewayFetchError.NoProviderBaseUrl: {
|
|
throw createError({
|
|
statusCode: 400,
|
|
message: 'Invalid provider URL',
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
const agentSearchConfig = topic.agent.config?.search;
|
|
let searchParam: false | { config: SearchTheWebConfig } = false;
|
|
|
|
if (agentSearchConfig?.enabled) {
|
|
searchParam = { config: { rerank: false, maxResults: agentSearchConfig.maxResults ?? 10 } };
|
|
|
|
if (agentSearchConfig.rerank && rerankConfig) {
|
|
const rerankModel = await db.query.models.findFirst({
|
|
where: {
|
|
id: rerankConfig.modelId,
|
|
userId,
|
|
},
|
|
with: {
|
|
provider: true,
|
|
}
|
|
});
|
|
|
|
if (rerankModel === undefined) {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: 'Invalid rerank model',
|
|
data: {
|
|
code: 'INVALID_RERANK_MODEL',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
|
|
const rerankProviderDetails = await getProviderDetails(rerankModel.provider, rerankConfig.providerApiKey, rerankModel);
|
|
if (!rerankProviderDetails.ok) {
|
|
switch (rerankProviderDetails.error) {
|
|
case GatewayFetchError.NoProviderApiKey: {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: `${rerankModel.provider.type} provider requires an API key`,
|
|
data: {
|
|
code: 'NO_RERANK_PROVIDER_API_KEY',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
case GatewayFetchError.NoProviderBaseUrl: {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: 'Invalid provider URL',
|
|
data: {
|
|
code: 'BAD_RERANK_PROVIDER_URL',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
const { gateway: rerankGateway } = rerankProviderDetails.data;
|
|
if (rerankGateway === null) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Invalid gateway',
|
|
data: {
|
|
code: 'INVALID_RERANK_GATEWAY',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
|
|
if (isRerankingProvider(rerankGateway.gateway)) {
|
|
searchParam.config = {
|
|
rerank: true,
|
|
maxResults: agentSearchConfig.maxResults ?? 10,
|
|
model: rerankGateway.gateway.reranking(rerankModel.externalId)
|
|
};
|
|
} else {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Invalid rerank model',
|
|
data: {
|
|
code: 'INVALID_RERANK_MODEL',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
let agentmessage = await db.transaction(async tx => {
|
|
const generationId = nanoid();
|
|
|
|
const [generation] = await tx.insert(generations).values({
|
|
id: generationId,
|
|
userId,
|
|
topicId,
|
|
modelId: model.externalId,
|
|
status: 'pending',
|
|
}).returning();
|
|
|
|
if (!generation) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Failed to insert generation',
|
|
message: 'Failed to insert generation',
|
|
});
|
|
}
|
|
|
|
const [agentmessage] = await tx.insert(messages).values({
|
|
id: nanoid(),
|
|
userId,
|
|
topicId,
|
|
content: null,
|
|
role: 'assistant',
|
|
generationId,
|
|
parentMessageId,
|
|
}).returning();
|
|
|
|
if (!agentmessage) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Failed to insert message',
|
|
message: 'Failed to insert message',
|
|
});
|
|
}
|
|
|
|
agentmessage.generationId = generation.id;
|
|
// @ts-ignore - doesnt exist on the type but yeah it does now
|
|
agentmessage.generation = generation;
|
|
|
|
return agentmessage;
|
|
});
|
|
const events = [{ type: 'MESSAGE_CREATED', payload: agentmessage }] as { type: string; payload: any }[];
|
|
|
|
let topicMessageTree = buildMessageTree(topic.messages);
|
|
|
|
if (parentMessageId) {
|
|
const parentMessage = await db.query.messages.findFirst({
|
|
where: {
|
|
id: parentMessageId,
|
|
}
|
|
});
|
|
if (!parentMessage) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Failed to find parent message',
|
|
message: 'Failed to find parent message',
|
|
});
|
|
}
|
|
|
|
await db.update(messages).set({ activeChildId: agentmessage.id }).where(eq(messages.id, parentMessageId));
|
|
events.push({ type: 'MESSAGE_UPDATED', payload: { id: parentMessageId, activeChildId: agentmessage.id } });
|
|
|
|
// we need to make the topic messages all the messages excluding the ones after the message we wish to regenerate
|
|
// and if parentMessageId is undefined, then excluding the last message
|
|
if (parentMessageId === undefined) {
|
|
topicMessageTree = topicMessageTree.slice(0, topicMessageTree.length - 1);
|
|
} else {
|
|
const parentMessageIdx = topicMessageTree.findIndex(m => m.id === parentMessageId);
|
|
if (parentMessageIdx === -1) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Failed to find parent message',
|
|
message: 'Failed to find parent message',
|
|
});
|
|
}
|
|
|
|
topicMessageTree = topicMessageTree.slice(0, parentMessageIdx);
|
|
}
|
|
}
|
|
|
|
topicEvents.emit(topicId, { type: 'MESSAGE_CREATED', payload: agentmessage });
|
|
|
|
const { gateway } = providerDetails.data;
|
|
if (gateway === null) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Invalid gateway',
|
|
data: {
|
|
code: 'INVALID_GATEWAY',
|
|
ok: false,
|
|
}
|
|
});
|
|
}
|
|
|
|
let logFile: fs.FileHandle | undefined;
|
|
let logMessage: ((message: string) => void) | undefined;
|
|
|
|
if (process.env.GENERATION_DEBUG) {
|
|
if (process.env.LOG_DIR) {
|
|
await fs.mkdir(process.env.LOG_DIR!, { recursive: true });
|
|
|
|
logFile = await fs.open(path.join(process.env.LOG_DIR!, `${Date.now()}-${agentmessage.generationId!}.log`), 'w');
|
|
logMessage = (message: string) => {
|
|
logFile!.write(message + '\n');
|
|
};
|
|
} else {
|
|
console.warn('Generation debug logging is enabled but LOG_DIR is not set');
|
|
}
|
|
}
|
|
|
|
const fileTokenSecret = process.env.BETTER_AUTH_SECRET!;
|
|
const topicMessages = marshallMessages(topic.agent, buildFocusedMessageTree(topicMessageTree), {
|
|
signFileUrl: (fileKey) => {
|
|
const { exp, sig } = generateFileToken(fileKey, fileTokenSecret);
|
|
return `exp=${exp}&sig=${sig}`;
|
|
},
|
|
});
|
|
if (topicMessages.ok === false) {
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Failed to marshall messages',
|
|
message: topicMessages.error,
|
|
});
|
|
}
|
|
|
|
event.waitUntil(
|
|
generateResponse(
|
|
agentmessage as MessageEntity,
|
|
{ gateway: gateway.gateway, model, parameters: args },
|
|
searchParam,
|
|
{
|
|
python: topic.agent.config?.tools?.python ?? false,
|
|
},
|
|
agentmessage.generationId!,
|
|
userId,
|
|
topicId,
|
|
topicMessages.data,
|
|
gateway.streamTransformer,
|
|
logMessage,
|
|
logFile
|
|
)
|
|
);
|
|
|
|
return {
|
|
ok: true
|
|
};
|
|
});
|
|
|
|
|
|
const INTERNAL_ERROR = 'An internal error occurred';
|
|
|
|
// todo message takes in variadics like console.log
|
|
const todo = (...args: any[]) => {
|
|
console.error('TODO', ...args);
|
|
};
|
|
|
|
const formatPartId = (partType: string, existingId: string) => {
|
|
return `veridian__part-${partType}-${existingId}-${nanoid()}`;
|
|
};
|
|
|
|
const formatToolCallId = (nativeId: string) => {
|
|
return `veridian__tool-${nativeId.slice(0, 16)}-${nanoid()}`;
|
|
};
|
|
|
|
const formatPythonValue = (value: unknown): string => {
|
|
if (value === undefined || value === null) {
|
|
return '';
|
|
}
|
|
if (typeof value === 'string') {
|
|
return value;
|
|
}
|
|
try {
|
|
return JSON.stringify(value, (_, v) => typeof v === 'bigint' ? v.toString() : v);
|
|
} catch {
|
|
return String(value);
|
|
}
|
|
};
|
|
|
|
const evalPython = async (code: string): Promise<string> => {
|
|
try {
|
|
let stdout = '';
|
|
const printCallback = (_: string, text: string) => {
|
|
stdout += text;
|
|
};
|
|
const m = new Monty(code);
|
|
const result = m.run({
|
|
printCallback,
|
|
limits: {
|
|
maxDurationSecs: 10,
|
|
maxMemory: 32 * 1024 * 1024,
|
|
maxRecursionDepth: 100,
|
|
},
|
|
});
|
|
|
|
const expressionOutput = formatPythonValue(result);
|
|
if (stdout && expressionOutput) {
|
|
return stdout.endsWith('\n')
|
|
? `${stdout}${expressionOutput}`
|
|
: `${stdout}\n${expressionOutput}`;
|
|
}
|
|
return stdout || expressionOutput || '';
|
|
} catch (error) {
|
|
if (error instanceof MontySyntaxError) {
|
|
return `SyntaxError: ${error.message}`;
|
|
}
|
|
if (error instanceof MontyRuntimeError) {
|
|
return error.display('traceback') || `RuntimeError: ${error.message}`;
|
|
}
|
|
if (error instanceof MontyTypingError) {
|
|
return error.display('concise') || `TypeError: ${error.message}`;
|
|
}
|
|
return error instanceof Error ? error.message : 'Python execution failed';
|
|
}
|
|
};
|
|
|
|
interface SearchTheWebRerankedConfig {
|
|
rerank: true;
|
|
maxResults: number;
|
|
model: RerankingModel;
|
|
}
|
|
|
|
interface SearchTheWebUnrankedConfig {
|
|
rerank?: false;
|
|
maxResults: number;
|
|
}
|
|
|
|
type SearchTheWebConfig = SearchTheWebRerankedConfig | SearchTheWebUnrankedConfig;
|
|
|
|
export const searchTheWeb = (config: SearchTheWebConfig) => {
|
|
return async (query: string) => {
|
|
const searchUrl = new URL(process.env.SEARCH_API_URL ?? process.env.SEARXNG_URL ?? "");
|
|
searchUrl.pathname = `${searchUrl.pathname.replace(/\/+$/, "")}/search`;
|
|
|
|
const results = await $fetch<{
|
|
results: Array<{
|
|
title: string;
|
|
url: string;
|
|
content: string;
|
|
}>;
|
|
}>(searchUrl.toString(), {
|
|
query: {
|
|
q: query,
|
|
format: 'json',
|
|
}
|
|
});
|
|
|
|
const sites = results.results.map((item: any) => ({
|
|
title: item.title,
|
|
link: item.url,
|
|
snippet: item.content,
|
|
}));
|
|
|
|
if (config.rerank) {
|
|
const { ranking } = await rerank({
|
|
model: config.model,
|
|
query,
|
|
documents: sites.map(site => site.snippet),
|
|
});
|
|
|
|
const ranked_sites = [];
|
|
|
|
for (let i = 0; i < ranking.length; i++) {
|
|
ranked_sites.push(sites[ranking[i]!.originalIndex]);
|
|
}
|
|
|
|
return ranked_sites.slice(0, config.maxResults || 10);
|
|
}
|
|
|
|
return sites.slice(0, config.maxResults || 10);
|
|
};
|
|
};
|
|
|
|
const fetchUrlTool = tool({
|
|
description: 'Fetches the content of a URL',
|
|
inputSchema: z.object({
|
|
url: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
content: z.string(),
|
|
}),
|
|
execute: async ({ url }) => {
|
|
const response = await fetch('https://api.firecrawl.dev/v2/scrape', {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': 'Bearer ' + process.env.FIRECRAWL_API_KEY,
|
|
},
|
|
body: JSON.stringify({
|
|
url,
|
|
}),
|
|
});
|
|
const content = await response.json();
|
|
return {
|
|
content: content.data.markdown,
|
|
};
|
|
},
|
|
});
|
|
|
|
const pythonTool = tool({
|
|
description: 'Executes Python in a sandbox. The value of the last expression is returned automatically (no print needed). stdout from print() is also included.',
|
|
inputSchema: z.object({
|
|
code: z.string().describe('Python code to run. Prefer a final expression over print(), e.g. `2 + 2` returns `4`.'),
|
|
}),
|
|
outputSchema: z.object({
|
|
output: z.string(),
|
|
}),
|
|
execute: async ({ code }) => {
|
|
const output = await evalPython(code);
|
|
return {
|
|
output,
|
|
};
|
|
},
|
|
});
|
|
|
|
async function generateResponse(
|
|
message: MessageEntity,
|
|
model: {
|
|
gateway: ModelGateway,
|
|
model: Model,
|
|
parameters?: Record<string, any>,
|
|
},
|
|
search: false | {
|
|
config: SearchTheWebConfig,
|
|
},
|
|
enabledTools: {
|
|
python?: boolean;
|
|
},
|
|
generationId: string,
|
|
userId: string,
|
|
topicId: string,
|
|
messages: ModelMessage[],
|
|
streamTransoforms: StreamTextTransform<{}> | StreamTextTransform<{}>[] | undefined,
|
|
log?: (message: string) => void,
|
|
logFile?: fs.FileHandle,
|
|
) {
|
|
const controller = new AbortController();
|
|
addPendingGeneration(generationId, controller);
|
|
|
|
let requestStart = undefined;
|
|
let ttft = undefined;
|
|
const activeParts = new Map<string, { id: string; accumulatedContent: string; providerOptions?: any }>();
|
|
const activeToolCalls = new Set<string>();
|
|
const nativeToDbToolCallId = new Map<string, string>();
|
|
|
|
const tools: Record<string, Tool> = {};
|
|
|
|
if (enabledTools.python) {
|
|
tools.python = pythonTool;
|
|
}
|
|
|
|
if (search) {
|
|
tools.search = tool({
|
|
description: 'Searches the web',
|
|
inputSchema: z.object({
|
|
query: z.string(),
|
|
}),
|
|
outputSchema: z.array(
|
|
z.object({
|
|
title: z.string(),
|
|
link: z.string(),
|
|
snippet: z.string(),
|
|
engine: z.string(),
|
|
})
|
|
),
|
|
execute: async ({ query }) => {
|
|
return await searchTheWeb(search.config)(query);
|
|
},
|
|
});
|
|
tools.fetchUrl = fetchUrlTool;
|
|
}
|
|
|
|
const response = streamText({
|
|
model: model.gateway(model.model.externalId),
|
|
messages,
|
|
allowSystemInMessages: true,
|
|
providerOptions: {
|
|
openrouter: {
|
|
debug: {
|
|
echo_upstream_body: true,
|
|
},
|
|
user: userId,
|
|
},
|
|
},
|
|
experimental_transform: streamTransoforms,
|
|
stopWhen: isLoopFinished(),
|
|
tools: model.model.capabilities.includes('tools') && Object.keys(tools).length > 0 ? tools : undefined,
|
|
onError: async (error: any) => {
|
|
// TODO: the docs say "The stream processing will pause until the callback promise is resolved." Suggesting that this error might not be fatal?
|
|
console.error('generation error', error);
|
|
log?.(error);
|
|
|
|
// Mark all active parts as finished
|
|
await Promise.all([...activeParts.values()].map(async part => {
|
|
await db.update(messageParts).set({ finished: true, lastUpdatedAt: new Date() }).where(eq(messageParts.id, part.id))
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-end',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
lastUpdatedAt: new Date(),
|
|
content: part.accumulatedContent,
|
|
}
|
|
})
|
|
}));
|
|
|
|
// Mark failed tool calls
|
|
await Promise.all([...activeToolCalls].map(async id => {
|
|
await db.update(toolCalls).set({
|
|
status: 'failed',
|
|
error: { type: ToolCallType.Text, value: 'Generation failed' }
|
|
}).where(eq(toolCalls.id, id))
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: id,
|
|
toolName: null,
|
|
error: {
|
|
type: ToolCallType.Text,
|
|
value: 'Generation failed'
|
|
}
|
|
}
|
|
})
|
|
}));
|
|
|
|
await db.update(generations).set({
|
|
status: 'failed',
|
|
error: typeof error === 'string' ? error : JSON.stringify(error)
|
|
}).where(eq(generations.id, generationId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'generation-failed',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
error: typeof error === 'string' ? error : JSON.stringify(error),
|
|
}
|
|
})
|
|
},
|
|
// onStepFinish: async (step) => {
|
|
// step.content.forEach(async (part) => {
|
|
// switch (part.type) {
|
|
// case 'text': {
|
|
// await db.insert(messageParts).values({
|
|
// userId,
|
|
// topicId,
|
|
// messageId: message.id,
|
|
// type: 'text',
|
|
// content: part.text,
|
|
// providerOptions: part.providerMetadata,
|
|
// finished: true,
|
|
// createdAt: new Date(),
|
|
// lastUpdatedAt: new Date(),
|
|
// });
|
|
// } break;
|
|
// }
|
|
// })
|
|
// },
|
|
abortSignal: controller.signal,
|
|
});
|
|
|
|
let curStepIdx = -1;
|
|
let key, type;
|
|
|
|
const pendingUpdates = new Map<string, NodeJS.Timeout>();
|
|
|
|
// since we arent streaming straight from the database, we can update less often
|
|
const TARGET_UPDATES_PER_SECOND = 10;
|
|
|
|
const scheduleUpdate = (key: string) => {
|
|
const part = activeParts.get(key);
|
|
if (!part || pendingUpdates.has(part.id)) return;
|
|
|
|
pendingUpdates.set(part.id, setTimeout(async () => {
|
|
const currentPart = activeParts.get(key);
|
|
if (!currentPart) return;
|
|
|
|
try {
|
|
await db.update(messageParts)
|
|
.set({
|
|
content: currentPart.accumulatedContent,
|
|
providerOptions: currentPart.providerOptions,
|
|
lastUpdatedAt: new Date(),
|
|
})
|
|
.where(eq(messageParts.id, currentPart.id));
|
|
} catch (e) {
|
|
console.warn('Failed to update message part', e);
|
|
}
|
|
pendingUpdates.delete(part.id);
|
|
}, 1000 / TARGET_UPDATES_PER_SECOND));
|
|
};
|
|
|
|
try {
|
|
for await (const token of response.fullStream) {
|
|
log?.(JSON.stringify(token, null, 2));
|
|
|
|
switch (token.type) {
|
|
case 'start': {
|
|
requestStart = performance.now();
|
|
break;
|
|
}
|
|
|
|
case 'start-step': {
|
|
curStepIdx++;
|
|
break;
|
|
}
|
|
|
|
case 'text-start':
|
|
case 'reasoning-start': {
|
|
const type = token.type.split('-')[0] as 'text' | 'reasoning';
|
|
const key = `${type}-${curStepIdx}`;
|
|
const nativeId = nanoid();
|
|
|
|
const [part] = await db.insert(messageParts).values({
|
|
id: formatPartId(type, nativeId),
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
type,
|
|
content: '',
|
|
finished: false,
|
|
createdAt: new Date(),
|
|
lastUpdatedAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!part) {
|
|
throw new Error('Failed to insert message part');
|
|
}
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeParts.set(key, { id: part.id, accumulatedContent: '' });
|
|
break;
|
|
}
|
|
|
|
case 'text-delta':
|
|
case 'reasoning-delta': {
|
|
if (ttft === undefined) {
|
|
ttft = performance.now() - requestStart!;
|
|
}
|
|
|
|
type = token.type.split('-')[0] as 'text' | 'reasoning';
|
|
key = `${type}-${curStepIdx}`;
|
|
const part = activeParts.get(key);
|
|
if (part === undefined) {
|
|
console.error('Received delta without a start');
|
|
break;
|
|
}
|
|
|
|
let shouldUpdate = false;
|
|
|
|
// TODO: we should potentially merge providerOptions, but for now, just overwrite them
|
|
if (token.providerMetadata !== undefined) {
|
|
shouldUpdate = true;
|
|
part.providerOptions = token.providerMetadata;
|
|
}
|
|
|
|
// OpenRouter sometimes puts [REDACTED] in thinking if reasoning is encrypted, so we need to remove it and hide it;
|
|
// do not trim or else we lose intentional whitespace and newlines potentially breaking the UI and having words comebined e.g. "the" "\n\n" "assistant" would become "theassistant"
|
|
const text = token.text.replaceAll('[REDACTED]', '');
|
|
if (text !== '') {
|
|
shouldUpdate = true;
|
|
part.accumulatedContent += token.text;
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
lastUpdatedAt: new Date(),
|
|
content: token.text,
|
|
}
|
|
})
|
|
}
|
|
|
|
if (shouldUpdate) {
|
|
scheduleUpdate(key);
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
case 'text-end':
|
|
case 'reasoning-end': {
|
|
type = token.type.split('-')[0];
|
|
key = `${type}-${curStepIdx}`;
|
|
const part = activeParts.get(key);
|
|
if (part === undefined) {
|
|
console.error('Received end without a start');
|
|
break;
|
|
}
|
|
|
|
activeParts.delete(key);
|
|
|
|
if (part.accumulatedContent === '' && !part.providerOptions) {
|
|
await db.delete(messageParts).where(eq(messageParts.id, part.id));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-delete',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
}
|
|
})
|
|
} else {
|
|
const [dbPart] = await db.update(messageParts)
|
|
.set({
|
|
content: part.accumulatedContent,
|
|
providerOptions: part.providerOptions,
|
|
finished: true,
|
|
lastUpdatedAt: new Date(),
|
|
})
|
|
.where(eq(messageParts.id, part.id)).returning();
|
|
|
|
if (!dbPart) {
|
|
throw new Error('Failed to update message part');
|
|
}
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-end',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
lastUpdatedAt: dbPart.lastUpdatedAt,
|
|
content: part.accumulatedContent,
|
|
}
|
|
})
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
case 'tool-input-start': {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.id;
|
|
const dbToolCallId = formatToolCallId(toolCallId);
|
|
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: dbToolCallId,
|
|
userId: userId,
|
|
toolName: token.toolName,
|
|
status: 'pending',
|
|
input: null,
|
|
output: null,
|
|
error: null,
|
|
createdAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!toolCall) {
|
|
throw new Error('Failed to insert tool call');
|
|
}
|
|
|
|
const [part] = await db.insert(messageParts).values({
|
|
id: formatPartId('tool-call', dbToolCallId),
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
type: 'tool-call',
|
|
content: null,
|
|
finished: false,
|
|
createdAt: new Date(),
|
|
lastUpdatedAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!part) {
|
|
throw new Error('Failed to insert message part');
|
|
}
|
|
|
|
// @ts-ignore - doesnt exist on the type but yeah it does now
|
|
part.toolCall = toolCall;
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeToolCalls.add(dbToolCallId);
|
|
|
|
activeParts.set(key, { id: part.id, accumulatedContent: '' });
|
|
} break;
|
|
case 'tool-call': {
|
|
let inputType: ToolCallType = ToolCallType.Text;
|
|
let inputValue: string = '';
|
|
|
|
switch (typeof token.input) {
|
|
case 'string':
|
|
inputType = ToolCallType.Text;
|
|
inputValue = token.input;
|
|
break;
|
|
case 'object':
|
|
inputType = ToolCallType.Json;
|
|
inputValue = JSON.stringify(token.input);
|
|
break;
|
|
default:
|
|
console.error('Unknown input type', token.input);
|
|
break;
|
|
}
|
|
|
|
const dbToolCallIdFromMap = nativeToDbToolCallId.get(token.toolCallId);
|
|
if (dbToolCallIdFromMap && activeToolCalls.has(dbToolCallIdFromMap)) {
|
|
const dbToolCallId = dbToolCallIdFromMap;
|
|
await db.update(toolCalls)
|
|
.set({
|
|
status: 'pending',
|
|
input: {
|
|
type: inputType,
|
|
value: inputValue,
|
|
},
|
|
})
|
|
.where(eq(toolCalls.id, dbToolCallId));
|
|
if (token.providerMetadata) await db.update(messageParts)
|
|
.set({
|
|
providerOptions: token.providerMetadata,
|
|
})
|
|
.where(eq(messageParts.toolCallId, dbToolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
toolName: token.toolName,
|
|
input: {
|
|
type: inputType,
|
|
value: inputValue,
|
|
},
|
|
}
|
|
})
|
|
} else {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.toolCallId;
|
|
const dbToolCallId = formatToolCallId(toolCallId);
|
|
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: dbToolCallId,
|
|
userId: userId,
|
|
toolName: token.toolName,
|
|
status: 'pending',
|
|
input: {
|
|
type: inputType,
|
|
value: inputValue,
|
|
},
|
|
output: null,
|
|
error: null,
|
|
createdAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!toolCall) {
|
|
throw new Error('Failed to insert tool call');
|
|
}
|
|
|
|
const [part] = await db.insert(messageParts).values({
|
|
id: formatPartId('tool-call', dbToolCallId),
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
providerOptions: token.providerMetadata,
|
|
type: 'tool-call',
|
|
content: null,
|
|
finished: false,
|
|
createdAt: new Date(),
|
|
lastUpdatedAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!part) {
|
|
throw new Error('Failed to insert message part');
|
|
}
|
|
|
|
// @ts-ignore - doesnt exist on the type but yeah it does now
|
|
part.toolCall = toolCall;
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeToolCalls.add(dbToolCallId);
|
|
|
|
activeParts.set(key, { id: part.id, accumulatedContent: '' });
|
|
}
|
|
} break;
|
|
case 'tool-result': {
|
|
let outputType: ToolCallType = ToolCallType.Text;
|
|
let outputValue: string = '';
|
|
const dbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
|
|
|
|
switch (typeof token.output) {
|
|
case 'string':
|
|
outputType = ToolCallType.Text;
|
|
outputValue = token.output;
|
|
break;
|
|
case 'object':
|
|
outputType = ToolCallType.Json;
|
|
outputValue = JSON.stringify(token.output);
|
|
break;
|
|
default:
|
|
console.error('Unknown output type', token.output);
|
|
if (dbToolCallId) {
|
|
await db.update(toolCalls).set({
|
|
status: 'failed',
|
|
error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
|
|
}).where(eq(toolCalls.id, dbToolCallId));
|
|
if (token.providerMetadata) await db.update(messageParts)
|
|
.set({
|
|
providerOptions: token.providerMetadata,
|
|
})
|
|
.where(eq(messageParts.toolCallId, dbToolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
toolName: token.toolName,
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
status: 'failed',
|
|
}
|
|
})
|
|
|
|
activeToolCalls.delete(dbToolCallId);
|
|
}
|
|
break;
|
|
}
|
|
|
|
if (dbToolCallId) {
|
|
await db.update(toolCalls).set({
|
|
status: 'completed',
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
}).where(eq(toolCalls.id, dbToolCallId));
|
|
if (token.providerMetadata) await db.update(messageParts)
|
|
.set({
|
|
providerOptions: token.providerMetadata,
|
|
})
|
|
.where(eq(messageParts.toolCallId, dbToolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
toolName: token.toolName,
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
status: 'completed',
|
|
}
|
|
})
|
|
|
|
activeToolCalls.delete(dbToolCallId);
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'tool-error': {
|
|
console.error('Tool error:', token);
|
|
|
|
let outputType: ToolCallType;
|
|
let outputValue: string;
|
|
|
|
switch (typeof token.error) {
|
|
case 'string':
|
|
outputType = ToolCallType.Text;
|
|
outputValue = token.error;
|
|
break;
|
|
case 'object':
|
|
outputType = ToolCallType.Json;
|
|
outputValue = JSON.stringify(token.error);
|
|
break;
|
|
default:
|
|
console.error('Unknown error type', token.error);
|
|
outputType = ToolCallType.Text;
|
|
outputValue = 'Tool returned invalid output';
|
|
break;
|
|
}
|
|
|
|
const existingDbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
|
|
if (existingDbToolCallId && activeToolCalls.has(existingDbToolCallId)) {
|
|
await db.update(toolCalls).set({
|
|
status: 'failed',
|
|
error: {
|
|
type: outputType as ToolCallType,
|
|
value: outputValue as string,
|
|
}
|
|
}).where(eq(toolCalls.id, existingDbToolCallId));
|
|
if (token.providerMetadata) await db.update(messageParts)
|
|
.set({
|
|
providerOptions: token.providerMetadata,
|
|
})
|
|
.where(eq(messageParts.toolCallId, existingDbToolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: existingDbToolCallId,
|
|
toolName: token.toolName,
|
|
error: {
|
|
type: outputType as ToolCallType,
|
|
value: outputValue as string,
|
|
},
|
|
}
|
|
})
|
|
} else {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.toolCallId;
|
|
const dbToolCallId = formatToolCallId(toolCallId);
|
|
nativeToDbToolCallId.set(toolCallId, dbToolCallId);
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: dbToolCallId,
|
|
userId: userId,
|
|
toolName: token.toolName,
|
|
status: 'failed',
|
|
input: null,
|
|
output: null,
|
|
error: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
createdAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!toolCall) {
|
|
throw new Error('Failed to insert tool call');
|
|
}
|
|
|
|
const [part] = await db.insert(messageParts).values({
|
|
id: formatPartId('tool-call', dbToolCallId),
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId: dbToolCallId,
|
|
providerOptions: token.providerMetadata,
|
|
type: 'tool-call',
|
|
content: null,
|
|
finished: false,
|
|
createdAt: new Date(),
|
|
lastUpdatedAt: new Date(),
|
|
}).returning();
|
|
|
|
if (!part) {
|
|
throw new Error('Failed to insert message part');
|
|
}
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeParts.set(key, { id: part.id, accumulatedContent: '' });
|
|
}
|
|
|
|
if (existingDbToolCallId) {
|
|
activeToolCalls.delete(existingDbToolCallId);
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'finish': {
|
|
let tps;
|
|
if (requestStart !== undefined && token.totalUsage.outputTokens !== undefined) {
|
|
const requestDuration = performance.now() - requestStart;
|
|
tps = token.totalUsage.outputTokens / (requestDuration / 1000);
|
|
}
|
|
|
|
await Promise.all([...activeParts.values()].map(async part => {
|
|
await db.update(messageParts)
|
|
.set({ finished: true, lastUpdatedAt: new Date() })
|
|
.where(eq(messageParts.id, part.id))
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-end',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
content: part.accumulatedContent,
|
|
}
|
|
})
|
|
}));
|
|
|
|
|
|
switch (token.finishReason) {
|
|
case 'error':
|
|
await db.update(generations).set({
|
|
status: 'failed',
|
|
error: INTERNAL_ERROR,
|
|
}).where(eq(generations.id, generationId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'generation-failed',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
error: INTERNAL_ERROR,
|
|
}
|
|
})
|
|
break;
|
|
case 'content-filter':
|
|
await db.update(generations).set({
|
|
status: 'failed',
|
|
error: 'Content was filtered',
|
|
}).where(eq(generations.id, generationId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'generation-failed',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
error: 'Content was filtered',
|
|
}
|
|
})
|
|
break;
|
|
}
|
|
|
|
await db.update(generations).set({
|
|
status: 'completed',
|
|
tokens: {
|
|
input: token.totalUsage.inputTokens,
|
|
cache: {
|
|
read: token.totalUsage.inputTokenDetails.cacheReadTokens,
|
|
write: token.totalUsage.inputTokenDetails.cacheWriteTokens
|
|
},
|
|
output: token.totalUsage.outputTokens,
|
|
thinking: token.totalUsage.outputTokenDetails.reasoningTokens,
|
|
ttft,
|
|
tps,
|
|
},
|
|
}).where(eq(generations.id, generationId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'generation-complete',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
tokens: {
|
|
input: token.totalUsage.inputTokens,
|
|
cache: {
|
|
read: token.totalUsage.inputTokenDetails.cacheReadTokens,
|
|
write: token.totalUsage.inputTokenDetails.cacheWriteTokens
|
|
},
|
|
output: token.totalUsage.outputTokens,
|
|
thinking: token.totalUsage.outputTokenDetails.reasoningTokens,
|
|
ttft,
|
|
tps,
|
|
},
|
|
}
|
|
})
|
|
|
|
// I hate you switch fallthroughs
|
|
break;
|
|
}
|
|
|
|
case 'abort': {
|
|
await Promise.all([
|
|
...[...activeParts.values()].map(async part => {
|
|
await db.update(messageParts)
|
|
.set({ finished: true, lastUpdatedAt: new Date() })
|
|
.where(eq(messageParts.id, part.id))
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'text-end',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
content: part.accumulatedContent,
|
|
}
|
|
})
|
|
}),
|
|
...[...activeToolCalls].map(async id => {
|
|
await db.update(toolCalls)
|
|
.set({ status: 'cancelled' })
|
|
.where(eq(toolCalls.id, id))
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-cancel',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: id,
|
|
}
|
|
})
|
|
}),
|
|
db.update(generations)
|
|
.set({ status: 'cancelled' })
|
|
.where(eq(generations.id, generationId)),
|
|
topicEvents.emit(topicId, {
|
|
type: 'generation-cancelled',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
}
|
|
})
|
|
]);
|
|
} break;
|
|
|
|
case 'file':
|
|
todo('file token type', token);
|
|
break;
|
|
case 'raw':
|
|
todo('raw token type', token);
|
|
break;
|
|
case 'source':
|
|
todo('source token type', token);
|
|
break;
|
|
case 'tool-approval-request':
|
|
todo('tool-approval-request token type', token);
|
|
break;
|
|
// typescript thinks this is not a real token type?
|
|
// case 'tool-output-denied':
|
|
// todo('tool-output-denied token type', token);
|
|
// break;
|
|
case 'error':
|
|
case 'finish-step':
|
|
case 'tool-input-delta':
|
|
case 'tool-input-end':
|
|
// handled or irrelevant
|
|
break;
|
|
}
|
|
}
|
|
} catch (error: any) {
|
|
console.error(error);
|
|
|
|
await db.update(generations)
|
|
.set({ status: 'failed', error: String(error) })
|
|
.where(eq(generations.id, generationId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'generation-failed',
|
|
payload: {
|
|
messageId: message.id,
|
|
generationId,
|
|
error: String(error),
|
|
}
|
|
})
|
|
} finally {
|
|
completeGeneration(generationId);
|
|
if (logFile !== undefined) logFile.close();
|
|
}
|
|
} |