message.ts•4.85 kB
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { extractJson } from './extract';
import type {
MessageData,
Part,
ToolRequestPart,
ToolResponsePart,
} from './model';
export type MessageParser<T = unknown> = (message: Message) => T;
/**
* Message represents a single role's contribution to a generation. Each message
* can contain multiple parts (for example text and an image), and each generation
* can contain multiple messages.
*/
export class Message<T = unknown> implements MessageData {
role: MessageData['role'];
content: Part[];
metadata?: Record<string, any>;
parser?: MessageParser<T>;
static parseData(
lenientMessage:
| string
| (MessageData & { content: string | Part | Part[]; role: string })
| MessageData,
defaultRole: MessageData['role'] = 'user'
): MessageData {
if (typeof lenientMessage === 'string') {
return { role: defaultRole, content: [{ text: lenientMessage }] };
}
return {
...lenientMessage,
content: Message.parseContent(lenientMessage.content),
};
}
static parse(
lenientMessage: string | (MessageData & { content: string }) | MessageData
): Message {
return new Message(Message.parseData(lenientMessage));
}
static parseContent(lenientPart: string | Part | (string | Part)[]): Part[] {
if (typeof lenientPart === 'string') {
return [{ text: lenientPart }];
} else if (Array.isArray(lenientPart)) {
return lenientPart.map((p) => (typeof p === 'string' ? { text: p } : p));
} else {
return [lenientPart];
}
}
constructor(message: MessageData, options?: { parser?: MessageParser<T> }) {
this.role = message.role;
this.content = message.content;
this.metadata = message.metadata;
this.parser = options?.parser;
}
/**
* Attempts to parse the content of the message according to the supplied
* output parser. Without a parser, returns `data` contained in the message or
* tries to parse JSON from the text of the message.
*
* @returns The structured output contained in the message.
*/
get output(): T {
return this.parser?.(this) || this.data || extractJson<T>(this.text);
}
toolResponseParts(): ToolResponsePart[] {
const res = this.content.filter((part) => !!part.toolResponse);
return res as ToolResponsePart[];
}
/**
* Concatenates all `text` parts present in the message with no delimiter.
* @returns A string of all concatenated text parts.
*/
get text(): string {
return this.content.map((part) => part.text || '').join('');
}
/**
* Concatenates all `reasoning` parts present in the message with no delimiter.
* @returns A string of all concatenated reasoning parts.
*/
get reasoning(): string {
return this.content.map((part) => part.reasoning || '').join('');
}
/**
* Returns the first media part detected in the message. Useful for extracting
* (for example) an image from a generation expected to create one.
* @returns The first detected `media` part in the message.
*/
get media(): { url: string; contentType?: string } | null {
return this.content.find((part) => part.media)?.media || null;
}
/**
* Returns the first detected `data` part of a message.
* @returns The first `data` part detected in the message (if any).
*/
get data(): T | null {
return this.content.find((part) => part.data)?.data as T | null;
}
/**
* Returns all tool request found in this message.
* @returns Array of all tool request found in this message.
*/
get toolRequests(): ToolRequestPart[] {
return this.content.filter(
(part) => !!part.toolRequest
) as ToolRequestPart[];
}
/**
* Returns all tool requests annotated with interrupt metadata.
* @returns Array of all interrupt tool requests.
*/
get interrupts(): ToolRequestPart[] {
return this.toolRequests.filter((t) => !!t.metadata?.interrupt);
}
/**
* Converts the Message to a plain JS object.
* @returns Plain JS object representing the data contained in the message.
*/
toJSON(): MessageData {
const out: MessageData = {
role: this.role,
content: [...this.content],
};
if (this.metadata) out.metadata = this.metadata;
return out;
}
}