rocket-ai/dist/index.d.mts

Simple AI Client that lets you access different LLMs in a unified way.

import { z } from 'zod';

type AiImageSize = "256x256" | "512x512" | "1024x1024" | "1792x1024" | "1024x1792" | null;

type AiMessage = {
    role: string;
    content: string;
};

type AiMessageUsage = {
    inputTokens: number;
    outputTokens: number;
    totalTokens: number;
};

type AiMessageResponse = {
    content: string;
    usage: AiMessageUsage;
};

type AiMessageResponseImage = {
    url: string;
    revisedPrompt: string;
};

declare enum AiModelType {
    Gpt4o = "gpt-4o",
    Gpt4oMini = "gpt-4o-mini",
    DallE3 = "dall-e3",
    GptTextToSpeech = "tts-1",
    O1Preview = "o1-preview",
    O1Mini = "o1-mini",
    Claude37SonnetLatest = "claude-3-7-sonnet-latest",
    Claude35HaikuLatest = "claude-3-5-haiku-latest",
    Gemini20FlashLatest = "gemini-2.0-flash-latest",
    Gemini15FlashLatest = "gemini-1.5-flash-latest",
    Gemini15ProLatest = "gemini-1.5-pro-latest",
    Llama33 = "fw-llama-3-3",
    DeepSeekV3 = "fw-deepseek-v3",
    DeepSeekR1 = "fw-deepseek-r1"
}

type AiVoice = 'alloy' | 'ash' | 'coral' | 'echo' | 'fable' | 'onyx' | 'nova' | 'sage' | 'shimmer';

type GenerateImageOptions = {
    model: string;
    prompt: string;
    size: string;
    n: number;
};

type InvokeOptions = {
    model: string;
    messages: any;
    systemPrompt: string;
    temperature?: number;
};

interface LlmGenerateImage {
    generateImage(model: string, prompt: string, size: AiImageSize, n: number): Promise<AiMessageResponseImage>;
}

interface LlmGenerateSpeech {
    generateSpeech(model: string, prompt: string, voice: AiVoice): Promise<string>;
}

interface LlmInvoke {
    invoke(model: string, messages: AiMessage[], systemPrompt: string, temperature?: number): Promise<AiMessageResponse>;
}

interface LlmStream {
    stream(model: string, messages: AiMessage[], systemPrompt: string): AsyncGenerator<AiMessageResponse, void, unknown>;
}

type StreamOptions = {
    model: string;
    messages: any;
    systemPrompt: string;
};

type Usage = {
    inputTokens: number;
    outputTokens: number;
    totalTokens: number;
};

/**
 * AiClient class provides methods to interact with various AI model providers.
 * It supports invoking models, streaming responses, generating images, and generating speech.
 */
