import { BaseToolkit } from "@langchain/core/tools";
import { VectorStoreQATool } from "../../../tools/vectorstore.js";
import { ZeroShotAgent } from "../../mrkl/index.js";
import { VECTOR_PREFIX, VECTOR_ROUTER_PREFIX } from "./prompt.js";
import { SUFFIX } from "../../mrkl/prompt.js";
import { LLMChain } from "../../../chains/llm_chain.js";
import { AgentExecutor } from "../../executor.js";
/**
 * Class representing a toolkit for working with a single vector store. It
 * initializes the vector store QA tool based on the provided vector store
 * information and language model.
 * @example
 * ```typescript
 * const toolkit = new VectorStoreToolkit(
 *   {
 *     name: "state_of_union_address",
 *     description: "the most recent state of the Union address",
 *     vectorStore: new HNSWLib(),
 *   },
 *   new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 }),
 * );
 * const result = await toolkit.invoke({
 *   input:
 *     "What did biden say about Ketanji Brown Jackson in the state of the union address?",
 * });
 * console.log(`Got output ${result.output}`);
 * ```
 */
export class VectorStoreToolkit extends BaseToolkit {
    constructor(vectorStoreInfo, llm) {
        super();
        Object.defineProperty(this, "tools", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "llm", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        const description = VectorStoreQATool.getDescription(vectorStoreInfo.name, vectorStoreInfo.description);
        this.llm = llm;
        this.tools = [
            new VectorStoreQATool(vectorStoreInfo.name, description, {
                vectorStore: vectorStoreInfo.vectorStore,
                llm: this.llm,
            }),
        ];
    }
}
/**
 * Class representing a toolkit for working with multiple vector stores.
 * It initializes multiple vector store QA tools based on the provided
 * vector store information and language model.
 */
export class VectorStoreRouterToolkit extends BaseToolkit {
    constructor(vectorStoreInfos, llm) {
        super();
        Object.defineProperty(this, "tools", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "vectorStoreInfos", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "llm", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.llm = llm;
        this.vectorStoreInfos = vectorStoreInfos;
        this.tools = vectorStoreInfos.map((vectorStoreInfo) => {
            const description = VectorStoreQATool.getDescription(vectorStoreInfo.name, vectorStoreInfo.description);
            return new VectorStoreQATool(vectorStoreInfo.name, description, {
                vectorStore: vectorStoreInfo.vectorStore,
                llm: this.llm,
            });
        });
    }
}
/** @deprecated Create a specific agent with a custom tool instead. */
export function createVectorStoreAgent(llm, toolkit, args) {
    const { prefix = VECTOR_PREFIX, suffix = SUFFIX, inputVariables = ["input", "agent_scratchpad"], } = args ?? {};
    const { tools } = toolkit;
    const prompt = ZeroShotAgent.createPrompt(tools, {
        prefix,
        suffix,
        inputVariables,
    });
    const chain = new LLMChain({ prompt, llm });
    const agent = new ZeroShotAgent({
        llmChain: chain,
        allowedTools: tools.map((t) => t.name),
    });
    return AgentExecutor.fromAgentAndTools({
        agent,
        tools,
        returnIntermediateSteps: true,
    });
}
/** @deprecated Create a specific agent with a custom tool instead. */
export function createVectorStoreRouterAgent(llm, toolkit, args) {
    const { prefix = VECTOR_ROUTER_PREFIX, suffix = SUFFIX, inputVariables = ["input", "agent_scratchpad"], } = args ?? {};
    const { tools } = toolkit;
    const prompt = ZeroShotAgent.createPrompt(tools, {
        prefix,
        suffix,
        inputVariables,
    });
    const chain = new LLMChain({ prompt, llm });
    const agent = new ZeroShotAgent({
        llmChain: chain,
        allowedTools: tools.map((t) => t.name),
    });
    return AgentExecutor.fromAgentAndTools({
        agent,
        tools,
        returnIntermediateSteps: true,
    });
}
