feat: impl extensions for langchain

Author: hlhr202

package-lock.json

@@ -47,14 +47,15 @@
     "@tensorflow/tfjs-node": "^4.2.0",
     "@types/node": "^18.15.5",
     "@types/semver": "^7.3.13",
-    "axios": "^1.3.5",
+    "axios": "*",
     "glob": "^9.3.4",
     "rimraf": "^4.4.1",
     "semver": "^7.3.8",
     "tsup": "^6.7.0",
     "tsx": "^3.12.6",
     "typescript": "^5.0.4",
-    "vitest": "^0.29.8"
+    "vitest": "^0.29.8",
+    "langchain": "^0.0.56"
   },
   "dependencies": {
     "@llama-node/cli": "0.0.27"
@@ -68,4 +69,4 @@
     "@llama-node/core": "0.0.27",
     "@llama-node/llama-cpp": "0.0.27"
   }
-}
+}

package.json

@@ -0,0 +1,29 @@
+import { AsyncCaller } from "langchain/dist/util/async_caller";
+import { Embeddings, type EmbeddingsParams } from "langchain/embeddings/base";
+import type { LLama } from "..";
+
+export class LLamaEmbeddings implements Embeddings {
+    caller: AsyncCaller;
+    llm: LLama;
+
+    constructor(params: EmbeddingsParams, llm: LLama) {
+        if ((params.maxConcurrency ?? 1) > 1) {
+            console.warn(
+                "maxConcurrency > 1 not officially supported for llama-node, use at your own risk"
+            );
+        }
+        this.caller = new AsyncCaller(params);
+        this.llm = llm;
+    }
+
+    embedDocuments(documents: string[]): Promise<number[][]> {
+        const promises = documents.map((doc) =>
+            this.llm.getDefaultEmbeddings(doc)
+        );
+        return Promise.all(promises);
+    }
+
+    embedQuery(document: string): Promise<number[]> {
+        return this.llm.getDefaultEmbeddings(document);
+    }
+}

src/extensions/langchain.ts

@@ -2,11 +2,11 @@ import { CompletionCallback } from "./llm";
 import type { LLM } from "./llm";
 
 export class LLama<
-    Instance,
-    LoadConfig,
-    LLMInferenceArguments,
-    LLMEmbeddingArguments,
-    TokenizeArguments
+    Instance = any,
+    LoadConfig = any,
+    LLMInferenceArguments = any,
+    LLMEmbeddingArguments = any,
+    TokenizeArguments = any
 > {
     llm: LLM<
         Instance,
@@ -48,6 +48,15 @@ export class LLama<
         }
     }
 
+    async getDefaultEmbeddings(text: string): Promise<number[]> {
+        if (!this.llm.getDefaultEmbedding) {
+            console.warn("getDefaultEmbedding not implemented for current LLM");
+            return [];
+        } else {
+            return this.llm.getDefaultEmbedding(text);
+        }
+    }
+
     async tokenize(content: TokenizeArguments): Promise<number[]> {
         if (!this.llm.tokenize) {
             console.warn("tokenize not implemented for current LLM");

src/index.ts

@@ -7,7 +7,7 @@ export interface LLM<
     LoadConfig,
     LLMInferenceArguments,
     LLMEmbeddingArguments,
-    LLMTokenizeArguments,
+    LLMTokenizeArguments
 > {
     readonly instance: Instance;
 
@@ -20,5 +20,7 @@ export interface LLM<
 
     getEmbedding?(params: LLMEmbeddingArguments): Promise<number[]>;
 
+    getDefaultEmbedding?(text: string): Promise<number[]>;
+
     tokenize?(content: LLMTokenizeArguments): Promise<number[]>;
 }

src/llm.d.ts

@@ -21,7 +21,13 @@ export interface TokenizeArguments {
 
 export class LLamaCpp
     implements
-        LLM<LLama, LoadConfig, LlamaInvocation, LlamaInvocation, TokenizeArguments>
+        LLM<
+            LLama,
+            LoadConfig,
+            LlamaInvocation,
+            LlamaInvocation,
+            TokenizeArguments
+        >
 {
     instance!: LLama;
 
@@ -82,6 +88,18 @@ export class LLamaCpp
         });
     }
 
+    async getDefaultEmbedding(text: string): Promise<number[]> {
+        return this.getEmbedding({
+            nThreads: 4,
+            nTokPredict: 1024,
+            topK: 40,
+            topP: 0.1,
+            temp: 0.1,
+            repeatPenalty: 1,
+            prompt: text,
+        });
+    }
+
     async tokenize(params: TokenizeArguments): Promise<number[]> {
         return new Promise<number[]>((res, rej) => {
             this.instance.tokenize(params.content, params.nCtx, (response) => {

src/llm/llama-cpp.ts

@@ -75,6 +75,18 @@ export class LLamaRS
         });
     }
 
+    async getDefaultEmbedding(text: string): Promise<number[]> {
+        return this.getEmbedding({
+            nThreads: 4,
+            numPredict: 1024,
+            topK: 40,
+            topP: 0.1,
+            temp: 0.1,
+            repeatPenalty: 1,
+            prompt: text,
+        });
+    }
+
     async tokenize(params: string): Promise<number[]> {
         return new Promise<number[]>((res) => {
             this.instance.tokenize(params, (response) => {

src/llm/llama-rs.ts

@@ -1,7 +1,8 @@
 import { defineConfig } from "tsup";
 
 export default defineConfig({
-    entry: ["src/index.ts", "src/llm/*.ts"],
+    entry: ["src/index.ts", "src/llm/*.ts", "src/extensions/*.ts"],
+    external: ["langchain"],
     target: ["es2015"],
     format: ["cjs", "esm"],
     dts: true,