add token use parameters

llm-tools · Jul 31, 2024 · e67d70d · e67d70d
1 parent 902a1da
commit e67d70d
Show file tree

Hide file tree

Showing 11 changed files with 797 additions and 757 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@llm-tools/embedjs",
-    "version": "0.0.93",
+    "version": "0.0.94",
     "description": "A NodeJS RAG framework to easily work with LLMs and custom datasets",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",
@@ -57,26 +57,26 @@
     "homepage": "https://github.com/llm-tools/embedjs#readme",
     "dependencies": {
         "@huggingface/inference": "^2.8.0",
-        "@langchain/anthropic": "^0.2.6",
+        "@langchain/anthropic": "^0.2.11",
         "@langchain/cohere": "^0.2.1",
-        "@langchain/community": "^0.2.20",
-        "@langchain/core": "^0.2.17",
-        "@langchain/google-vertexai": "^0.0.20",
-        "@langchain/mistralai": "^0.0.26",
-        "@langchain/openai": "^0.2.4",
+        "@langchain/community": "^0.2.22",
+        "@langchain/core": "^0.2.18",
+        "@langchain/google-vertexai": "^0.0.21",
+        "@langchain/mistralai": "^0.0.27",
+        "@langchain/openai": "^0.2.5",
         "axios": "^1.7.2",
         "compute-cosine-similarity": "^1.1.0",
         "confluence.js": "^1.7.4",
         "csv-parse": "^5.5.6",
-        "debug": "^4.3.5",
+        "debug": "^4.3.6",
         "html-to-text": "^9.0.5",
-        "langchain": "^0.2.10",
+        "langchain": "^0.2.12",
         "md5": "^2.3.0",
         "mime": "^4.0.4",
         "office-text-extractor": "^3.0.3",
-        "sitemapper": "^3.1.8",
+        "sitemapper": "^3.2.9",
         "stream-mime-type": "^2.0.0",
-        "usetube": "^2.0.2",
+        "usetube": "^2.2.7",
         "uuid": "^10.0.0",
         "youtube-transcript": "^1.2.1"
     },
@@ -94,15 +94,15 @@
         "@types/jest": "29.5.12",
         "@types/md5": "^2.3.5",
         "@types/ms": "^0.7.34",
-        "@types/node": "^20.14.11",
+        "@types/node": "^22.0.2",
         "@types/usetube": "^2.1.2",
         "@types/webidl-conversions": "^7.0.3",
         "@types/yargs-parser": "^21.0.3",
-        "@typescript-eslint/eslint-plugin": "^7.16.1",
-        "@typescript-eslint/parser": "^7.16.1",
+        "@typescript-eslint/eslint-plugin": "^8.0.0",
+        "@typescript-eslint/parser": "^8.0.0",
         "chromadb": "^1.8.1",
-        "cohere-ai": "^7.11.0",
-        "eslint": "^8.57.0",
+        "cohere-ai": "^7.11.1",
+        "eslint": "^9.8.0",
         "eslint-config-prettier": "^9.1.0",
         "eslint-plugin-prettier": "^5.2.1",
         "hnswlib-node": "^3.0.0",
@@ -113,28 +113,28 @@
         "prettier": "^3.3.3",
         "rimraf": "^6.0.1",
         "ts-jest": "29.2.3",
