@dataclouder/nest-vertex 0.0.65 → 0.0.67

package/controllers/llm/llm-adapter.controller.js

@@ -37,6 +37,10 @@ let VertexAdapterLLMController = VertexAdapterLLMController_1 = class VertexAdap
         return result;
     }
     async continueConversation(conversation) {
+        if (conversation.returnJson) {
+            const result = await this.llmAdapterService.chatAndExtractJson(conversation);
+            return result;
+        }
         return this.llmAdapterService.chat(conversation);
     }
 };

package/models/adapter.models.d.ts

@@ -38,6 +38,7 @@ export declare class ChatMessageDict {
     content: string;
     role: ChatRole;
     metadata?: any;
+    json?: any;
 }
 export interface GeneratedTextAdapter {
     text: string;

package/models/adapter.models.js

@@ -61,6 +61,7 @@ class ChatMessageDict {
     content;
     role;
     metadata;
+    json;
 }
 exports.ChatMessageDict = ChatMessageDict;
 class ImageGenAdapterResponse {

package/package.json

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

package/services/adapter-llm.service.d.ts

@@ -8,10 +8,7 @@ export declare class LLMAdapterService implements IModelAdapterMethods {
     generate(request: ChatLLMRequestAdapter): Promise<any>;
     private getDefaultModel;
     chat(request: ChatLLMRequestAdapter): Promise<ChatMessageDict>;
-    chatAndExtractJson(request: ChatLLMRequestAdapter): Promise<{
-        json: any;
-        metadata: any;
-    }>;
+    chatAndExtractJson(request: ChatLLMRequestAdapter): Promise<ChatMessageDict>;
     private getQualityModel;
     describeImage(request: DescribeImageRequestAdapter): Promise<GeneratedTextAdapter>;
 }

package/services/adapter-llm.service.js

@@ -59,7 +59,7 @@ let LLMAdapterService = LLMAdapterService_1 = class LLMAdapterService {
             if (json === null) {
                 throw new Error('No JSON found in the initial response.');
             }
-            return { json, metadata: response.metadata };
+            return { ...response, json };
         }
         catch (initialError) {
             const warningMessage = `Initial JSON extraction failed: ${initialError.message}. Attempting recovery.`;
@@ -80,7 +80,7 @@ let LLMAdapterService = LLMAdapterService_1 = class LLMAdapterService {
                 }
                 response.metadata.recovered = true;
                 this.logger.log('Successfully recovered JSON on second attempt.');
-                return { json: fixedJson, metadata: response.metadata };
+                return { ...response, json: fixedJson };
             }
             catch (recoveryError) {
                 this.logger.error(`Recovery attempt also failed: ${recoveryError.message}`, recoveryError.stack);

package/services/google-genai.service.js

@@ -37,20 +37,14 @@ let GoogleGenaiService = GoogleGenaiService_1 = class GoogleGenaiService {
                 }, null);
                 if (balancedKey?.key) {
                     this.logger.debug(`⚖️ Using balanced key: ${balancedKey.id})  ${balancedKey.name} - ${balancedKey.key} `);
-                    return {
-                        client: new genai_1.GoogleGenAI({ apiKey: balancedKey.key }),
-                        balancedKey: balancedKey,
-                    };
+                    return { client: new genai_1.GoogleGenAI({ apiKey: balancedKey.key }), balancedKey: balancedKey };
                 }
             }
             catch (error) {
                 this.logger.error(`Failed to get balanced key: ${error.message}`);
                 const randomKey = process.env.GEMINI_API_KEY;
                 console.log(`⚖️ Using random key ::::::::: ${randomKey} 🚨🚨🚨🚨🚨🚨`);
-                return {
-                    client: new genai_1.GoogleGenAI({ apiKey: randomKey }),
-                    balancedKey: null,
-                };
+                return { client: new genai_1.GoogleGenAI({ apiKey: randomKey }), balancedKey: null };
             }
             if (!this.defaultApiKey) {
                 throw new nest_core_1.AppException({

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

@@ -2,10 +2,17 @@ import { HttpService } from '@nestjs/axios';
 import { AvailableKeyResult, ModelType, TierType } from '../models/key-balancer.models';
 export interface PostQueryKeyRequest {
     provider: string;
+    service?: string;
     tierType?: TierType;
     aiType?: ModelType;
     keyType?: string;
-    service?: string;
+}
+export interface KeyFailureDTO {
+    provider: string;
+    key: string;
+    error: string;
+    service: string;
+    ttlSeconds: number;
 }
 export declare class KeyBalancerClientService {
     private readonly httpService;
@@ -15,6 +22,6 @@ export declare class KeyBalancerClientService {
     getBestKey(keyRequest: PostQueryKeyRequest, token: any): Promise<AvailableKeyResult>;
     getKeys(): Promise<AvailableKeyResult[]>;
     private getUnavailableServiceResult;
-    recordFailedRequest(keyId: string, error: any, modelType: ModelType, model: string, ttlSeconds?: number): Promise<void>;
+    recordFailedRequest(failure: KeyFailureDTO): Promise<void>;
     updateUsage(keyUsageId: string, usageCount: number): Promise<any>;
 }

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

@@ -43,7 +43,7 @@ let KeyBalancerClientService = KeyBalancerClientService_1 = class KeyBalancerCli
             keyRequest,
         };
         const keyBalancerHost = process.env.KEY_BALANCER_HOST || 'http://localhost:8085';
-        const url = `${keyBalancerHost}/api/key-usage/best-key-balanced`;
+        const url = `${keyBalancerHost}/api/key-usage/get-best-key-balanced`;
         try {
             const { data } = await (0, rxjs_1.firstValueFrom)(this.httpService.post(url, finalRquest));
             return data;
@@ -73,11 +73,11 @@ let KeyBalancerClientService = KeyBalancerClientService_1 = class KeyBalancerCli
             errorDescription: 'Failed to fetch key from key-balancer API.',
         };
     }
-    async recordFailedRequest(keyId, error, modelType, model, ttlSeconds) {
+    async recordFailedRequest(failure) {
         const keyBalancerHost = process.env.KEY_BALANCER_HOST;
-        const url = `${keyBalancerHost}/api/key-balancer/redis/key/failed`;
-        const { data } = await (0, rxjs_1.firstValueFrom)(this.httpService.post(url, { keyId, error, modelType, model, ttlSeconds }));
-        this.logger.warn(`Failed to get LLM key for model ${model} from the external service.`, data);
+        const url = `${keyBalancerHost}/api/key-usage/update-key-failure`;
+        const { data } = await (0, rxjs_1.firstValueFrom)(this.httpService.post(url, failure));
+        this.logger.warn(`Failed to get LLM key for model ${failure.service} from the external service.`, data);
         return;
     }
     async updateUsage(keyUsageId, usageCount) {

package/services/vertex-gemini-chat.service.d.ts

@@ -11,6 +11,7 @@ export declare class GeminiChatService {
     private mapToGeminiRole;
     private formatMessagesToContent;
     chat(messages: MessageLLM[], model?: string, keyType?: TierType): Promise<ChatMessageDict>;
+    private _handleChatError;
     chatStream(messages: MessageLLM[], model?: string, keyType?: TierType): Promise<AsyncIterable<ChatMessageDict>>;
     listModels(): Promise<Record<string, string>[]>;
     getDefaultQualityModel(quality: EModelQuality): string;

package/services/vertex-gemini-chat.service.js

@@ -57,81 +57,103 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
         if (!messages || messages.length === 0) {
             return { content: '', role: adapter_models_1.ChatRole.Assistant, metadata: { finishReason: 'NO_INPUT' } };
         }
-        const { client, balancedKey } = await this.googleGenaiService.getGoogleGenAIClient({
-            model: model,
-            keyTierType: keyType || key_balancer_models_1.TierType.FREE_TIER,
-        });
-        try {
-            const systemMessage = messages
-                .filter(msg => msg.role === 'system')
-                .map(msg => msg.content)
-                .join('\n');
-            let lastMessage = messages[messages.length - 1];
-            if (lastMessage.role !== 'user' || typeof lastMessage.content !== 'string') {
-                messages.push({ role: adapter_models_1.ChatRole.User, content: 'please say something to start/continue conversation' });
-                lastMessage = messages[messages.length - 1];
-            }
-            const formattedContents = this.formatMessagesToContent(messages);
-            this.logger.debug(`Sending request to Gemini model 🦾 ${model} ♊️  with ${formattedContents.length} content parts and system instruction: ${!!systemMessage}`);
-            const response = await client.models.generateContent({
-                model: model,
-                contents: formattedContents,
-                config: { systemInstruction: systemMessage },
-            });
-            if (!response.text) {
-                if (response?.promptFeedback?.blockReason) {
-                    throw new nest_core_1.AppException({
-                        error_message: 'Gemini Bloqueo la solicitud: ' + response.promptFeedback.blockReason,
-                        explanation: 'Lo sentimos a gemini no le gusta esta conversación, y simplemente no responde, intenta reiniciar la conversación y ser muy respetuoso.',
-                    });
+        const maxRetries = 2;
+        for (let attempt = 1; attempt <= maxRetries; attempt++) {
+            const { client, balancedKey } = await this.googleGenaiService.getGoogleGenAIClient({ model: model, keyTierType: keyType || key_balancer_models_1.TierType.FREE_TIER });
+            try {
+                const systemMessage = messages
+                    .filter(msg => msg.role === 'system')
+                    .map(msg => msg.content)
+                    .join('\n');
+                let lastMessage = messages[messages.length - 1];
+                if (lastMessage.role !== 'user' || typeof lastMessage.content !== 'string') {
+                    messages.push({ role: adapter_models_1.ChatRole.User, content: 'please say something to start/continue conversation' });
+                    lastMessage = messages[messages.length - 1];
                 }
-                throw new nest_core_1.AppException({
-                    error_message: 'Gemini response is null or undefined',
-                    explanation: 'Probablemente Google bloqueo esta solicitud. intenta reiniciar la conversación y ser muy respetuoso.',
+                const formattedContents = this.formatMessagesToContent(messages);
+                this.logger.debug(`Sending request to Gemini model 🦾 ${model} ♊️  (Attempt ${attempt}/${maxRetries})`);
+                const response = await client.models.generateContent({
+                    model: model,
+                    contents: formattedContents,
+                    config: { systemInstruction: systemMessage },
                 });
-            }
-            this.logger.debug(`Received Gemini response text. ${response?.text.slice(0, 50).replace(/\n/g, '')} ...`);
-            const responseText = response?.text ?? '';
-            const tokens = {
-                input: response?.usageMetadata?.promptTokenCount,
-                output: response?.usageMetadata?.candidatesTokenCount,
-                total: response?.usageMetadata?.totalTokenCount,
-            };
-            return {
-                content: responseText,
-                role: adapter_models_1.ChatRole.Assistant,
-                metadata: { tokens },
-            };
-        }
-        catch (error) {
-            this.logger.error(`First Error Gemini Gemini {${balancedKey?.key} chat.sendMessage failed: ${error.message}`, client['apiClient']);
-            if (error instanceof nest_core_1.AppException) {
-                throw error;
-            }
-            if (balancedKey?.id) {
-                if (error.status === 429) {
-                    this.logger.error(`Gemini chat.sendMessage failed: Rate limited ${model} - con clave ${balancedKey?.name} ${balancedKey?.key}`);
-                    await this.keyBalancer.recordFailedRequest(balancedKey.id, error, key_balancer_models_1.ModelType.LLM, model, 15);
+                if (!response.text) {
+                    if (response?.promptFeedback?.blockReason) {
+                        throw new nest_core_1.AppException({
+                            error_message: 'Gemini Bloqueo la solicitud: ' + response.promptFeedback.blockReason,
+                            explanation: 'Lo sentimos a gemini no le gusta esta conversación, y simplemente no responde, intenta reiniciar la conversación y ser muy respetuoso.',
+                        });
+                    }
                     throw new nest_core_1.AppException({
-                        error_message: `Muchas solicitudes ${model} - con clave ${balancedKey?.key} Rate limited`,
-                        explanation: error.message,
+                        error_message: 'Gemini response is null or undefined',
+                        explanation: 'Probablemente Google bloqueo esta solicitud. intenta reiniciar la conversación y ser muy respetuoso.',
                     });
                 }
-                if (error.status === 503) {
-                    this.logger.error(`Gemini chat.sendMessage failed: Service Unavailable ${model} - ${balancedKey?.name} con clave ${balancedKey?.key}`);
-                    await this.keyBalancer.recordFailedRequest(balancedKey.id, error, key_balancer_models_1.ModelType.LLM, model, 40);
-                    throw new nest_core_1.AppException({
-                        error_message: `Problema de Google, No puede procesar la solicitud, invalidar ${balancedKey?.name} clave ${balancedKey?.key} por unos segundos.`,
-                        explanation: error.message,
-                    });
+                this.logger.debug(`Received Gemini response text. ${response?.text.slice(0, 50).replace(/\n/g, '')} ...`);
+                const responseText = response?.text ?? '';
+                const tokens = {
+                    input: response?.usageMetadata?.promptTokenCount,
+                    output: response?.usageMetadata?.candidatesTokenCount,
+                    total: response?.usageMetadata?.totalTokenCount,
+                };
+                return {
+                    content: responseText,
+                    role: adapter_models_1.ChatRole.Assistant,
+                    metadata: { tokens },
+                };
+            }
+            catch (error) {
+                this.logger.warn(`Attempt ${attempt} failed for model ${model}. Error: ${error.message}`);
+                await this._handleChatError(error, balancedKey, model, attempt < maxRetries);
+                if (attempt >= maxRetries) {
+                    this.logger.error(`All ${maxRetries} attempts failed for model ${model}. Rethrowing final error.`);
+                    throw error;
                 }
-                await this.keyBalancer.recordFailedRequest(balancedKey.id, error, key_balancer_models_1.ModelType.LLM, model);
             }
-            if (error instanceof nest_core_1.AppException) {
+        }
+        throw new Error('Chat failed after all retry attempts.');
+    }
+    async _handleChatError(error, balancedKey, model, isRetryable) {
+        this.logger.error(`Error during Gemini chat for model ${model} with key ${balancedKey?.key}: ${error.message}`, isRetryable ? undefined : error.stack);
+        if (error instanceof nest_core_1.AppException) {
+            if (balancedKey?.key) {
+                await this.keyBalancer.recordFailedRequest({
+                    key: balancedKey.key,
+                    provider: 'google',
+                    service: model,
+                    error: error.message,
+                    ttlSeconds: 0,
+                });
+            }
+            if (!isRetryable)
                 throw error;
+            return;
+        }
+        let ttlSeconds = 0;
+        let errorMessage = `An unexpected error occurred with model ${model}.`;
+        let explanation = error.message;
+        if (balancedKey?.id) {
+            switch (error.status) {
+                case 429:
+                    ttlSeconds = 15;
+                    errorMessage = `Rate limited for model ${model} with key ${balancedKey?.key}.`;
+                    this.logger.warn(`Gemini chat failed: Rate limited ${model} - key ${balancedKey?.name}`);
+                    break;
+                case 503:
+                    ttlSeconds = 40;
+                    errorMessage = `Service unavailable for model ${model}. Key ${balancedKey?.name} will be invalidated for a few seconds.`;
+                    this.logger.warn(`Gemini chat failed: Service Unavailable ${model} - key ${balancedKey?.name}`);
+                    break;
+                default:
+                    break;
             }
-            this.logger.error(`Gemini chat.sendMessage failed: ${error.message}`, error.stack);
-            throw new Error(`Failed to get Gemini chat completion: ${error.message}`);
+            await this.keyBalancer.recordFailedRequest({ key: balancedKey.key, provider: 'google', service: model, error: error.message, ttlSeconds });
+        }
+        if (!isRetryable) {
+            throw new nest_core_1.AppException({
+                error_message: errorMessage,
+                explanation: explanation,
+            });
         }
     }
     async chatStream(messages, model = this.defaultModel, keyType) {
@@ -174,7 +196,13 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
             }
             catch (error) {
                 if (balancedKey?.id) {
-                    await this.keyBalancer.recordFailedRequest(balancedKey.id, error, key_balancer_models_1.ModelType.LLM, model);
+                    await this.keyBalancer.recordFailedRequest({
+                        key: balancedKey.key,
+                        provider: 'google',
+                        service: model,
+                        error: error.message,
+                        ttlSeconds: 0,
+                    });
                 }
                 this.logger.error(`Gemini generateContentStream failed: ${error.message}`, error.stack);
                 throw new Error(`Failed to get Gemini chat stream completion: ${error.message}`);
@@ -227,6 +255,9 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
             try {
                 const fixPrompt = `I cant extract the json from the following text, something is wrong, please analyze the following text, correct any JSON formatting errors invalid characters or whatever thing is wrong, change or complete the text if needed, and return ONLY the valid JSON object or array. Do not include any explanatory text before or after the JSON itself. JSON to fix: ${responseText}`;
                 const fixResponse = await this.chat([{ role: adapter_models_1.ChatRole.User, content: fixPrompt }], model, keyType);
+                if (!fixResponse) {
+                    throw new Error('Recovery chat call returned undefined.');
+                }
                 this.logger.debug(`Recovery response content: ${fixResponse.content}`);
                 const fixedJson = (0, llm_models_1.extractJsonFromResponse)(fixResponse.content);
                 if (fixedJson === null) {
@@ -246,6 +277,12 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
     }
     async chatAndExtractJson(messages, model = gemini_models_1.GeminiModels.Gemini2_5Lite, keyType) {
         const response = await this.chat(messages, model, keyType);
+        if (!response) {
+            throw new nest_core_1.AppException({
+                error_message: 'chatAndExtractJson received no response from chat()',
+                explanation: 'The chat service did not return a response, so JSON could not be extracted.',
+            });
+        }
         const json = await this._extractJsonWithRecovery(response.content, model, keyType);
         response.json = json;
         return response;
@@ -290,7 +327,13 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
         }
         catch (error) {
             if (balancedKey?.id) {
-                await this.keyBalancer.recordFailedRequest(balancedKey.id, error, key_balancer_models_1.ModelType.LLM, visionModelLLM);
+                await this.keyBalancer.recordFailedRequest({
+                    key: balancedKey.key,
+                    provider: 'google',
+                    service: visionModelLLM,
+                    error: error.message,
+                    ttlSeconds: 0,
+                });
             }
             this.logger.error(`Error in describeImageByUrl vision call: ${error.message}`, error.stack);
             throw new Error(`Failed to get Gemini vision completion: ${error.message}`);
@@ -317,6 +360,12 @@ let GeminiChatService = GeminiChatService_1 = class GeminiChatService {
         }
         else {
             const response = await this.chat(conversation.messages, modelName, tierType);
+            if (!response) {
+                throw new nest_core_1.AppException({
+                    error_message: 'chatWithConversation received no response from chat()',
+                    explanation: 'The chat service did not return a response.',
+                });
+            }
             const endTime = Date.now();
             const processTime = (endTime - startTime) / 1000;
             return {