327 lines
14 KiB
TypeScript
327 lines
14 KiB
TypeScript
import type { AssistantContent, FilePart, ImagePart, JSONValue, ModelMessage, ToolContent } from "ai";
|
|
import { ToolCallType } from "../drizzle/schema";
|
|
import { Err, Ok, type Result } from "../types/result";
|
|
import type { Message, MessageEntity } from "../app/composables/useChat";
|
|
import type { Agent } from "../app/composables/useAgents";
|
|
import { resolveFileUrl } from "./url";
|
|
|
|
export const getActiveMessage = <T extends { id: string; deleted?: boolean | null; activeChildId?: string | null; children: (T | undefined)[] }>(
|
|
message: T
|
|
): T | null => {
|
|
if (message.activeChildId && message.children.length > 0) {
|
|
const child = message.children.find(c => c?.id === message.activeChildId);
|
|
if (child) {
|
|
return child;
|
|
}
|
|
}
|
|
|
|
if (message.deleted === true) {
|
|
if (message.children.length > 0 && message.children[0]) {
|
|
return message.children[0];
|
|
}
|
|
return null;
|
|
}
|
|
|
|
return message;
|
|
};
|
|
|
|
export const getMessageCount = (message: {
|
|
deleted?: boolean | null;
|
|
children: unknown[];
|
|
}): number => {
|
|
if (message.deleted === true) {
|
|
return message.children.length;
|
|
}
|
|
|
|
return message.children.length + 1;
|
|
};
|
|
|
|
export const getCurrentChildIndex = (message: {
|
|
deleted?: boolean | null;
|
|
activeChildId?: string | null;
|
|
children: ({ id: string } | undefined)[];
|
|
}): number => {
|
|
if (message.deleted === true) {
|
|
if (!message.activeChildId) {
|
|
return 0;
|
|
}
|
|
const idx = message.children.findIndex(c => c?.id === message.activeChildId);
|
|
return idx >= 0 ? idx : 0;
|
|
}
|
|
|
|
if (!message.activeChildId) {
|
|
return 0;
|
|
}
|
|
const idx = message.children.findIndex(c => c?.id === message.activeChildId);
|
|
return idx >= 0 ? idx + 1 : 0;
|
|
};
|
|
|
|
export const getNextActiveChildId = (
|
|
message: {
|
|
deleted?: boolean | null;
|
|
activeChildId?: string | null;
|
|
children: ({ id: string } | undefined)[];
|
|
},
|
|
direction: -1 | 1
|
|
): string | null => {
|
|
const allItems = message.deleted === true
|
|
? message.children
|
|
: [null, ...message.children];
|
|
|
|
const currentIdx = allItems.findIndex(item => {
|
|
if (!item) {
|
|
return !message.activeChildId;
|
|
}
|
|
return item.id === message.activeChildId;
|
|
});
|
|
|
|
const validCurrentIdx = currentIdx >= 0 ? currentIdx : 0;
|
|
const newIdx = Math.max(0, Math.min(allItems.length - 1, validCurrentIdx + direction));
|
|
const newItem = allItems[newIdx];
|
|
|
|
return newItem?.id ?? null;
|
|
};
|
|
|
|
export const buildMessageTree = (flatMessages: MessageEntity[]) => {
|
|
const messagesMap = new Map(flatMessages.map(m => [m.id, { ...m, children: [] as MessageEntity[] }]));
|
|
const roots = [] as Message[];
|
|
|
|
for (const msg of messagesMap.values()) {
|
|
if (msg.parentMessageId && messagesMap.has(msg.parentMessageId)) {
|
|
messagesMap.get(msg.parentMessageId)!.children.push(msg);
|
|
} else {
|
|
roots.push(msg);
|
|
}
|
|
}
|
|
return roots.filter(root => !root.deleted || root.children.length > 0);
|
|
};
|
|
|
|
export const buildFocusedMessageTree = (messages: Readonly<Message[]>): MessageEntity[] => {
|
|
const focusedMessageTree: MessageEntity[] = [];
|
|
|
|
for (const message of messages) {
|
|
const activeMessage = getActiveMessage(message);
|
|
if (activeMessage) {
|
|
focusedMessageTree.push(activeMessage);
|
|
}
|
|
}
|
|
|
|
return focusedMessageTree;
|
|
};
|
|
|
|
export interface MarshallOptions {
|
|
signFileUrl?: (fileKey: string) => string;
|
|
}
|
|
|
|
export const marshallMessages = (agent: Agent, messages: Readonly<MessageEntity[]>, opts?: MarshallOptions): Result<ModelMessage[], string> => {
|
|
const marshalledMessages: ModelMessage[] = [];
|
|
|
|
if (agent && agent.systemPrompt) {
|
|
marshalledMessages.push({
|
|
role: 'system',
|
|
content: agent.systemPrompt,
|
|
});
|
|
}
|
|
|
|
for (const message of messages) {
|
|
switch (message.role) {
|
|
case 'user': {
|
|
const attachments = message.attachments.map(attachment => {
|
|
let url = resolveFileUrl(attachment.file.url);
|
|
|
|
if (opts?.signFileUrl && attachment.file.url.startsWith('/api/files/')) {
|
|
const fileKey = attachment.file.url.slice('/api/files/'.length);
|
|
const token = opts.signFileUrl(fileKey);
|
|
url = `${url}?${token}`;
|
|
}
|
|
|
|
if (attachment.file.mimeType.startsWith('image/')) {
|
|
return {
|
|
type: 'image',
|
|
image: new URL(url),
|
|
};
|
|
}
|
|
|
|
return {
|
|
type: 'file',
|
|
data: new URL(url),
|
|
filename: attachment.file.name,
|
|
mediaType: attachment.file.mimeType,
|
|
};
|
|
}) as (FilePart | ImagePart)[];
|
|
|
|
let messageDate = new Date(message.createdAt);
|
|
const prompt = `[${messageDate.toDateString()} ${messageDate.toLocaleTimeString()}]: ${message.content!}`;
|
|
|
|
marshalledMessages.push({
|
|
role: 'user',
|
|
content: attachments ? [
|
|
{
|
|
type: 'text',
|
|
text: `[${messageDate.toDateString()} ${messageDate.toLocaleTimeString()}]: ${message.content!}`
|
|
},
|
|
...attachments,
|
|
] : prompt,
|
|
});
|
|
break;
|
|
}
|
|
case 'assistant':
|
|
let assistantPart: AssistantContent = [];
|
|
let toolParts: ToolContent = [];
|
|
|
|
for (const part of (message.parts || [])) {
|
|
if (!part) return Err('Part is undefined');
|
|
|
|
switch (part.type) {
|
|
case 'text':
|
|
case 'reasoning': {
|
|
if (toolParts.length > 0) {
|
|
marshalledMessages.push({
|
|
role: 'assistant',
|
|
content: assistantPart,
|
|
});
|
|
marshalledMessages.push({
|
|
role: 'tool',
|
|
content: toolParts,
|
|
});
|
|
toolParts = [];
|
|
assistantPart = [];
|
|
}
|
|
|
|
if (part.providerOptions || part.content) assistantPart.push({
|
|
type: part.type,
|
|
text: part.content || '',
|
|
providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
|
|
})
|
|
break;
|
|
}
|
|
case 'tool-call': {
|
|
if (part.toolCall === null) return Err('Tool call is null');
|
|
|
|
if (part.toolCall.status === 'pending') {
|
|
return Err('Marshalling tool call that is still pending. This is likely a UI bug if this happens.');
|
|
}
|
|
|
|
let inputValue: string | object = '';
|
|
|
|
switch (part.toolCall.input!.type) {
|
|
case ToolCallType.Text:
|
|
inputValue = part.toolCall.input!.value;
|
|
break;
|
|
case ToolCallType.Json:
|
|
if (typeof part.toolCall.input!.value === 'string') {
|
|
inputValue = JSON.parse(part.toolCall.input!.value);
|
|
} else {
|
|
inputValue = part.toolCall.input!.value;
|
|
}
|
|
break;
|
|
}
|
|
|
|
assistantPart.push({
|
|
type: 'tool-call',
|
|
toolCallId: part.toolCall.id!,
|
|
toolName: part.toolCall.toolName!,
|
|
input: inputValue,
|
|
providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
|
|
})
|
|
|
|
if (part.toolCall.status === 'failed') {
|
|
let failureType: 'error-text' | 'error-json';
|
|
let failureValue: string | JSONValue;
|
|
|
|
if (part.toolCall.error === null || part.toolCall.error === undefined) {
|
|
failureType = 'error-text';
|
|
failureValue = 'An unknown error occurred';
|
|
} else {
|
|
switch (part.toolCall.error!.type) {
|
|
case ToolCallType.Text:
|
|
failureType = 'error-text';
|
|
failureValue = part.toolCall.error!.value;
|
|
break;
|
|
case ToolCallType.Json:
|
|
failureType = 'error-json';
|
|
if (typeof part.toolCall.error!.value === 'string') {
|
|
failureValue = JSON.parse(part.toolCall.error!.value);
|
|
} else {
|
|
failureValue = part.toolCall.error!.value;
|
|
}
|
|
break;
|
|
}
|
|
|
|
failureType = 'error-json';
|
|
// failureValue = JSON.stringify(part.toolCall.error!.value);
|
|
if (typeof part.toolCall.error!.value === 'string') {
|
|
failureValue = JSON.parse(part.toolCall.error!.value);
|
|
} else {
|
|
failureValue = part.toolCall.error!.value;
|
|
}
|
|
}
|
|
|
|
toolParts.push({
|
|
type: 'tool-result',
|
|
toolCallId: part.toolCall.id!,
|
|
toolName: part.toolCall.toolName!,
|
|
// @ts-expect-error - This is a type error, because typescript cant provie that the value must be a string when the type is error-text
|
|
output: {
|
|
type: failureType,
|
|
value: failureValue,
|
|
},
|
|
providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
|
|
});
|
|
break;
|
|
}
|
|
|
|
if (part.toolCall.status === 'completed') {
|
|
let outputType: 'text' | 'json';
|
|
let outputValue: string;
|
|
|
|
switch (part.toolCall.output!.type) {
|
|
case ToolCallType.Text:
|
|
outputType = 'text';
|
|
outputValue = part.toolCall.output!.value;
|
|
break;
|
|
case ToolCallType.Json:
|
|
outputType = 'json';
|
|
if (typeof part.toolCall.output!.value === 'string') {
|
|
outputValue = JSON.parse(part.toolCall.output!.value);
|
|
} else {
|
|
outputValue = part.toolCall.output!.value;
|
|
}
|
|
break;
|
|
}
|
|
|
|
toolParts.push({
|
|
type: 'tool-result',
|
|
toolCallId: part.toolCall.id!,
|
|
toolName: part.toolCall.toolName!,
|
|
output: {
|
|
type: outputType,
|
|
value: outputValue,
|
|
},
|
|
providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
|
|
});
|
|
break;
|
|
}
|
|
} break;
|
|
default:
|
|
return Err(`Unknown part type: ${part.type}`);
|
|
}
|
|
}
|
|
|
|
if (assistantPart.length > 0) marshalledMessages.push({
|
|
role: 'assistant',
|
|
content: assistantPart,
|
|
});
|
|
if (toolParts.length > 0) marshalledMessages.push({
|
|
role: 'tool',
|
|
content: toolParts,
|
|
});
|
|
break;
|
|
default:
|
|
return Err(`Unknown message role: ${message.role}`);
|
|
}
|
|
}
|
|
|
|
return Ok(marshalledMessages);
|
|
};
|