1263 lines
46 KiB
TypeScript
1263 lines
46 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 { glob } from 'glob';
|
|
import { type ModelMessage, streamText, type StreamTextTransform, tool } from "ai";
|
|
import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
|
|
import { topicEvents } from "~~/server/utils/events";
|
|
import { nanoid } from "nanoid";
|
|
import { spawn } from "child_process";
|
|
import { type Model } from "~/composables/useModels";
|
|
import { eq } from "drizzle-orm";
|
|
import { buildFocusedMessageTree, buildMessageTree, marshallMessages } from "~~/utils/message";
|
|
import path from "path";
|
|
|
|
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(),
|
|
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, 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',
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
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
|
|
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({ focusedIndex: (parentMessage.focusedIndex ?? 0) + 1 }).where(eq(messages.id, parentMessageId));
|
|
events.push({ type: 'MESSAGE_UPDATED', payload: { focusedIndex: (parentMessage.focusedIndex ?? 0) + 1 } });
|
|
|
|
// 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 topicMessages = marshallMessages(topic.agent, buildFocusedMessageTree(topicMessageTree));
|
|
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 },
|
|
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);
|
|
throw new Error('TODO');
|
|
};
|
|
|
|
const evalPython = async (code: string) => {
|
|
// This is the wrapper logic from above, minified or stored as a string
|
|
// Or you can save the wrapper script to a file and call that.
|
|
const wrapper = `
|
|
import ast
|
|
import sys
|
|
code = sys.stdin.read()
|
|
tree = ast.parse(code)
|
|
last_node = tree.body[-1] if tree.body else None
|
|
namespace = {}
|
|
if len(tree.body) > 1:
|
|
exec(compile(ast.Module(tree.body[:-1], []), "<ast>", "exec"), namespace)
|
|
if isinstance(last_node, ast.Expr):
|
|
res = eval(compile(ast.Expression(last_node.value), "<ast>", "eval"), namespace)
|
|
if res is not None: print(res)
|
|
elif last_node:
|
|
exec(compile(ast.Module([last_node], []), "<ast>", "exec"), namespace)
|
|
`.trim();
|
|
|
|
return new Promise<string>((resolve, reject) => {
|
|
const child = spawn('python3', ['-c', wrapper]);
|
|
|
|
let output = '';
|
|
let errorOutput = '';
|
|
|
|
child.stdout.on('data', (data) => {
|
|
output += data.toString();
|
|
});
|
|
|
|
child.stderr.on('data', (data) => {
|
|
errorOutput += data.toString();
|
|
});
|
|
|
|
child.on('close', (exitCode) => {
|
|
if (exitCode !== 0) {
|
|
reject(errorOutput || `Exit code ${exitCode}`);
|
|
} else {
|
|
resolve(output.trim());
|
|
}
|
|
});
|
|
|
|
// Send the agent's code to the wrapper via stdin
|
|
child.stdin.write(code);
|
|
child.stdin.end();
|
|
});
|
|
};
|
|
|
|
// TODO: obviously come up with a better way for the user to define their own tools
|
|
const { listDirectoryTool, globTool, readFileTool, readFilesTool, fetchUrlTool, pythonTool, bashTool } = {
|
|
listDirectoryTool: tool({
|
|
description: 'Lists the contents of a directory',
|
|
inputSchema: z.object({
|
|
path: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
files: z.array(z.object({ name: z.string(), type: z.string() })),
|
|
}),
|
|
execute: async ({ path }) => {
|
|
const rawFiles = await fs.readdir(path, { withFileTypes: true });
|
|
const files = rawFiles.map((file) => ({
|
|
name: file.name,
|
|
type: file.isFile() ? 'file' : 'directory',
|
|
}));
|
|
|
|
return {
|
|
files,
|
|
};
|
|
},
|
|
}),
|
|
globTool: tool({
|
|
description: 'Lists files matching a glob pattern',
|
|
inputSchema: z.object({
|
|
pattern: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
files: z.array(z.object({ path: z.string(), type: z.string() })),
|
|
}),
|
|
execute: async ({ pattern }) => {
|
|
const rawFiles = await glob(pattern, { withFileTypes: true });
|
|
const files = rawFiles.map((file) => ({
|
|
path: file.parentPath + '/' + file.name,
|
|
type: file.isFile() ? 'file' : 'directory',
|
|
}));
|
|
|
|
return {
|
|
files,
|
|
};
|
|
},
|
|
}),
|
|
readFileTool: tool({
|
|
description: 'Reads the contents of a file',
|
|
inputSchema: z.object({
|
|
path: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
path: z.string(),
|
|
content: z.string(),
|
|
}),
|
|
execute: async ({ path }) => {
|
|
const file = await fs.readFile(path);
|
|
return {
|
|
path,
|
|
content: file.toString(),
|
|
};
|
|
},
|
|
}),
|
|
readFilesTool: tool({
|
|
description: 'Reads the contents of multiple files',
|
|
inputSchema: z.object({
|
|
paths: z.array(z.string()).describe('a list of file paths to read'),
|
|
}),
|
|
outputSchema: z.object({
|
|
files: z.array(
|
|
z.object({
|
|
path: z.string(),
|
|
content: z.string(),
|
|
}),
|
|
),
|
|
}),
|
|
execute: async ({ paths }) => {
|
|
const files = await Promise.all(
|
|
paths.map(async (path) => {
|
|
const file = await fs.readFile(path);
|
|
return {
|
|
path: path,
|
|
content: file.toString(),
|
|
};
|
|
}),
|
|
);
|
|
|
|
return {
|
|
files,
|
|
};
|
|
},
|
|
}),
|
|
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(url);
|
|
const content = await response.text();
|
|
return {
|
|
content,
|
|
};
|
|
},
|
|
}),
|
|
pythonTool: tool({
|
|
description: 'Executes a Python code snippet',
|
|
inputSchema: z.object({
|
|
code: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
output: z.string(),
|
|
}),
|
|
execute: async ({ code }) => {
|
|
const output = await evalPython(code);
|
|
return {
|
|
output,
|
|
};
|
|
},
|
|
}),
|
|
bashTool: tool({
|
|
description: 'Executes a Bash command',
|
|
inputSchema: z.object({
|
|
code: z.string(),
|
|
}),
|
|
outputSchema: z.object({
|
|
output: z.string(),
|
|
}),
|
|
execute: async ({ code }) => {
|
|
const { exec } = await import('child_process');
|
|
const { promisify } = await import('util');
|
|
const execAsync = promisify(exec);
|
|
|
|
async function runCommand(command: string) {
|
|
const { stdout, stderr } = await execAsync(command);
|
|
if (stderr) {
|
|
console.error(`Error: ${stderr}`);
|
|
return stderr;
|
|
} else {
|
|
return stdout;
|
|
}
|
|
}
|
|
|
|
const output = await runCommand(code);
|
|
return {
|
|
output,
|
|
};
|
|
},
|
|
}),
|
|
}
|
|
|
|
async function generateResponse(
|
|
message: MessageEntity,
|
|
model: {
|
|
gateway: ModelGateway,
|
|
model: Model,
|
|
parameters?: Record<string, any>,
|
|
},
|
|
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>();
|
|
|
|
// TODO: somehow let the user turn on and off tools
|
|
const tools = {
|
|
// writeFile: tool({
|
|
// inputSchema: z.object({
|
|
// path: z.string(),
|
|
// content: z.string(),
|
|
// }),
|
|
// outputSchema: z.object({
|
|
// success: z.boolean(),
|
|
// }),
|
|
// execute: async ({ path, content }) => {
|
|
// await fs.writeFile(path, content);
|
|
// return {
|
|
// success: true,
|
|
// };
|
|
// }
|
|
// }),
|
|
listDirectory: listDirectoryTool,
|
|
glob: globTool,
|
|
readFile: readFileTool,
|
|
readFiles: readFilesTool,
|
|
// fetchUrl: fetchUrlTool,
|
|
python: pythonTool,
|
|
bash: bashTool,
|
|
};
|
|
|
|
console.log({ messages });
|
|
|
|
const response = streamText({
|
|
model: model.gateway(model.model.externalId),
|
|
messages,
|
|
providerOptions: {
|
|
openrouter: {
|
|
debug: {
|
|
echo_upstream_body: true,
|
|
},
|
|
user: userId,
|
|
}
|
|
},
|
|
experimental_transform: streamTransoforms,
|
|
// a little trick that makes it so that the stream doesnt stop because of tool calls, and will continue an unbounded amount of time and steps
|
|
stopWhen: [],
|
|
tools: model.model.capabilities.includes('tools') ? 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,
|
|
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 [part] = await db.insert(messageParts).values({
|
|
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,
|
|
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,
|
|
content: part.accumulatedContent,
|
|
}
|
|
})
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
case 'tool-input-start': {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.id;
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: toolCallId,
|
|
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({
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId,
|
|
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
|
|
part.toolCall = toolCall;
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
partId: part.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeToolCalls.add(toolCallId);
|
|
|
|
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;
|
|
}
|
|
|
|
if (activeToolCalls.has(token.toolCallId)) {
|
|
await db.update(toolCalls)
|
|
.set({
|
|
status: 'pending',
|
|
input: {
|
|
type: inputType,
|
|
value: inputValue,
|
|
},
|
|
})
|
|
.where(eq(toolCalls.id, token.toolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
toolName: token.toolName,
|
|
input: {
|
|
type: inputType,
|
|
value: inputValue,
|
|
},
|
|
}
|
|
})
|
|
} else {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.toolCallId;
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: toolCallId,
|
|
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({
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
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
|
|
part.toolCall = toolCall;
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-start',
|
|
payload: {
|
|
messageId: message.id,
|
|
part,
|
|
}
|
|
})
|
|
|
|
activeToolCalls.add(toolCallId);
|
|
|
|
activeParts.set(key, { id: part.id, accumulatedContent: '' });
|
|
}
|
|
} break;
|
|
case 'tool-result': {
|
|
let outputType: ToolCallType = ToolCallType.Text;
|
|
let outputValue: string = '';
|
|
|
|
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);
|
|
await db.update(toolCalls).set({
|
|
status: 'failed',
|
|
error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
|
|
}).where(eq(toolCalls.id, token.toolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
toolName: token.toolName,
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
status: 'failed',
|
|
}
|
|
})
|
|
|
|
activeToolCalls.delete(token.toolCallId);
|
|
break;
|
|
}
|
|
|
|
await db.update(toolCalls).set({
|
|
status: 'completed',
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
}).where(eq(toolCalls.id, token.toolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
toolName: token.toolName,
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
status: 'completed',
|
|
}
|
|
})
|
|
|
|
activeToolCalls.delete(token.toolCallId);
|
|
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;
|
|
}
|
|
|
|
if (activeToolCalls.has(token.toolCallId)) {
|
|
await db.update(toolCalls).set({
|
|
status: 'failed',
|
|
error: {
|
|
type: outputType as ToolCallType,
|
|
value: outputValue as string,
|
|
}
|
|
}).where(eq(toolCalls.id, token.toolCallId));
|
|
|
|
await topicEvents.emit(topicId, {
|
|
type: 'tool-call-delta',
|
|
payload: {
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
toolName: token.toolName,
|
|
error: {
|
|
type: outputType as ToolCallType,
|
|
value: outputValue as string,
|
|
}
|
|
}
|
|
})
|
|
} else {
|
|
key = `tool-call-${curStepIdx}`;
|
|
|
|
const toolCallId = token.toolCallId;
|
|
|
|
const [toolCall] = await db.insert(toolCalls).values({
|
|
id: toolCallId,
|
|
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({
|
|
userId,
|
|
topicId,
|
|
messageId: message.id,
|
|
toolCallId: token.toolCallId,
|
|
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: '' });
|
|
}
|
|
|
|
activeToolCalls.delete(token.toolCallId);
|
|
break;
|
|
}
|
|
|
|
case 'finish': {
|
|
let tps;
|
|
if (ttft !== undefined && token.totalUsage.outputTokens !== undefined) {
|
|
const tokenStreamStart = requestStart! + ttft;
|
|
// this is the *real* request duration, excluding the
|
|
// TTFT
|
|
const requestDuration = performance.now() - tokenStreamStart;
|
|
|
|
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();
|
|
}
|
|
} |