import { AppConfig, ChatOptions, GenerationConfig } from "./config";
import { ChatInterface, GenerateProgressCallback, InitProgressCallback, InitProgressReport } from "./types";
import { ChatCompletionRequestBase, ChatCompletionRequestStreaming, ChatCompletionRequestNonStreaming, ChatCompletion, ChatCompletionMessageParam, ChatCompletionChunk } from "./openai_api_protocols/index";
/**
 * Message kind used by worker
 */
type RequestKind = ("return" | "throw" | "reload" | "generate" | "runtimeStatsText" | "interruptGenerate" | "unload" | "resetChat" | "initProgressCallback" | "generateProgressCallback" | "getMaxStorageBufferBindingSize" | "getGPUVendor" | "forwardTokensAndSample" | "chatCompletionNonStreaming" | "chatCompletionStreamInit" | "chatCompletionStreamNextChunk" | "customRequest");
interface ReloadParams {
    localIdOrUrl: string;
    chatOpts?: ChatOptions;
    appConfig?: AppConfig;
}
interface GenerateParams {
    input: string | Array<ChatCompletionMessageParam>;
    streamInterval?: number;
    genConfig?: GenerationConfig;
}
interface ResetChatParams {
    keepStats: boolean;
}
interface GenerateProgressCallbackParams {
    step: number;
    currentMessage: string;
}
interface ForwardTokensAndSampleParams {
    inputIds: Array<number>;
    curPos: number;
    isPrefill: boolean;
}
interface ChatCompletionNonStreamingParams {
    request: ChatCompletionRequestNonStreaming;
}
interface ChatCompletionStreamInitParams {
    request: ChatCompletionRequestStreaming;
}
export interface CustomRequestParams {
    requestName: string;
    requestMessage: string;
}
type MessageContent = GenerateProgressCallbackParams | ReloadParams | GenerateParams | ResetChatParams | ForwardTokensAndSampleParams | ChatCompletionNonStreamingParams | ChatCompletionStreamInitParams | CustomRequestParams | InitProgressReport | string | null | number | ChatCompletion | ChatCompletionChunk | void;
/**
 * The message used in exchange between worker
 * and the main thread.
 */
export interface WorkerMessage {
    kind: RequestKind;
    uuid: string;
    content: MessageContent;
}
/**
 * Worker handler that can be used in a WebWorker
 *
 * @example
 *
 * // setup a chat worker handler that routes
 * // requests to the chat
 * const chat = new ChatModule();
 * cont handler = new ChatWorkerHandler(chat);
 * onmessage = handler.onmessage;
 */
export declare class ChatWorkerHandler {
    protected chat: ChatInterface;
    protected chatCompletionAsyncChunkGenerator?: AsyncGenerator<ChatCompletionChunk, void, void>;
    constructor(chat: ChatInterface);
    handleTask<T extends MessageContent>(uuid: string, task: () => Promise<T>): Promise<void>;
    onmessage(event: MessageEvent): void;
}
interface ChatWorker {
    onmessage: any;
    postMessage: (message: any) => void;
}
/**
 * A client of chat worker that exposes the chat interface
 *
 * @example
 *
 * const chat = new webllm.ChatWorkerClient(new Worker(
 *   new URL('./worker.ts', import.meta.url),
 *   {type: 'module'}
 * ));
 */
export declare class ChatWorkerClient implements ChatInterface {
    worker: ChatWorker;
    private initProgressCallback?;
    private generateCallbackRegistry;
    private pendingPromise;
    constructor(worker: any);
    setInitProgressCallback(initProgressCallback: InitProgressCallback): void;
    protected getPromise<T extends MessageContent>(msg: WorkerMessage): Promise<T>;
    reload(localIdOrUrl: string, chatOpts?: ChatOptions, appConfig?: AppConfig): Promise<void>;
    getMaxStorageBufferBindingSize(): Promise<number>;
    getGPUVendor(): Promise<string>;
    generate(input: string | Array<ChatCompletionMessageParam>, progressCallback?: GenerateProgressCallback, streamInterval?: number, genConfig?: GenerationConfig): Promise<string>;
    runtimeStatsText(): Promise<string>;
    interruptGenerate(): void;
    unload(): Promise<void>;
    resetChat(keepStats?: boolean): Promise<void>;
    forwardTokensAndSample(inputIds: Array<number>, curPos: number, isPrefill: boolean): Promise<number>;
    /**
     * Every time the generator is called, we post a message to the worker asking it to
     * decode one step, and we expect to receive a message of `ChatCompletionChunk` from
     * the worker which we yield. The last message is `void`, meaning the generator has nothing
     * to yield anymore.
     */
    chatCompletionAsyncChunkGenerator(): AsyncGenerator<ChatCompletionChunk, void, void>;
    chatCompletion(request: ChatCompletionRequestNonStreaming): Promise<ChatCompletion>;
    chatCompletion(request: ChatCompletionRequestStreaming): Promise<AsyncIterable<ChatCompletionChunk>>;
    chatCompletion(request: ChatCompletionRequestBase): Promise<AsyncIterable<ChatCompletionChunk> | ChatCompletion>;
    onmessage(event: any): void;
}
export {};
//# sourceMappingURL=web_worker.d.ts.map