usage handler calc for azure openai

pele-ai · Aug 4, 2024 · a919aa4 · a919aa4
1 parent acbf37b
commit a919aa4
Show file tree

Hide file tree

Showing 6 changed files with 9,854 additions and 7,983 deletions.
diff --git a/...ponents/nodes/agents/ConversationalRetrievalToolAgent/ConversationalRetrievalToolAgent.ts b/...ponents/nodes/agents/ConversationalRetrievalToolAgent/ConversationalRetrievalToolAgent.ts
@@ -129,7 +129,7 @@ class ConversationalRetrievalToolAgent_Agents implements INode {
         if (isStreamable) {
             const handler = new CustomChainHandler(options.socketIO, options.socketIOClientId)
             const customerId = nodeData.inputs?.vars?.customerId ?? '';
-            const usageHandler = new UsageHandler(customerId);
+            const usageHandler = new UsageHandler(customerId, nodeData?.inputs?.model?.modelName);
             res = await executor.invoke({ input }, { callbacks: [loggerHandler, handler, usageHandler, ...callbacks] })
             if (res.sourceDocuments) {
                 options.socketIO.to(options.socketIOClientId).emit('sourceDocuments', flatten(res.sourceDocuments))

diff --git a/packages/components/nodes/chatmodels/ChatOpenAI/ChatOpenAI.ts b/packages/components/nodes/chatmodels/ChatOpenAI/ChatOpenAI.ts
@@ -182,7 +182,7 @@ class ChatOpenAI_ChatModels implements INode {
             modelName,
             openAIApiKey,
             streaming: streaming ?? true,
-            stream_options: { include_usage: true }
+            // stream_options: { include_usage: true }
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)

diff --git a/packages/components/nodes/vectorstores/Milvus/Milvus.ts b/packages/components/nodes/vectorstores/Milvus/Milvus.ts
@@ -247,7 +247,7 @@ const similaritySearchVectorWithScore = async (query: number[], k: number, vecto
             anns_field: vectorStore.vectorField,
             topk: k.toString(),
             metric_type: vectorStore.indexCreateParams.metric_type,
-            params: vectorStore.indexSearchParams
+            params: '' + vectorStore.indexSearchParams
         },
         output_fields: outputFields,
         vector_type: DataType.FloatVector,

diff --git a/packages/components/package.json b/packages/components/package.json
@@ -35,6 +35,7 @@
         "@google/generative-ai": "^0.15.0",
         "@huggingface/inference": "^2.6.1",
         "@langchain/anthropic": "^0.2.1",
+        "@langchain/aws": "^0.0.6",
         "@langchain/cohere": "^0.0.7",
         "@langchain/community": "^0.2.17",
         "@langchain/exa": "^0.0.5",
@@ -80,6 +81,7 @@
         "graphql": "^16.6.0",
         "html-to-text": "^9.0.5",
         "ioredis": "^5.3.2",
+        "js-tiktoken": "^1.0.12",
         "jsdom": "^22.1.0",
         "jsonpointer": "^5.0.1",
         "langchain": "^0.2.11",

diff --git a/packages/components/src/usageHandler.ts b/packages/components/src/usageHandler.ts
@@ -1,39 +1,56 @@
 import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
-import { Run } from '@langchain/core/tracers/base';
+import { LLMResult } from "@langchain/core/outputs";
+import { Serialized } from "@langchain/core/load/serializable";
+import { encodingForModel, Tiktoken } from "js-tiktoken";
 import axios from "axios";
-import { BaseMessage } from "@langchain/core/messages";
-import { LLMResult, Generation } from "@langchain/core/outputs";
-import { AIMessageChunk } from "@langchain/core/messages";
-type Message = BaseMessage | Generation | string;
-
 
 export class UsageHandler extends BaseCallbackHandler {
 
     name = "usage_handler" as const
     customerId: string;
+    encoding: Tiktoken;
 
-    constructor(customerId: string) {
+    constructor(customerId: string, model: string) {
         super();
         this.customerId = customerId;
+        //@ts-ignore
+        this.encoding = encodingForModel(model);
+    }
+
+    handleLLMStart(_llm: Serialized, prompts: string[]) {
+        try {
+            const prompt = prompts[0];
+            const tokens = this.encoding?.encode(prompt);
+            if (tokens) {
+                const url = `${process.env.USAGE_URL}/usage`;
+                axios.post(url, {
+                    customerId: this.customerId,
+                    usage: {
+                        input_tokens: tokens.length
+                    }
+                });
+            }
+        } catch (e) {
+            console.log(e);
+        }
     }
 
-    handleLLMEnd(output: LLMResult, runId: string) {
+    handleLLMEnd(output: LLMResult) {
         try {
-            //@ts-ignore
-            const message: AIMessageChunk = output?.generations[0][0].message;
-            const { usage_metadata } = message?.lc_kwargs || {};
-            if (usage_metadata && this.customerId) {
+            const prompt = output?.generations[0][0].text;
+            const tokens = this.encoding?.encode(prompt);
+
+            if (tokens) {
                 const url = `${process.env.USAGE_URL}/usage`;
-                console.log('calling usage: ', url);
-                if (url) {
-                    axios.post(url, {
-                        customerId: this.customerId,
-                        usage: usage_metadata
-                    });
-                }
+                axios.post(url, {
+                    customerId: this.customerId,
+                    usage: {
+                        output_tokens: tokens.length
+                    }
+                });
             }
-        } catch (err) {
-            console.log("usage handler error:" + err);
+        } catch (e) {
+            console.log(e);
         }
     }
 }