@retrivora-ai/rag-engine 2.2.9 → 2.3.1

package/src/core/VectorPlugin.ts

@@ -7,6 +7,7 @@ import { IngestDocument, ChatResponse } from '../types';
 import { ChatMessage } from '../types';
 import { LicenseVerifier } from './LicenseVerifier';
 import { wrapError, RetrivoraErrorCode } from '../exceptions';
+import { FreeTierLimitsGuard } from './FreeTierLimitsGuard';
 
 /**
  * VectorPlugin — main orchestrator class.
@@ -78,6 +79,13 @@ export class VectorPlugin {
     );
   }
 
+  private estimateInputTokens(message: string, history: ChatMessage[] = []): number {
+    const allContent =
+      message.length +
+      history.reduce((acc, m) => acc + ((m as any)?.content?.length || (m as any)?.text?.length || 0), 0);
+    return Math.ceil(allContent / 4);
+  }
+
   /**
    * Run a chat query.
    */
@@ -88,6 +96,14 @@ export class VectorPlugin {
       throw wrapError(err, 'CONFIGURATION_ERROR');
     }
     try {
+      const inputTokens = this.estimateInputTokens(message, history);
+      const tokenCheck = (FreeTierLimitsGuard as any).checkRequestAllowed({
+        operationType: 'chat',
+        inputTokens,
+      });
+      if (!tokenCheck.allowed && tokenCheck.reason?.toLowerCase().includes('token')) {
+        throw wrapError(new Error(tokenCheck.reason), 'RATE_LIMITED');
+      }
       return await this.pipeline.ask(message, history, namespace);
     } catch (err) {
       const msg = String(err);
@@ -95,6 +111,9 @@ export class VectorPlugin {
       if (msg.includes('Embed') || msg.includes('embed')) {
         defaultCode = 'EMBEDDING_FAILED';
       }
+      if (msg.toLowerCase().includes('token limit')) {
+        defaultCode = 'RATE_LIMITED';
+      }
       throw wrapError(err, defaultCode);
     }
   }
@@ -109,6 +128,14 @@ export class VectorPlugin {
       throw wrapError(err, 'CONFIGURATION_ERROR');
     }
     try {
+      const inputTokens = this.estimateInputTokens(message, history);
+      const tokenCheck = (FreeTierLimitsGuard as any).checkRequestAllowed({
+        operationType: 'chat_stream',
+        inputTokens,
+      });
+      if (!tokenCheck.allowed && tokenCheck.reason?.toLowerCase().includes('token')) {
+        throw wrapError(new Error(tokenCheck.reason), 'RATE_LIMITED');
+      }
       const stream = this.pipeline.askStream(message, history, namespace);
       for await (const chunk of stream) {
         yield chunk;
@@ -119,6 +146,9 @@ export class VectorPlugin {
       if (msg.includes('Embed') || msg.includes('embed')) {
         defaultCode = 'EMBEDDING_FAILED';
       }
+      if (msg.toLowerCase().includes('token limit')) {
+        defaultCode = 'RATE_LIMITED';
+      }
       throw wrapError(err, defaultCode);
     }
   }