-        "typescript": "^5.5.3",
-        "vectordb": "^0.7.1",
+        "typescript": "^5.5.4",
+        "vectordb": "^0.8.0",
         "weaviate-ts-client": "^2.2.0"
     },
     "peerDependencies": {
         "@azure/cosmos": "^4.0.0",
         "@pinecone-database/pinecone": "^2.2.2",
         "@qdrant/js-client-rest": "^1.10.0",
         "chromadb": "^1.8.1",
-        "cohere-ai": "^7.11.0",
+        "cohere-ai": "^7.11.1",
         "hnswlib-node": "^3.0.0",
         "ioredis": "^5.4.1",
         "lmdb": "^3.0.12",
         "mongodb": "^6.8.0",
-        "vectordb": "^0.7.1",
+        "vectordb": "^0.8.0",
         "weaviate-ts-client": "^2.2.0"
     },
     "overrides": {
         "@pinecone-database/pinecone": "^2.2.2",
         "weaviate-ts-client": "^2.2.0",
         "hnswlib-node": "^3.0.0",
-        "vectordb": "^0.7.1"
+        "vectordb": "^0.8.0"
     },
     "peerDependenciesMeta": {
         "@pinecone-database/pinecone": {

diff --git a/src/core/rag-application.ts b/src/core/rag-application.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 
 import { BaseDb } from '../interfaces/base-db.js';
 import { BaseLoader } from '../interfaces/base-loader.js';
-import { AddLoaderReturn, Chunk, InsertChunkData, LoaderChunk, Message } from '../global/types.js';
+import { AddLoaderReturn, Chunk, InsertChunkData, LoaderChunk, QueryResponse } from '../global/types.js';
 import { DynamicLoader, LoaderParam } from './dynamic-loader-selector.js';
 import { RAGApplicationBuilder } from './rag-application-builder.js';
 import { DEFAULT_INSERT_BATCH_SIZE } from '../global/constants.js';
@@ -358,11 +358,7 @@ export class RAGApplication {
      * the LLM model with the provided query template, user query, context, and conversation history. The
      * `sources` property is an array of strings representing unique sources used to generate the LLM response.
      */
-    public async query(
-        userQuery: string,
-        conversationId?: string,
-        customContext?: Chunk[],
-    ): Promise<Extract<Message, { actor: 'AI' }>> {
+    public async query(userQuery: string, conversationId?: string, customContext?: Chunk[]): Promise<QueryResponse> {
         if (!this.model) {
             throw new Error('LLM Not set; query method not available');
         }

diff --git a/src/global/types.ts b/src/global/types.ts
@@ -67,3 +67,18 @@ export type Conversation = {
     conversationId: string;
     entries: Message[];
 };
+
+export type ModelResponse = {
+    result: string;
+    tokenUse?: {
+        inputTokens: number;
+        outputTokens: number;
+    };
+};
+
+export type QueryResponse = Extract<Message, { actor: 'AI' }> & {
+    tokenUse: {
+        inputTokens: number | 'UNKNOWN';
+        outputTokens: number | 'UNKNOWN';
+    };
+};
diff --git a/src/interfaces/base-model.ts b/src/interfaces/base-model.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 import { v4 as uuidv4 } from 'uuid';
 
 import { BaseConversation } from './base-conversations.js';
-import { Chunk, Message, SourceDetail } from '../global/types.js';
+import { Chunk, Message, ModelResponse, QueryResponse, SourceDetail } from '../global/types.js';
 
 export abstract class BaseModel {
     private readonly baseDebug = createDebugMessages('embedjs:model:BaseModel');
@@ -35,7 +35,7 @@ export abstract class BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         conversationId: string = 'default',
-    ): Promise<Extract<Message, { actor: 'AI' }>> {
+    ): Promise<QueryResponse> {
         const conversation = await BaseModel.conversations.getConversation(conversationId);
         this.baseDebug(`${conversation.entries.length} history entries found for conversationId '${conversationId}'`);
 
@@ -48,20 +48,26 @@ export abstract class BaseModel {
         });
 
         // Run LLM implementation in subclass
-        const result = await this.runQuery(system, userQuery, supportingContext, conversation.entries);
+        const response = await this.runQuery(system, userQuery, supportingContext, conversation.entries);
 
         const uniqueSources = this.extractUniqueSources(supportingContext);
         const newEntry: Message = {
             id: uuidv4(),
             timestamp: new Date(),
-            content: result,
+            content: response.result,
             actor: 'AI',
             sources: uniqueSources,
         };
 
         // Add AI response to history
         await BaseModel.conversations.addEntryToConversation(conversationId, newEntry);
-        return newEntry;
+        return {
+            ...newEntry,
+            tokenUse: {
+                inputTokens: response.tokenUse?.inputTokens ?? 'UNKNOWN',
+                outputTokens: response.tokenUse?.outputTokens ?? 'UNKNOWN',
+            },
+        };
     }
 
     private extractUniqueSources(supportingContext: Chunk[]): SourceDetail[] {
@@ -89,5 +95,5 @@ export abstract class BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string>;
+    ): Promise<ModelResponse>;
 }
diff --git a/src/models/anthropic-model.ts b/src/models/anthropic-model.ts
@@ -3,7 +3,7 @@ import { ChatAnthropic } from '@langchain/anthropic';
 import { HumanMessage, AIMessage, SystemMessage } from '@langchain/core/messages';
 
 import { BaseModel } from '../interfaces/base-model.js';
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 
 export class Anthropic extends BaseModel {
     private readonly debug = createDebugMessages('embedjs:model:Anthropic');
@@ -24,7 +24,7 @@ export class Anthropic extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const pastMessages: (AIMessage | SystemMessage | HumanMessage)[] = [
             new SystemMessage(
                 `${system}. Supporting context: ${supportingContext.map((s) => s.pageContent).join('; ')}`,
@@ -44,6 +44,13 @@ export class Anthropic extends BaseModel {
         this.debug('Executing anthropic model with prompt -', userQuery);
         const result = await this.model.invoke(pastMessages);
         this.debug('Anthropic response -', result);
-        return result.content.toString();
+
+        return {
+            result: result.content.toString(),
+            tokenUse: {
+                inputTokens: result.response_metadata.usage.input_tokens,
+                outputTokens: result.response_metadata.usage.output_tokens,
+            },
+        };
     }
 }
diff --git a/src/models/huggingface-model.ts b/src/models/huggingface-model.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 import { HuggingFaceInference } from '@langchain/community/llms/hf';
 
 import { BaseModel } from '../interfaces/base-model.js';
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 
 export class HuggingFace extends BaseModel {
     private readonly debug = createDebugMessages('embedjs:model:HuggingFace');
@@ -36,7 +36,7 @@ export class HuggingFace extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const pastMessages = [system];
         pastMessages.push(`Data: ${supportingContext.map((s) => s.pageContent).join('; ')}`);
 
@@ -57,6 +57,9 @@ export class HuggingFace extends BaseModel {
         this.debug(`Executing hugging face '${this.model.model}' model with prompt -`, userQuery);
         const result = await this.model.invoke(finalPrompt);
         this.debug('Hugging response -', result);
-        return result;
+
+        return {
+            result,
+        };
     }
 }
diff --git a/src/models/mistral-model.ts b/src/models/mistral-model.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 import { ChatMistralAI } from '@langchain/mistralai';
 import { AIMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
 
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 import { BaseModel } from '../interfaces/base-model.js';
 
 export class Mistral extends BaseModel {
@@ -27,7 +27,7 @@ export class Mistral extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const pastMessages: (AIMessage | SystemMessage | HumanMessage)[] = [new SystemMessage(system)];
         pastMessages.push(
             new SystemMessage(`Supporting context: ${supportingContext.map((s) => s.pageContent).join('; ')}`),
@@ -46,6 +46,9 @@ export class Mistral extends BaseModel {
         this.debug('Executing mistral model with prompt -', userQuery);
         const result = await this.model.invoke(pastMessages);
         this.debug('Mistral response -', result);
-        return result.content.toString();
+
+        return {
+            result: result.content.toString(),
+        };
     }
 }
diff --git a/src/models/ollama-model.ts b/src/models/ollama-model.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 import { Ollama as ChatOllamaAI } from '@langchain/community/llms/ollama';
 import { AIMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
 
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 import { BaseModel } from '../interfaces/base-model.js';
 
 export class Ollama extends BaseModel {
@@ -22,7 +22,7 @@ export class Ollama extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const pastMessages: (AIMessage | SystemMessage | HumanMessage)[] = [new SystemMessage(system)];
         pastMessages.push(
             new SystemMessage(`Supporting context: ${supportingContext.map((s) => s.pageContent).join('; ')}`),
@@ -41,6 +41,9 @@ export class Ollama extends BaseModel {
         this.debug(`Executing ollama model ${this.model} with prompt -`, userQuery);
         const result = await this.model.invoke(pastMessages);
         this.debug('Ollama response -', result);
-        return result.toString();
+
+        return {
+            result: result.toString(),
+        };
     }
 }
diff --git a/src/models/openai-model.ts b/src/models/openai-model.ts
@@ -3,7 +3,7 @@ import { ChatOpenAI, ClientOptions } from '@langchain/openai';
 import { HumanMessage, AIMessage, SystemMessage } from '@langchain/core/messages';
 
 import { BaseModel } from '../interfaces/base-model.js';
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 
 export class OpenAi extends BaseModel {
     private readonly debug = createDebugMessages('embedjs:model:OpenAi');
@@ -38,7 +38,7 @@ export class OpenAi extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const pastMessages: (AIMessage | SystemMessage | HumanMessage)[] = [new SystemMessage(system)];
         pastMessages.push(
             new SystemMessage(`Supporting context: ${supportingContext.map((s) => s.pageContent).join('; ')}`),
@@ -57,6 +57,13 @@ export class OpenAi extends BaseModel {
         this.debug('Executing openai model with prompt -', userQuery);
         const result = await this.model.invoke(pastMessages);
         this.debug('OpenAI response -', result);
-        return result.content.toString();
+
+        return {
+            result: result.content.toString(),
+            tokenUse: {
+                inputTokens: result.response_metadata.tokenUsage.promptTokens,
+                outputTokens: result.response_metadata.tokenUsage.completionTokens,
+            },
+        };
     }
 }
diff --git a/src/models/vertexai-model.ts b/src/models/vertexai-model.ts
@@ -2,7 +2,7 @@ import createDebugMessages from 'debug';
 import { ChatVertexAI } from '@langchain/google-vertexai';
 import { AIMessage, HumanMessage, SystemMessage } from '@langchain/core/messages';
 
-import { Chunk, Message } from '../global/types.js';
+import { Chunk, Message, ModelResponse } from '../global/types.js';
 import { BaseModel } from '../interfaces/base-model.js';
 
 export class VertexAI extends BaseModel {
@@ -19,7 +19,7 @@ export class VertexAI extends BaseModel {
         userQuery: string,
         supportingContext: Chunk[],
         pastConversations: Message[],
-    ): Promise<string> {
+    ): Promise<ModelResponse> {
         const systemString =
             system + '\n' + `Supporting context: ${supportingContext.map((s) => s.pageContent).join('; ')}`;
         const pastMessages: (AIMessage | SystemMessage | HumanMessage)[] = [new SystemMessage(systemString)];
@@ -37,6 +37,9 @@ export class VertexAI extends BaseModel {
         this.debug('Executing VertexAI model with prompt -', userQuery);
         const result = await this.model.invoke(pastMessages);
         this.debug('VertexAI response -', result);
-        return result.content.toString();
+
+        return {
+            result: result.content.toString(),
+        };
     }
 }