@dataclouder/nest-vertex 0.0.55 → 0.0.56

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dataclouder/nest-vertex",
-  "version": "0.0.55",
+  "version": "0.0.56",
   "description": "NestJS Vertex AI library for Dataclouder",
   "author": "dataclouder",
   "license": "MIT",

package/services/key-balancer-api.service.d.ts

@@ -16,4 +16,5 @@ export declare class KeyBalancerClientService {
     getKeys(): Promise<AvailableKeyResult[]>;
     private getUnavailableServiceResult;
     recordFailedRequest(keyId: string, error: any, modelType: ModelType, model: string, ttlSeconds?: number): Promise<void>;
+    updateUsage(keyUsageId: string, usageCount: number): Promise<any>;
 }

package/services/key-balancer-api.service.js

@@ -74,6 +74,19 @@ let KeyBalancerClientService = KeyBalancerClientService_1 = class KeyBalancerCli
         this.logger.warn(`Failed to get LLM key for model ${model} from the external service.`, data);
         return;
     }
+    async updateUsage(keyUsageId, usageCount) {
+        const keyBalancerHost = process.env.KEY_BALANCER_HOST || 'http://localhost:8085';
+        const url = `${keyBalancerHost}/api/key-usage/update-usage`;
+        try {
+            const { data } = await (0, rxjs_1.firstValueFrom)(this.httpService.post(url, { keyUsageId, usageCount }));
+            this.logger.verbose(`Usage updated for keyUsageId: ${keyUsageId} with usageCount: ${usageCount}`);
+            return data;
+        }
+        catch (error) {
+            this.logger.error(`Failed to update usage for keyUsageId: ${keyUsageId}`, error.stack);
+            throw error;
+        }
+    }
 };
 exports.KeyBalancerClientService = KeyBalancerClientService;
 exports.KeyBalancerClientService = KeyBalancerClientService = KeyBalancerClientService_1 = __decorate([

package/services/vertex-tts.service.js

@@ -32,6 +32,7 @@ let VertextTtsService = class VertextTtsService extends base_tts_service_1.BaseT
             keyType: 'service_account',
         }, token);
         if (selectedKey.key) {
+            this.keyBalancerClient.updateUsage(selectedKey.id, input.text.length);
             this.logger.verbose(` Balanced key Found : ${selectedKey.name} ${selectedKey.id}`);
             try {
                 const credentials = JSON.parse(selectedKey.key);