declare class AiClient {
    private clients?;
    private modelProviderMap?;
    private structuredOutputSchema?;
    private options;
    /**
     * Constructs an instance of AiClient.
     * @param clients - Optional custom clients for AI providers.
     * @param modelProviderMap - Optional custom mapping of models to providers.
     */
    constructor(clients?: any | undefined, modelProviderMap?: any | undefined);
    /**
     * Gets the provider name by model type.
     * @param model - The model type.
     * @returns The provider name.
     * @throws Will throw an error if no provider is found for the model.
     */
    _getProviderByModel(model: string): any;
    /**
     * Sets the structured output schema for the client.
     * @param schema - The Zod schema for structured output.
     * @returns The AiClient instance.
     */
    withStructuredOutput(schema: z.ZodSchema<any>): AiClient;
    /**
     * Sets options for invoking or streaming.
     * @param options - The options to set.
     * @returns The AiClient instance.
     */
    setOptions(options: Partial<InvokeOptions & StreamOptions>): AiClient;
    /**
     * Invokes an AI model with the given options.
     * @param options - The options for invoking the model.
     * @returns A promise that resolves to the AI message response.
     */
    invoke(options: InvokeOptions): Promise<AiMessageResponse>;
    /**
     * Streams responses from an AI model with the given options.
     * @param options - The options for streaming responses.
     * @returns A promise that resolves to an async generator of AI message responses.
     */
    stream(options: StreamOptions): Promise<AsyncGenerator<AiMessageResponse, void, unknown>>;
    /**
     * Prepares options by adding the structured output schema to the system prompt if available.
     * @param options - The options to prepare.
     * @returns The prepared options.
     */
    private prepareOptions;
    /**
     * Invokes or streams responses from an AI model based on the method specified.
     * @param options - The options for invoking or streaming.
     * @param method - The method to use ('invoke' or 'stream').
     * @returns A promise that resolves to the AI message response or an async generator of AI message responses.
     * @throws Will throw an error if the provider is not configured.
     */
    private _invokeOrStream;
    /**
     * Generates an image using the specified options.
     * @param options - The options for generating the image.
     * @returns A promise that resolves to the AI message response image.
     * @throws Will throw an error if the provider is not configured.
     */
    generateImage(options: GenerateImageOptions): Promise<AiMessageResponseImage>;
    /**
     * Generates speech using the specified model, messages, and voice.
     * @param model - The model type.
     * @param messages - The messages to convert to speech.
     * @param voice - The voice to use for speech generation.
     * @returns A promise that resolves to the generated speech as a string.
     * @throws Will throw an error if the provider is not configured.
     */
    generateSpeech(model: string, messages: string, voice: AiVoice): Promise<string>;
}

/**
 * The Agent class is responsible for managing the interaction with the AI client,
 * handling tool registration, and executing tasks based on user input.
 */
declare class Agent {
    private client;
    private toolRegistry;
    private promptTemplate;
    private usageCounter;
    /**
     * Constructs an Agent instance.
     * @param {string} name - The name of the agent.
     * @param {string} [systemPrompt] - An optional system prompt to initialize the agent with.
     */
    constructor(name: string, systemPrompt?: string);
    /**
     * Registers a list of tools with the agent.
     * @param {Array<any>} tools - An array of tools to register.
     */
    registerTools(tools: Array<any>): void;
    /**
     * Executes a task based on the provided input.
     * @param {string} input - The input string for the task.
     * @returns {Promise<AiMessageResponse>} - The final response from the agent.
     */
    executeTask(input: string): Promise<AiMessageResponse>;
}

declare class Tool {
    name: string;
    description: string;
    queryFormat: z.ZodType<any>;
    func: Function;
    constructor(func: Function, options: {
        name: string;
        description: string;
        queryFormat: z.ZodType<any>;
    });
    execute(...args: any[]): Promise<any>;
    getQueryFormat(): string;
}

declare class ToolRegistry {
    private tools;
    registerTools(tools: Tool[]): void;
    callTool(toolName: string, ...args: any[]): Promise<any>;
    getToolInfo(toolName: string): {
        name: string;
        description: string;
        queryFormat: any;
    };
    getAllTools(): Record<string, {
        name: string;
        description: string;
        queryFormat: any;
    }>;
}

declare function tool(func: Function, options: {
    name: string;
    description: string;
    queryFormat: z.ZodType<any>;
}): Tool;

declare class StructuredOutputTemplate {
    static getFormatInstructions(schema: z.ZodSchema<any>): string;
}

declare class BasicPromptTemplate {
    private template;
    constructor(name: string);
    addTemplateSection(name: string, content: string): void;
    getTemplate(): string;
}

export { Agent, AiClient, type AiImageSize, type AiMessage, type AiMessageResponse, type AiMessageResponseImage, type AiMessageUsage, AiModelType, type AiVoice, BasicPromptTemplate, type GenerateImageOptions, type InvokeOptions, type LlmGenerateImage, type LlmGenerateSpeech, type LlmInvoke, type LlmStream, type StreamOptions, StructuredOutputTemplate, Tool, ToolRegistry, type Usage, tool };