import agui from '@ag-ui/core';

interface ImageData {
    format: string;
    bytes: string;
}
type Role = agui.Role;
type SystemMessage = agui.SystemMessage;
type DeveloperMessage = agui.DeveloperMessage;
type ToolCall = agui.ToolCall;
type ToolResult = agui.ToolMessage & {
    toolName?: string;
};
type AIMessage = agui.AssistantMessage & {
    generativeUI?: (props?: any) => any;
    agentName?: string;
    state?: any;
    image?: ImageData;
};
type UserMessage = agui.UserMessage & {
    image?: ImageData;
};
type Message = AIMessage | ToolResult | UserMessage | SystemMessage | DeveloperMessage;

export { AIMessage, DeveloperMessage, ImageData, Message, Role, SystemMessage, ToolCall, ToolResult, UserMessage };