agsamantha/node_modules/@langchain/community/dist/chat_models/webllm.js

137 lines
4.5 KiB
JavaScript
Raw Normal View History

2024-10-02 15:15:21 -05:00
import { SimpleChatModel, } from "@langchain/core/language_models/chat_models";
import { AIMessageChunk } from "@langchain/core/messages";
import { ChatGenerationChunk } from "@langchain/core/outputs";
import * as webllm from "@mlc-ai/web-llm";
/**
* To use this model you need to have the `@mlc-ai/web-llm` module installed.
* This can be installed using `npm install -S @mlc-ai/web-llm`.
*
* You can see a list of available model records here:
* https://github.com/mlc-ai/web-llm/blob/main/src/config.ts
* @example
* ```typescript
* // Initialize the ChatWebLLM model with the model record.
* const model = new ChatWebLLM({
* model: "Phi-3-mini-4k-instruct-q4f16_1-MLC",
* chatOptions: {
* temperature: 0.5,
* },
* });
*
* // Call the model with a message and await the response.
* const response = await model.invoke([
* new HumanMessage({ content: "My name is John." }),
* ]);
* ```
*/
export class ChatWebLLM extends SimpleChatModel {
static lc_name() {
return "ChatWebLLM";
}
constructor(inputs) {
super(inputs);
Object.defineProperty(this, "engine", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "appConfig", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "chatOptions", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "temperature", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "model", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
this.appConfig = inputs.appConfig;
this.chatOptions = inputs.chatOptions;
this.model = inputs.model;
this.temperature = inputs.temperature;
this.engine = new webllm.MLCEngine({
appConfig: this.appConfig,
});
}
_llmType() {
return "web-llm";
}
async initialize(progressCallback) {
if (progressCallback !== undefined) {
this.engine.setInitProgressCallback(progressCallback);
}
await this.reload(this.model, this.chatOptions);
}
async reload(modelId, newChatOpts) {
await this.engine.reload(modelId, newChatOpts);
}
async *_streamResponseChunks(messages, options, runManager) {
const messagesInput = messages.map((message) => {
if (typeof message.content !== "string") {
throw new Error("ChatWebLLM does not support non-string message content in sessions.");
}
const langChainType = message._getType();
let role;
if (langChainType === "ai") {
role = "assistant";
}
else if (langChainType === "human") {
role = "user";
}
else if (langChainType === "system") {
role = "system";
}
else {
throw new Error("Function, tool, and generic messages are not supported.");
}
return {
role,
content: message.content,
};
});
const stream = await this.engine.chat.completions.create({
stream: true,
messages: messagesInput,
stop: options.stop,
logprobs: true,
});
for await (const chunk of stream) {
// Last chunk has undefined content
const text = chunk.choices[0].delta.content ?? "";
yield new ChatGenerationChunk({
text,
message: new AIMessageChunk({
content: text,
additional_kwargs: {
logprobs: chunk.choices[0].logprobs,
finish_reason: chunk.choices[0].finish_reason,
},
}),
});
await runManager?.handleLLMNewToken(text);
}
}
async _call(messages, options, runManager) {
const chunks = [];
for await (const chunk of this._streamResponseChunks(messages, options, runManager)) {
chunks.push(chunk.text);
}
return chunks.join("");
}
}