@langchain/community
Version:
Third-party integrations for LangChain.js
1 lines • 11.6 kB
Source Map (JSON)
{"version":3,"file":"sagemaker_endpoint.cjs","names":["LLM","fields: SageMakerEndpointInput","SageMakerRuntimeClient","prompt: string","options: this[\"ParsedCallOptions\"]","runManager?: CallbackManagerForLLMRun","InvokeEndpointCommand","InvokeEndpointWithResponseStreamCommand","GenerationChunk"],"sources":["../../src/llms/sagemaker_endpoint.ts"],"sourcesContent":["import {\n InvokeEndpointCommand,\n InvokeEndpointWithResponseStreamCommand,\n SageMakerRuntimeClient,\n SageMakerRuntimeClientConfig,\n} from \"@aws-sdk/client-sagemaker-runtime\";\nimport { CallbackManagerForLLMRun } from \"@langchain/core/callbacks/manager\";\nimport { GenerationChunk } from \"@langchain/core/outputs\";\nimport {\n type BaseLLMCallOptions,\n type BaseLLMParams,\n LLM,\n} from \"@langchain/core/language_models/llms\";\n\n/**\n * A handler class to transform input from LLM to a format that SageMaker\n * endpoint expects. Similarily, the class also handles transforming output from\n * the SageMaker endpoint to a format that LLM class expects.\n *\n * Example:\n * ```\n * class ContentHandler implements ContentHandlerBase<string, string> {\n * contentType = \"application/json\"\n * accepts = \"application/json\"\n *\n * transformInput(prompt: string, modelKwargs: Record<string, unknown>) {\n * const inputString = JSON.stringify({\n * prompt,\n * ...modelKwargs\n * })\n * return Buffer.from(inputString)\n * }\n *\n * transformOutput(output: Uint8Array) {\n * const responseJson = JSON.parse(Buffer.from(output).toString(\"utf-8\"))\n * return responseJson[0].generated_text\n * }\n *\n * }\n * ```\n */\nexport abstract class BaseSageMakerContentHandler<InputType, OutputType> {\n contentType = \"text/plain\";\n\n accepts = \"text/plain\";\n\n /**\n * Transforms the prompt and model arguments into a specific format for sending to SageMaker.\n * @param {InputType} prompt The prompt to be transformed.\n * @param {Record<string, unknown>} modelKwargs Additional arguments.\n * @returns {Promise<Uint8Array>} A promise that resolves to the formatted data for sending.\n */\n abstract transformInput(\n prompt: InputType,\n modelKwargs: Record<string, unknown>\n ): Promise<Uint8Array>;\n\n /**\n * Transforms SageMaker output into a desired format.\n * @param {Uint8Array} output The raw output from SageMaker.\n * @returns {Promise<OutputType>} A promise that resolves to the transformed data.\n */\n abstract transformOutput(output: Uint8Array): Promise<OutputType>;\n}\n\nexport type SageMakerLLMContentHandler = BaseSageMakerContentHandler<\n string,\n string\n>;\n\n/**\n * The SageMakerEndpointInput interface defines the input parameters for\n * the SageMakerEndpoint class, which includes the endpoint name, client\n * options for the SageMaker client, the content handler, and optional\n * keyword arguments for the model and the endpoint.\n */\nexport interface SageMakerEndpointInput extends BaseLLMParams {\n /**\n * The name of the endpoint from the deployed SageMaker model. Must be unique\n * within an AWS Region.\n */\n endpointName: string;\n /**\n * Options passed to the SageMaker client.\n */\n clientOptions: SageMakerRuntimeClientConfig;\n /**\n * Key word arguments to pass to the model.\n */\n modelKwargs?: Record<string, unknown>;\n /**\n * Optional attributes passed to the InvokeEndpointCommand\n */\n endpointKwargs?: Record<string, unknown>;\n /**\n * The content handler class that provides an input and output transform\n * functions to handle formats between LLM and the endpoint.\n */\n contentHandler: SageMakerLLMContentHandler;\n streaming?: boolean;\n}\n\n/**\n * The SageMakerEndpoint class is used to interact with SageMaker\n * Inference Endpoint models. It uses the AWS client for authentication,\n * which automatically loads credentials.\n * If a specific credential profile is to be used, the name of the profile\n * from the ~/.aws/credentials file must be passed. The credentials or\n * roles used should have the required policies to access the SageMaker\n * endpoint.\n */\nexport class SageMakerEndpoint extends LLM<BaseLLMCallOptions> {\n lc_serializable = true;\n\n static lc_name() {\n return \"SageMakerEndpoint\";\n }\n\n get lc_secrets(): { [key: string]: string } | undefined {\n return {\n \"clientOptions.credentials.accessKeyId\": \"AWS_ACCESS_KEY_ID\",\n \"clientOptions.credentials.secretAccessKey\": \"AWS_SECRET_ACCESS_KEY\",\n \"clientOptions.credentials.sessionToken\": \"AWS_SESSION_TOKEN\",\n };\n }\n\n endpointName: string;\n\n modelKwargs?: Record<string, unknown>;\n\n endpointKwargs?: Record<string, unknown>;\n\n client: SageMakerRuntimeClient;\n\n contentHandler: SageMakerLLMContentHandler;\n\n streaming: boolean;\n\n constructor(fields: SageMakerEndpointInput) {\n super(fields);\n\n if (!fields.clientOptions.region) {\n throw new Error(\n `Please pass a \"clientOptions\" object with a \"region\" field to the constructor`\n );\n }\n\n const endpointName = fields?.endpointName;\n if (!endpointName) {\n throw new Error(`Please pass an \"endpointName\" field to the constructor`);\n }\n\n const contentHandler = fields?.contentHandler;\n if (!contentHandler) {\n throw new Error(\n `Please pass a \"contentHandler\" field to the constructor`\n );\n }\n\n this.endpointName = fields.endpointName;\n this.contentHandler = fields.contentHandler;\n this.endpointKwargs = fields.endpointKwargs;\n this.modelKwargs = fields.modelKwargs;\n this.streaming = fields.streaming ?? false;\n this.client = new SageMakerRuntimeClient(fields.clientOptions);\n }\n\n _llmType() {\n return \"sagemaker_endpoint\";\n }\n\n /**\n * Calls the SageMaker endpoint and retrieves the result.\n * @param {string} prompt The input prompt.\n * @param {this[\"ParsedCallOptions\"]} options Parsed call options.\n * @param {CallbackManagerForLLMRun} runManager Optional run manager.\n * @returns {Promise<string>} A promise that resolves to the generated string.\n */\n /** @ignore */\n async _call(\n prompt: string,\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun\n ): Promise<string> {\n return this.streaming\n ? await this.streamingCall(prompt, options, runManager)\n : await this.noStreamingCall(prompt, options);\n }\n\n private async streamingCall(\n prompt: string,\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun\n ): Promise<string> {\n const chunks = [];\n for await (const chunk of this._streamResponseChunks(\n prompt,\n options,\n runManager\n )) {\n chunks.push(chunk.text);\n }\n return chunks.join(\"\");\n }\n\n private async noStreamingCall(\n prompt: string,\n options: this[\"ParsedCallOptions\"]\n ): Promise<string> {\n const body = await this.contentHandler.transformInput(\n prompt,\n this.modelKwargs ?? {}\n );\n const { contentType, accepts } = this.contentHandler;\n\n const response = await this.caller.call(() =>\n this.client.send(\n new InvokeEndpointCommand({\n EndpointName: this.endpointName,\n Body: body,\n ContentType: contentType,\n Accept: accepts,\n ...this.endpointKwargs,\n }),\n { abortSignal: options.signal }\n )\n );\n\n if (response.Body === undefined) {\n throw new Error(\"Inference result missing Body\");\n }\n return this.contentHandler.transformOutput(response.Body);\n }\n\n /**\n * Streams response chunks from the SageMaker endpoint.\n * @param {string} prompt The input prompt.\n * @param {this[\"ParsedCallOptions\"]} options Parsed call options.\n * @returns {AsyncGenerator<GenerationChunk>} An asynchronous generator yielding generation chunks.\n */\n async *_streamResponseChunks(\n prompt: string,\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun\n ): AsyncGenerator<GenerationChunk> {\n const body = await this.contentHandler.transformInput(\n prompt,\n this.modelKwargs ?? {}\n );\n const { contentType, accepts } = this.contentHandler;\n\n const stream = await this.caller.call(() =>\n this.client.send(\n new InvokeEndpointWithResponseStreamCommand({\n EndpointName: this.endpointName,\n Body: body,\n ContentType: contentType,\n Accept: accepts,\n ...this.endpointKwargs,\n }),\n { abortSignal: options.signal }\n )\n );\n\n if (!stream.Body) {\n throw new Error(\"Inference result missing Body\");\n }\n\n for await (const chunk of stream.Body) {\n if (chunk.PayloadPart && chunk.PayloadPart.Bytes) {\n const text = await this.contentHandler.transformOutput(\n chunk.PayloadPart.Bytes\n );\n yield new GenerationChunk({\n text,\n generationInfo: {\n ...chunk,\n response: undefined,\n },\n });\n await runManager?.handleLLMNewToken(text);\n } else if (chunk.InternalStreamFailure) {\n throw new Error(chunk.InternalStreamFailure.message);\n } else if (chunk.ModelStreamError) {\n throw new Error(chunk.ModelStreamError.message);\n }\n }\n }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAyCA,IAAsB,8BAAtB,MAAyE;CACvE,cAAc;CAEd,UAAU;AAmBX;;;;;;;;;;AAgDD,IAAa,oBAAb,cAAuCA,0CAAwB;CAC7D,kBAAkB;CAElB,OAAO,UAAU;AACf,SAAO;CACR;CAED,IAAI,aAAoD;AACtD,SAAO;GACL,yCAAyC;GACzC,6CAA6C;GAC7C,0CAA0C;EAC3C;CACF;CAED;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA,YAAYC,QAAgC;EAC1C,MAAM,OAAO;AAEb,MAAI,CAAC,OAAO,cAAc,OACxB,OAAM,IAAI,MACR,CAAC,6EAA6E,CAAC;EAInF,MAAM,eAAe,QAAQ;AAC7B,MAAI,CAAC,aACH,OAAM,IAAI,MAAM,CAAC,sDAAsD,CAAC;EAG1E,MAAM,iBAAiB,QAAQ;AAC/B,MAAI,CAAC,eACH,OAAM,IAAI,MACR,CAAC,uDAAuD,CAAC;EAI7D,KAAK,eAAe,OAAO;EAC3B,KAAK,iBAAiB,OAAO;EAC7B,KAAK,iBAAiB,OAAO;EAC7B,KAAK,cAAc,OAAO;EAC1B,KAAK,YAAY,OAAO,aAAa;EACrC,KAAK,SAAS,IAAIC,0DAAuB,OAAO;CACjD;CAED,WAAW;AACT,SAAO;CACR;;;;;;;;;CAUD,MAAM,MACJC,QACAC,SACAC,YACiB;AACjB,SAAO,KAAK,YACR,MAAM,KAAK,cAAc,QAAQ,SAAS,WAAW,GACrD,MAAM,KAAK,gBAAgB,QAAQ,QAAQ;CAChD;CAED,MAAc,cACZF,QACAC,SACAC,YACiB;EACjB,MAAM,SAAS,CAAE;AACjB,aAAW,MAAM,SAAS,KAAK,sBAC7B,QACA,SACA,WACD,EACC,OAAO,KAAK,MAAM,KAAK;AAEzB,SAAO,OAAO,KAAK,GAAG;CACvB;CAED,MAAc,gBACZF,QACAC,SACiB;EACjB,MAAM,OAAO,MAAM,KAAK,eAAe,eACrC,QACA,KAAK,eAAe,CAAE,EACvB;EACD,MAAM,EAAE,aAAa,SAAS,GAAG,KAAK;EAEtC,MAAM,WAAW,MAAM,KAAK,OAAO,KAAK,MACtC,KAAK,OAAO,KACV,IAAIE,yDAAsB;GACxB,cAAc,KAAK;GACnB,MAAM;GACN,aAAa;GACb,QAAQ;GACR,GAAG,KAAK;EACT,IACD,EAAE,aAAa,QAAQ,OAAQ,EAChC,CACF;AAED,MAAI,SAAS,SAAS,OACpB,OAAM,IAAI,MAAM;AAElB,SAAO,KAAK,eAAe,gBAAgB,SAAS,KAAK;CAC1D;;;;;;;CAQD,OAAO,sBACLH,QACAC,SACAC,YACiC;EACjC,MAAM,OAAO,MAAM,KAAK,eAAe,eACrC,QACA,KAAK,eAAe,CAAE,EACvB;EACD,MAAM,EAAE,aAAa,SAAS,GAAG,KAAK;EAEtC,MAAM,SAAS,MAAM,KAAK,OAAO,KAAK,MACpC,KAAK,OAAO,KACV,IAAIE,2EAAwC;GAC1C,cAAc,KAAK;GACnB,MAAM;GACN,aAAa;GACb,QAAQ;GACR,GAAG,KAAK;EACT,IACD,EAAE,aAAa,QAAQ,OAAQ,EAChC,CACF;AAED,MAAI,CAAC,OAAO,KACV,OAAM,IAAI,MAAM;AAGlB,aAAW,MAAM,SAAS,OAAO,KAC/B,KAAI,MAAM,eAAe,MAAM,YAAY,OAAO;GAChD,MAAM,OAAO,MAAM,KAAK,eAAe,gBACrC,MAAM,YAAY,MACnB;GACD,MAAM,IAAIC,yCAAgB;IACxB;IACA,gBAAgB;KACd,GAAG;KACH,UAAU;IACX;GACF;GACD,MAAM,YAAY,kBAAkB,KAAK;EAC1C,WAAU,MAAM,sBACf,OAAM,IAAI,MAAM,MAAM,sBAAsB;WACnC,MAAM,iBACf,OAAM,IAAI,MAAM,MAAM,iBAAiB;CAG5C;AACF"}