import { Bedrock, FoundationModelSummary } from "@aws-sdk/client-bedrock";
import { BedrockRuntime, ConverseRequest, ConverseResponse, ConverseStreamOutput } from "@aws-sdk/client-bedrock-runtime";
import { AwsCredentialIdentity, Provider } from "@aws-sdk/types";
import { AbstractDriver, AIModel, Completion, CompletionChunkObject, DataSource, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, ImageGeneration, PromptOptions, PromptSegment, TrainingJob, TrainingOptions } from "@llumiverse/core";
import { NovaMessagesPrompt } from "@llumiverse/core/formatters";
export interface BedrockModelCapabilities {
    name: string;
    canStream: boolean;
}
export interface BedrockDriverOptions extends DriverOptions {
    /**
     * The AWS region
     */
    region: string;
    /**
     * The bucket name to be used for training.
     * It will be created if does not already exist.
     */
    training_bucket?: string;
    /**
     * The role ARN to be used for training
     */
    training_role_arn?: string;
    /**
     * The credentials to use to access AWS
     */
    credentials?: AwsCredentialIdentity | Provider<AwsCredentialIdentity>;
}
export type BedrockPrompt = NovaMessagesPrompt | ConverseRequest;
export declare class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockPrompt> {
    static PROVIDER: string;
    provider: string;
    private _executor?;
    private _service?;
    private _service_region?;
    constructor(options: BedrockDriverOptions);
    getExecutor(): BedrockRuntime;
    getService(region?: string): Bedrock;
    protected formatPrompt(segments: PromptSegment[], opts: PromptOptions): Promise<BedrockPrompt>;
    static getExtractedExecution(result: ConverseResponse, _prompt?: BedrockPrompt): CompletionChunkObject;
    static getExtractedStream(result: ConverseStreamOutput, _prompt?: BedrockPrompt): CompletionChunkObject;
    requestTextCompletion(prompt: ConverseRequest, options: ExecutionOptions): Promise<Completion>;
    extractRegion(modelString: string, defaultRegion: string): string;
    private getCanStream;
    protected canStream(options: ExecutionOptions): Promise<boolean>;
    requestTextCompletionStream(prompt: ConverseRequest, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>>;
    preparePayload(prompt: ConverseRequest, options: ExecutionOptions): ConverseRequest;
    requestImageGeneration(prompt: NovaMessagesPrompt, options: ExecutionOptions): Promise<Completion<ImageGeneration>>;
    startTraining(dataset: DataSource, options: TrainingOptions): Promise<TrainingJob>;
    cancelTraining(jobId: string): Promise<TrainingJob>;
    getTrainingJob(jobId: string): Promise<TrainingJob>;
    validateConnection(): Promise<boolean>;
    listTrainableModels(): Promise<AIModel<string>[]>;
    listModels(): Promise<AIModel[]>;
    _listModels(foundationFilter?: (m: FoundationModelSummary) => boolean): Promise<AIModel[]>;
    generateEmbeddings({ text, image, model }: EmbeddingsOptions): Promise<EmbeddingsResult>;
}
//# sourceMappingURL=index.d.ts.map