UNPKG

langchain

Version:
131 lines (130 loc) 4.59 kB
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({ 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, }); }