import { PromptMixin, PromptTemplate, ModuleRecord } from '@llamaindex/core/prompts';
import { QueryType, QueryBundle } from '@llamaindex/core/query-engine';
import { ToolMetadata, LLM } from '@llamaindex/core/llms';
import { BaseOutputParser } from '@llamaindex/core/schema';

/**
 * Top level types to avoid circular dependencies
 */

/**
 * StructuredOutput is just a combo of the raw output and the parsed output.
 */
interface StructuredOutput<T> {
    rawOutput: string;
    parsedOutput: T;
}
type ToolMetadataOnlyDescription = Pick<ToolMetadata, "description">;

interface SingleSelection {
    index: number;
    reason: string;
}
type SelectorResult = {
    selections: SingleSelection[];
};
type MetadataType = string | ToolMetadataOnlyDescription;
declare abstract class BaseSelector extends PromptMixin {
    select(choices: MetadataType[], query: QueryType): Promise<SelectorResult>;
    abstract _select(choices: ToolMetadataOnlyDescription[], query: QueryType): Promise<SelectorResult>;
}

type Answer = {
    choice: number;
    reason: string;
};

type SingleSelectPrompt = PromptTemplate<[
    "context",
    "query",
    "numChoices"
]>;
type MultiSelectPrompt = PromptTemplate<[
    "contextList",
    "query",
    "maxOutputs",
    "numChoices"
]>;

/**
 * A selector that uses the LLM to select a single or multiple choices from a list of choices.
 */
declare class LLMMultiSelector extends BaseSelector {
    llm: LLM;
    prompt: MultiSelectPrompt;
    maxOutputs: number;
    outputParser: BaseOutputParser<StructuredOutput<Answer[]>>;
    constructor(init: {
        llm: LLM;
        prompt?: MultiSelectPrompt;
        maxOutputs?: number;
        outputParser?: BaseOutputParser<StructuredOutput<Answer[]>>;
    });
    _getPrompts(): {
        prompt: MultiSelectPrompt;
    };
    _updatePrompts(prompts: {
        prompt: MultiSelectPrompt;
    }): void;
    protected _getPromptModules(): ModuleRecord;
    /**
     * Selects a single choice from a list of choices.
     * @param choices
     * @param query
     */
    _select(choices: ToolMetadataOnlyDescription[], query: QueryBundle): Promise<SelectorResult>;
    asQueryComponent(): unknown;
}
/**
 * A selector that uses the LLM to select a single choice from a list of choices.
 */
declare class LLMSingleSelector extends BaseSelector {
    llm: LLM;
    prompt: SingleSelectPrompt;
    outputParser: BaseOutputParser<StructuredOutput<Answer[]>>;
    constructor(init: {
        llm: LLM;
        prompt?: SingleSelectPrompt;
        outputParser?: BaseOutputParser<StructuredOutput<Answer[]>>;
    });
    _getPrompts(): Record<string, SingleSelectPrompt>;
    _updatePrompts(prompts: Record<string, SingleSelectPrompt>): void;
    /**
     * Selects a single choice from a list of choices.
     * @param choices
     * @param query
     */
    _select(choices: ToolMetadataOnlyDescription[], query: QueryBundle): Promise<SelectorResult>;
    asQueryComponent(): unknown;
    protected _getPromptModules(): {};
}

declare const getSelectorFromContext: (isMulti?: boolean) => BaseSelector;

export { BaseSelector, LLMMultiSelector, LLMSingleSelector, getSelectorFromContext };
export type { SelectorResult, SingleSelection };